nowait-reasoning-optimizer

NOWAIT 推論オプティマイザー

論文「Wait, We Don't Need to 'Wait'! Removing Thinking Tokens Improves Reasoning Efficiency」（Wang et al., 2025）の NOWAIT 手法を実装します。

概要

NOWAIT は、学習を必要とせず推論時に介入する手法で、生成中に自己内省トークン（例: "Wait"、"Hmm"、"Alternatively"）を抑制します。これにより、モデルの有用性を損なうことなく、思考連鎖（CoT）の軌跡長を 27〜51% 削減します。

使用シーン

• 限られた計算資源で R1 スタイルの推論モデルをデプロイする場合
• 本番システムの推論レイテンシーを削減したい場合
• 推論タスクのトークンコストを最適化したい場合
• 冗長な CoT 出力を効率化する必要がある場合

対応モデル

モデルシリーズ	タイプ	トークン削減率
QwQ-32B	RL ベース	16〜31%
Phi4-Reasoning-Plus	RL ベース	23〜28%
Qwen3-32B	RL ベース	13〜16%
Kimi-VL-A3B	マルチモーダル	40〜60%
QvQ-72B-Preview	マルチモーダル	20〜30%

重要: NOWAIT は RL ベースのモデルで最も効果を発揮します。蒸留モデル（Qwen3-4B/8B/14B）では、内省トークンを抑制すると性能が低下します。

クイックスタート

1. 基本的な実装

from scripts.nowait_processor import NOWAITLogitProcessor

# Initialize processor for your model's tokenizer
processor = NOWAITLogitProcessor(tokenizer)

# Use during generation
outputs = model.generate(
    inputs,
    logits_processor=[processor],
    max_new_tokens=32768
)

2. 抑制されるキーワード

完全なリストは references/keywords.md を参照してください。主要なキーワードは以下のとおりです。

wait, alternatively, hmm, but, however, check, 
double-check, maybe, verify, again, oh, ah

仕組み

1. キーワードの初期化: 経験的分析から内省キーワードを特定します
2. トークンバリアントへの展開: キーワードを語彙内のすべてのトークンバリアントに対応付けます（例: "wait" → " wait"、"Wait"、" Wait"、".wait"、"WAIT"）
3. 推論中の抑制: デコード時に内省トークンのロジットを大きな負の値に設定します

Logits (Before)         Logits (After)
Wait     0.8     →     Wait     -inf
First    0.6     →     First    0.6
Hmm      0.5     →     Hmm      -inf
Let      0.4     →     Let      0.4

主な知見

効果の理由

• NOWAIT は自己内省を完全になくすわけではなく、不要な「待機」的推論をスキップするようモデルを導きます
• モデルは重要な判断ポイントでは、依然として必要な検証を行います
• より直線的で分かりやすい推論経路になります

RL モデルと蒸留モデルの比較

モデルタイプ	NOWAIT の効果	推奨
RL ベース（QwQ、Phi4、Qwen3-32B）	精度は安定し、トークンを大幅に削減	✅ 推奨
蒸留（Qwen3-4B/8B/14B）	難しいタスクで精度が低下	⚠️ 注意して使用

蒸留モデルは学習データの CoT 構造に強く依存しているため、内省トークンを除去すると推論パターンが乱れます。

統合例

HuggingFace Transformers

from transformers import AutoModelForCausalLM, AutoTokenizer
from scripts.nowait_processor import NOWAITLogitProcessor

model = AutoModelForCausalLM.from_pretrained("Qwen/QwQ-32B")
tokenizer = AutoTokenizer.from_pretrained("Qwen/QwQ-32B")

processor = NOWAITLogitProcessor(tokenizer)

response = model.generate(
    tokenizer(prompt, return_tensors="pt").input_ids,
    logits_processor=[processor],
    max_new_tokens=32768,
    do_sample=True,
    temperature=0.7
)

vLLM

from vllm import LLM, SamplingParams
from scripts.nowait_processor import get_nowait_bad_words_ids

llm = LLM(model="Qwen/QwQ-32B")
bad_words_ids = get_nowait_bad_words_ids(llm.get_tokenizer())

sampling_params = SamplingParams(
    max_tokens=32768,
    bad_words_ids=bad_words_ids
)

期待される結果

タスクタイプ	元のトークン数	NOWAIT のトークン数	削減率
数学（AIME）	15,000	10,500	30%
視覚的 QA（MMMU）	2,900	1,450	50%
動画 QA（MMVU）	1,700	1,250	27%

制限事項

• CoT のオーバーヘッドがすでに小さい非常に単純な問題では、効果が限定的です
• 蒸留モデルは、難易度の高いタスクで精度が低下する場合があります
• 分野によっては、モデルごとのキーワード調整が必要になる場合があります

参考資料

• 論文: arXiv:2506.08343v2
• キーワードの完全なリスト: references/keywords.md
• 実装: scripts/nowait_processor.py