TEMPO:通过 EM 风格的交替优化扩展大推理模型的测试时训练 TEMPO: Scaling Test-time Training for Large Reasoning Models
用 EM 算法统一测试时训练,交替校准评论器与精炼策略来突破奖励漂移与多样性坍缩。
前置知识
测试时训练(Test-time Training, TTT)
指模型在推理阶段利用无标签或带标签的测试样本对自身参数进行在线更新,从而突破离线训练的能力上限。在 LRM 推理场景中,TTT 通常以 self-rewarding 强化学习形式出现,例如 TTRL 用多数投票、EMPO 用熵作为伪奖励。
TEMPO 的核心切入点是 TTT 在 LRM 上为何会快速饱和与多样性坍缩,不理解 TTT 的初衷和典型实现就无法体会 TEMPO 在 E-step 上补齐的关键差异。
PPO(Proximal Policy Optimization)
一种信赖域风格的策略梯度算法,通过 clip 目标函数限制新旧策略的更新幅度以稳定训练。在 LLM RLHF / RLVR 中,actor 与 critic 通常分别用 LLM 和 value head 实现。
TEMPO 用 PPO 完成 RLVR 初始化,并在 M-step 沿用 token 级策略梯度框架;若不熟悉 PPO 的 actor-critic 拆分,难以理解 critic $V_\phi$ 在 TEMPO 中如何作为可微奖励提供者。
EM 算法与 ELBO
Expectation-Maximization 在隐变量模型上交替做 E-step(估计隐变量后验)和 M-step(最大化期望似然),其优化目标是 evidence lower bound(ELBO)。ELBO 与真实目标只差一个非负 KL 项。
TEMPO 把 TTT 重新表述为 EM:critic 重校准=E-step(估计正确响应的后验),policy refinement=M-step(最大化期望对数似然),这是整篇论文的理论骨架。
Pass@K 与多样性坍缩
Pass@K 衡量 K 次独立采样中至少命中正确答案的比例,反映模型输出空间的多样性。多样性坍缩指 RL 训练后模型只生成单一思路,使 pass@K 随训练下降,是 self-rewarding TTT 的典型失败模式。
TEMPO 的核心卖点之一是同时提升 avg@K 和 pass@K,需要这两个指标才能理解作者用什么样的实证证据证明自己解决了 diversity collapse。
研究动机
近期的测试时训练(TTT)方法,例如 TTRL、EMPO、Theta-Evolve 等,承诺让 LRM 在推理阶段通过 self-rewarding 强化学习继续提升,但实际上它们在 AIME 这类高难基准上都会很快进入饱和期:在 OLMO3-7B-Base 上 TTRL 跑完只能让 AIME24 从 33.0% 提到 40.8%,继续迭代反而在 pass@8 上从 56.1 崩到 45.6;EMPO 也几乎与之重叠(41.6% / 43.3)。问题根源被作者归结为:没有 ground-truth 时,self-generated 奖励(多数投票、熵、self-consistency)天然被模型的初始能力所封顶;随着策略对自己偏好的答案越来越自信,奖励信号会与「真实正确性」持续脱钩,形成自我强化闭环,导致 reward drift、性能 plateau 与多样性坍缩同时发生。
本文的目标是论文希望在不依赖测试时刻 ground-truth 的前提下,仍然让 TTT 在长时间尺度上稳定提升 LRM 的推理表现,并同时维持输出多样性。具体目标有三:一是把已有 self-rewarding 方法的失败原因诊断清楚;二是给出一个可以同时利用无标签测试集 $D_u$ 和有标签数据集 $D_L$ 的统一框架;三是让该框架在数学、逻辑、STEM 等多领域推理任务上同时跑出 avg@K 与 pass@K 的双赢。
与已有工作不同的是,现有 TTT 工作的普遍做法是只更新 policy 而不更新 reward model,相当于把整个 EM 算法退化成了纯 M-step。TEMPO 的独特切入角度是显式补齐缺失的 E-step——周期性地在有标签数据 $D_L$ 上重新训练一个 token-level 的 critic $V_\phi$,让 critic 持续提供 grounded 的连续奖励,而不是用单调的多数投票或熵这样的启发式信号。这样 reward signal 与真实正确性之间的 KL gap 会被持续压低,policy 在数百步测试时训练下仍能稳步爬升。
核心方法
TEMPO(Test-time Expectation-Maximization Policy Optimization)把 TTT 在 LRM 上重新建模为 EM 优化。首先在带标签数据 $D_L$(如 DAPO-Math-17K)上用标准 RLVR/PPO 完成 actor $\pi_\theta$ 与 critic $V_\phi$ 的初始化;然后进入测试阶段,在一个无标签题目集合 $D_u$(例如 AIME 24/25、BBH、ZebraLogic)上做交替优化——E-step 在 $D_L$ 上最小化 critic 的 MSE 损失 $\mathcal{L}_{critic}=\mathbb{E}_{(x,y,I)\sim D_L}\|V_\phi(x,y_t)-I\|_2^2$ 让 critic 重新对齐真实正确性;M-step 在 $D_u$ 上采样 $y\sim\pi_\theta$,用 critic 末端值 $R=V_\phi(x,y_T)$ 作为整条轨迹的回报,用中间 token-level value $V_\phi(x,y_{1:t})$ 作为 baseline,得到 advantage $A_t=R-V_\phi(x,y_{1:t})$ 并跑带 clip 的策略梯度 $\mathcal{L}_{policy}=-\mathbb{E}_{x\in D_u,y\sim\pi_\theta}\sum_t A_t\log\pi_\theta(y_t|x,y_{<t})$。两条数据流(labeled 与 unlabeled)在训练图中呈「批萨饼式」交替。
与已有 TTT 方法的本质区别在于:TTRL 用二元多数投票 $q(y|x)\propto\mathbf{1}(y\in\mathcal{Y}_{majority})\pi_{\theta_0}(y|x)$,EMPO 用熵或 self-consistency 作为奖励,两者都是冻结的、heuristic 的奖励来源,本质上只执行 M-step,等价于 EM 的退化版本。TEMPO 的创新点是引入一个可学习的 critic $V_\phi$ 同时负责两件事:作为 $P(Correct|y,x)$ 的可微近似提供连续质量分数,作为后验估计的载体被周期性 E-step 校准。于是 auxiliary distribution $q(y|x)\propto V_\phi(x,y_T)\pi_{\theta_0}(y|x)$ 比二元投票平滑得多,且不会因为 policy 演化而漂移。这一设计让 ELBO 中的 $KL(q\|P(y|x,Correct))$ 在持续训练下不爆炸,是 TEMPO 能在 350 步以后仍在上升的根本原因。
方法步骤详情
完整流程对应 Algorithm 1:(1) 输入为 $D_L$、$D_u$、初始 $\pi_{\theta_0}$、初始 $V_{\phi_0}$、总迭代轮数 $N$;(2) Stage 1:在 $D_L$ 上跑 PPO 完成 RLVR 初始化,得到 RLVR-ready 的 actor 和 critic;(3) 进入交替循环,对每个 $k=1,\dots,N$:E-step 从 $D_L$ 中采样 $(x,y,I)$,按式 6 更新 $\phi$ 使 $V_\phi$ 重新逼近正确答案;M-step 从 $D_u$ 中采样 $x\sim D_u$、生成 $y\sim\pi_\theta$,按式 11 计算 token-level advantage $A_t=V_\phi(x,y_T)-V_\phi(x,y_{1:t})$,再按式 12 最小化带 clip 的策略梯度损失更新 $\theta$;(4) 输出 $\pi_{\theta_N}$。关键超参:math 实验 batch=256、mini-batch=64(≤8B),14B 用 128/32,最大响应长度 16K,sequence-clip dual ratio $3\times 10^{-4}$ 与 $5\times 10^{-4}$(沿用 GSPO 设置)。
技术新颖性
技术新颖性体现在三层。第一层是把 TTT-LRM 严格映射到 EM 并推导出 ELBO,从而把各类 self-rewarding 方法统一解释为 $KL(q\|P)$ 失控的退化情况,给后续工作提供了诊断框架。第二层是让 critic 同时承担「奖励生成者」和「后验估计器」两个角色,并通过周期性的 RLVR 风格 MSE 重训让 critic 与 actor 同步演化,避免 RL 中常见的 reward hacking 与 distributional drift。第三层是实现层面把 $V_\phi(x,y_{1:t})$ 当作 token-level baseline 而不是 GAE-style 多步回报,巧妙地把 critic 既当 critic 又当 value baseline 用,工程上不需要额外维护 TD(λ) 链路。
实验结果
TEMPO 在 5 个数学基准(AIME 24/25/26、Beyond AIME、OlymMath)以及 4 个通用推理基准(BBH、AGI Eval、ZebraLogic、GPQA-Diamond)上一致碾压 RLVR、TTRL、EMPO 三类基线,并保持多样性。数学方向最显眼的数据是 Qwen3-14B-Base:AIME24 avg@16 从 Zero-RL 的 42.3% 提升到 TEMPO 的 65.8%(绝对 +23.5 pp),AIME25 从 37.1% 提到 44.6%(+7.5 pp),pass@8 分别从 69.1/59.0 提到 73.3/60.0;OLMO3-7B-Base 的 AIME24 也从 33.0% 提到 51.1%(+18.1 pp)、AIME25 从 26.3% 提到 37.0%(+10.7 pp)。同时 TTRL/EMPO 在 pass@8 上出现明显多样性退化,例如 Qwen3-14B 在 AIME25 上 TTRL pass@8 跌到 45.8、EMPO 跌到 46.7,而 TEMPO 维持 60.0。Figure 4 还显示 Qwen3-14B 训练到 350 步仍未饱和,说明论文报告的数字其实是保守下界。Figure 5 给出更直接的证据——同一个已经收敛 192 步的 OLMO3,继续用有标签 PPO 几乎不动,而切换到 TEMPO 在 200 步内额外拿到 15+ avg@16 准确率。Figure 6 的消融把 critic 冻结后性能在约 100 步开始 plateau,证实 E-step 是不可省略的。通用推理上 OLMO3-7B 在 BBH、AGI Eval、ZebraLogic、GPQA-Diamond 上分别 +21.4、+24.5、+12.9、+10.5 pp,多项超过 7B 级专用模型如 General-Reasoner-7B(65.6/63.6/25.9/35.1)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME 2024 (Qwen3-14B-Base) | avg@16 | 65.8% | Zero-RL (PPO) 42.3% | +23.5 pp,同时 pass@8 = 73.3 (vs 69.1) |
| AIME 2025 (Qwen3-14B-Base) | avg@16 | 44.6% | Zero-RL (PPO) 37.1% | +7.5 pp, pass@8 60.0 (vs 59.0),TTRL/EMPO 同期 pass@8 仅 45.8/46.7 |
| AIME 2024 (OLMO3-7B-Base) | avg@16 | 51.1% | Zero-RL (PPO) 33.0% | +18.1 pp,pass@8 从 56.1 提到 61.6 |
| AIME 2025 (OLMO3-7B-Base) | avg@16 | 37.0% | Zero-RL (PPO) 26.3% | +10.7 pp |
| BigBenchHard (OLMO3-7B) | Avg@1 | 68.2 | Zero-RL (PPO) 46.8 | +21.4,超越 General-Reasoner-7B (65.6) |
| AGI Eval (OLMO3-7B) | Avg@1 | 62.4 | Zero-RL (PPO) 37.9 | +24.5 |
| ZebraLogic (Qwen3-8B) | Avg@1 | 33.9 | Zero-RL (PPO) 25.7 | +8.2 |
| GPQA-Diamond (OLMO3-7B) | Avg@8 | 32.4 | Zero-RL (PPO) 21.9 | +10.5,pass@8 从 62.1 提到 69.4 |
局限与改进
作者在第 6 节明确承认四点限制。第一,actor + critic 双模型带来额外的显存和算力开销,对比单模型的 TTRL 部署成本明显更高;第二,critic 重校准依赖有标签数据集 $D_L$,且 $D_L$ 的规模与分布会影响 critic 在 $D_u$ 上的泛化能力——这一限制在 $D_L$ 与 $D_u$ 偏移较大时会暴露;第三,实验仅覆盖数学、逻辑和 STEM 推理,未验证代码生成、agentic 等场景;第四,EM 视角虽提供原则化分析,但并未给出形式化的交替优化收敛保证,需要未来工作补足。我自己观察到的隐含限制包括:(i) TEMPO 的最大响应长度被截在 16K,对于更长链路的问题可能仍受上下文限制;(ii) Table 1 中 Qwen3-14B 在 Beyond AIME 与 OlymMath 上 pass@K 出现负增长(-3.7、-1.4),说明在已经较强的模型上 TTT 仍可能牺牲少量多样性换取精度;(iii) 全程依赖 gpt-oss-120b 作为通用域评分裁判,引入了额外模型依赖与潜在偏差。
独立分析的弱点
从工程和理论两侧看,TEMPO 都还有可改进之处。(1) 双模型 actor+critic 在 14B 规模上已把 batch 砍半(256→128),进一步扩展到 32B 或 MoE 时显存压力会更突出;改进方向是引入 LoRA 或共享 backbone 的轻量 critic 来降低内存。(2) E-step 调度频率是隐式超参,论文未给出系统化敏感性实验,建议未来加入按 KL 漂移自适应触发 E-step 的机制,避免人为周期性带来的额外噪声。(3) 在已经很强的 Qwen3-14B 上,Beyond AIME 的 pass@K 仍为负,说明 reward 重新加权可能在低质量样本密度大的子任务上误删部分合理但低频解,需要在 critic 训练阶段加入更长 horizon 的 TD 目标或优势归一化来缓解。(4) 当前 critic 是单标量 value,没有建模答案的多模态分布,因此对多解类推理任务(如有多种等价解法的题目)的覆盖仍然有限,可考虑 outcome-conditioned 或 token-prior 形式的扩展 critic。
未来方向
作者在结论中给出三条方向:证明交替优化的形式化收敛性、把框架拓展到 agentic / 多轮任务、研究 critic 校准频率与算力之间的 trade-off。基于论文成果,我自己还能想到几条延伸:(i) 用更便宜的 discriminative 替代品(如 process reward model 而非 outcome critic)替换 MSE-trained critic,进一步压低 GPU 成本;(ii) 把 EM 视角用到 RLHF 后期,让 critic 在人为偏好数据 + 模型自生成数据上交替更新,缓解 reward hacking;(iii) 把 TEMPO 与 self-consistency decoding、多数投票推理融合,让 pass@K 与 avg@K 同时进入 closed-loop 优化;(iv) 探索在 $D_L$ 极小甚至纯 synthetic-labeled 条件下的可行性,例如用基础模型自身 judge 充当弱监督以缩小对 labeled set 的依赖。
复现评估
代码与模型已通过论文首页的 Project Page、GitHub、HuggingFace 链接开源,说明至少在 repo 层做到了可访问。复现所需资源:实验覆盖三个 base 模型(OLMO3-7B、Qwen3-8B、Qwen3-14B),按 14B 模型 batch=128、max length 16K、TRL 风格 GRPO/PPO 训练,粗略估计算力需求在数百到上千张 H100/A100 量级;而 Figure 4 中把 Qwen3-14B 训练到 350 步的报告需要更长时间。$D_L$ 为公开数据集 DAPO-Math-17K(math 实验)和 Dolci-RL-Zero-General(通用域),$D_u$ 与评测集(AIME 24/25/26、Beyond AIME、OlymMath、BBH、AGI Eval、ZebraLogic、GPQA-Diamond)均为公开。复现难度主要是 GPU 资源和分布式 RL 工程经验,数据与算法层面对开源用户友好。
论文图表
已收敛的 OLMO3 在继续做有标签 PPO(蓝)几乎不动,改为 TEMPO TTT(橙)后 200 步内额外拿到 15+ avg@16。
把 TTT 与继续在同一分布上做监督强化学习直接对比,是回答 why bother TTT 最干净的消融实验。
伪代码按 Stage 1 初始化 + Stage 2 交替循环展开:每轮先在 $D_L$ 上做 critic 重校准(Eq. 6),再到 $D_u$ 上做策略更新(Eq. 12),共循环 $N$ 轮。
把整篇论文的形式化叙述压缩为可执行的工程流程,是读者复现时唯一需要的步骤清单。