TEMPO: Scaling Test-time Training for Large Reasoning Models

1. Problem Statement

Large Reasoning Models (LRMs) excel at complex tasks (e.g., Olympiad math, logic puzzles) by allocating extensive computation at inference time. However, their parameters remain static after training, preventing them from learning from unlabeled test instances.

Test-time training (TTT) aims to update model parameters on unlabeled test data to extend capabilities beyond the original training distribution. Current TTT methods (e.g., TTRL, EMPO, Theta-Evolve) rely on self-generated reward signals (entropy, majority voting, self-consistency) to refine policies without ground-truth labels. These methods suffer from two critical failure modes:

1. Performance Plateaus: As the model evolves, self-reward signals become self-reinforcing and overestimate the quality of narrow reasoning patterns, causing performance to saturate quickly.
2. Diversity Collapse: In the pursuit of higher average scores, models converge to a single dominant reasoning pattern, suppressing valid alternative solutions and degrading $pass@k$ metrics.

The root cause is the lack of external calibration; the reward signal drifts as the policy changes, creating a degenerate optimization loop.

2. Methodology: TEMPO

The authors propose TEMPO (Test-time Expectation-Maximization Policy Optimization), a framework that decouples reward generation from policy optimization using an alternating Actor-Critic design formalized through the Expectation-Maximization (EM) algorithm.

Theoretical Foundation

The objective is to maximize the expected log-probability of generating a correct response, $J(\theta) = \mathbb{E}_x[\log P(\text{Correct}|x; \theta)]$. Since correctness is unobserved on unlabeled test data ($D_u$), the authors derive an Evidence Lower Bound (ELBO):
$$J(\theta) \geq \mathcal{L}(q, \theta) = \sum_{x \in D_u} \sum_{y} q(y|x) \log \frac{P(\text{Correct}|y, x)\pi_\theta(y|x)}{q(y|x)}$$
where $q(y|x)$ is an auxiliary distribution approximating the posterior of correct responses.

The Alternating Procedure

TEMPO alternates between two steps to tighten the ELBO:

1. E-Step (Critic Recalibration):

   • Goal: Estimate the posterior distribution $q(y|x)$ (i.e., the probability that a response is correct).
   • Mechanism: A Critic model ($V_\phi$) is periodically trained on a small, labeled dataset ($D_L$) to predict token-level correctness.
   • Function: The critic provides a grounded, external reward signal. It approximates the posterior by reweighting the model's current outputs: $q(y|x) \propto V_\phi(x, y_T) \pi_{\theta_0}(y|x)$.
   • Significance: This step prevents reward drift by anchoring the signal to ground truth.

2. M-Step (Policy Refinement):

   • Goal: Update the Actor policy ($\pi_\theta$) to maximize the lower bound given the fixed $q(y|x)$.
   • Mechanism: The actor generates trajectories on unlabeled test questions ($D_u$). The policy is updated using Policy Gradient (PPO-style) where the advantage is calculated based on the critic's value predictions.
   • Advantage Calculation: $A_t = R - V_\phi(x, y_{1:t})$, where $R = V_\phi(x, y_T)$ is the final reward. This encourages tokens leading to high-value outcomes.

Algorithm Flow

1. Initialize Actor and Critic via standard RLVR on labeled data.
2. Loop:
   • E-Step: Sample from $D_L$, update Critic parameters $\phi$.
   • M-Step: Sample from $D_u$, generate responses, compute advantages using the updated Critic, and update Actor parameters $\theta$.

3. Key Contributions

1. TEMPO Framework: A scalable TTT framework that achieves sustained performance gains by interleaving policy refinement with periodic critic recalibration, avoiding the diversity collapse and plateaus of prior self-rewarding methods.
2. Unified EM Perspective: The paper provides a theoretical unification of TTT methods, characterizing existing approaches (TTRL, EMPO) as degenerate EM variants that omit the crucial E-step (posterior recalibration). This identifies the missing E-step as the root cause of scalability failures.
3. Empirical Validation: Extensive experiments across model families (OLMO3, Qwen3) and scales (7B to 14B) demonstrate that TEMPO significantly outperforms baselines in both accuracy and output diversity across math, logic, and STEM domains.

4. Experimental Results

The authors evaluated TEMPO on AIME 2024/2025/2026, OlymMath, BigBenchHard, AGI Eval, and ZebraLogic.

• Scalability (Math Reasoning):
  • OLMO3-7B: Improved AIME 2024 accuracy from 33.0% to 51.1% (+18.1 points).
  • Qwen3-14B: Improved AIME 2024 accuracy from 42.3% to 65.8% (+23.5 points).
  • Unlike baselines (TTRL, EMPO) which plateaued after initial gains, TEMPO showed a consistent upward trajectory over 350+ training steps.
• Diversity Preservation:
  • Baselines suffered from diversity collapse (e.g., TTRL's $pass@8$ dropped significantly as $avg@16$ rose).
  • TEMPO maintained high $pass@k$ scores (e.g., Qwen3-14B achieved 73.3% pass@8 on AIME 2024 vs. 56.7% for TTRL), proving it preserves multiple valid reasoning paths.
• Generalization:
  • TEMPO generalized to non-math domains. On BBH and AGI Eval, OLMO3-7B saw gains of +21.4 and +24.5 points respectively, outperforming specialized frontier models.
• Ablation Studies:
  • Frozen Critic: Removing the E-step (keeping the critic static) caused performance to plateau after ~100 steps, confirming the necessity of periodic recalibration.
  • Supervised Continuation: Continuing training on labeled data alone yielded negligible gains, proving that novel test-time problems are required for further improvement.

5. Significance

• Theoretical Insight: The paper reframes Test-time Training not just as a heuristic engineering trick, but as an incomplete EM algorithm. It establishes that external calibration (the E-step) is mathematically necessary to prevent the evidence lower bound from diverging.
• Practical Impact: TEMPO demonstrates that additional test-time compute can be effectively converted into sustained capability gains, breaking the "ceiling" imposed by offline training and self-rewarding loops.
• Robustness: By maintaining a grounded critic, TEMPO solves the "self-reinforcement trap," enabling models to explore diverse, high-quality reasoning paths rather than collapsing into a single mode.

Conclusion: TEMPO represents a significant step forward in making Large Reasoning Models truly adaptive, offering a principled, scalable, and diversity-preserving approach to test-time learning.