[論文レビュー] Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection

Yu Bai, Fan Chen+3 more|arXiv (Cornell University)|Jun 7, 2023

Adversarial Robustness in Machine Learning被引用数 9

ひとこと要約

要約: 本論文は、トランスフォーマーが広範な文脈内学習アルゴリズム（例：最小二乗法、リッジ、ラッソ、GLM）を実装し、異なるタスクに適応するための文脈内アルゴリズム選択を実行できることを、理論的保証と実証的検証とともに示す総合的な理論を構築する。

ABSTRACT

Neural sequence models based on the transformer architecture have demonstrated remarkable \emph{in-context learning} (ICL) abilities, where they can perform new tasks when prompted with training and test examples, without any parameter update to the model. This work first provides a comprehensive statistical theory for transformers to perform ICL. Concretely, we show that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, learning generalized linear models, and gradient descent on two-layer neural networks, with near-optimal predictive power on various in-context data distributions. Using an efficient implementation of in-context gradient descent as the underlying mechanism, our transformer constructions admit mild size bounds, and can be learned with polynomially many pretraining sequences. Building on these ``base'' ICL algorithms, intriguingly, we show that transformers can implement more complex ICL procedures involving \emph{in-context algorithm selection}, akin to what a statistician can do in real life -- A \emph{single} transformer can adaptively select different base ICL algorithms -- or even perform qualitatively different tasks -- on different input sequences, without any explicit prompting of the right algorithm or task. We both establish this in theory by explicit constructions, and also observe this phenomenon experimentally. In theory, we construct two general mechanisms for algorithm selection with concrete examples: pre-ICL testing, and post-ICL validation. As an example, we use the post-ICL validation mechanism to construct a transformer that can perform nearly Bayes-optimal ICL on a challenging task -- noisy linear models with mixed noise levels. Experimentally, we demonstrate the strong in-context algorithm selection capabilities of standard transformer architectures.

研究の動機と目的

• transformers がパラメータ更新なしでインコンテキスト実装できる標準MLアルゴリズムを示す。
• Expressive power、インコンテキスト予測、事前学習サンプル効率のエンドツーエンド理論を提供する。
• インコンテキストアルゴリズム選択の仕組みを導入・分析する（事前ICLテストと事後ICL検証）。
• 単一のトランスフォーマーがタスクとデータ分布間でベースICLアルゴリズムを適応的に選択できることを示す。

提案手法

• インコンテキストリッジ回帰と最小二乗法のトランスフォーマーベース実装を構築する。
• インコンテキストGLMと凸リスク最小化へ拡張する。
• トランスフォーマー内の効率的なインコンテキスト勾配降下メカニズムを開発する。
• 事後ICL検証と事前ICLテストという2つのアルゴリズム選択メカニズムを証明する。
• 多項式個数のシーケンスから複数のICLタスクを学習できる事前学習結果を提供する。
• ノイズのある設定でベイズ最適ICLに近づくことができることを理論的・実証的に検証する。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーは文脈内で標準MLアルゴリズムの広範なクラス（例：最小二乗法、リッジ、ラッソ、GLM）を実装できるか。
• RQ2 単一のトランスフォーマーがシーケンス全体で異なる基盤ICLアルゴリズムを選択する仕組みは何か。
• RQ3 正確な文脈内学習のために必要な層・ヘッドの複雑さと重みノルムをどう境界付けるか。
• RQ4 多様なタスクで訓練されたトランスフォーマーは困難なノイズ環境でほぼベイズ最適ICLを示すか。
• RQ5 事前学習はトランスフォーマーの文脈内学習とアルゴリズム選択能力にどのように影響するか。

主な発見

• トランスフォーマーは文脈内リッジ回帰と最小二乗法を誤差保証とともに近似可能。
• トランスフォーマーはGLMの凸リスク最小化と2層ネットワークの文脈内勾配降下を実装できる。
• 事前ICLテストと事後ICL検証という2つの一般的なアルゴリズム選択メカニズムが可能で、タスク/アルゴリズムの適応的選択を実現する。
• ノイズのある線形モデルで混合ノイズレベルを用いた事後ICL検証によりベイズ最適ICLに近づく単一トランスフォーマーを達成できる。
• 多項式的に多くの訓練シーケンスから複数のICLタスクを学習できる事前学習をサポート。
• 実験的結果は代表的なタスクでの強力な文脈内アルゴリズム選択能力を裏付ける。