[論文レビュー] Uncovering mesa-optimization algorithms in Transformers

Johannes von Oswald, Eyvind Niklasson+10 more|arXiv (Cornell University)|Sep 11, 2023

Topic Modeling被引用数 5

ひとこと要約

この論文は、自己回帰的に訓練されたトランスフォーマーが推論中に勾配ベースのメサ最適化を暗黙的に行っていることを示しており、内部の目的関数がコンテキストから構築され、予測を生成するために最適化されている。著者らはこのプロセスを逆設計し、コンテキスト内で最適化問題を明示的に解く新しいメサレイヤーを導入し、few-shot なコンテキスト内学習および系列モデリングタスクにおける性能向上を示した。

ABSTRACT

Transformers have become the dominant model in deep learning, but the reason for their superior performance is poorly understood. Here, we hypothesize that the strong performance of Transformers stems from an architectural bias towards mesa-optimization, a learned process running within the forward pass of a model consisting of the following two steps: (i) the construction of an internal learning objective, and (ii) its corresponding solution found through optimization. To test this hypothesis, we reverse-engineer a series of autoregressive Transformers trained on simple sequence modeling tasks, uncovering underlying gradient-based mesa-optimization algorithms driving the generation of predictions. Moreover, we show that the learned forward-pass optimization algorithm can be immediately repurposed to solve supervised few-shot tasks, suggesting that mesa-optimization might underlie the in-context learning capabilities of large language models. Finally, we propose a novel self-attention layer, the mesa-layer, that explicitly and efficiently solves optimization problems specified in context. We find that this layer can lead to improved performance in synthetic and preliminary language modeling experiments, adding weight to our hypothesis that mesa-optimization is an important operation hidden within the weights of trained Transformers.

研究の動機と目的

• 自己回帰的に訓練されたトランスフォーマーが、少数ショット学習のためのメタ学習を受けていないにもかかわらず、フォワードパス内に内部的な最適化プロセス（メサ最適化）を実装しているかどうかを調査すること。
• 単純な系列モデリングタスクを逆設計し、自己注意機構がどのように内部の訓練データを構築し、目的関数を最適化しているかを解明すること。
• このようなメサ最適化が、モデルサイズに依存せず、標準的な自己回帰的トランスフォーマーにおいてコンテキスト内学習能力を可能にしているかどうかを評価すること。
• 内部の最小二乗最適化問題をコンテキスト内で明示的かつ効率的に解く新しい自己注意メカニズム、すなわちメサレイヤーを設計・評価すること。
• 合成的および初期の言語モデリングタスクにおいてメサレイヤーの性能をテストし、コンテキスト内学習の向上可能性を評価すること。

提案手法

• 単純な系列タスクで訓練された自己回帰的トランスフォーマーを逆設計し、内部的に訓練データを構築し、目的関数を最適化するメカニズムを同定すること。
• 初期の自己注意レイヤーが、入力トークンをグループ化・複製することで内部データセットを形成し、暗黙的に目的関数を定義していることを同定すること。
• より深いレイヤーが、これらの内部目的関数を勾配ベースで最適化して予測を生成していることを確認すること。
• 最小二乗最適化問題を閉形式解により直接解くことで勾配ステップを避ける新しい自己注意メカニズム、すなわちメサレイヤーを提案すること。
• メサレイヤーの訓練を安定化させるために、キーとクエリの正規化を導入し、収束性と性能を向上させること。
• GPT-2風のアーキテクチャにおいて標準の自己注意レイヤーをメサレイヤーに置き換え、The Pile データセットを用いて標準的な自己回帰的言語モデリング設定で評価すること。

実験結果

リサーチクエスチョン

• RQ1 自己回帰的に訓練されたトランスフォーマーは、メタ学習を受けていなくても、推論中に勾配ベースのメサ最適化を暗黙的に実装しているか？
• RQ2 トランスフォーマー内の最適化プロセスは、コンテキストからどのように内部目的関数を構築し、解いているかを逆設計で解明できるか？
• RQ3 標準的な自己回帰的トランスフォーマーが、メサ最適化のおかげでどの程度、コンテキスト内少数ショット学習能力を示すか？
• RQ4 メサレイヤーのような新しい自己注意メカニズムは、コンテキスト内で最適化問題を明示的に解くことで、標準的な自己注意よりも系列モデリングで優れた性能を発揮できるか？
• RQ5 標準的な注意機構をメサレイヤーに置き換えることで、言語モデリングタスクにおけるコンテキスト内学習性能が向上するか？

主な発見

• 自己回帰的に訓練されたトランスフォーマーは、二段階のメサ最適化プロセスを実装している：(i) 注意メカニズムによる内部訓練データの構築、および (ii) 予測生成のための内部目的関数の勾配ベース最適化。
• 訓練済みトランスフォーマーのフォワードパスには、文脈から構築された目的関数を対象に勾配降下アルゴリズムが暗黙的に存在し、標準的な大規模言語モデルにおいてもメサ最適化が実現していることを示している。
• 標準的な自己回帰損失で訓練された、小型で単純なトランスフォーマーでさえ、モデルサイズに依存せず、強力なコンテキスト内少数ショット学習能力を示している。
• プロンプトチューニングは、これらのモデルにおけるコンテキスト内学習性能を向上させており、大規模言語モデルにおけるその役割を模倣している。
• 最小二乗最適化問題を直接解くことで、標準のソフトマックスおよび線形自己注意トランスフォーマーを上回る性能を示した、提案されたメサレイヤーは合成的系列タスクで優れた性能を発揮した。
• 初期の言語モデリング実験では、標準的な注意機構をメサレイヤーに置き換えることで、The Pile データセットにおける性能向上と強力なコンテキスト内学習能力が得られた。