Transformerにおけるメサ最適化アルゴリズムの発見

要旨

Transformerは深層学習において支配的なモデルとなっているが、その優れた性能の理由は十分に理解されていない。本研究では、Transformerの強力な性能は、メサ最適化（モデルの順伝播中に実行される学習プロセス）に対するアーキテクチャ的なバイアスに起因すると仮説を立てる。このプロセスは以下の2つのステップから構成される：(i) 内部的な学習目標の構築、(ii) 最適化を通じたその対応する解の発見。この仮説を検証するため、単純な系列モデリングタスクで訓練された一連の自己回帰型Transformerを逆解析し、予測生成を駆動する基盤となる勾配ベースのメサ最適化アルゴリズムを明らかにした。さらに、学習された順伝播最適化アルゴリズムが、教師ありfew-shotタスクを解くために即座に再利用可能であることを示し、メサ最適化が大規模言語モデルの文脈内学習能力の基盤となっている可能性を示唆する。最後に、文脈内で指定された最適化問題を明示的かつ効率的に解く新しいセルフアテンション層、メサ層を提案する。この層が合成データおよび予備的な言語モデリング実験において性能向上をもたらすことを確認し、メサ最適化が訓練済みTransformerの重み内に隠された重要な操作であるという仮説に重みを加える。

English

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.

Transformerにおけるメサ最適化アルゴリズムの発見

Uncovering mesa-optimization algorithms in Transformers

要旨

Support