LYNX: 信頼度制御推論のための動的エグジット学習

要旨

大規模推論モデルは、長い思考連鎖を生成することで複雑なタスクにおいて高い性能を発揮するが、しばしば「過剰思考」に陥る。すなわち、正答に十分な情報が得られた後も推論を継続してしまう。これは推論時の計算資源を浪費し、精度を損なう可能性がある。既存の早期終了手法は、追加のサンプリングやヒューリスティクスを用いたデコード操作、補助的な検証モデルへの依存、あるいは事後解析パイプラインとしてのみ機能するものが多く、形式的保証を欠いていた。本研究では、モデル自身の隠れ状態の認識を信頼度制御された停止判断に変換するオンライン早期終了機構LYNXを提案する。LYNXは、生成中に自然発生する推論の手がかり（例：「うーん」「待てよ」）に終了判断を付与し、強制終了による教師信号を用いてこれらの手がかりトークンにおける隠れ状態に軽量なプローブを訓練する。さらに、得られたスコアを分割コンフォーマル予測で包むことで、早期終了に対する分布非依存的な制御を実現する。重要な点は、このプローブを汎用数学コーパスで一度訓練・較正するだけで、ベンチマークやデコード温度、非数学タスクにわたって変更なく再利用できることである。1.5Bから32Bパラメータにわたる3つのモデルファミリーにおいて、ベースモデルごとに数学的に訓練された単一のプローブが、精度と効率性の優れたトレードオフを実現した。GSM8Kでは、ベースライン精度を維持または向上させつつトークン数を40-65%削減。MATH-500では最大12ポイントの精度向上と約35-60%のトークン削減。AIME 2024ではベースライン精度を維持しつつ50%超のトークン節約。非数学ベンチマークであるCommonsenseQAでは、ゼロショット転移によりわずかな精度向上と最大70%のトークン削減を達成した。最先端の早期終了手法と比較して、LYNXは競争力のあるあるいは優れたパレートフロンティアを提供しつつ、完全なオンライン動作、推論時の代理モデル不要、ユーザー調整可能な明示的な信頼保証を実現している。

English

Large reasoning models achieve strong performance on complex tasks by generating extended chains of thought, but they often "overthink": continuing to reason long after they have enough information to answer correctly. This wastes inference-time compute and can hurt accuracy. Existing attempts to stop early either manipulate decoding with extra sampling and heuristics, rely on auxiliary verifier models, or operate only as post-hoc analysis pipelines without formal guarantees. We introduce LYNX, an online early-exit mechanism that turns a model's own hidden-state awareness into confidence-controlled stopping decisions. LYNX attaches exit decisions to naturally occurring reasoning cues (e.g., "hmm", "wait") during generation, trains a lightweight probe on hidden states at those cue tokens using supervision from forced exits, and wraps the resulting scores in split conformal prediction to obtain distribution-free control over premature exits. Crucially, we train and calibrate this probe once on a generic mathematical corpus and reuse it unchanged across benchmarks, decoding temperatures, and even non-mathematical tasks. Across three model families spanning 1.5B to 32B parameters, a single mathematically trained probe per base model yields strong accuracy--efficiency tradeoffs. On GSM8K, LYNX matches or improves baseline accuracy while reducing tokens by 40--65\%; on MATH-500 it improves accuracy by up to 12 points with roughly 35--60\% fewer tokens; on AIME 2024 it recovers baseline accuracy with more than 50\% token savings; and on CommonsenseQA, a non-math benchmark, it transfers zero-shot with modest accuracy gains and up to 70\% fewer tokens. Compared to state-of-the-art early-exit methods, LYNX offers competitive or superior Pareto frontiers while remaining fully online, requiring no proxy models at inference, and providing explicit, user-tunable confidence guarantees.

LYNX: 信頼度制御推論のための動的エグジット学習

LYNX: Learning Dynamic Exits for Confidence-Controlled Reasoning

要旨

Support