論文「To See the Unseen: on the Generalization Ability of Transformers in Symbolic Reasoning」の技術的サマリー

本論文は、デコーダ型トランスフォーマーモデルが、トレーニング中に観測されたことのないトークン（記号）を含む文脈内推論（In-Context Learning）において、抽象的な記号推論を行う能力を調査したものです。特に、論理推論タスクにおける「未観測トークンへの一般化」の失敗要因と、その解決策を理論的・実験的に明らかにしています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳述します。

1. 問題定義

現代の大規模言語モデル（LLM）は、数学や科学などの形式的推論タスクで高い性能を示しますが、入力トークンの名前や数値を変更するなどの単純な変化に対して脆弱であることが知られています。
本研究は、**「トレーニングデータに含まれない任意の新しいトークンを記号として扱って、文脈内のアルゴリズムを実行できるか」**という点に焦点を当てています。

• タスク: 命題論理（Propositional Logic）の問題解決。ルール（例：「A なら B」）と事実（例：「A」）が与えられ、クエリ（例：「B か？」）の真偽を推論する。
• 課題: 既存の研究（Boix-Adserà et al., 2024 など）では、モデルが未観測トークンのコピー（生成）に失敗することが示されました。しかし、その背後にあるメカニズム（なぜコピーが難しいのか）は完全には解明されていませんでした。

2. 主要な発見とメカニズム：(un)embedding の崩壊（Collapse）

本研究の核心的な発見は、**「未観測トークンの埋め込み（embedding）とアン埋め込み（unembedding）が、トレーニング過程でほぼ同じベクトルに収束（崩壊）する」**という現象です。

• 現象: トレーニングが進むと、モデルが観測したことのないトークンの重みベクトルが、互いに非常に類似した方向に収束します。
• 理論的根拠: L2 正則化（重み減衰）付きの勾配降下法（SGD）において、ラベルがデータに存在しない場合、そのラベルに対応する重みベクトル間の距離が減少することを証明しました（Lemma 4.1）。LayerNorm の存在により入力が有界になることが、この崩壊を加速させます。
• 影響: 埋め込みとアン埋め込みが共有されている場合、異なる未観測変数（例：U1 と U2）が区別できなくなり、モデルはそれらを同一視してしまいます。これが、未観測トークンのコピーや生成を困難にしている主要なメカニズムです。

3. 手法と提案ソリューション

この「埋め込みの崩壊」を克服し、未観測トークンへの一般化を実現するために、以下の 3 つの要素を組み合わせたアプローチを提案・検証しました。

1. コピー対応アーキテクチャの導入（Copy Attention）:
   • 既存の Boix-Adserà et al. (2024) の提案を改良し、入力トークンの埋め込みを直接出力にコピーするための「コピーアテンション」ヘッドを追加しました。これにより、MLP を経由せずに入力トークンを正確にコピーする経路を確保します。
2. データの多様性（Data Diversity）:
   • トレーニングデータにおける記号（述語）の種類の多さを増やすことで、モデルが特定のトークンに依存しない構造を学習できるようにしました。
3. 埋め込みの凍結またはリセット（Freezing / Active Forgetting）:
   • 凍結: トレーニング中に埋め込み行列を固定し、初期状態（ほぼ直交しており区別しやすい状態）を維持します。
   • アクティブ・フォージティング（Active Forgetting）: 定期的に（例：100 ステップごと）埋め込みを再初期化することで、崩壊を防ぎます。

4. 実験結果

合成データ（命題論理タスク）と大規模オープンウェイトモデル（Gemma 3）を用いた実験で以下の結果が得られました。

合成データ実験

• 標準トランスフォーマー: 未観測トークン（クエリのみ、またはすべての述語）が含まれる場合、精度は著しく低下します。モデルは未観測トークンを生成せず、既知のトークンで置き換えるか無視します。
• コピーアテンションのみの効果: 単一の未観測トークンの場合、多様性を高めればある程度機能しますが、すべての変数が未観測の場合は失敗します（埋め込み崩壊のため）。
• 凍結・リセットの併用:
  • 埋め込みを凍結した場合、コピーアテンションの有無にかかわらず、高い多様性があれば未観測トークンへの一般化が成功しました。
  • アクティブ・フォージティング（トレーニング初期または定期的にリセット）も、コピーアテンションと組み合わせることで、未観測トークンへの一般化を可能にしました。
  • 単なる「コピーアテンション」や「データ多様性」だけでは不十分であり、「コピー経路の確保」と「埋め込みの崩壊防止」の両方が必要であることが示されました。

Gemma 3 モデルへの適用

• 未使用トークンの崩壊: Gemma 3 ファミリー（1B〜27B）には、下流タスク用に意図的に未使用に設定された 99 個のトークンが存在します。これらのトークンの埋め込みを分析したところ、使用済みトークンに比べてコサイン類似度が非常に高く（崩壊している）、特にモデルサイズが小さいほど顕著でした。
• ファインチューニングへの影響: 未使用トークンを記号として用いて命題論理タスクでファインチューニングを行ったところ、既知の記号（英単語やアルファベット）に比べて学習が著しく遅い（90% 精度到達に 5000 ステップ必要 vs 500 ステップ）ことが確認されました。これは埋め込みの相関が初期値として不適切であることを示唆しています。

5. 意義と結論

• 理論的貢献: 記号推論における未観測トークンへの失敗が、単なる「コピー能力の欠如」だけでなく、**最適化過程における重みの崩壊（Collapse）**に起因することを初めて理論的・実証的に解明しました。
• 既存研究の統合: Boix-Adserà et al. (2024) が提案したアーキテクチャ変更の限界と、Anand et al. (2025) が提案した「アクティブ・フォージティング」の有効性を、埋め込み崩壊という共通のメカニズムで説明しました。
• 実用的示唆: 大規模モデル（Gemma 3）においても同様の現象が観測されており、下流タスクでのファインチューニングにおいて、意図的に未使用に設定されたトークンを利用する際の注意点（初期値の相関による学習の遅延）を指摘しました。
• 将来の展望: 埋め込みの崩壊を防ぐための戦略（凍結、リセット、アーキテクチャ変更）は、LLM が真の抽象推論を行うための重要な手がかりとなります。

本論文は、LLM の推論能力の限界を「表面的なパターンマッチング」だけでなく、**重み空間の幾何学的な性質（崩壊）**という観点から理解し、それを克服するための具体的な設計指針を提供した点で重要です。