[論文レビュー] Transformers learn to implement preconditioned gradient descent for in-context learning

Kwangjun Ahn, Xiang Cheng+2 more|arXiv (Cornell University)|Jun 1, 2023

Stochastic Gradient Optimization Techniques被引用数 6

ひとこと要約

この論文は、ランダムな線形回帰インスタンスで訓練されたトランスフォーマーが、損失関数の局所的構造を分析することで、データ分布やデータ不足に起因する分散に適応するプリコンディショニング付き勾配降下法を学習することを示している。1層のアテンション層では、グローバル・オプティマルが、データ分布とデータ不足に起因する分散に適応する1ステップのプリコンディショニング付き勾配降下に対応する。より深いトランスフォーマーでは複数ステップの反復が実装され、臨界点はGD++を含む適応的最適化アルゴリズムと一致する。

ABSTRACT

Several recent works demonstrate that transformers can implement algorithms like gradient descent. By a careful construction of weights, these works show that multiple layers of transformers are expressive enough to simulate iterations of gradient descent. Going beyond the question of expressivity, we ask: Can transformers learn to implement such algorithms by training over random problem instances? To our knowledge, we make the first theoretical progress on this question via an analysis of the loss landscape for linear transformers trained over random instances of linear regression. For a single attention layer, we prove the global minimum of the training objective implements a single iteration of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the input distribution but also to the variance induced by data inadequacy. For a transformer with $L$ attention layers, we prove certain critical points of the training objective implement $L$ iterations of preconditioned gradient descent. Our results call for future theoretical studies on learning algorithms by training transformers.

研究の動機と目的

• ランダムな問題インスタンスへの訓練を通じて、トランスフォーマーが勾配ベースの最適化アルゴリズムを学習できるかどうかを調査すること。
• ランダムな線形回帰インスタンスで訓練された線形トランスフォーマーの損失関数の局所的構造を分析し、最適化アルゴリズムが訓練からどのように生じるかを理解すること。
• パラメータ空間におけるグローバル・ミニマおよび臨界点の構造を特徴づけ、それらがどの最適化アルゴリズムを実装しているかを同定すること。
• トランスフォーマーの理論的表現力と、イン・コンテキスト学習における実際の挙動のギャップを、とくに勾配ベースの手法に関して埋める。
• 理論的発見を実験的に検証し、学習された臨界点がGD++のような既知の適応的最適化アルゴリズムと一致することを示すこと。

提案手法

• 非ソフトマックスアテンション機構を用いて、ランダムな等方的線形回帰インスタンスで訓練された1層線形トランスフォーマーの損失関数の局所的構造を分析する。
• 訓練目的関数のグローバル・ミニマが、データ分布およびデータ不足に起因する分散に適応するプリコンディショニング付き勾配降下法の1ステップに対応することを証明する。
• パラメータ空間にスパarsity条件を導入し、kステップの適応的勾配ベース最適化アルゴリズムに制限された探索を可能にし、複数層トランスフォーマーにおける臨界点の特徴づけを可能にする。
• より深いトランスフォーマー（L層）において、特定の臨界点がデータに依存するプリコンディショニング付きのLステップのプリコンディショニング付き勾配降下法を実装することを示す。
• スパarsity条件を緩和して全パラメータ空間を解析し、勾配降下ステップと線形変換を組み合わせて条件数をさらに改善する、新しい勾配ベースのアルゴリズムを実装する臨界点を同定する。
• 理論的発見を実験的に検証するため、学習された重みを可視化し、トランスフォーマー予測子のテスト損失を標準的な最適化ベースライン（GD、プリコンディショニング付きGD、OLS）と比較する。

(a) $\operatorname{Dist}(\Sigma^{1/2}A_{0}\Sigma^{1/2},I)$

実験結果

リサーチクエスチョン

• RQ1 ランダムな線形回帰インスタンスで訓練されたトランスフォーマーは、非凸最適化を通じて勾配降下法を学習できるか？
• RQ2 ランダムな回帰問題で訓練された1層線形トランスフォーマーのグローバル・ミニマが実装する最適化アルゴリズムは何か？
• RQ3 より深いトランスフォーマー（L層）における臨界点は、反復的最適化アルゴリズムとどのように対応するか？
• RQ4 パラメータのスパarsity制約を解除すると、学習されたアルゴリズムはどのように変化するか？
• RQ5 理論的に導出された臨界点は、訓練済みトランスフォーマーで観察される実際の挙動と一致するか？

主な発見

• 1層線形トランスフォーマーのグローバル・ミニマは、入力データの分布およびデータ不足に起因する分散に適応するプリコンディショニング付き勾配降下法の1ステップを実装する。
• スパarsity条件を課した2層トランスフォーマーにおいて、グローバル・ミニマは適応的ステップサイズを有する勾配降下に対応し、効果的にkステップの適応的最適化プロセスを実装する。
• より深いトランスフォーマー（L層）では、訓練目的関数の特定の臨界点が、データに依存するプリコンディショニング付きのLステップの勾配降下法を実装する。
• スパarsity制約を解除した場合、ある臨界点は勾配降下法と線形変換を組み合わせて条件数をさらに改善する、新しいアルゴリズムに対応する。データ共分散が等方的である場合、これはGD++アルゴリズムと一致する。
• 実験的検証により、3層トランスフォーマーの学習された重みが理論的停留点と一致し、その点での目的関数値は0に近く、グローバル・オプティマルであることが示唆される。
• テスト損失の比較により、トランスフォーマーが学習した予測子は3ステップのプリコンディショニング付き勾配降下法と同等の性能を示し、理論的整合性と実際の最適化挙動の一致を検証した。

(b) $\operatorname{Dist}(\Sigma^{1/2}A_{1}\Sigma^{1/2},I)$