⚡ electrical engineering

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss

This paper introduces Double Preconditioning (DoPr), a novel optimization paradigm that combines gradient-wise and activation-wise preconditioning to mitigate test-time feedback error accumulation in autoregressive and generative tasks, significantly improving downstream performance even when validation loss remains unchanged.

Original authors: Thomas T. Zhang, Alok Shah, Yifei Zhang, Vincent Zhang, Nikolai Matni, Max Simchowitz

Published 2026-06-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Thomas T. Zhang, Alok Shah, Yifei Zhang, Vincent Zhang, Nikolai Matni, Max Simchowitz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Practice vs. Performance" Gap

Imagine you are teaching a robot to walk across a room.

The Training Phase: You show the robot a video of a human walking perfectly. You tell the robot, "Copy this exact step." The robot practices on this video over and over. It gets really good at copying the video. In machine learning terms, this is minimizing the Validation Loss (getting a low score on the practice test).
The Real World (Test-Time): Now, you turn the robot loose. It has to walk across the room on its own. If it takes a tiny step slightly off-center, the next step has to compensate for that error. If it takes another tiny step off-center, the errors start to pile up. By the time it reaches the other side, it might be stumbling or falling over.

The paper calls this Test-Time Feedback (TTF). It's the phenomenon where a model that looks perfect on the practice data (the video) fails miserably in the real world because small mistakes compound into big disasters when the model relies on its own previous predictions.

The Old Way: Focusing Only on the Scoreboard

For a long time, engineers have built "optimizers" (the engines that teach the AI) with one goal: Make the practice score as low as possible as fast as possible.

Think of these optimizers like a coach who only cares about the robot's form while watching the video. The coach says, "Great job, you matched the video perfectly!" But the coach ignores the fact that the robot learned to walk in a way that is very fragile. If the robot has to walk on its own, that "perfect" form breaks down immediately.

The paper argues that low practice scores do not guarantee good real-world performance. In fact, chasing the lowest practice score can sometimes make the robot learn "bad habits" that only work when the teacher is holding its hand.

The New Solution: Double Preconditioning (DoPr)

The authors propose a new way to train these models called Double Preconditioning (DoPr). They describe it as a "plug-in" upgrade for existing training engines.

To understand how it works, imagine the robot's brain is made of layers of gears.

The First Gear (Activation Preconditioning): The paper says the old optimizers were biased. They learned to be perfect in directions where the data was easy, but they ignored directions where the data was messy or "stretched out."
- The Analogy: Imagine trying to learn to drive on a perfectly flat, empty highway (easy data). You get great at that. But if you try to drive on a bumpy, winding mountain road (messy data), you crash.
- The Fix: The first part of DoPr forces the robot to learn uniformly. It makes the robot practice driving on both the flat highway and the bumpy mountain road equally. It ensures the robot learns the shape of the road correctly, not just the easy parts. This is called Activation Preconditioning.
The Second Gear (Gradient Preconditioning): Learning on the bumpy road is hard and slow. If you just force the robot to learn that way, it might get stuck or move too slowly.
- The Analogy: This is like adding a turbocharger to the car. It helps the car move fast and smoothly without losing control.
- The Fix: The second part of DoPr uses the standard, fast optimizers (like Adam or Muon) to speed things up.

The Magic: DoPr combines these two. It first "straightens out" the learning path so the robot learns the right features (the bumpy road), and then it uses the turbocharger to make it fast.

The Surprising Result

Here is the most interesting part of the paper:

When the authors tested DoPr, they found that the robots didn't necessarily get better scores on the practice test (Validation Loss). Sometimes, the practice score was even slightly worse than before!

However, when they sent the robots out into the real world (Test-Time):

Robotics: The robots walked further and completed tasks more often.
Language Models: The AI wrote better code and solved math problems more accurately, even though it didn't predict the next word in the training text as perfectly as the old models.

The Takeaway

The paper concludes that we have been looking at the wrong scoreboard. We have been obsessed with how well a model minimizes its training error, but that doesn't mean it will succeed when it has to act on its own.

Double Preconditioning is a tool that teaches models to be robust and learn "uniformly" across all types of data, rather than just memorizing the easy parts. It's a way to build AI that doesn't just look good in the classroom, but actually performs well in the real world.

In short: DoPr stops the AI from "cheating" on the practice test by learning the right lessons for the real exam, even if it means the practice score doesn't look quite as shiny.

Technical Summary: Double Preconditioning (DoPr)

1. Problem Statement: Test-Time Feedback (TTF) and the Loss-Performance Gap

The paper addresses a critical disconnect in modern deep learning applications where models are trained on one-step supervised objectives (e.g., next-token prediction, per-step imitation) but deployed via sequential rollout along their own predictions. The authors term this phenomenon Test-Time Feedback (TTF).

In TTF settings—such as autoregressive language modeling, flow-based generative modeling, and robot policy learning—small errors in the model's predictions compound over time. This causes the distribution of states encountered during deployment to diverge from the training distribution (a phenomenon known as TTF shift).

The core problem identified is that standard optimizers (e.g., Adam, Muon, Shampoo) are designed to minimize validation loss ( $L_{val}$ ) under the training distribution. However, the paper argues that minimizing $L_{val}$ does not necessarily minimize the downstream task reward ( $R_{test}$ ) or the idealized loss ( $L_{ideal}$ ) defined over the model's own deployment distribution. Specifically, standard optimizers may induce imperfect feature learning where the model learns features that are optimal for the training distribution but fail to generalize under the feedback dynamics of TTF, leading to degraded downstream performance even when validation loss converges well.

2. Methodology: Double Preconditioning (DoPr)

To mitigate TTF shift without modifying data, architecture, or the training objective, the authors propose Double Preconditioning (DoPr). DoPr is a "plug-in" optimization framework that combines two distinct preconditioning strategies:

Activation Preconditioning (AP): A layer-wise preconditioner that uses the empirical covariance of layer activations ( $\hat{\Sigma}_z \approx E[zz^\top]$ ) to precondition the gradient.
- Mechanism: The raw gradient $G$ is multiplied by the inverse of the activation covariance matrix: $M = G \cdot \hat{\Sigma}_z^{-1}$ .
- Theoretical Motivation: The authors demonstrate that when layer activations are non-isotropic (ill-conditioned), standard gradient descent suffers from biased feature learning. AP corrects this bias, encouraging uniform feature learning across all directions in the state space. This ensures that the model learns features that are robust even in directions that are underweighted by the training loss but critical for preventing error accumulation under TTF.
- Invariance Principle: The paper establishes that AP induces invariance to affine transformations of the input distribution. If the input distribution is transformed affinely, the AP update rule ensures the layer outputs remain invariant, stabilizing the optimization trajectory against distribution shifts.
Gradient Preconditioning (GP): A standard gradient-based preconditioner (e.g., Adam, Muon, Signum) applied to the AP-preconditioned gradient.
- Role: While AP addresses the geometry of feature learning, it can be numerically unstable on its own. The GP component stabilizes the training process, accelerates convergence, and provides momentum.
- Integration: DoPr applies AP first to the gradient, then passes the result to the GP. This allows DoPr to leverage the stability and hyperparameter scaling rules of established GPs (like Adam or Muon) while inheriting the TTF-mitigating properties of AP.

Algorithmic Flow:
For a layer with weights $W$ and activations $z$ :

Compute batch gradient $G = \nabla_W \mathcal{L}$ .
Compute empirical activation covariance $\hat{\Sigma}_z$ .
Apply AP: $M = G \cdot (\hat{\Sigma}_z + \gamma \text{tr}(\hat{\Sigma}_z)I)^{-1}$ .
Apply GP: $D = \text{GP}(M, \dots)$ .
Update weights: $W \leftarrow W - \eta D$ .

3. Key Contributions

Theoretical Insight: The paper rigorously connects TTF shift to imperfect feature learning. It proves (via Proposition 3.2 and 4.1) that standard optimizers can converge to solutions with low validation loss but poor subspace alignment of features, which exacerbates TTF.
DoPr Framework: Introduction of a novel optimization paradigm that decouples directionality (handled by AP to ensure uniform feature learning) from adaptivity/stability (handled by GP).
Drop-in Compatibility: DoPr is designed as a drop-in replacement for existing optimizers. It does not require changes to the model architecture or data curation.
Hyperparameter Transfer: The authors show that because DoPr retains the GP component, existing hyperparameter scaling rules (e.g., $\mu$ P scaling for learning rates and weight decay) derived for base GPs like Adam or Muon transfer directly to DoPr, facilitating large-scale training.

4. Experimental Results

The authors evaluate DoPr across continuous control, robotics, and language modeling tasks. Key findings include:

Continuous Control (Humanoid-v5): DoPr consistently improves terminal reward compared to baseline optimizers (AdamW, Muon, Signum, AdaMuon). Notably, this improvement in downstream performance does not consistently correlate with improvements in training or validation loss; in some cases, DoPr achieves higher rewards with equal or slightly worse loss convergence.
Robot Policy Learning (Robomimic): On challenging tasks like "Tool-Hang" and "Transport," DoPr variants outperform baselines in success rates. Again, the gains in task success are not accompanied by uniform improvements in training loss.
Language Models (Mathematical Reasoning):
- 3B SFT: DoPr-AdamW achieves higher peak accuracy on GSM8K across various learning rates compared to AdamW, despite having comparable or higher token-level training loss.
- 8B SFT: In larger-scale fine-tuning, DoPr-AdamW demonstrates a stronger correlation between validation loss and downstream performance (GSM8K, MATH-500) compared to baseline AdamW. Baseline AdamW models often showed improved validation loss but degraded task performance at higher learning rates, whereas DoPr maintained robust performance.
- Cross-Task Conflict: DoPr reduces conflict between tasks; optimal learning rates for GSM8K-CoT were also near-optimal for MATH-500, a property not observed in baseline AdamW.

5. Significance and Claims

The paper claims that optimization is a new design axis for mitigating error accumulation in TTF settings, distinct from data curation or architecture design.

Decoupling Loss and Performance: The primary significance is the demonstration that improving downstream performance in TTF settings does not require (and may not even accompany) improvements in validation loss. This challenges the conventional wisdom that minimizing training loss is the sole proxy for model quality.
Feature Learning as a Lever: The work posits that "uniform feature learning" (ensured by AP) is the critical factor for robustness under TTF, rather than the speed of loss convergence.
Practical Utility: DoPr offers a practical, low-overhead intervention (a "plug-in") that improves model robustness in real-world deployment scenarios where models must roll out their own predictions.

The authors conclude by asking whether there is still an "optimizer free-lunch" available, suggesting that the decoupling of directionality and adaptivity in optimizers remains a fertile ground for research, particularly for tasks where the training objective diverges from the deployment reality.