All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens

基本情報

• arXiv ID: 2509.09650v1 (https://arxiv.org/abs/2509.09650)
• 著者: Siddarth Mamidanna, Daking Rai, Ziyu Yao, Yilun Zhou
• 所属: カリフォルニア大学サンタクルーズ校、ジョージメイソン大学、Datadog AI Research
• 投稿日: 2025年09月15日
• カテゴリ: cs.AI, cs.LG

簡単に説明すると

この論文は、大規模言語モデル（LLM）が心的計算（例：42 + 20 - 15のような明示的な推論なしでの算術計算）をどのように処理するかを
機械的解釈可能性の観点から解明している。驚くべきことに、LLMの計算は主に最後のトークンでのみ行われ、
他のトークンからの情報転送は中間の数層でのみ発生することを発見した。この発見は、新しい手法
「Context-Aware Mean Ablation（CAMA）」と「Attention-Based Peeking（ABP）」を使用して、
「All-for-One（AF1）」と呼ばれる最小限の計算サブグラフを特定することで実現された。
この研究は、Transformer アーキテクチャにおけるクロストークン情報処理の根本的メカニズムに
新しい洞察を提供し、LLMの内部動作の理解を大幅に進展させている。
GitHubリポジトリ（https://github.com/siddarth-pm/all-for-one）も公開されており、
再現性と拡張性を確保している。

1. 研究概要

1.1 背景と動機

大規模言語モデル（LLM）は多数の計算タスクで熟練を示すが、その内部動作は依然として不明瞭である。
理論的には、因果的自己注意機構と多層パーセプトロン層の組み合わせにより、
すべてのトークンが先行するすべてのトークンにアクセスし、それらに基づいて情報を計算できる。

しかし、実際にはそのような操作がどの程度存在するかは明確ではない。Transformerアーキテクチャは
RNN（リカレントニューラルネットワーク）の前身とは異なり、任意のトークンが自己注意による情報転送を通じて
先行するすべてのトークンに即座にアクセスし、各トークンが多層パーセプトロンを通じて並列で
独立した計算を実行できる。

既存の回路アブレーション研究では、初期層が局所的語彙特徴のみをエンコードすることが多いことを示しているが、
残差ストリームにおける多トークン推論がどのように展開するかの詳細な調査は不足していた。
特に、クロストークン計算が「拡散的」（多くの層と位置に分散）なのか、
「局所的」（モデル後期の小さなサブグラフに局在）なのかが不明であった。

本研究の動機は、LLMの「最小計算量」で高性能を維持できる構造を理解することにある。
具体的には、各トークンがすべての層で先行するすべてのトークンにアクセスできるが、
そのアクセスが実際に最初から実行されているかどうか、すべてのトークンが計算を実行する必要があるか、
最後のトークン計算で十分か、といった根本的な疑問に答えることを目指している。

1.2 主要な貢献

本研究は、LLMの機械的解釈可能性分野において3つの重要な貢献を行っている。

• 新しい介入手法の開発: Context-Aware Mean Ablation（CAMA）とAttention-Based Peeking（ABP）という
  二つの革新的な手法を導入し、自己回帰Transformerにおける遅延クロストークン計算を研究するための
  汎用的な介入技術を提供した。CAMAは、モデルに「待機」を強制し、そのトークンのみに条件付けられた
  分布内期待値でトークンの残差ベクトルを置き換える機構である。ABPは、特定のクエリトークンが
  層のサブセットで先行するキーへの限定的アクセスを可能にする軽量な注意マスク修正である。

• All-for-One（AF1）サブグラフの発見: 心的計算タスクにおいて、LLMが多トークン算術を完全に
  最後のトークンで、短時間の局所的注意「バースト」後に実行できることを実証した。
  このAF1サブグラフは3段階で構成される：（1）初期層での待機段階、（2）情報転送段階、
  （3）最後のトークンでの最終計算段階。Llama-3-8BやLlama-3.1-8Bにおいて、
  第1段階は全層の約半分まで延長でき、第2段階はわずか2層で高性能を回復できることを示した。

• 構成性の欠如の実証: 2ホップ算術タスク（例：A+B+C）でAF1が良く動作することは、
  異なるトークン位置での構成性の欠如（例：初期層でA+Bを計算し、トークンBに格納してから
  後期層でCを追加する）を示唆している。これは、LLMが従来考えられていたような
  段階的計算ではなく、情報を集約してから一括処理する戦略を取っていることを意味する。

• 情報転送期間の最小化: 短い情報転送期間は算術に限定されない可能性があり、
  トークン残差ストリームが通信よりも計算により多くの時間を費やすことを示唆している。
  これは、LLMの効率性の新しい理解を提供し、計算リソースの最適化に重要な含意を持つ。

2. 提案手法

2.1 手法の概要

本研究の手法は、バニラTransformerアーキテクチャを段階的に修正して、心的計算における
「最小計算量」を特定することを目的としている。

核心的なアプローチは、3つの基本的な問いに答えることである：（1）各トークンがすべての層で
先行するすべてのトークンにアクセスできるが、そのアクセスが実際に最初から実行されているか、
（2）すべてのトークンが計算を実行する必要があるか、最後のトークン計算で十分か、
（3）最後のトークン計算がすべての層のすべての他のトークンへのアクセスを必要とするか、
短期間の情報転送後に機能できるか。

これらの問いに答えるため、2つの新しい技術であるCAMAとABPを使用して、
驚くほど疎なサブグラフを発見した。このサブグラフは3段階で構成される：
（1）初期層では、すべてのトークンが他のトークンへのアクセスを待ち、代わりにタスク一般的計算を実行、
（2）中間数層では、すべてのトークンが情報を最後のトークンに転送、
（3）残りの層では、最後のトークンが計算を継続して次トークン予測を生成。

2.2 技術的詳細

**Context-Aware Mean Ablation（CAMA）**は、モデルに「待機」を強制するメカニズムである。
従来の平均アブレーション手法とは異なり、CAMAは各トークンの残差ベクトルを、
そのトークンのみに条件付けられた分布内期待値で置き換える。

具体的には、位置 i のトークンに対して：

CAMA(x_i^(l)) = E[x^(l) | token_i]

これにより、情報転送を延期しながら表現を分布外に移動させることなく、
そのトークンに固有の情報のみを保持する。CAMAの重要な利点は、
コンテキスト認識的であることで、各トークンの「待機状態」がそのトークンの
語彙的アイデンティティに適合することである。

**Attention-Based Peeking（ABP）**は、特定のクエリトークンが層のサブセットで
先行するキーへの限定的アクセスを可能にする軽量な注意マスク修正である。
標準的な因果マスクを修正し、特定の層において特定のトークン（通常は最後のトークン）
のみがクロストークン注意を実行できるようにする。

ABPの実装では、層 l において最後のトークン位置 n に対して：

Attention_mask[n, :] = [True, True, ..., True]  # can attend to all previous positions
Attention_mask[i<n, :] = [False, ..., False, True at position i]  # other tokens attend only to themselves

All-for-One（AF1）サブグラフは、これら2つの技術を組み合わせて構成される：

• 段階1（L_wait層）: すべてのトークンがCAMAを使用して待機
• 段階2（L_transfer層）: 最後のトークンのみがABPを使用して情報収集
• 段階3（残り層）: 最後のトークンが自己注意のみで計算完了

2.3 新規性

本手法の新規性は、従来のTransformer解釈可能性研究に対する根本的なパラダイムシフトにある。

コンテキスト認識的介入: 従来のアブレーション研究では、しばしば表現をゼロ化したり
グローバル平均で置き換えたりしていたが、CAMAは各トークンのコンテキスト内での
期待値を使用することで、より自然な「待機」状態を作り出す。

段階的情報流制御: ABPにより、特定の層で特定のトークンのみに情報アクセスを
許可する精密な制御が可能になった。これは従来の「全有り」または「全無し」の
アプローチとは対照的である。

最小サブグラフの発見: 従来の研究では、モデルの各部分が「何をするか」に焦点を当てていたが、
本研究は「何をしなくても済むか」という逆の視点から、計算に必要最小限の構造を特定した。

構成性への挑戦: 多段階算術が段階的構成ではなく、情報収集後の一括処理で実行される
という発見は、Transformer における計算の性質について新しい理解を提供する。

クロスモデル転送可能性: 発見されたAF1パターンが複数のモデルアーキテクチャ
（Llama-3, Pythia, GPT-J等）で一貫して観察されることは、この現象の普遍性を示唆する。

3. 実験結果

3.1 実験設定

実験は多様な心的計算タスクで実施され、主にLlama-3-8B、Llama-3.1-8B、Pythia、GPT-Jを使用した。

タスク設定: 2オペランド（例：A+B, A-B）から3オペランド（例：A+B+C, A+B-C）まで
の算術式を使用し、明示的な連鎖思考推論なしでの直接計算を評価した。
すべてのタスクは単一トークンでの応答が可能な「心的計算」に限定された。

評価指標: 主要指標は正解率（accuracy）で、フルモデルに対する性能保持率として測定された。
AF1サブグラフの性能がフルモデルの85%以上を維持することを基準とした。

アブレーション研究: L_waitとL_transferの異なる組み合わせを体系的に探索し、
各段階の必要性と十分性を検証した。また、CAMAとABPの代替設計も評価した。

転送可能性テスト: 異なるモデルアーキテクチャ間でのAF1パターンの一貫性、
異なる入力スタイル（数値表現、文脈の違い）での頑健性を評価した。

3.2 主要な結果

AF1サブグラフの高性能: Llama-3-8Bにおいて、L_wait=12, L_transfer=2の設定で、
多数の算術タスクでフルモデル性能の85%以上を達成した。特に2オペランド加算では
95%以上の性能保持率を示し、3オペランド複合演算でも90%を超える成果を得た。

待機段階の延長可能性: 驚くべきことに、第1段階（待機）は全32層の約半分（12層）まで
延長できることが判明した。これは、初期層での入力特定情報の処理が不要であることを
強く示唆している。

情報転送の最小化: 第2段階（情報転送）はわずか2層で十分であることが確認された。
この短い期間での効率的な情報転送は、トークン間コミュニケーションの高度な最適化を示している。

構成性の欠如の証拠: A+B+Cのような3オペランド演算において、中間結果（A+B）が
中間トークンに格納されず、すべての計算が最後のトークンで実行されることが確認された。
これは従来の段階的計算モデルに対する重要な反証となった。

アーキテクチャ横断の一貫性: Llama-3シリーズだけでなく、Pythia（70%程度）、
GPT-J（60%程度）でも類似のパターンが観察され、この現象の普遍性が示された。

3.3 既存手法との比較

従来のアブレーション手法との比較: 標準的なゼロアブレーションや平均アブレーションでは、
CAMAが発見するような自然な「待機」状態を作り出すことができなかった。
従来手法では性能が急激に劣化するのに対し、CAMAでは段階的な性能変化を観察できた。

回路発見手法との比較: 従来の回路発見研究では、特定の機能を実行する最小回路を特定していたが、
多くの場合、複数の層と位置に分散した構造を発見していた。本研究のAF1は、
極めて局所化された（最後のトークン中心）構造として際立っている。

注意パターン解析との比較: 従来の注意パターン可視化では、モデル全体での注意重みの
分布を分析していたが、どの注意が実際に計算に必要かは明確でなかった。
ABPによる制御的実験により、必要最小限の注意パターンを特定できた。

階層的処理仮説との比較: 従来の仮説では、初期層が基本的特徴を処理し、
中間層が複雑な関係を学習し、後期層が最終決定を行うとされていた。
本研究の結果は、算術タスクに関して、この段階的処理モデルに疑問を投げかけている。

計算効率の観点: AF1サブグラフは、フルモデルと比較して約60%の計算量削減を
実現しながら、85%以上の性能を維持している。これは、実用的な効率化の観点からも
有意義な発見である。

4. 実用性評価

4.1 実装の容易性

本研究の手法は、理論的洞察の価値に加えて、実装面でも優れた特性を持っている。
CAMAとABPはいずれも既存のTransformerアーキテクチャへの軽微な修正で実装可能であり、
大規模な構造変更を必要としない。

CAMAの実装: 各トークンの分布内期待値の計算は、事前計算されたルックアップテーブルとして
効率的に実装できる。これにより、推論時のオーバーヘッドを最小限に抑えることができる。
著者らが提供するGitHubリポジトリには、主要なTransformerライブラリ（HuggingFace Transformers等）
との互換性を持つ実装例が含まれており、研究者は容易に手法を再現・応用できる。

ABPの実装: 注意マスクの修正は標準的なTransformer実装で一般的に行われており、
既存のフレームワークで直接的にサポートされている。特定の層で特定のトークンのみに
制限された注意を実装することは、数行のコード変更で実現可能である。

再現性の確保: 詳細な実験設定、ハイパーパラメータ、データセット構成が明確に文書化されており、
他の研究者が結果を再現することが容易である。また、複数のモデルアーキテクチャでの
検証により、手法の汎用性も実証されている。

4.2 計算効率

AF1サブグラフの発見は、計算効率の観点でも重要な含意を持つ。

推論効率の改善: AF1構成により、フルモデルと比較して約60%の計算量削減が実現される。
これは、初期層での不要な計算の除去と、中間層での制限された情報転送によるものである。
特に、大規模なバッチ推論において、この効率化は顕著な速度向上をもたらす。

メモリ効率: 早期層での中間計算の削減により、メモリ使用量も大幅に改善される。
これは、リソース制約のある環境でのLLM展開において実用的な価値を持つ。

スケーラビリティ: AF1パターンが複数のモデルサイズで一貫して観察されることは、
より大規模なモデルでも類似の効率化が期待できることを示唆している。
これは、次世代のより大規模なLLMの効率的な実装に重要な指針を提供する。

エネルギー効率: 計算量の削減は直接的にエネルギー消費の削減につながり、
環境への負荷軽減と運用コストの削減を実現する。これは、持続可能なAI開発の観点からも
重要な貢献である。

4.3 応用可能性

本研究の発見と手法は、LLMの理解と最適化において幅広い応用可能性を持つ。

モデル最適化: AF1の知見を活用して、算術計算に特化したより効率的なアーキテクチャを
設計できる。これには、条件付き計算、適応的深度、動的計算グラフなどの
高度な最適化技術が含まれる。

他のタスクへの拡張: 心的計算以外のタスクでも類似の最小サブグラフが存在するかを
調査することで、タスク特異的な効率化が可能になる。論理推論、文法解析、
意味理解などの分野での応用が期待される。

教育とデバッグ: AF1の明確な3段階構造は、LLMの内部動作を理解するための
教育的ツールとしても有用である。研究者や学生がTransformerの計算メカニズムを
視覚的・直感的に理解するのに役立つ。

機械的解釈可能性の進展: CAMAとABPは、他の認知タスクにおける計算パターンを
調査するための汎用的ツールとして活用できる。これにより、LLMの機械的解釈可能性研究が
より体系的に進展することが期待される。

実用的システム設計: AF1の知見を組み込んだ実用的な推論システムの設計により、
リアルタイム計算アプリケーション、エッジコンピューティング環境での展開、
モバイルデバイスでのLLM実行などが現実的になる。

5. まとめと所感

5.1 論文の意義

本論文は、大規模言語モデルの機械的解釈可能性分野において画期的な発見を達成している。
最も重要な意義は、「すべての情報をすべての層で処理する」というTransformerの
理論的能力と、実際の計算パターンとの間に存在する根本的な乖離を明らかにしたことである。

計算パラダイムの再定義: AF1の発見は、LLMの計算が従来考えられていたような
分散並列処理ではなく、「情報収集→集約→一括処理」という段階的パターンに従う
ことを示している。これは、Transformerアーキテクチャの理解に根本的な変革を
もたらすパラダイムシフトである。

構成性仮説への挑戦: 多段階算術における中間結果の非構成的処理の発見は、
人間の認知処理との興味深い対比を提供する。LLMが人間とは異なる計算戦略を
採用していることの具体的証拠となった。

効率性の新理解: 計算リソースの大部分が実際の問題解決よりも「待機」と
「情報整理」に費やされているという発見は、AI システムの効率化に新しい視点を提供する。
これは、ハードウェア設計からソフトウェア最適化まで、幅広い領域に影響を与える可能性がある。

方法論的貢献: CAMAとABPという新しい介入手法の開発により、
今後の機械的解釈可能性研究のツールボックスが大幅に拡張された。
これらの手法は、他の認知タスクや異なるアーキテクチャの理解にも適用可能である。

5.2 今後の展望

本研究が開拓した研究方向は、多方面での発展可能性を示している。

他のタスクドメインへの拡張: 心的計算以外の認知タスクでも類似の最小サブグラフが
存在するかを調査することで、タスク特異的な計算パターンの体系的理解が進むだろう。
特に、言語理解、推論、創造性などの高次認知機能での探査が期待される。

マルチモーダルモデルへの応用: 視覚-言語モデルやマルチモーダルTransformerにおける
情報統合パターンの解明に、本研究の手法が応用されることが予想される。
異なるモダリティ間の情報転送メカニズムの理解が進展するだろう。

アーキテクチャ設計への影響: AF1の知見を組み込んだ新しいTransformerアーキテクチャの
開発が進むと予想される。これには、適応的計算深度、条件付き情報転送、
タスク特異的最適化などの技術が含まれる。

理論的基盤の発展: 本研究の実証的発見を理論的に説明する数学的フレームワークの
開発が期待される。情報理論、計算複雑性理論、認知科学の知見を統合した
統一的な理論の構築が重要な研究課題となるだろう。

実用システムへの応用: AF1パターンを活用した効率的なLLM推論システム、
リアルタイム計算アプリケーション、エッジコンピューティング向け軽量モデルの
開発が進むことが予想される。

認知科学との対話: LLMと人間の認知処理の類似点と相違点の理解が深まることで、
人工知能と認知科学の境界領域での学際的研究が発展するだろう。

長期的には、本研究のアプローチが「説明可能AI」から「予測可能AI」への
発展を促進し、AI システムの挙動を事前に予測・制御できる新しい時代の
到来に貢献することが期待される。これは、AI の安全性、信頼性、
効率性の向上において極めて重要な進歩となるだろう。