RiemannInfer: リーマン幾何学でトランスフォーマーの推論を改善する

なぜ日常のAI利用者に重要なのか

現代のチャットボットやAIアシスタントは数学の問題を解いたり、エッセイを書いたり、医療関連の話題を説明したりできますが、それらがどのように結論に達しているのか、あるいはより確実に推論させるにはどうすればよいかはまだ十分に分かっていません。本論文は「RiemannInfer」と呼ばれる新しい手法を紹介します。これは大規模言語モデル（LLM）の内部活動を曲がった幾何学的な面上の運動として扱う視点です。この見方はシステムの思考のより直感的な像を提供するだけでなく、推論をより速く、より正確にします。

AIの思考を幾何学的な旅に変える

GPT-4やLlamaのようなLLMでは、文中の各単語は高次元のベクトルで表現され、層状の「アテンション」が単語同士の影響の強さを決めます。著者らは、これらの隠れ状態をモデルの言語理解を符号化した全体的な形状を持つ広大な空間の点として見ることができると観察します。推論をテキスト上の確率計算の連続として扱う代わりに、それを経路探索問題として再解釈します：モデルは初期状態（質問）から最終状態（回答）へこの空間を通って移動します。曲面の数学であるリーマン幾何学の道具を用いて、アテンションのパターンがこの内部風景をどのように曲げ、伸ばすかを捉える曲がった多様体を構築します。

全体像を失わずに複雑さを圧縮する

LLMの生の内部空間は非常に巨大なので、RiemannInferの最初のステップは本質的構造を保ちながらそれを縮小することです。著者らは点のつながり方を扱う位相幾何学の手法と、広く使われる次元削減アルゴリズムUMAPを組み合わせます。次元を減らす前に、隠れ状態の点群の「形状」を解析して、重要な接続パターンが圧縮で保たれることを確認します。その結果、トークン間の重要な関係――どの単語が互いに強く注意を払っているかといった関係――が大まかに保存された低次元空間が得られます。このコンパクトな幾何学マップにより、距離、最短経路、曲率といった精密な計算が現実的になります。

アテンションから曲がった地図を構築する

核となる革新はアテンション重みを幾何学的な距離の尺度に変換することです。モデルが二つのトークン間で強いアテンションを向けると、RiemannInferはそれらを多様体上で近い点として扱い、弱いアテンションではより遠くに置きます。これらの関係から、長さや角度を決める数学的ルールである計量を定義し、それを使って測地線（直線の曲がった等価物）や空間がどれだけ鋭く曲がるかを示す曲率を計算します。マルチヘッドアテンションは自然に複数の計量の混合となり、文法や意味など言語構造の異なる側面をそれぞれ捉えます。この構成により、モデルの決定は情報が濃い場所や希薄な場所を反映した起伏のある風景を通る特定の経路を選ぶこととして解釈できます。

少ない労力で済む推論経路を計画する

多様体が構築されると、著者らは推論を質問から回答へ「容易な」経路を見つける問題として再定式化します――経路に沿って行う総仕事量を最小化するものです。山登りのアナロジーを借りると、急でギザギザのルートを登るには滑らかな尾根道を辿るより多くのエネルギーが必要です。LLMの文脈では、曲率が急勾配の役割を果たし、モデルの推論作業は経路に沿って内部の不確実性がどれだけ減るかに対応します。測地線や曲率の近似式と効率的なグラフ探索アルゴリズム（ダイクストラ法）を組み合わせることで、RiemannInferはこの仕事を最小化するルートを迅速に特定し、実質的にモデルをより効率的な思考の連鎖へ導きます。

実際のモデルで実験が示すこと

著者らはRiemannInferをGPT-4o、Llama-3-405B、DeepSeek-V2-400Bなどの最先端LLMで、GSM8K、MATH500、StrategyQA、AGIEvalといった難易度の高い数学・推論ベンチマーク上でテストしました。いずれの場合も、これらのモデルをRiemannInferフレームワークでラップすると精度が数パーセント向上しました――性能の最前線では小さな差ですが重要な改善です――しかも速度は維持されるかわずかに改善されました。単純で純粋に線形な手法との比較では、隠れ状態の曲がった幾何学を無視すると性能が著しく低下することが示され、多様体的視点の重要性が強調されます。

大局：AIの推論に物理的な感触を与える

一般読者への要点は、著者らが大規模言語モデルの不透明な内部動作をより具体的なものに変えたことです：良い推論は滑らかで労力の少ない道に従うことに対応する、曲がった地形を横断する旅として。アテンションのパターンや推論のステップを距離、曲率、仕事といった幾何学的・物理的概念に基づかせることで、RiemannInferは結果を改善する実用的手段と、AIと連続空間の物理学との間の概念的架け橋の両方を提供します。現行の方法は近似的で改善すべき点も多いものの、この枠組みは、幾何学と物理の言葉でAIの思考過程を分析・最適化・場合によっては設計する将来のシステムへの道を示しています。

引用: Mao, R., Zhang, Z., Yang, M. et al. RiemannInfer: improving transformer inference through Riemannian geometry. Sci Rep 16, 6636 (2026). https://doi.org/10.1038/s41598-026-37328-x

キーワード: 大規模言語モデル, 幾何学的ディープラーニング, リーマン多様体, アテンション機構, 推論効率