AIの再現性を取り戻す：Thinking Machines Labが挑む非決定性問題の打破
AI、特に大規模言語モデル（LLM）がしばしば直面する課題の一つに、同一の入力に対して毎回異なる応答を返す「非決定性（nondeterminism）」があります。これは科学研究や企業利用では信頼性や再現性を揺るがす重大な問題ですが、TMLはこれを単なる制約ではなく「解決可能な技術課題」と見なし、研究を進めています。
1. 非決定性とは何か？その原因を探る
1-1. 非決定性の現れ方
LLMに同じ質問を複数回投げかけると、答えが微妙に異なるのは珍しくありません。これは一般に「サンプリング」の過程で生じるものであり、「温度パラメータ」を0にして理論上もっとも確率の高いトークンを選ぶようにしても、実際には非決定性が残ります。
1-2. GPUカーネルの“数値演算”が鍵
多くのエンジニアや研究者は、「並行処理と浮動小数点演算（floating-point）の非結合性（非可換性／非結合性）が原因」と考えます。たとえば、(a + b) + c ≠ a + (b + c) のように、計算順序によって結果が異なる可能性があるためです。
2. 根本原因：バッチ不変性の欠如
2-1. カーネルが同じ順序で計算できない理由
Horace Heらは、非決定性の真の原因は「バッチサイズなどの条件によって計算順序や実装戦略が変わってしまうこと」にあると指摘します。このため、同じ入力でもサーバの負荷やバッチ処理の状況によって結果が変わってしまうのです。
2-2. バッチ不変性とは？
「バッチ不変性（batch invariance）」とは、バッチサイズやその要素の位置に関係なく、各計算結果が一貫して得られる性質を指します。しかし、実際にはこの性質が崩れることで非決定性が発生します。
2-3. 具体的に問題となる演算
RMSNorm（正規化演算）
行列乗算（MatMul）
アテンション機構
これらは「削減（reduction）」を含むため、バッチサイズや分割処理の方法によって計算順序や結果が変化し得ます。
3. 解決へのアプローチ：バッチ不変なカーネル設計
3-1. 一貫した削減順序を維持する設計
例えば、RMSNormでは、各バッチ要素を専用のGPUコアで処理し、一貫した順序で削減する方法を取ることで、バッチ不変性が維持可能です。
3-2. 行列乗算とアテンションの工夫
行列乗算では、データ並列戦略（data-parallel）を採用し、削減を一貫した順序で行いつつ、パフォーマンス低下を最小限に抑える工夫が求められます。
アテンションでは、特に「KVキャッシュ」など処理の分割の影響を受けやすいため、「固定サイズの分割戦略（fixed-size split）」によってバッチサイズに関係なく順序を保つ必要があります。
3-3. 実装と成果
vLLM＋FlexAttentionをベースに、バッチ不変なカーネルを使うモードを組み込んだ実験では、サンプリング結果が完全に一致する（すなわち「真の再現性（完全決定性）」を達成）ことが確認されました。
加えて、性能面ではデフォルトモードより若干遅いものの、「致命的な遅延」には至らず、実用範囲にあることも示されています。
4. 実用面へのインパクトと今後の展望
4-1. 強化学習（RL）への恩恵
再現可能な応答により、強化学習におけるオフポリシー問題が緩和され、訓練と推論の整合性が向上します。実際、この決定性を確保することで「真のオンポリシーRL」が実現でき、訓練プロセスがより滑らかになると報告されています。
4-2. 透明性とオープンリサーチへの姿勢
TMLは新たなブログシリーズ「Connectionism」を開始し、研究成果の定期的な公開とコードの共有を通じて、公開性の高い研究文化を志向しています。
4-3. 製品化と社会実装への期待
創業者のMira Murati氏は、近い将来「研究者やスタートアップに役立つプロダクト」をリリースする意向を示しており、この決定性技術が新しいAIサービスの基盤技術となる可能性が高まっています。
結論
Thinking Machines Labが示した「Defeating Nondeterminism in LLM Inference」は、LLMの“予測できない応答”という難題に対し、GPUカーネルの設計にまで踏み込んでバッチ不変性を徹底することで、決定性を回復し再現性を高めるという、非常に具体的かつ実践的なアプローチです。
この研究は、AIを研究目的や業務利用といった場面でさらに信頼性の高いツールへと進化させる礎となるでしょう。今後のTMLの取り組みと、その成果から目が離せません。
