Training Transformers as a Universal Computer

以下は、論文「Training Transformers as a Universal Computer」の詳細な技術的サマリーです。

1. 問題定義

チェーン・オブ・トークン（CoT）を拡張したトランスフォーマーが計算的に万能（チューリング機械をシミュレート可能）であるという理論的研究はあるものの、表現力（理論的に表現可能なもの）と学習性（勾配ベースの最適化を通じて実際に学習できるもの）の間には大きな隔たりがあります。

現在の課題は以下の通りです：

• 汎化の失敗: モデルは長さ汎化（例：訓練で見たことのない長い数の加算）や構成的汎化（プリミティブの新しい組み合わせを解くこと）に苦戦しています。
• コンテキストのボトルネック: 標準的な CoT は、タスクの時間計算量 $T$ に比例する $O(T)$ のコンテキスト長を必要とします。これにより、限定されたコンテキストウィンドウ内での、長く複雑なプログラムの訓練および推論が非現実的になります。
• 万能な学習の欠如: 以前のニューラルインタプリタの訓練試行は、特定の事前設計された関数や専門的なアーキテクチャに依存することが多く、標準的なトランスフォーマーが一般的でチューリング完全な言語を解釈することを学習できることを示すには至っていませんでした。

ここで扱われる中心的な問いは、標準的なトランスフォーマーを訓練して万能なコンピュータとして機能させ、未見のプログラムや長さに汎化させることは可能か、というものです。

2. 手法

著者らは、新しいプログラミング言語、メモリ効率の良い実行スケッフォールド、および特定の訓練レジームを組み合わせたフレームワークを提案します。

A. MicroPy: 計算可能な万能言語

著者らは、簡素化されつつもチューリング完全なオブジェクト指向プログラミング言語であるMicroPyを設計しました。

• 意味論: これは、検索（retrieval）と書き換え（rewriting）のルールという、小さく有限なセットに依存しています。
  • 検索: 文脈（スタックフレームまたは状態）から関数定義、変数バインディング、または属性値を検索すること。
  • 書き換え: 現在の状態に基づき、プリミティブ（例：IF、SEQ、LookupAttr）をより単純な式に還元すること。
• 構造: プログラムは手続き定義と式から構成されます。実行には、式をステップバイステップで還元し、可変状態（オブジェクトの属性）とスタックフレーム（局所変数のため）を管理するプロセスが含まれます。
• 目的: これらの有限なルールを学習することで、モデルは任意の MicroPy プログラムを解釈することを学習し、実質的に万能なインタプリタとして機能します。

B. PENCIL スケッフォールド：限定コンテキスト実行

コンテキスト長のボトルネック（$O(T)$）に対処するため、著者らはPENCIL（Yang et al., 2025 による）を利用します。

• メカニズム: PENCIL は、完了した中間計算を回収する還元ルールを導入します。サブ計算（スタックフレーム）が完了すると、中間の推論ステップは消去され、最終結果のみが保持されます。
• 計算量: これにより、コンテキスト長は時間計算量 $O(T)$ ではなく、空間計算量 $O(S)$ によって制限されます。
• アナロジー: これはプログラミング言語におけるスタックの規律のように機能します。関数が戻ると、その局所的なスタックフレームがポップされ、長時間実行されるプログラムであってもコンテキストウィンドウのサイズを管理可能な範囲に保ちます。

C. 訓練データ生成

モデルは、ランダムに生成された MicroPy プログラムから生成された実行トレースを用いて訓練されます。

• 2 つのサンプラー:
  1. プログラムレベルサンプラー: 確率的文脈自由文法を通じて、多様な構文構造を生成します。
  2. プランサンプラー: 目標とする「プラン」（スタック構造）をサンプリングし、それを実現するプログラムをコンパイルすることで、稀なランタイム構成（特定のスタックスケルトンと検索パターン）のカバレッジを確保します。
• 監督: 訓練目的は、ステップごとの実行トレースにおける次のトークン予測です。モデルは、現在のコンテキストに基づいて次の書き換えステップ（検索または局所還元）を予測します。
• 制約: 訓練プログラムは多様性を確保しつつモデルを圧倒しないよう、トレース行数を 128 行に制限しますが、評価でははるかに長いプログラムでテストされます。

D. モデルアーキテクチャ

• アーキテクチャ: RoPE（回転位置埋め込み）を備えたデコーダ専用トランスフォーマー（5950 万パラメータ）。
• 最適化: 埋め込みには AdamW を、2 次元行列パラメータには Muon を使用します。
• 入力: 手続き定義 + 評価する式 + 現在の状態/スタック。
• 出力: 実行トレースの次のステップ（トークンごと）。

3. 主要な貢献

1. MicroPy 言語設計: 検索と書き換えルールの有限セットに基づく意味論を持つチューリング完全言語の作成。これにより、構成的学習が促進されます。
2. 万能学習の実証的証拠: 万能言語のランダムなトレースのみで訓練された標準的なトランスフォーマーが、一度も見たことのない複雑な人間作成のプログラムを実行することに汎化できることを初めて実証しました。
3. PENCIL の統合: PENCIL スケッフォールドの成功した適用により、固定されたコンテキストウィンドウ内で、訓練時に見たものよりも最大60 倍長いトレース長を持つプログラムをトランスフォーマーが実行できるようにしました。
4. 完全な汎化: ビット操作、算術、論理解決に関する保持されたタスクにおいて 100% の精度を達成し、強力な構成的および長さ汎化を実証しました。

4. 実験結果

モデルは、訓練データよりも実質的に長く、意味的に有意義な保持された人間作成の MicroPy プログラムで評価されました。

• 評価されたタスク:
  • ビット操作: リンクドリスト内のビットのコピーと反転。
  • 算術: ビット文字列表現を使用した二進加算と乗算。
  • 論理: SAT 検証と SAT 解決（バックトラッキング探索を使用）。
• 性能:
  • 精度: 試験されたコンテキストウィンドウ内のすべてのタスクで、モデルは100% のトークンレベル精度を達成しました。
  • 長さ汎化: 訓練プログラムは 128 行に制限されていましたが、モデルは最大7,552 行のトレース長を持つプログラムを正常に実行しました。
  • 分布のギャップ: 訓練プログラムの分布とテストスイートの分布の間にはほぼゼロの重なりがあり、モデルが特定のプログラムを暗記したのではなく、実行のルールを学習したことを確認しました。
• 演算カバレッジ: 訓練ジェネレータは、テストスイートで使用されたすべての演算タイプ（例：If、LookupVar、TailApp）を網羅しており、モデルが訓練中にすべての必要なプリミティブを目にしていたことを保証しました。

5. 意義と含意

• 万能計算は学習可能: この論文は、計算的万能性がトランスフォーマーの理論的な性質に留まらず、適切な言語における標準的な次のトークン予測訓練を通じて獲得可能な能力であることを示す強力な実証的証拠を提供します。
• 構成的性への帰納的バイアス: 結果は、PENCIL などのコンテキスト管理スケッフォールドと組み合わされたトランスフォーマーが、句構造と構成的意味論に適した帰納的バイアスを有することを示唆しています。これはプログラミング言語と自然言語の両方の構造を反映しています。
• ギャップの埋め合わせ: この研究は、I/O トレースを学習する「ニューラルコンピュータ」と形式的インタプリタの間のギャップを埋め、単一のモデルが形式的でチューリング完全な言語を実行することを学習できることを示しています。
• 将来の方向性: 著者らは、プログラミング言語インタプリタの学習を可能にする同じメカニズムが、トランスフォーマーが自然言語タスクにおいてこれほど効果的である理由を説明する可能性があると示唆しています。両者とも構成的意味論に依存しているためです。今後の研究では、これをより豊かな言語や、明示的な実行トレースなしの推論へと拡張することが考えられます。

要約すると、この論文は、実行が効率的にコンテキストを管理するようにスケッフォールドされている場合、標準的なトランスフォーマーを訓練して、チューリング完全な言語の任意のプログラムを完全な精度で実行可能な万能コンピュータとして機能させることができることを実証しています。