Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures

要約

トランスフォーマー、畳み込みニューラル ネットワーク、グラフ ニューラル ネットワークなど、多くの最新のニューラル ネットワーク アーキテクチャのコア コンポーネントは、$\textit{weight-sharing}$ を使用して線形層として表現できます。
二次最適化手法であるクロネッカー因子近似曲率 (K-FAC) は、ニューラル ネットワークのトレーニングを高速化し、それによって計算コストを削減することが期待されています。
ただし、現時点では、一般的なアーキテクチャ、特に線形重み共有層を備えたアーキテクチャに適用するフレームワークはありません。
この研究では、K-FAC の 2 つの種類、$\textit{expand}$ と $\textit{reduce}$ を動機付ける線形重み共有層の 2 つの異なる設定を特定します。
それぞれの設定で重みを共有する深層線形ネットワークに対して、それらが正確であることを示します。
特に、K-FAC-reduce は一般に K-FAC-expand より高速であり、これを利用して Wide ResNet の限界尤度を最適化することで自動ハイパーパラメータ選択を高速化します。
最後に、グラフ ニューラル ネットワークとビジョン トランスフォーマーの両方をトレーニングするためにこれら 2 つの K-FAC バリエーションを使用した場合、これら 2 つの K-FAC バリエーションにはほとんど違いがないことがわかります。
ただし、どちらのバリエーションも、一次リファレンス実行のステップ数の $50$ ～ $75\%$ で固定の検証メトリクス目標に到達することができ、これは実時間の同等の改善につながります。
これは、K-FAC を最新のニューラル ネットワーク アーキテクチャに適用できる可能性を強調しています。

要約(オリジナル)

The core components of many modern neural network architectures, such as transformers, convolutional, or graph neural networks, can be expressed as linear layers with $\textit{weight-sharing}$. Kronecker-Factored Approximate Curvature (K-FAC), a second-order optimisation method, has shown promise to speed up neural network training and thereby reduce computational costs. However, there is currently no framework to apply it to generic architectures, specifically ones with linear weight-sharing layers. In this work, we identify two different settings of linear weight-sharing layers which motivate two flavours of K-FAC — $\textit{expand}$ and $\textit{reduce}$. We show that they are exact for deep linear networks with weight-sharing in their respective setting. Notably, K-FAC-reduce is generally faster than K-FAC-expand, which we leverage to speed up automatic hyperparameter selection via optimising the marginal likelihood for a Wide ResNet. Finally, we observe little difference between these two K-FAC variations when using them to train both a graph neural network and a vision transformer. However, both variations are able to reach a fixed validation metric target in $50$-$75\%$ of the number of steps of a first-order reference run, which translates into a comparable improvement in wall-clock time. This highlights the potential of applying K-FAC to modern neural network architectures.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google