[論文レビュー] Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models

Chenyang Zhang, Qingyue Zhao+2 more|arXiv (Cornell University)|Mar 24, 2026

Advanced Graph Neural Networks被引用数 0

ひとこと要約

要約の直接回答: 1層のトランスフォーマーと位置情報のみの自己注意を用いることで、勾配降下法によって教師モデルの広いクラスを証明的に学習できることを示し、教師のパラメータ収束と良好なOOD一般化を達成し、上限と下限が一致する。

ABSTRACT

Transformers have achieved great success across a wide range of applications, yet the theoretical foundations underlying their success remain largely unexplored. To demystify the strong capacities of transformers applied to versatile scenarios and tasks, we theoretically investigate utilizing transformers as students to learn from a class of teacher models. Specifically, the teacher models covered in our analysis include convolution layers with average pooling, graph convolution layers, and various classic statistical learning models, including a variant of sparse token selection models [Sanford et al., 2023, Wang et al., 2024] and group-sparse linear predictors [Zhang et al., 2025]. When learning from this class of teacher models, we prove that one-layer transformers with simplified "position-only'' attention can successfully recover all parameter blocks of the teacher models, thus achieving the optimal population loss. Building upon the efficient mimicry of trained transformers towards teacher models, we further demonstrate that they can generalize well to a broad class of out-of-distribution data under mild assumptions. The key in our analysis is to identify a fundamental bilinear structure shared by various learning tasks, which enables us to establish unified learning guarantees for these tasks when treating them as teachers for transformers.

研究の動機と目的

• 教師モデルのクラスから学習する際の理論的保証を通じたトランスフォーマーの理解を動機づける。
• CNN、GCN、スパーストークン選択、グループスパース予測子を教師モデルとして包含する統一的なバイリニア構造フレームワークを定義する。
• この設定で勾配降下法によって訓練された1層トランスフォーマーに対する収束と一般化保証を確立する。
• 合成データと実データの実験を通じて、理論が観測される訓練ダイナミクスと注意パターンと整合することを示す。

提案手法

• バイリニア構造を持つ形式 f* (X) = sigma(V* X S*) の教師モデルを定義し、さまざまな具現化（CNN、GCN、STS、GSLP）を含む。
• 位置情報のみの自己注意を用いた簡略化された1層トランスフォーマーを採用する：TF(Z; WV; WKQ) = sigma(WV X S) ここで S は学習されるアテンションスコア。
• WV と WKQ のゼロ初期化とガウス分布入力 X に対する母集団損失での勾配降下法による訓練と反復更新式（3.3)-(3.4) の導出。
• 真の成分 (V*, S*) への収束を理論的に分析し、過剰損失とパラメータ収束の厳密な境界を提供（定理3.1）」。
• 軽度のモーメント仮定の下でのOOD一般化境界を拡張（定理3.2）。
• CNN、GCN、STS、GSLP のタスクやMNISTベースの設定を含む合成データ・実データ実験で検証。

実験結果

リサーチクエスチョン

• RQ1 勾配降下法で訓練された1層トランスフォーマーは、 Bilinear 構造を持つ広範な教師モデルのすべてのパラメータブロックを回復できるのか。
• RQ2 学習されたアテンションスコアと値行列の収束速度はどの程度で、反復回数とともに過剰損失はどうスケールするのか。
• RQ3 訓練分布を超えたOODデータへ学習済みトランスフォーマーは一般化するのか。
• RQ4 CNN、GCN、スパーストークン選択、グループスパース線形予測子の教師モデルに対して理論結果はどのように現れるのか。
• RQ5 経験的実験は予測されたパラメータ・損失ダイナミクスと注意パターンを反映しているのか。

主な発見

• 1層トランスフォーマーの勾配降下法による訓練は、教師の値行列 V* とソフトマックススコア S* を正確な収束保証とともに回復できる。
• アテンションスコアは基底 S* へ収束し、収束率は ||S(T)−S*||F = Theta(D^{5/2} / (||V*|| sqrt(eta T)))。
• 値行列は基底 V* へ収束し、収束率は ||WV^(T)−V*||F = Theta(D^2 sqrt(K/(eta T)))。
• 過剰損失 L(WV^(T); WKQ^(T)) − L_opt は Theta(K D^4 /(eta T)) によって上下に境界付けられる。
• フレームワークはCNN（平均プーリング）、正則グラフ上のGCN、スパーストークン選択、グループスパース線形予測子など、多様な教師モデルを含み、合成実験で収束する損失と整合した注意パターンを示すことを理論と一致させて確認。
• OOD一般化境界は、境界付き二次モーメントの下で訓練されたトランスフォーマーのOOD損失が教師のOOD損失から epsilon 内にあり、堅牢な一般化を確立している。