Adaptive Nucleus Truncation for Long-Form Reasoning

本論文は、多様なタスクや生成予算において長文推論モデルの性能を安定させ、大幅に向上させるために、トークンの切り捨て閾値を動的に調整するエントロピー条件付きメカニズムであるAdaptive Nucleus Truncation Sampling (ANTS) を導入する。

要約

あなたは、非常に賢いが少し散漫な生徒に、複雑な問題の解き方を教えているところだと想像してください。あなたはプロンプト（問い）を与え、自分の思考プロセスをステップ・バイ・ステップで書き出すよう求めています。

問題は、生徒が書き進めるにつれて、話が逸れてしまうことです。彼らは無関係な詳細に気を取られたり、同じことを繰り返したり、行き止まりへと続く誤った道に迷い込んだりすることがあります。AIの世界では、これを「ドリフト（漂流）」や「不安定性」と呼びます。

この論文では、AIが特に長い回答を書かなければならない時に、脱線せずに軌道を維持するための新しいツールであるANTS（Adaptive Nucleus Truncation Sampling：適応型核截断サンプリング）を紹介しています。

以下に、その仕組みをシンプルな概念に分解して説明します。

1. 問題点：「固定フィルター」対「スマートフィルター」

AIが、次に進むべき数千もの可能性のある経路（単語）が交差する場所に立っていると想像してください。

• 従来の手法： 従来のツールは、固定された門番のように機能します。「状況がどうであれ、上位50%の経路のみを通す」と命じます。
  • 欠点： 時には、非常に集中する必要がある場合（数学の問題など）もあり、広い門はノイズを入れすぎてしまいます。また、時には創造性が必要な場合（物語の執筆など）もあり、狭い門は良いアイデアを切り捨ててしまいます。固定された門は、状況に応じて判断を変えることができません。
• ANTSによる解決策： ANTSは、スマートで適応的なガイドとして機能します。固定された門の代わりに、現在の状況を見て、「生徒がいま、どの程度混乱しているか？」を問いかけます。
  • 生徒が答えに対して非常に確信を持っている場合（混乱度が低い場合）、ガイドは門を狭めて集中力を維持させます。
  • 生徒が確信を持てていない場合（混乱度が高い場合）、ガイドは門を広げて、より多くの選択肢を探索できるようにします。

2. 秘訣：「ロジット」と「エントロピー」

これらの決定を下すために、ANTSは2つの特別なツールを使用します。

• ロジット（生のスコア）： ほとんどのAIツールは、最終的な「確率」（パーセンテージのようなもの）を見ます。しかし、本論文は、これはフィルターやリサイズが行われた後の「加工された写真」を見ているようなものだと主張しています。ANTSは、フィルタリングが行われる前の**生のスコア（ロジット）**を見ます。これは、料理される前の「生の食材」を見るようなものであり、AIが実際に「最善であると考えている」単語の、より鮮明な姿を捉えることができます。
• エントロピー（混乱計）： ANTSは「エントロピー」を測定します。これは、基本的には、その特定の瞬間におけるAIの混乱度や不確実性の尺度です。ANTSはこの計器を使用して、門をどの程度広く開けるかを決定します。

3. セーフティネット：「フォールバック・アーム」

これがこの発明における最も重要な部分です。
想像してみてください。スマートなガイド（ANTS）が、あまりにも親切になろうとしすぎて、経路を切り捨てすぎてしまい、生徒が立ち往生したり、デタラメな内容を生成し始めたりしたとします。

• フォールバック（退避策）： ANTSには、特別な「緊急ボタン」（フォールバック・アームと呼ばれます）があります。もしガイドが、経路を切り捨てすぎることが状況を悪化させていると気づいた場合、即座にボタンを押し、経路の切り捨てを完全に停止することができます。これにより、元の、フィルターのない方法へと戻ります。
• なぜ重要か： 昔のやり方では、もしフィルターが厳しすぎると、AIはどんどん悪化していくだけでした。しかし、ANTSがあれば、システムは「いつ厳格なのをやめて、自由に戻すべきか」を「学習」することができ、トレーニングのプロセスを安定させることができます。

4. 結果：会話が長くなるほど、より良くなる

研究者たちは、異なる「予算（AIが生成できる単語数の制限）」を用いて、大規模なAIモデルでテストを行いました。

• 短い予算（8K単語）： 結果はまちまちでした。例えばコーディングのようなタスクでは、ANTSを使用した方がAIの性能がむしろ低下しました。非常に限られたスペースで作業する場合、許可する単語を厳選しすぎると、最終的な結果を損なう可能性があるようです。
• 長い予算（16Kおよび32K単語）： ここでこそ、ANTSが真価を発揮します。許容される長さが増えるにつれて、ANTSは大幅に向上しました。
  • 指示への追従性： 複雑なルールに従って長時間記述するよう求められた際、ANTSはAIがルールを忘れたり、とりとめもなく喋り続けたりするのを防ぎました。
  • 数学と論理： 難解な数学の問題において、ANTSはAIが間違ったステップを「幻覚（ハルシネーション）」として生成することを防ぎ、より高いスコアをもたらしました。
  • 「Codeforces」のひねり： 興味深いことに、コーディングのタスクについては、短い長さでは性能が悪かったものの、長い長さでは驚異的な成果を出しました。これは、複雑なコーディングにおいては、正しいものに落ち着く前に多くのアイデアを探索する自由が必要であるが、それには十分なスペースが必要であることを示唆しています。

大きな教訓

この論文は、サンプリング手法（AIが次の単語を選ぶ方法）を、単にオン・オフできる単純な設定として扱うべきではないと主張しています。代わりに、以下の要素に基づいて挙動を変化させるダイナミックなコントローラーとして扱うべきです。

1. 回答にどれだけの長さが必要か。
2. その瞬間にAIがどの程度混乱しているか。
3. 現在の戦略が機能しているのか、それともリセットするために「緊急ボタン」を押す必要があるのか。

要するに、ANTSは、AIに対し、「いつ集中すべきか」「いつ創造的になるべきか」、そして「いつ『賢く』振る舞うのをやめて、自然に流れるままにすべきか」を知る方法を教えるシステムであり、長い会話の途中で迷子にならないように保証するものなのです。

技術概要

技術要約：長文推論のための適応型核型切断（Adaptive Nucleus Truncation）

問題提起

長文形式の言語モデルによる推論において、生成予算（generation budget）は、多段階の推論、自己洗練、およびツール利用をサポートするのに十分な大きさになることがよくあります。これらのレジームにおいて、デコーディング・サンプラーは単なる後処理の詳細ではなく、到達可能な推論パスと訓練軌道の安定性を決定する重要な制御メカニメントとなります。既存の切断手法（top-$p$、min-$p$、固定されたtop-$n_\sigma$など）は、無制限のサンプリングよりも改善されていますが、以下のような重大な限界があります。

• 固定された閾値： top-$p$やmin-$p$のような手法は、エントロピー、タスクの難易度、訓練段階、または生成予算の変化に適応できない固定された確率閾値に依存しています。
• 確率空間の歪み： 確率ベースの手法は、ソフトマックスおよび温度スケーリングの後に作用するため、確率正規化による歪みを継承します。その結果、同じ潜在的なロジット幾何学であっても、温度が変化すると異なる候補セットが生じる可能性があります。
• 適応性の欠如： 単一の固定された切断パラメータ（例：top-$n_\sigma$における固定された$n$）では、鋭い枝刈りを必要とする決定的なコンテキストと、広いサポートを必要とする曖昧なコンテキストの両方に同時に対応することはできません。さらに、不安定な強化学習（RL）フェーズにおいては、過度な切断は訓練の不安定性（エントロピー、KLダイバージェンス、および勾配ノルムのオーバーシュート）を招く可能性があります。

手法：適応型核型切断サンプリング（ANTS）

著者らは、**適応型核型切断サンプリング（Adaptive Nucleus Truncation Sampling: ANTS）**を提案しています。これは、固定されたデコーディング・ヒューリスティックとしてのtop-$n_\sigma$サンプリングを、適応的で予算を考慮したロールアウト制御メカニズムへと拡張したものです。ANTSは以下の3つのステージで動作します。

1. ロジット空間におけるサポート推定

確率ベースの手法とは異なり、ANTSは温度適用前のロジット空間において候補セットを選択します。語彙ごとの標準偏差$\sigma(\ell_t)$に基づき、最大ロジット$\ell_{t, \text{max}}$の周囲に近傍$N_t(n)$を定義します。
$$N_t(n) = \{ v \in V \mid \ell_{t,v} > \max_{u \in V} \ell_{t,u} - n \cdot \sigma(\ell_t) \}$$
この選択は、正のロジット再スケーリング（温度スケーリング）に対して不変であり、温度の変化に関わらず候補セットが一貫していることを保証します。

2. エントロピー条件付き閾値設定

切断近傍の幅$n_t$は固定されておらず、モデルの局所的な不確実性に基づいて適応します。これは以下のように計算されます。
$$n_t = n_0 + \gamma \cdot H(p^{(0)}_t)$$
ここで、$H(p^{(0)}_t)$はユニット温度分布のエントロピーです。高エントロピーのコンテキスト（曖昧な場合）はより多くの選択肢を保持し（大きな$n_t$）、低エントロピーのコンテキスト（決定的な場合）はより鋭く枝刈りされます。

3. オンライン・トンプソン・サンプリング・コントローラー

最適な切断強度$\gamma$を動的に決定するために、ANTSは有限の腕（arms）に対するトンプソン・サンプリング・バンディット・コントローラーを採用しています。

• 有限の腕： 対数間隔の$\gamma$値（$\gamma_k = 10^{\eta_k}$）を持つ$K$個の腕。
• フォールバック腕： $K+1$番目のクリティカルな腕として$\gamma_{K+1} = +\infty$を設定しており、これは実質的に切断を無効化（$N_t = V$）し、ベースラインのサンプラーを復元します。
• 報酬信号： コントローラーは、各腕によって誘導される分布のエントロピーに基づく内在的報酬を使用します。十分なサポートを保持している（平均より高いエントロピーを持つ）腕には、より大きな事後更新が行われます。
• 安定化メカニズム： フォールバック腕は「安全な脱出路」として機能します。RLのロールアウト中に切断が安全ではなくなった場合（KLダイバージェンスや勾配ノルムなどの指標がドリフトした場合）、コントローラーはフォールバック腕を選択して訓練を安定化させるように学習できます。

主な貢献

1. 適応型ロジット空間サンプラー： 候補セットの選択と、セット内における確率的変動を分離した、温度不変なサンプラーとしてのANTSの定式化。
2. エントロピー駆動型制御： 切断強度をオンラインで適応させ、訓練の不安定性を防ぐためのフォールバック腕を含む、エントロピー条件付きトンプソン・コントローラーの導入。
3. 生成予算スケーリングの評価： 8K、16K、および32Kの生成予算にわたる包括的な評価を行い、性能向上が一般に予算とともにスケールすることを実証。
4. 予算を考慮した切断： 短い予算でのコード生成における特定の失敗モードを特定し、生成予算およびタスクタイプと結合して学習される必要がある切断ポリシーの必要性を提示。

実験結果

本手法は、33B合計 / 4Bアクティブのスパース混合エキスパート（MoE）推論モデルを用いて評価されました。

• 予算によるスケーリング： パーセンテージベースのベンチマークにおける平均性能は、8Kで**+1.9ポイント、16Kで+3.8**、32Kで**+5.2**ポイント向上しました。
• 指示遵守（IFBench）： 最も強い改善を示し、32Kにおいて「Loose」で**+10.5**、「Strict」で**+10.8**ポイント向上しました。ANTSは、制約に違反する冗長または脱線的な継続へとドリフトすることを防ぎました。
• 数学的推論（AIME 2025）： 32Kにおいて**+7.0**ポイント向上し、その利得は8Kから32Kにかけてほぼ倍増しました。
• コード生成（Codeforces）： 予算との決定的な相互作用が明らかになりました。8Kにおいて、ANTSはより多くのソリューション・トークンを生成しているにもかかわらず、ベースラインを**-59 ELO下回りました。しかし、16Kおよび32Kではこの差を逆転させ、それぞれ+230および+212**のELO向上を達成しました。これは、切断が厳密に制約されたコードのレジームでは有害であるが、実行可能なプログラム空間がより大きい場合には有益であることを示唆しています。
• 科学的QA（GPQA）および知識（MMLU Pro）： ほぼ同等または緩やかな向上にとどまり、精密な単一回答の選択を必要とするタスクと比較して、ロジット空間の切断が少ないレバレッジしか提供しないことを示しました。
• トークン割り当て： ANTSはトークンの使用パターンを変化させました。中程度の予算では、「思考（thought）」トークンを増加させました。32Kでは、精度を維持しながら思考トークンを減少させており、これは単に長い連鎖を生成するのではなく、不要な継続を抑制していることを示唆しています。

意義と主張

本論文は、サンプラー設計を単なる固定されたデコーディング・ハイパーパラメータとしてではなく、推論時のスケーリングおよびロールアウト制御の不可欠な構成要素として扱うべきであると主張しています。

• 安定性： フォールバック腕の包含は本手法の成功の核心であり、固定された手法では逃れることのできない不安定な切断状態からシステムを回復させることを可能にします。
• コンテキストへの感度： 結果は、最適なサンプリング戦略がコンテキスト依存であり、生成予算、タスクタイプ（例：コード vs 数学）、および訓練状態によって変化することを実証しています。
• 効率性： ANTSはしばしばベースラインよりも早い段階で性能バンドに到達し、低品質または不安定なロールアウトの頻度を減らすことで、RL設定における実効的な計算効率を向上させます。

著者らは、ANTSがすべてのタスクにおいて既存のサンプラーを普遍的に圧倒するわけではないものの、特に長文の推論や指示遵守において、有用な推論パスを除去することなくノイズの多いテールを抑制することが極めて重要となる場面において、強力な生成予算スケーリングパターンを示すと結論付けています。