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

Abstract: Many modern applications of deep learning involve training a neural network via a one-step prediction loss (e.g., $L^2$ regression, cross-entropy), but deploy the network by rolling out along its own predictions. Key examples include autoregressive language modeling, flow-based generative modeling, and robot policy learning. It is well-documented that these settings induce a phenomenon we call test-time feedback (TTF): the mismatch between the training/validation loss and downstream metrics of interest, such as task success rate and generation quality, which grows with task length. While data curation, architecture, and objective design have been proposed to combat train-test shift in TTF settings, this paper proposes optimization as a new design axis to mitigate error accumulation. Specifically, we introduce a new optimization paradigm called double-preconditioning (DoPr) uniquely tailored to the challenges of TTF. DoPr combines gradient-wise preconditioning, as in Adam and Muon, with activation-wise preconditioning (AP), such as in KFAC. We show that the addition of AP yields a drop-in intervention for increasing downstream model performance across a range of TTF settings. Interestingly, these gains in test-time performance do not consistently accompany improvements in validation loss, opening new questions about how to properly evaluate models trained with one-step supervised objectives.

• The paper introduces DoPr, an optimizer framework that wraps a gradient preconditioner with an activation-based preconditioner to mitigate test-time feedback shift.
• It demonstrates that conventional optimizers can achieve low validation loss while failing in deployment due to non-uniform feature learning and error amplification.
• Empirical results in robotics and language model fine-tuning show that DoPr consistently improves terminal rewards and downstream accuracies.

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

Introduction and Motivation

Deploying deep learning models in sequential decision-making contexts such as language generation and robotics introduces a fundamental mismatch between standard supervised training objectives and the natural downstream performance metrics. This paper introduces a rigorous framework for understanding test-time feedback (TTF) shift, the phenomenon where rolling out models on their own predictions induces input distributions differing from training, leading to compounding errors. Popular optimizers, especially those based solely on gradient statistics (e.g., Adam, Muon, Shampoo), are shown to often yield non-uniform feature learning that exacerbates these TTF-induced distribution shifts. The authors propose Double Preconditioning (DoPr), a drop-in optimizer framework that wraps any gradient preconditioner (GP) with a preceding activation-based preconditioner (AP), and systematically analyze its statistical and empirical properties.

Test-Time Feedback and Feature Learning Pathologies

In TTF settings (autoregressive language modeling, imitation learning in robotics), errors in model outputs recursively alter the input distribution, compounding even minor inaccuracies. The supervised training objective is thus poorly aligned with the model’s eventual deployment environment.

Figure 1: TTF shift arises as models are trained on per-step supervision but deployed autoregressively, amplifying errors in directions underweighted by the training loss.

The paper formally characterizes this compounding error via a behavior cloning framework in finite-horizon MDPs, defining TTF shift as the divergence between the state distributions under the demonstrator policy (training) and learned policy (deployment). Crucially, even models that achieve superior in-distribution validation loss can yield worse downstream trajectories if their feature learning is anisotropic or misaligned with sensitive directions of environment feedback.

Figure 2: Standard optimizers promote non-uniform feature learning, making models highly sensitive to feedback-induced distributional shift.

The analysis leverages minimal instances (linear dynamical systems) to show that layerwise non-isotropy in activations leads to pathological error amplification under sequential deployment. A constructed case explicitly demonstrates that lower validation loss does not guarantee superior test-time performance due to the accumulation of error in underrepresented subspaces.

Figure 3: Feedback amplifies distribution shift layerwise—errors at early layers irreversibly distort all downstream activations under TTF.

Empirically, models trained with conventional GPs can display rapid convergence of validation loss but fail to recover ground-truth feature subspaces, while their deployment performance degrades.

Figure 4: GP optimizers may accelerate validation loss without improving feature subspace alignment, resulting in poor TTF robustness unless AP is integrated.

The Double Preconditioning (DoPr) Framework

DoPr inserts a layerwise AP step prior to applying the GP of choice, fundamentally altering the trajectory of parameter updates. For a given layer, AP constructs the empirical activation covariance and preconditions the gradient matrix accordingly before passing it to the base optimizer (Adam, Muon, etc.). This operation can be interpreted as transforming the optimization geometry so that feature learning proceeds as if the activations were isotropic, thus neutralizing coordinate-dependent biases that standard GPs cannot correct.

Critically, DoPr is invariant to affine transformations of layer activations: affinely reparameterized networks subject to AP updates will follow identical optimization paths, a property not satisfied by vanilla SGD or Adam.

Figure 5: DoPr induces invariance to input affine transformations, whereas SGD optimization paths diverge under activation reparameterization.

The framework generalizes to convolutions and attention by derivations rooted in this invariance principle. Importantly, AP does not accelerate loss convergence per se but induces uniform improvement across sensitive directions, mitigating pathological error propagation under TTF.

Theoretical Consequences and Scaling Laws

The theoretical results include precise update forms for AP in both feedforward networks and more complex layers. A key proposition (adapted from [zhang2023meta]) demonstrates that single-step AP contracts feature subspace distance monolithically, while gradient descent's contraction is non-uniform (proportional to the spectrum of the activation covariance).

DoPr’s modular construction allows hyperparameter scaling rules, such as those prescribed by $\mu$P, to be inherited directly from the base GP. This property is empirically validated:

Figure 6: DoPr-AdamW inherits $\mu$P scaling rules; learning-rate and weight decay transfer seamlessly, demonstrating width-independent hyperparameter stability.

Empirical Evaluation Across Domains

Continuous Control and Robotics

Extensive experiments on state-based imitation learning benchmarks (e.g., Humanoid-v5, Half-Cheetah-v5) demonstrate that DoPr applied to AdamW, Muon, Signum, or AdaMuon consistently achieves higher terminal rewards, even though its training/validation losses are not always superior to the underlying GP.

Figure 7: Across optimizers, DoPr variants attain higher terminal rewards without monotonically improving train/val loss.

On high-dimensional, pixel-based robotic tasks (Robomimic), DoPr again elevates best-task success rates for both AdamW and Muon baselines.

Figure 8: Best rollout success rates on Robomimic tasks: DoPr-optimized agents outperform GP-only agents.

LLM Fine-Tuning

LLM supervised fine-tuning for math reasoning on GSM8K and MATH-500 tasks provides a compelling illustration: DoPr-AdamW and DoPr-Muon deliver higher downstream accuracies across learning rates and, counterintuitively, these improvements do not require lower final training losses.

Figure 9: DoPr-AdamW produces higher peak GSM8K accuracy versus AdamW alone across the learning rate sweep; this benefit does not trace training loss.

At larger scales (Llama 8B), DoPr not only improves absolute accuracy, but downstream metrics (GSM8K, CoT, MATH-500) correlate better with held-out validation NLLs relative to the baseline, indicating robust, reliable optimization alignment.

Figure 10: On Llama-3.1-8B SFT, DoPr-AdamW ensures best downstream accuracy tracks validation NLL, eliminating counterintuitive collapse in metrics observed for AdamW alone.

Practical Implications and Limitations

DoPr offers a lightweight, optimizer-agnostic recipe for improving robustness to deployment distribution shift with minimal changes to existing codebases. Hyperparameter tuning is simplified given inheritance from base optimizer rules. The computational and memory overhead is moderate, driven mainly by additional activation covariance computations and Cholesky decompositions per layer.

Practically, DoPr is especially advantageous for applications where the natural deployment trajectory diverges from IID assumptions: robotics, generative models, and system identification under covariate drift.

Theoretically, this work establishes that acceleration of in-distribution learning is neither necessary nor sufficient for downstream reliability in sequential contexts; the geometry of feature learning, and specifically its uniformity, is a critical axis of improvement orthogonal to standard GP benefits.

Conclusion

The paper identifies a structural limitation of widely used adaptive optimizers in TTF regimes and introduces Double Preconditioning as a theoretically principled and empirically validated intervention. By aligning feature learning isotropically, DoPr provides systematic gains in distributionally sensitive metrics that are untouched by in-distribution validation loss. The modularity and invariance properties facilitate broad integration with scaling law frameworks and emerging optimizer technologies. These insights suggest future directions in optimizer design—not only for further decoupling feature learning from loss acceleration, but also for extending deployment robustness in continually adaptive or open-ended learning systems.