[ICLR 2025 投稿] One-Token Verification：仅需一个 Token，让 LLM 推理正确性“一眼便知”

TL;DR

在 LLM 推理（尤其是数学推理）中，我们通常生成多个分支进行投票。但如何快速、精准地踢掉那些“一本正经胡说八道”的路径？本文提出的 One-Token Verification (OTV) 给出了优雅的方案：通过在序列末尾插入一个特殊的 [ToT] token，利用 LoRA 探测 KV Cache，模型就能在一次前向传播中告诉你当前推理路径的“靠谱程度”。该方法在 AIME 竞赛题目上表现惊艳，且比传统验证器节省了 90% 的计算资源。

核心定位：从“事后审查”到“实时监控”

在 OpenAI o1 开启的“强化推理”时代，Test-time Scaling（测试时计算缩放） 成了提升性能的关键。传统的做法是生成 N 个完整预测，再用一个 PRM（过程奖励模型）挨个打分。

• 痛点：PRM 往往是一个独立的、巨大的模型，推理极慢；且往往要等模型说完了才能打分，如果模型中途就走歪了，剩下的 token 全是浪费。
• OTV 的直觉：模型其实“知道”自己在心虚。与其训练一个外部法官，不如在模型内部装一个探测器，实时读取它推理时的“心跳”（KV Cache）。

核心机制：LoRA 探测器与 KV Cache 的化学反应

OTV 的精髓在于其非侵入性与高效性：

1. LoRA 门控机制：作者通过 LoRA 引入了一个验证分支，但这个分支只有在检测到 [ToT] token 时才会激活。平时模型依然是那个纯粹的 Reasoner，不会被带偏。
2. KV Cache 探测：传统的 Probe 只看最后一层的隐藏状态（损耗很大）。OTV 通过 Cross-Attention 让 [ToT] token 能够回溯之前所有层、所有位置的 KV Cache，获取最完整的信息流。
3. 线性斜率（Linear Ramp）伪标签：由于过程标签（Process-level labels）获取成本极高，作者巧妙地使用结果标签（Outcome labels）构造了一个从 0.5 到 1（或 0）平滑过渡的伪评分轨迹进行训练。

图 1：OTV 概念图。通过复用 KV Cache，模型在生成过程中随时可以插入 [ToT] 进行自我审核。

并行训练：这才是 Transformer 的味道

为了不拖慢训练，作者设计了一种三角掩码（Triangular Mask）机制。在一个前向传播中，可以在序列的每一个位置都插入一个 [ToT] 探针，同时计算出整条路径上每个点的正确性预测值。这种设计完美契合了 Transformer 的并行优势。

实验结果：以小博大的胜利

在 AIME24/25 数学竞赛任务中，OTV 的表现令人印象深刻：

• 胜过外部基准：即便与规模更大的外部奖励模型（如 AceMath-RM-7B）相比，集成在 4B 或 8B 模型内部的 OTV 依然取得了更高的加权投票（Weighted Majority Voting）准确率。
• 效率飞跃：通过 OTV 引导的早期终止，模型一旦发现得分暴跌就会果断弃路。实验显示，在不损失性能的前提下，推理 Token 总量节省了约 90%。

表 1：在 AIME 基准测试中，OTV 在各种规模的 Qwen3 和 DAPO 模型上均取得了 SOTA 结果。

深度可视化：正确与错误的“分水岭”

通过观察置信度轨迹（图 2），我们可以清晰地看到：

• 正确路径（红线）：随着推理步步深入，OTV 的信心值稳步攀升，呈现清晰的上升趋势。
• 错误路径（绿线）：即使模型在“一本正经”地写步骤，OTV 预测的置信度也会在关键逻辑点骤降或处于低位水平。 这种**区分度（Separability）**正是 OTV 能够实现高效剪枝的根本原因。

图 2：正确与错误路径的置信度动态演化，注意红绿线条之间的明显鸿沟。

局限性与展望

尽管 OTV 表现优异，但其目前重度依赖于最终结果的伪标签。未来的方向包括引入更精细的“不确定性”建模（例如区分模型是“真不会”还是“算错了”），以及将此机制推广到代码纠错、长程规划等对过程正确性要求极高的场景中。

总结

OTV 告诉我们：LLM 其实比我们想象中更了解自己的错误。通过一个简单的 [ToT] token 和 LoRA，我们就能解锁模型自带的“自我反思”天赋。这为未来的推理模型提供了一个极其轻量且高效的工程化方向。