[論文レビュー] Transformers learn through gradual rank increase

Enric Boix-Adserà, Etai Littwin+3 more|arXiv (Cornell University)|Jun 12, 2023

Neural Networks and Applications被引用数 5

ひとこと要約

この論文は、視覚および言語トランスフォーマーが訓練中に段階的学習ダイナミクスを示すことを示しており、訓練済み重みと初期重みの差のランクが段階的に徐々に増加することを明らかにしている。理論的には、初期化が小さいおよび重み行列が対角であるという仮定のもとで、勾配フロー訓練は、各段階でランクが最大で1ずつ増加する。実験的にも、この現象は簡略化仮定がなくても成立し、実世界の視覚および言語モデルで観察された。

ABSTRACT

We identify incremental learning dynamics in transformers, where the difference between trained and initial weights progressively increases in rank. We rigorously prove this occurs under the simplifying assumptions of diagonal weight matrices and small initialization. Our experiments support the theory and also show that phenomenon can occur in practice without the simplifying assumptions.

研究の動機と目的

• トランスフォーマーが単純なニューラルネットワークで観察されるような段階的学習ダイナミクスを示すかどうかを調査すること。
• 初期化が小さく、重み行列が対角であるという簡略化仮定のもとで、トランスフォーマーにおける段階的学習の理論的基盤を確立すること。
• 理論的結果を、簡略化仮定を満たさない現実的設定で実験的に検証すること。
• 観察されたダイナミクスと、LoRAのような効率的な微調整手法との関係を調査すること。

提案手法

• 理論的分析では、対角的注意ヘッド重みをもつ簡略化されたトランスフォーマー変種を用い、勾配フローを適用することで、2つの重みベクトルの要素ごとの積に依存するネットワークに問題を還元する。
• 関数 $ f_{\mathsf{NN}}(\mathbf{x};\mathbf{u},\mathbf{v}) = h(\mathbf{x}; \mathbf{u} \odot \mathbf{v}) $ の形をとるネットワークに対して一般定理を証明し、初期化が小さい場合、訓練が離散的な段階に分けられ、各段階で $ \mathbf{u} \odot \mathbf{v} $ に最大で1つの新たな顕著な要素が追加されることを示す。
• 理論的結果を注意ヘッドに適用するため、積 $ \mathbf{W}_K \mathbf{W}_Q^\top $ を対角行列の積として表現し、有効重み行列のランク増加を分析可能にする。
• 実験では、初期化時と訓練後の $ \mathbf{W}_K \mathbf{W}_Q^\top $ のスペクトルを比較し、重み差のランクを測定することで段階的変化を追跡する。
• CIFAR-10およびImageNetで訓練されたViTモデルを用い、対角性や初期化の小ささを仮定しない注意ヘッド重みを分析する。
• 訓練ステップごとの重み差の正規化スペクトルの可視化を含み、段階的ランク増加を示す。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーは、学習関数の複雑さが段階的に増加するような段階的学習ダイナミクスを示すか？
• RQ2 初期化が小さく、重み行列が対角であるという仮定のもとで、トランスフォーマーにおける段階的学習ダイナミクスを厳密に証明できるか？
• RQ3 理論的仮定が満たされない状況でも、重み差のランク増加が段階的に継続するか？
• RQ4 観察された段階的ダイナミクスと、LoRAのような効率的な微調整手法との関係は何か？

主な発見

• 初期化が小さく、重み行列が対角であるという仮定のもとで、トランスフォーマーにおける勾配フロー訓練は離散的な段階に分けられ、各段階は $ \Theta(\log(1/\alpha)) $ の時間にわたり、有効重みベクトル $ \mathbf{u} \odot \mathbf{v} $ に最大で1つの新たな顕著な要素が追加される。
• 理論的分析により、鞍点からの脱出ダイナミクスに起因し、訓練済み重みと初期重みの差のランクが、各段階で最大で1ずつ段階的に増加することが証明された。
• ImageNetで訓練されたViTモデルの実験では、訓練に伴い、訓練済みと初期の $ \mathbf{W}_K \mathbf{W}_Q^\top $ 行列の差のランクが段階的に増加しており、スペクトルプロットから明確な段階的挙動が観察された。
• 理論的仮定（非対角的重み、大きな初期化）が満たされない標準的な訓練でも、視覚および言語モデルの両方で段階的ランク増加が依然として観察された。
• CIFAR-10よりもImageNetのような大規模データセットの方が、この現象が顕著に現れ、データの複雑さが段階的ダイナミクスに影響を与える可能性を示唆している。
• 結果からLoRA微調整との関連性が示唆され、両者とも低ランク更新を伴うため、観察された段階的ダイナミクスがこのような手法の有効性を説明する可能性がある。