[論文レビュー] In-Context Convergence of Transformers

Huang Yu, Yuan Cheng+1 more|arXiv (Cornell University)|Oct 8, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用数 3

ひとこと要約

本稿は、線形関数のコンテキスト内学習のためのソフトマックスアテンションを備えた1層Transformerにおける勾配降下法による学習ダイナミクスの最初の理論的分析を提供する。バランス型およびアンバランス型の特徴量分布の両方において有限時間収束を確立し、支配的特徴量が最初に学習され、その後に低頻度特徴量が学習される段階的学習プロセスを明らかにした。これは、競合するアテンション重みダイナミクスの新しい分析に基づく。

ABSTRACT

Transformers have recently revolutionized many domains in modern machine learning and one salient discovery is their remarkable in-context learning capability, where models can solve an unseen task by utilizing task-specific prompts without further parameters fine-tuning. This also inspired recent theoretical studies aiming to understand the in-context learning mechanism of transformers, which however focused only on linear transformers. In this work, we take the first step toward studying the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent in order to in-context learn linear function classes. We consider a structured data model, where each token is randomly sampled from a set of feature vectors in either balanced or imbalanced fashion. For data with balanced features, we establish the finite-time convergence guarantee with near-zero prediction error by navigating our analysis over two phases of the training dynamics of the attention map. More notably, for data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process, where the transformer first converges to a near-zero prediction error for the query tokens of dominant features, and then converges later to a near-zero prediction error for the query tokens of under-represented features, respectively via one and four training phases. Our proof features new techniques for analyzing the competing strengths of two types of attention weights, the change of which determines different training phases.

研究の動機と目的

• 勾配降下法による学習が行われるソフトマックスベースのTransformerのコンテキスト内学習における学習ダイナミクスを理解すること。
• 一部の特徴量が低頻度に出現するアンバランスな特徴量分布が存在する状況における収束行動を分析すること。
• 競合するアテンション重みが駆動する段階的転移によって、学習プロセスがどのように進行するかを特定すること。
• バランス型およびアンバランス型の両設定において、有限時間収束の保証を確立すること。

提案手法

• K個の特徴量ベクトルから確率{pk}k=1Kでトークンを抽出する構造化データモデル上で、ソフトマックスアテンションを備えた1層Transformerを勾配降下法で学習する。
• コンテキスト内学習を、入力-ラベルペアを含むプロンプトシーケンスとクエリトークンを用いた線形回帰タスクとしてモデル化する。
• バランス型特徴量に対しては、2段階の解析を導入：急速なアテンション整合化の後、損失最小化が続く。
• 「目的」アテンション重みと「非目的」アテンション重みの間の相互作用を追跡する新しい証明技法を提案し、明確に区別された学習段階を特定する。
• 集中不等式と再帰的バインドを用いて、勾配降下法下でのアテンションマップの進化を制御する。
• 確率的議論を用いて、多項式(K)のサンプルサイズおよび誤差許容度のlog(1/ε)依存性において、高確率での収束を示す。

実験結果

リサーチクエスチョン

• RQ1 ソフトマックスベースのTransformerが勾配降下法の学習ダイナミクスを経てコンテキスト内学習を行う際、どのように収束するのか？
• RQ2 目的特徴量と非目的特徴量の間のアテンション重みの競合が、学習段階をどのように形作るのか？
• RQ3 なぜアンバランスなデータ分布においても、低頻度特徴量に対してコンテキスト内学習が成功するのか？
• RQ4 バランス型およびアンバランス型の両特徴量分布において、有限時間収束を保証できるのか？
• RQ5 特徴量がアンバランスな場合、学習ダイナミクスにおける明確な学習段階は何か？

主な発見

• バランス型特徴量では、2段階に分けた収束が達成される：急速なアテンション整合化の後、損失最小化が続く。予測誤差はほぼゼロに近づく。
• アンバランス型特徴量では、段階的収束が観察される。支配的特徴量が速やかに学習され、その後に4つの明確な学習段階を経て低頻度特徴量が学習される。
• 支配的特徴量の予測誤差は、勾配降下ステップ数O(log(ε^(-1/2)) / (ηε))で収束し、誤差はεで抑えられる。
• 適切な特徴量のアテンションスコアは、十分な反復後、高確率1 - e^(-Ω(poly(K)))でほぼ1に集中する。
• 解析により、時間経過に伴い「目的」アテンション重みと「非目的」アテンション重みの優位性が変化し、これが学習段階を定義することが明らかになった。
• 競合するアテンション重みダイナミクスに基づくこの証明技法は一般化可能であり、他のTransformer関連問題への応用も可能である。