TEMPO: Scaling Test-time Training for Large Reasoning Models
| Authors | Qingyang Zhang et al. |
| Year | 2026 |
| HF Upvotes | 33 |
| arXiv | 2604.19295 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
Test-time training (TTT) adapts model parameters on unlabeled test instances during inference time, which continuously extends capabilities beyond the reach of offline training. Despite initial gains, existing TTT methods for LRMs plateau quickly and do not benefit from additional test-time compute. Without external calibration, the self-generated reward signal increasingly drifts as the policy model evolves, leading to both performance plateaus and diversity collapse. We propose TEMPO, a TTT framework that interleaves policy refinement on unlabeled questions with periodic critic recalibration on a labeled dataset. By formalizing this alternating procedure through the Expectation-Maximization (EM) algorithm, we reveal that prior methods can be interpreted as incomplete variants that omit the crucial recalibration step. Reintroducing this step tightens the evidence lower bound (ELBO) and enables sustained improvement. Across diverse model families (Qwen3 and OLMO3) and reasoning tasks, TEMPO improves OLMO3-7B on AIME 2024 from 33.0% to 51.1% and Qwen3-14B from 42.3% to 65.8%, while maintaining high diversity.
Engineering Breakdown
Plain English
TEMPO addresses a critical bottleneck in test-time training (TTT) for large reasoning models: existing methods plateau quickly and fail to benefit from additional compute at inference time. The core problem is that self-generated reward signals drift as the policy evolves, causing both performance stagnation and mode collapse where the model stops exploring diverse solutions. The authors propose TEMPO, which alternates between refining the policy on unlabeled test questions and recalibrating a critic model on labeled data, formalized through the Expectation-Maximization framework. This approach reveals that prior TTT methods are incomplete variants missing the critic recalibration step, enabling continuous scaling of test-time compute for reasoning models.
Core Technical Contribution
The key insight is modeling TTT as an incomplete EM algorithm where the E-step refines the policy on unlabeled test instances and the M-step recalibrates the critic on labeled validation data. This formalization reveals that prior methods like online policy gradient adaptation omit the critic recalibration, leading to reward signal drift and early plateaus. TEMPO's novelty lies in the periodic recalibration mechanism that prevents reward function degradation as the policy distribution shifts during test-time adaptation. By treating the problem through the EM lens, the authors provide both theoretical justification for the alternating procedure and empirical evidence that this structure enables scaling test-time compute without performance collapse.
How It Works
TEMPO operates in a two-phase cycle during inference on each test instance or batch of unlabeled questions. In the policy refinement phase (E-step), the model performs gradient updates using self-generated reward signals from the current critic, allowing the reasoning policy to adapt to the specific test distribution. Periodically, the system switches to the recalibration phase (M-step) where it freezes policy updates and retrains the critic model on a small labeled validation set, preventing reward signal drift from accumulated policy changes. The critic retraining uses the new policy distribution to generate fresh training data, ensuring the reward signal stays calibrated to the evolving model behavior. This alternating procedure repeats across multiple test-time steps, with the frequency of critic recalibration as a key hyperparameter. The framework maintains a circular buffer of labeled data for efficient critic updates without requiring external labeled test data for each instance.
Production Impact
In production systems, TEMPO enables inference-time model improvement that continues scaling with additional compute budget, eliminating the hard ceiling seen in current TTT approaches. This is particularly valuable for high-stakes reasoning tasks (e.g., math, code generation, scientific problem-solving) where spending extra compute at inference time to improve solution quality is acceptable. The main implementation change is adding a critic recalibration loop to your inference pipeline, requiring you to maintain a labeled validation dataset (typically 1-5% of training data) separate from test data. The trade-off is increased latency and memory during inference—each recalibration pass adds overhead proportional to the critic model size and validation set size, potentially adding 10-30% to per-sample inference time depending on recalibration frequency. Integration complexity is moderate: you need to orchestrate policy gradient updates and critic training, manage circular buffer state, and monitor reward signal statistics to detect drift.
Limitations and When Not to Use This
TEMPO requires a labeled validation set for critic recalibration, making it unsuitable for pure zero-shot deployment where no labeled data is available for the target distribution. The method assumes the critic can be meaningfully retrained on small validation sets—if your task has high label variance or the critic is severely undertrained, recalibration may introduce noise rather than calibration. Computational cost during inference is non-trivial; for models with large critics or tight latency budgets, the recalibration overhead may be prohibitive. The paper doesn't thoroughly address failure modes when the policy diverges far from the training distribution before the first recalibration, or how to set recalibration frequency optimally across different task types and model scales.
Research Context
This work builds on a growing literature in test-time adaptation and online learning for LLMs, particularly extending recent progress in test-time scaling for reasoning models. It connects to prior TTT methods that use self-training and online policy optimization but identifies and fixes a fundamental limitation—reward signal decay—that explains empirical plateaus in those approaches. The EM formalization provides a principled theoretical lens for understanding and designing TTT procedures, opening new research directions around optimal recalibration schedules and critic architecture choices. TEMPO's results likely benchmark on reasoning-heavy tasks like MATH, AIME, or code generation where test-time compute scaling has shown promise, advancing the state-of-the-art in continuous model improvement during inference.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::