適応型並列推論:効率的推論スケーリングの次なるパラダイム

要点

適応型並列推論(APR)は、新しいパラダイムで、Large Language Models(LLMs)が推論時に並列操作と逐次操作の間で計算資源を動的に割り当てることを可能にします。モデルが独立した推論スレッドを生成するタイミングとそれらをどのように調整するかを決定できるようにすることで、APRはエンドツーエンドのレイテンシを削減し、長い逐次推論チェーンに伴う「context-rot」を防止します。

逐次推論スケーリングの問題

逐次推論は探索量に比例して線形にスケールし、主に3つのボトルネックを引き起こします:

  1. Context-Rot:モデルがより多くの中間探索パスを生成するにつれ、有用な情報と妨害情報を区別することが難しくなり、性能が低下します。
  2. Latency:生成時間は推論の長さに比例して増加し、複雑なタスクでは数十分から数時間の待ち時間になることがあります。
  3. Compute Intensity:出力シーケンス長に沿ってスケールすると、推論が遅くなり信頼性が低下します。

並列推論は、モデルが複数のスレッドを独立して(互いのコンテキストに依存せず)かつ同時に(同時に実行)探索できるようにすることで、これらの課題に対処します。

固定並列性から適応制御へ

従来の並列推論手法は存在しましたが、通常はモデル外部から課された構造に依存していました。APRはこの制御をモデル自身にシフトします。

非適応的アプローチの制限

  • Simple Fork-and-Join(Self-consistency、BoN):軌道が独立にサンプリングされるため、冗長な計算が発生しがちです。
  • Heuristic-based Structured Search(Tree/Graph/Skeleton of Thoughts、MCTS):分解戦略に関する事前知識が必要で、すべての問題に対して知られているわけではありません。
  • Fixed-Stage Variants(ParaThinker、GroupThink、Hogwild! Inference):問題が恩恵を受けるかどうかに関わらず、特定の並列構造を課します。

APR の利点

適応型並列推論は、モデルが問題の複雑さに基づいて並列化のレベルを決定できるようにします。これにより、3つの主要な利点が得られます:

  • No Domain-Specific Heuristics:モデルは強化学習(RL)を通じて一般的な分解戦略を学習し、同時自己検証のような新たなパターンを発見します。
  • Reduced Redundancy:モデルは分岐前にユニークで重複しないサブタスクを生成でき、Best-of-N(BoN)とは異なります。
  • Dynamic Allocation:モデルは単純な問題については並列化しないことを選択でき、不要な並列化のオーバーヘッドを回避します。

適応並列性のための推論システム

適応的な並列ブランチの実行は通常、フォーク・ジョイン設計に従います:モデルは問題を同時実行のサブタスクに分割し、最終的な回答に結合します。主な技術的課題は KV キャッシュレベルで結果を集約することです。

エンジン変更アプローチ(Multiverse)

Multiverse、Parallel-R1、NPR のような手法は、推論エンジンを変更し、独立スレッドからの非連続 KV キャッシュメモリブロックをコピーして結合します。

  • Trade-offs:これにより2回目のプリフィルを回避できますが、システムの脆弱性(例:不正なポインタやキャッシュの追い出し)や位置エンコーディングの分布シフトが生じ、動作を合わせるために大規模な訓練と修正されたアテンションマスクが必要になります。

エンジン非依存アプローチ(ThreadWeaver)

ThreadWeaver はオーケストレーションをクライアント側に移します。クライアントは独立ブランチからのテキスト出力を連結し、エンジンは最終結論の KV キャッシュを生成するために2回目のプリフィルを実行します。

  • Trade-offs:プリフィル時に多少の計算冗長性が生じますが、プリフィルはデコードよりもはるかにコストが低いため、効率的です。エンジンの改変を回避し、標準的な因果注意と互換性があるため、逐次モデルへの適応が容易になります。

訓練と報酬設計

モデルに並列制御フロートークンの使用を教えるには、デモンストレーションデータと特定のインセンティブ構造の両方が必要です。

デモンストレーションと SFT

Supervised Fine-Tuning(SFT)は、並列制御フローの構文を教えるために使用されます。SFT が新たな推論能力を誘発するのか、既存の能力を特定のトークン構文に合わせるだけなのかについては議論が続いています。

報酬メカニズム

結果の正確性だけでは並列化を促すには不十分なことが多く、研究者はさまざまな報酬設計を検討しています:

  • Structure-only Rewards:スレッド数や構造の正しい使用に報酬を与えると、モデルが無意味なスレッドを生成して簡単に「ゲーム」できてしまいます。
  • Critical Path Rewards:レイテンシ最適化のために、報酬はクリティカルパス(因果的に依存するトークンの最長シーケンス)に焦点を当てます。ThreadWeaver は $1 - L_{\mathrm{critical}} / L_{\mathrm{total}}$ の報酬を使用し、クリティカルパスが全トークンの割合として小さくなるほど報酬が増加します。
  • Correctness Gating:非効率だが誤ったパスに報酬が与えられないように、最終回答が正しい場合にのみ並列化報酬が付与されます:$R = \mathbf{1}(\text{Correctness}) + \mathbf{1}(\text{Correctness}) \times (\text{parallelization metric})$。

パフォーマンスと未解決の質問

APR の評価は目的によって異なります。Multiverse と ThreadWeaver は精度とレイテンシを優先し、NPR は Genuine Parallelism Rate に焦点を当て、Parallel-R1 は APR を RL 後の性能向上のための中間訓練探索スキャフォールドとして使用します。

残された課題

  • Inference vs. Training:推論時の並列化が一貫して精度を向上させるのか、あるいは主な価値が訓練時の探索スキャフォールドにあるのかは不明です。
  • Stability:並列化報酬が緩和されると、モデルは逐次推論に戻りがちです。
  • Hardware Awareness:将来的な訓練では、利用可能な計算予算に基づいてハードウェアを意識した並列化判断が可能になるかもしれません。
  • Depth of Parallelism:現在の構造はフラットであり、Recursive Language Models(RLMs)を用いた再帰的並列化(深さ > 1)の探索が次のステップとして考えられます。

Sources