- LLMが数学オリンピック級の問題を解ける一方で、単純な足し算や数独ですら外部ツールなしでは正確に実行できないという限界を乗り越えるため、トランスフォーマ内部に実際のコンピュータを構築するアプローチ
- 任意のCコードをトークンに変換し、モデル自体が数百万ステップの実行トレースを直接生成しながら、CPUで毎秒3万トークン以上の速度でストリーミング可能
- 中核技術はアテンションヘッドの次元を2Dに制限し、凸包（convex hull）ベースの幾何学的検索へ変換することで、線形時間スキャンを対数時間に置き換えること
- WebAssemblyインタプリタをトランスフォーマの重みに実装し、外部ツール呼び出しなしでモデルのデコーディングループ内でプログラム実行が透過的に行われる
- 将来的にはプログラムを重みに直接コンパイルし、学習された表現とコンパイル済みアルゴリズムを単一の計算基盤に統合するAIシステムへ拡張できる可能性を示す
LLMの計算上の限界
- 最先端のLLMは国際数学オリンピック金メダル級の問題を解いたり、未解決の数学・科学問題に挑んだりできる一方で、純粋な計算タスクでは依然として失敗する
- 基本的な足し算でも誤りが起き、Sudoku-Benchのようなベンチマークでは外部支援なしでの解答成功率が非常に低い
- 現在、このギャップを埋める回避策は大きく2つある
- ツール使用（Tool use）: モデルがコードを書き、外部インタプリタがそれを実行して結果を返す
- エージェントオーケストレーション: 外部ループが中間状態を保存し、作業を分解してモデルを繰り返し呼び出す
- こうしたアプローチは有用だが、実際の計算能力がモデルの外部に存在するという根本的な限界を際立たせる
- たとえるなら、人間が飛行機を作ったからといって人間自身が飛べるようになったわけではないのと同じ状況
- モデルは計算について推論したりツールを調整したりできても、自分で実行できなければ真のコンピュータではない
トランスフォーマをコンピュータにした方法
- トランスフォーマアーキテクチャがチューリングマシンをシミュレートできるという理論的普遍性は複数の研究で示されているが、理論上の表現力が実用的な実行効率を保証するわけではない
- 純粋に理論的な計算モデルではなく、現代的なRAMコンピュータをトランスフォーマ内部に実装し、各命令を最大5トークンに対応付けた
- しかし、より深い問題はデコーディング過程そのものにある
- 標準的な自己回帰デコーディングでは、各ステップで伸び続ける履歴全体と相互作用する
- KVキャッシュは過去のキー／値の再計算を避けるが、キャッシュサイズに比例するアテンションコストは依然として発生する
- この限界を解決するために、アテンションヘッドの次元を2Dに制限し、実行スタイルのトレースに対する効率的なデコーディング経路を導入した
- 主要な検索／更新演算を系列長に対して対数時間で実行できる
- これにより、数百万ステップのプログラムをトランスフォーマ内部で実行できる
計算の意味 — ツール使用 vs モデル内実行
- 従来のツール使用方式: モデルが
python -c "print(3+5)" のようなコードを出力 → 外部インタプリタが実行 → 結果をトークン列に注入
- このシステムの方式: WebAssemblyインタプリタをトランスフォーマの重みに実装
- WebAssemblyは高速で決定論的な実行のための低水準命令セットであり、C/C++の汎用コンパイルターゲットでもある
- 3 + 5 を計算するとき、モデルはwasm命令（
i32.const, i32.add）を出力した後、高速デコーディングモードに切り替えてステップごとの実行トレースを直接生成する
- 主な違い
- ツール使用は不透明（opaque）: モデルが制御を渡し、ブラックボックスの答えを受け取る
- モデル内実行は透明（transparent）: すべての中間段階がトレースに現れ、モデルが自身のデコーディングループを離れない
数独デモ — 正確な長期計算のストレステスト
- 学習されたニューラルネットワークの手法は易しい数独では良好な性能を示すが、難しい問題では完全に失敗する
- 自己回帰モデルは制約充足問題に本質的に不向きだという説明が一般的だが、このシステムは完全に自己回帰型でありながら100%の正確性を達成した
- 実際のボトルネックは自己回帰パラダイムそのものではなく、難しい数独が非常に長い実行トレースを必要とし、標準アテンションでは長いコンテキスト生成が非実用的になることにある
- コンパイル済みの完全な数独ソルバをトランスフォーマ内部で実行し、学習済みヒューリスティックではなく実際のアルゴリズムをステップごとに実行する
- Arto Inkalaの「世界で最も難しい数独」を3分以内に正確に解いた
- コンパイル済みソルバが正しければトランスフォーマの実行も正しいという一般的な保証を提供する
計算のエンコード方式 — 追記専用トレース
- 自己回帰トランスフォーマを、自身の履歴の中に住む機械になぞらえる
- 従来のコンピュータは編集可能なメモリを更新するが、トランスフォーマにはその概念がない
- 固定されたプロンプト（入力／プログラム）と、増え続けるトレース（生成トークン）から構成される
- 追記専用ノートの比喩
- 最初の数行が入力（プロンプト）で、その後の各行が計算の次のステップを記録する
- 過去の行は修正できず、各ステップで参照できる過去の位置は少数に限られる
- 多くのアルゴリズムは「各ステップで固定された少数の過去位置だけを参照する、追記専用トレース」として表現できる
- このシステムでモデルが生成するトークンは、仮想マシンの命令ポインタ、メモリ／スタック操作、算術、制御フロー、出力などの進化する状態を表現する
- アテンションヘッドは共有された1次元配列のように機能し、各トークンが1つのインデックスに値を書き込み、別のインデックスから値を読む書き込み後読み出しプリミティブを提供する
- アテンションは**累積和（cumulative sums）**も計算でき、命令ポインタやスタック深度などをデルタ増分の累積和として追跡できる
中核となる解放鍵: 指数的に高速なアテンション
- 実際のコンピュータは命令ごとにほぼ一定の作業量で小さな状態を更新するが、標準的なトランスフォーマはt番目のデコーディングステップで長さtのプレフィックスと相互作用し、総コストが二次的に増大する
- HullKVCache の導入によってこの二次的な爆発を解決した
- 実ベンチマークでは、HullKVCacheは毎秒31,037トークン（6,747行/秒）、標準KVCacheは毎秒272トークン（59行/秒）で、約114倍の差があった
- ステップごとにΘ(t)時間かかる代わりに、O(log t) 時間でアテンション検索を実行する
-
2Dアテンションの高速経路
- トランスフォーマ全般を加速したり新しいアーキテクチャを導入したりするのではなく、通常のトランスフォーマでヘッド次元を2Dに制限する扱いやすいパラメータ化に焦点を当てている
d_model = 36, n_heads = 18 とし、ヘッドあたりちょうど2次元、7層、標準の nn.MultiheadAttention、ゲート付きフィードフォワードネットワークを使用
- カスタムアテンションカーネルや疎マスクを使わず、純粋なバニラPyTorch で構成
- 層数、ヘッド数、埋め込みサイズは任意に大きく設定できるため、モデル全体が小さい必要はない
n_heads = d_model / 2 として、より多くのヘッドを使う
-
2Dアテンションの幾何学的視点
- アテンション機構では、過去トークンが2Dキー・ベクトル kⱼ と値 vⱼ を提供し、現在ステップがクエリ q を形成して、内積が最大のキーの値を返す（hardmaxアテンション）
- これは計算幾何学の古典的な**「支持点（supporting point）クエリ」**と正確に同じで、方向 q が与えられたときに凸包上でその方向に最も遠い点を見つける問題に対応する
- トークン生成と同時に凸包データ構造を維持することで、過去全ポイントに対して対数時間検索が可能になる
- メモリ／スタック操作では「インデックス i に最後に保存された値」を取得する必要があり、これが2Dキーを必要とする理由である
- インデックス j を2Dキー
kⱼ = (2j, -j²) として保存し、方向 q = (i, 1) でクエリすると、二次項 -j² が正確に一致したものだけが勝つようにするペナルティとして働く
- チューリング完全性には2Dアテンションで十分であり、このブログが示すようにRAMコンピュータ全体を表現できる
今後の計画
-
より豊かなアテンション機構
- 現在の実装はhardmaxアテンションを用いているが、これは本質的な制約ではない
- k-sparse softmaxアテンションで近似可能: 入れ子になった凸包から上位k個のキーを検索し、それらに対してのみsoftmaxを実行することで、コストは
O(k + log n) になる
- この幾何学的高速経路は実行器構造に限られず、原理的には2Dヘッドを持つあらゆるトランスフォーマのデコーディング時間を高速化できる
- 3Dヘッドにも3D凸包を通じて自然に拡張できるが、さらに高次元になると効率は急激に低下する
-
2Dヘッドで大規模モデルを学習
- 2Dヘッドのパラメータ化は本質的に小規模というわけではなく、より多くのヘッドと層を使って標準トランスフォーマに近い総パラメータ数を維持できる
- 活用形態はいくつか考えられる
- より遅く汎用的なモデルと組み合わせる専用高速経路
- 単一システム内の高速／低速ハイブリッドアーキテクチャ
- 2Dモデルがトークンを高速に提案し、通常のアテンションモデルが検証する推測デコーディング（speculative decoding）
- 実行トレースはフォワードパスの一部なので、全過程が**微分可能（differentiable）**である。これは外部ツールとは根本的に異なり、より大きなモデルへ直接統合できる学習可能な計算基盤になる
-
プログラムを重みにコンパイル
- 現在はモデルが重みにエンコードされたインタプリタを学習しているが、重み生成コンパイル機構を拡張すれば、任意のプログラムをトークン列なしで直接重みにコンパイルできる
- 重みそのものがソフトウェアのデプロイ先（deployment target）となり、モデルがソフトウェアのような振る舞いを学習するのではなく、コンパイル済みプログラムのロジックを内部回路の一部として含むことになる
- ロジックを重みにコンパイルできるなら、勾配降下法だけがモデルを変更する唯一の方法ではなくなる。構造、アルゴリズム、保証をネットワークに直接挿入する別の経路が生まれる
- 最終的には、経験から学ぶだけでなく、自身の重みを修正・拡張して内部機械を自ら書き換えるシステムへ発展できる可能性がある
-
AIシステムをソフトウェアのように成長させる
- 現代のソフトウェアエコシステムがモジュール、抽象化、再利用可能なコンポーネントを蓄積しながら進化してきたように、AIシステム内部でも新しい計算能力が段階的に追加される同様の過程が可能である
- 新しい機能を、システム全体を再学習することなく既存回路へ接続できる
- 将来のAIシステムはソフトウェアを単に使うのではなく、ソフトウェアを内部に含み、学習された表現とコンパイル済みアルゴリズムを単一の計算基盤に統合する
1件のコメント
Hacker Newsの意見
単なる計算の実行よりもはるかに興味深いアプローチだ
モデルはトークン数の対数に比例する動的アテンション切り替えを実行できる
これにより、テキストで表現されたレジスタやスタックを追跡しながら、プログラム実行を模倣できる
もしLLMが「集中モード」に切り替わって非常に高速にトークンを生成できるなら、多数の仮説を探索・整理する推論段階を高速化できるはずだ
論文では、高速経路と低速経路を組み合わせたハイブリッド構造や、speculative executionモデルとして活用できる可能性が示されている
最初は「なぜこんなことをする必要があるのか？」と思ったが、今では訓練のブートストラップという観点で見ている
たとえば正確度80%のエキスパートシステムをモデルに内蔵し、その結果を学習データとして使って精度を上げられる
さまざまなタスクの訓練コストを下げるほど、AI競争への参入障壁は下がる
softmax attentionと違って、average-hard attentionはキーとクエリに対して微分不可能である
straight-through推定で補正しても、逆伝播速度は向上しない
人間の脳は計算能力に優れているわけではない
10桁の掛け算にも時間がかかる。こうした計算は論理ゲートのほうがはるかに効率的だ
だとすれば、LLMが直接計算するよりALUのような外部モジュールを呼び出すほうがよいのではないかと思う
AIが書いたような文体だが、肝心のメッセージが不明瞭だ
モデルが外部システムの代わりに内部でプログラムを実行することの利点が、速度・逆伝播・ベンチマークのどこにあるのかはっきりしない
一部にはsymbolic logicが必須だと考える人もいるが、こうした試みは単に興味深い実験と見なせる
これは外部ツール呼び出しとは根本的に異なる
また**O(k + log n)**複雑度のデコードコストを持ち、Sudokuのような問題を100%の正確さで解けるなら十分意味がある
こうした方式をMoEスタイルで組み合わせたり、WASMベースのVM内蔵のように多様なインタプリタを試したりできるだろう
モデルの内部計算経路にツールを統合するというアイデアは、解釈可能性の観点から非常に興味深い
単純なTransformerでこれほどの効率を出せる点には驚かされる
潜在力はあるが、現状では実用性は低い
重みや「コンパイラ」ツールが公開されておらず、実験が難しい
それでも、事前定義された計算プリミティブをLLMに内蔵するという発想は依然として有用かもしれない
この一文が核心だ:
「実行トレースが順伝播の一部であるため、計算そのものを通じて勾配を伝播できる」
つまり、外部ツール呼び出しとは違って訓練可能な計算基盤になる
また訓練データや損失関数の設計も不明瞭だ
とはいえ、ツール呼び出しがバッチ効率を壊すという点を考えると、内部計算サブネットを通すことで大規模な効率改善が可能かもしれない
ただしそのサブネットは必ずしもTransformerである必要はなく、GPUに最適化された非学習レイヤーでも十分だと思われる
論文は肝心な部分を隠している
アテンションヘッド次元を2に制限すると、対数時間複雑度でトークンの検索・更新ができる
しかし、なぜこの戦略が「コードトークン」にだけ適用されるのかは不明だ
WASMをターゲットにしている点も効率面では疑問が残る
たとえばself-attentionを「lookup table」と表現するのは不正確だ
コード例では
d_model = 36, n_heads = 18として1ヘッドあたり2次元を構成していたが、それでもなお不明瞭であるSudokuソルバをTransformerの重みにどうコンパイルしたのか、具体的な説明がない
直接コードから重みへコンパイルしたのか、それとも学習で獲得したのかが不明だ
興味深いが、「なぜわざわざこうするのか」という疑問は残る
人間の脳もチューリングマシンをシミュレートできるが遅い。だから外部ツールを使う
モデルも同様に、外部ツールを使うほうが効率的ではないかと思う
Elixirのような言語を内蔵して、より短いコードを実行することもできるだろう
モデルが実行中にコードを修正したり、デバッグ能力を持ったりできるという発想だ