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

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** ($\text{DoPr}$) uniquely tailored to the challenges of TTF. $\text{DoPr}$ combines gradient-wise preconditioning, as in $\text{Adam}$ and $\text{Muon}$, with activation-wise preconditioning ($\text{AP}$), such as in $\text{KFAC}$. We show that the addition of $\text{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.