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

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

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

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

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

arXiv:2606.06418v1 [cs.LG] 4 Jun 2026

Thomas T. Zhanga,† Alok Shaha,† Yifei Zhangb,† Vincent Zhanga,† Nikolai Matnia,‡ Max Simchowitzb,c,‡ a

University of Pennsylvania b Carnegie Mellon University c Amazon FAR

†

Equal contribution

‡

Equal advising

Figure 1: Standard optimizers, while effective at accelerating validation loss convergence, may induce poor feature learning. This can exacerbate distribution shift due to test-time feedback (TTF), the growing compounding errors as the model is deployed along its own predictions, ultimately leading to degraded downstream performance. We propose Double Preconditioning (DoPr) as a plug-in approach, where we apply a particular preconditioner to the layerwise gradient to encourage more “uniform” feature learning before passing it to a more standard optimizer. This reduces susceptibility to TTF, thereby improving downstream performance. a b

{ttz2, alokshah, vsyzhang, nmatni}@seas.upenn.edu {yifeizh3,msimchow}@andrew.cmu.edu

Code: https://tinyurl.com/3kfuhmpf Date: June 5, 2026 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.

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

Introduction

Recent years have seen a resurgence in the design of novel optimizers for deep learning, particularly via the use of “matrix-shaped” normalization or preconditioning. These optimizers adjust gradient updates by a “preconditioning” matrix to reweight the descent direction. The successes highlight improved convergence time on training or validation loss, such as Shampoo [Gupta et al., 2018, Anil et al., 2020, Shi et al., 2023] winning the AlgoPerf optimizer competition [MLCommons, 2024] as measured by an overall “holdout error per unit compute” metric, and Muon [Jordan et al., 2024] coming to prominence in greatly accelerating low “time to target validation loss” on NanoGPT pretraining and seeing adoption in training production-scale language models [Team et al., 2025]. These successes implicitly assume that low generalization error on the training objective implies good downstream performance, e.g. pass@k in the context of language reasoning. However, it has been widely noted that validation loss may not track natural performance metrics. The two diverge under a regime we call test-time feedback (TTF), ubiquitous in contemporary deep learning, in which one trains on per-step supervised objectives, while deployment involves rolling out multiple steps according to the model’s own predictions. Key examples include: (1) robot policy learning where per-step (or per-chunk) imitation of expert actions along expert trajectories, and deployment involves rolling out using the learned policy’s actions and (2) autoregressive LLM training in which models are trained via next-token prediction, yet deployed via autoregressive sampling. In these settings, sequential deployment can cause the distribution of inputs (e.g., context tokens in language modeling) to diverge from that of the training distribution, leading to distribution shift driven solely by the model’s own prediction errors. This “TTF shift” is widely attested in both robotics [Ross et al., 2011, Simchowitz et al., 2025] and in long-horizon language generation (e.g., [Bengio et al., 2015, Ranzato et al., 2015, Song et al., 2023]). Motivated by these findings and recent achievements in deep-learning optimizer design, we ask: Is there a design space for lightweight, generally applicable interventions on deep learning optimizers that target improved downstream performance? Contributions. First, we establish a rigorous connection between distribution shift induced by TTF and imperfect feature learning under today’s popular optimizers that precondition updates solely using gradient statistics (e.g. Kingma and Ba [2015], Jordan [2025]). Specifically, we observe that when activations/inputs at each layer are non-isotropic, feature learning suffers, which induces errors that can disproportionately drive distribution shift under TTF. Moreover, these cannot be remedied by updates at later layers (Proposition 3.2). Second, based on the above insights, we propose the Double Preconditioning (DoPr) framework, which couples an activation-based preconditioner (AP) using layer activations to compute a preconditioning matrix (e.g. Martens et al. [2018]) with more standard gradient-based preconditioners (GP), such as in Adam [Kingma and Ba, 2015, Loshchilov, 2017] or Muon [Jordan et al., 2024]. AP encourages more uniformly accurate feature learning, mitigating the TTF distribution shift described above, whilst benefiting from the training speed and stability afforded by GP. Whereas past variants of AP in the literature were introduced as standalone optimizers, DoPr shows how to incorporate AP as an drop-in modification to any GP optimizer, such as Adam and Muon, providing a stable training recipe and produces models with strong downstream performance. We develop an invariance principle under which DoPr updates can be systematically derived for any {architecture, GP } pairing, e.g., convolution and self-attention layers. Further, optimal hyperparameters can be reliably predicted by popular GP scaling heuristics [Yang and Littwin, 2023]. Finally, we carefully evaluate the capabilities of DoPr across a range of continuous-control, robotics, and language generation tasks. We find that DoPr consistently improves downstream performance— measured via natural, task-specific metrics across numerous applications— by intervening purely on the optimizer, without additional modifications to the data, training objective, or architecture. 2

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

Notation. We denote vector quantities by bold lower-case, and matrix/tensor quantities by bold upper-case. We use ⊙ to denote element-wise (Hadamard) product, ⊗ for Kronecker product, and vec(·) the column-major vectorization operator. We reserve W to denote weights in their tensor shape, and θ = vec(W) be the vectorized version. We use E[ f (x)] to denote the expectation of f (x). We reserve Σ for uncentered covariance matrices, e.g. Σx = E[xx⊤ ]. We denote layer-indexing with subscripts and (k) (optimizer) iterates with superscripts, e.g., Wℓ . Lastly, we use the index shorthand [L] = {1, . . . , L}.

Learning under Test-Time Feedback

We study the problem of learning under test-time feedback (TTF) — where our goal is to learn a model that is deployed in feedback with its own generations (as in language modeling) or with the environment (as in robotics). We formalize TTF as a problem of behavior cloning [Pomerleau, 1988] in a Markov Decision Process (MDP). Consider a finite-horizon MDP with states s t ∈ S, actions a t ∈ A, transitions s t+1 ∼ P(s t , a t ), and total horizon T . Models or policies π are maps from S → ∆(A), and we let Eπ [·] denote the expectation over actions a t ∼ π(s t ) and state transitions s t+1 ∼ P(s t , a t ). In behavior cloning, we learn a policy πθ : S → ∆(A) via demonstration sequences (s1 , a1 , . . . , s T , a T ). Given a per-example training loss Ltrain (π; s, a) measuring the distance between π(s) and a, and pairs of states and corresponding actions (s, a) (blue denoting demonstrator/data distribution), we train a learned policy πθ by minimizing P Ltrain (πθ ) = (s,a)∈data Ltrain (πθ ; s, a). (2.1) In the infinite data regime, and assuming (s1 , a1 , . . . ) are collected by a demonstrator πdemo deployed in the MDP, minimizing Eq. (2.1) approximately minimizes validation loss P T Lval (πθ ) = Eπdemo T1 t=1 Ltrain (πθ ; s t , a t ) . (2.2) When there is no risk of confusion, we will abbreviate Ltrain = L and Lval = L . Subsequently, πθ is deployed in the MDP, and we evaluate its expected reward under the distribution of states (ŝ1 , â1 , . . . , ŝ T , â T ) induced by πθ under a trajectory-wise reward function R(·): Rtest (πθ ) = Eπθ R(ŝ1 , â1 , . . . , ŝ T , â T ) . (2.3) Language modeling. Language modeling with tokens x t can be cast as an instance of TTF [Ouyang et al., 2022, Foster et al., 2024], where the states s t = x 1:t denote the current context, actions a t = x t+1 are the next token, and the dynamics are induced by concatenation: s t+1 = x 1:t+1 . The language model then produces tokens x t+1 ∼ πθ (· | x 1:t ). Next-token prediction uses the cross entropy loss Ltrain (πθ ; s t = x 1:t , a t = x t+1 ) = − log πθ (x t+1 | x 1:t ). Typical choices of reward include success rate or pass@k [Chen et al., 2021]. Robotic behavior cloning. In robotics, s corresponds to the robot and environment state, and a to the robot actions. In practice, s is replaced by robot observations (e.g., pixels, tactile), and actions a may be short sequences or “action-chunks”. Earlier works parameterize Gaussian policies πθ (· | s) = N(µθ (s), σ2 I), motivating an L2 training loss Ltrain (πθ ; s, a) = ∥µθ (s) − a∥2 whereas modern works use generative model parameterizations (e.g. flows or diffusion [Chi et al., 2023, Pan et al., 2025]). Typical rewards can include task success (e.g., an object was successfully moved to a desired location) or dense locomotive rewards (e.g., how far the robot has traversed).

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

Figure 2: Many settings involve using per-step supervised objectives along training sequences. However, due to rolling out along the model’s own predictions, mismatches between directions salient for Lval (πθ ) versus Rtest (πθ ) cause TTF shift. Hypothetically, instead optimizing for directions salient for Lideal (πθ )—often unavailable for offline training—would induce smaller TTF shift.

Optimizer Preconditioning for Minimizing Ltrain in TTF. In our TTF formulation, Ltrain is a standard supervised-learning loss. We consider an optimization algorithm alg that produces iterates {θ (k) }k≥1 in order to minimize Ltrain . Examples include mini-batch stochastic gradient descent (SGD), Adam (W), Shampoo [Gupta et al., 2018, Shi et al., 2023], and the recently popularized Muon optimizer [Jordan et al., 2024]. Most popular deep learning optimizers can be viewed as gradient-based preconditioners (GP), where the update direction is shaped purely using gradient information. Many prominent members therein, such as Adam, Shampoo, Muon etc., are P specifically structured approximations (e.g., diagonal, layer-wise) of the AdaGrad gradient-covariance k≥1 ∇Ltrain (θ (k) )∇Ltrain (θ (k) )⊤ [Duchi et al., 2011], and admit specific interpretations as gradient whiteners or normalizers. We provide full discussion in Appendix A. Past work has extensively studied how the choice of optimization improves performance only on the validation loss Lval , both in terms of final iterate limk→∞ Lval (θ (k) ), or the convergence rate, i.e., how quickly Lval (θ (k) ) approaches its limit. In this work, we consider a different problem statement: Problem Statement 1 How do we design an optimizer to minimize Ltrain , so as to maximize downstream performance of the resulting model πθ when evaluated at test-time, via Rtest (πθ )?

Distribution Shift under Test-Time Feedback

Whereas the validation loss Lval evaluates training loss under the distribution of states induced by πdemo , test-time reward Rtest considers the distribution of states under πθ . Due to sequential deployment in feedback with the MDP (e.g. through autoregressive generation in language modeling or with the environment in robotics), these two distributions do not agree (Figure 2). We call their consequent difference test-time feedback (TTF) shift. Definition 3.1 (TTF Shift). Let πdemo be the data-collection policy, and πθ a learned policy. Given a probability distance/divergence D (e.g. KL, Wasserstein Distance, etc. . . ), we define the TTF shift via π π π Dt (πθ ) := D(Ps tdemo , Pŝ θ ), where Pπ s denotes the marginal distribution of s t ∼ P . t

When TTF shift occurs, the distribution used to measure Lval , Ps tdemo , and that used for Rtest , Pŝ θ , t differ, which has been shown in the literature to degrade performance Rtest . The effects of test-time distribution shifts has been extensively studied in the robotics and sequence-modeling literature [Ross et al., 2011, Bengio et al., 2015, Simchowitz et al., 2025]; see Appendix A for a full account. Importantly, in contrast to general distribution shifts, TTF shift is due to model error propagated through sequential deployment. To reduce both TTF shift and its effects on downstream reward, past work has suggested that one strive to minimize an idealized validation loss with respect to actions from πdemo , but under 4

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

the state distribution of the learned policy πθ [Ross and Bagnell, 2010]: PT Lideal (πθ ) := Eπθ [ T1 t=1 Ltrain (πθ ; ŝ t , a⋆t )], ŝ t ∼ Pπθ , a⋆t ∼ πdemo (· | ŝ t ).

(3.1)

In particular, prior work [Ross and Bagnell, 2010, Foster et al., 2024] show that not only is Lideal a good surrogate for Rtest , but that minimizing Lideal would suffice to reduce TTF by implicitly correcting πθ towards πdemo when it veers off distribution (Figure 2(c)). However, minimizing Lideal in practice is often challenging, due to its circular dependence on πθ ’s own distribution and the necessity of collecting πdemo actions on πθ states. Adaptive or multi-stage data collection, such as the seminal DAgger algorithm [Ross et al., 2011], attempt to bridge this gap. However, we are often only given data from Pπdemo (e.g., finetuning corpora or teleoperator demonstrations). Hence, we can refine Problem Statement 1: Problem Statement 2 Can we use optimization preconditioning to encourage minimization of Lideal by mitigating TTF shift, even if we are only given data from πdemo ?

3.1

TTF Shift and Imperfect Feature Learning

We consider linear dynamical systems (LDS) as a minimal example of TTF. We consider a LDS with state s ∈ Rn , actions a ∈ Rm , state transitions s t+1 = As t + Ba t + w t , where w t ∼ N(0, Σw ) is process noise and A, B are matrices of appropriate shape. We assume that the demonstrator πdemo (s) = K⋆ s and imitator πθ (s) = Kθ s are both state feedback policies described by matrices in Rm×n , and the latter uses an L2 training loss: Ltrain (πθ ; s, a) = ∥Kθ s − a∥22 = ∥(K⋆ − Kθ )s∥22 = Ltrain (πθ ; s), where the last equality is redundant. Lastly, we consider a reward R(ŝ1 , â1 , . . . , ŝ T , â T ) = P T notes that the a argument PT − T1 t=1 Ltrain (πθ ; ŝ t ) = − T1 t=1 ∥(K⋆ − Kθ )ŝ t ∥2 . In this setup, we can describe all relevant quantities in closed form:1 P t−1 Proposition 3.1. Define the state-covariance matrix Γ t (K) = s=0 (A + BK)s Σw (A + BK)s⊤ , and Γ̄ t (K) = Pt πdemo π 1 = N(0, Γ t (K⋆ )), Pŝ θ = N(0, Γ t (Kθ )), and t s=1 Γ s (K). Then, Ps t t

Lval (πθ ) = (K⋆ − Kθ )Γ̄ T (K⋆ ) /2 F ,

(3.2)

2 . F

(3.3)

−Rtest (πθ ) = Lideal (πθ ) = (K⋆ − Kθ )Γ̄ T (Kθ ) /2 1

If D denotes the Wasserstein-2 distance, we have Dt (πθ ) = ∥Γ t (K⋆ ) /2 − Γ t (Kθ ) /2 ∥2F . 1

Validation loss need not track downstream performance due to TTF. Proposition 3.1 characterizes TTF shift in terms of the state-covariance matrices induced by K⋆ , via Γ t (K⋆ ) (resp. Kθ via Γ t (Kθ )). These same covariances weight the errors K⋆ − Kθ in both Lval and Lideal . Thus, TTF and downstream performance depends not only on error magnitudes, but in what directions the error between policies K⋆ − Kθ lies with respect to covariances. In Appendix B.5.1, we construct two linear policies πi (s) = Kπi s, i ∈ {1, 2} such that Lval (π1 ) ≪ Lval (π2 ), but Lideal (π1 ) ≫ Lideal (π2 ): Kπ1 − K⋆ has error along small eigenspaces of Γ T (K⋆ ), but through feedback between Kπ1 and the dynamics, these errors lead to large eigenvalues of Γ T (Kπ1 ), driving up Lideal . Understanding TTF from Feature Learning. To see how feature learning influences TTF, we overparameterize the linear learner policy πθ as a product of linear layers πθ (s) = Kθ s = Fθ Gθ s, where Fθ ∈ Rm×d , Gθ ∈ Rd×n ; similarly write K⋆ = F⋆ G⋆ , and assume rank(K⋆ ) = d for simplicity. A necessary 1

See Appendix B.5 for proofs and further details of the results in this section.

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

Figure 3: A depiction of how test-time feedback exacerbates distribution shift: errors in the network’s predictions affect the ensuing states, which changes their distribution away from the one seen in training. Changes in the state distribution (red) also affect the quality of the learned features (depicted as blue, to purple, to red) at intermediate layers (e.g., Proposition 3.2). Errors propagate layerwise, further exacerbating TTF. Notably, the learning signal (e.g., the loss gradient ∇Lval ) only supervises the predictions on the training distribution (blue).

condition for globally optimal imitation of πdemo (·) is accurately recovering the (row-)span of G⋆ , e.g., measured by a subspace distance: dist(G, G⋆ ) := PG P⋆⊥ op ,

(3.4)

where PG , P⋆⊥ are orthogonal projections onto rowsp(PG ), rowsp(G⋆ )⊥ . Vanilla GD, by minimizing Lval , always ensures that (Gθ − G⋆ )s is small along directions s which align with large eigenvalues of Γ̄ T (K⋆ ). However, large subspace distance dist(Gθ , G⋆ ) means that there exist states s for which G⋆ s is non-zero, but Gθ s vanishes. Effectively, Gθ “zeros out” certain directions of state-space, and therefore Fθ Gθ will as well. Errors in these directions can compound, driving up TTF, and in the worst case, no adjustment of later layers can mitigate the errors made in Gθ : Proposition 3.2 (Poor Feature Learning causes TTF Shift, Informal). There exists (F1 , G1 ), (F2 , G2 ), (F⋆ , G⋆ ) and linear dynamics A, B, Σw such that Lval (F1 G1 ) ≪ Lval (F2 G2 ), but dist(G1 , G⋆ ) > dist(G2 , G⋆ ), and for any possible “last-layer” Fθ , we have Lideal (F2 G2 ) ≪ Lideal (Fθ G1 ). Proposition 3.2 demonstrates that a model that appears better measured by in-distribution validation loss can hide significantly poorer feature learning, and can thus suffer vastly worse TTF, even assuming the downstream layers can be retrained. A natural question would be whether one can encourage better feature learning, as a separate consideration from validation loss performance. Key Takeaways 1. In TTF settings, distribution shift is inevitable: future predictions are rolled out along a model’s own (imperfect) predictions. 2. Therefore, under TTF, what is crucial is not necessarily the training objective, but rather the error directions exacerbated under feedback dynamics, see Proposition 3.1. 3. TTF further exposes a mismatch between in-distribution loss convergence and “feature learning.” Poor convergence on the latter can further exacerbate TTF layer-wise.

Double Preconditioning (DoPr)

Toward addressing the pathologies of TTF, we propose the double preconditioning (DoPr) framework, which be summarized as applying layer-wise an activation-covariance preconditioner (AP) onto the gradient, then passing the AP-gradient into a gradient preconditioner (GP) of choice, such as in Adam or Muon. For concreteness, consider an L-layer feedforward network: fθ (x) = W L φ(W L−1 · · · φ(W1 x) · · · ), 6

(4.1)

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

where θ is the concatenation of θ ℓ = vec(Wℓ ), ℓ ∈ [L]. Denote the intermediate activations zℓ ≜ φ(Wℓ−1 zℓ−1 · · · ), z1 ≜ x and pre-activations hℓ = Wℓ zℓ . We omit the layer index ℓ when it is clear from context. The core double preconditioning update can be summarized as: ˆ W L ( fθ ) Σ̂−1 , Σ̂z ≈ E[zz⊤ ] M=∇ z D = GP(G), W

(AP)

← W − ηD,

(DoPr)

ˆ W is the minibatch gradient, and Σ̂z the batch empirical uncentered covariance of z. We display where ∇ the core algorithm in Algorithm 1, and discuss derivations for general architectures in Section 4.2. Full details and practical features are described in Appendix B. Whereas much of the prior literature has focused on the role of gradient pre-conditioned optimization for stabilizing training and improving validation loss, we will show that DoPr’s use of activation preconditioning mitigates TTF, and therefore improve downstream model performance independent of validation loss performance. Gradient- v.s. Activation- preconditioning. We reAlgorithm 1: Double Preconditioning (Feedformark that activation preconditioning (AP) and many ward Layer, layer index ℓ suppressed) gradient preconditioners (GP’s) are motivated as different approximations to different curvature preconditionInput: learning rate η, weight decay λ, ers; see Appendix A for an extensive account. Prior work damping γ has modeled GP’s (e.g., Adam, Muon) as performing for k = 1 to K do (layer-wise) steepest descent with respect to a salient Sample batch B, |B| = n. norm ∥·∥ [Bernstein and Newhouse, 2024, Pethick et al., c(θ (k−1) ) G ← ∇W L P 2025]: Wnext = W − η arg max∥G∥≤1 〈G, ∇W L 〉.2 Vern (k−1) (k−1) ⊤ Σ̂z ← 1n i=1 zi zi sions of AP have appeared in prior literature as a stan−1 M ← G · Σ̂z + γ tr(Σ̂z )I (AP) dalone optimizer (see e.g., Amid et al. [2022], Benz∗∗ D ← GP(M, kwargs) (DoPr) ing [2022], Zhang et al. [2024]); we provide a full (k) (k−1) W ← (1 − ηλ)W − ηD account in Appendix A. Rather than establishing the end for convergence/acceleratory properties of AP, we demonstrate how it is directly relevant to TTF. However, just as vanilla gradient descent can suffer from numerical instability in deep networks, the AP update requires stabilization to converge meaningfully on modern taskloads. This motivates the combination of AP with GP, yielding DoPr. In the steepest descent GP framing described above, we can express DoPr by keeping the induced norm constraint ∥G∥ ≤ 1 while replacing the Euclidean inner product with the local “activation-covariance metric” induced by the activation covariance matrix: −1

Wnext = W − η arg max 〈G, ∇W L 〉Σ̂−1 = W − η arg max〈G, ∇W L · Σ̂z 〉. ∥G∥≤1

∥G∥≤1

For example, setting ∥·∥ = ∥·∥∞ as the entry-wise ℓ∞ -norm yields a Sign-Descent AP update (∼Adam), and ∥·∥RMS−RMS recovers a Spectral-Descent AP update (∼Muon).

4.1

Theoretical Motivation: AP Mitigates TTF via Isotropic Feature Learning

Proposition 3.2 therefore reveals that non-uniform feature learning at a given layer (i.e., larger dist(G1 , G⋆ )) can induce greater TTF shift and reduce downstream performance, regardless of what is learned at subsequent layers (the choice of F1 ). We now argue that AP optimizers exhibit more uniform feature learning by correcting a bias in GD that occurs when inputs are non-isotropic, i.e. Eπdemo [s t s⊤ t ] 2

Practical features like momentum and weight decay can be incorporated with slight modifications [Pethick et al., 2025].

6 × 100 4 × 100

SGD Adam DoPr-SGD

3 × 100 2 × 100

100 0

250

500

750

Subspace Distance

Validation Loss

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

0.75

SGD Adam DoPr-SGD

0.50

0.25 0

1000 1250 1500 1750 2000

250

500

750

Steps

1000 1250 1500 1750 2000

Steps

Figure 4: Mismatch between validation loss and feature learning. In-distribution validation loss converges and is even accelerated by GP. However, unless AP is applied, the feature subspace distance (3.4) converges poorly, which can exacerbate TTF (Proposition 3.2). Details in Appendix C.2.

is ill conditioned [Collins et al., 2021, Zhang et al., 2024]. Taken together, these findings imply that AP optimizers have the potential to mitigate TTF shift and its consequences. next,GD

Proposition 4.1 (Informal, adapted from Zhang et al. [2024]). Let η > 0 denote the step size, Gθ be next,AP Gθ after a step of full-batch gradient step on Ltrain (Fθ Gθ ) with respect to Gθ , and let Gθ denote the same with the AP update. Then, setting Σs = Γ̄ T (K⋆ ), we have: next,GD

P⋆⊥ = (Gθ − ηFθ ⊤ Fθ Gθ Σs )P⋆⊥ + ηFθ ⊤ F⋆ G⋆ Σs P⋆⊥

next,AP

P⋆⊥ = (I − ηFθ ⊤ Fθ )Gθ P⋆⊥ .

Gθ

(4.2)

Consequently, omitting some technical details, when cond(Γ̄ T (K⋆ )) ≫ 1, next,GD

dist(Gθ

, G⋆ ) ≈ dist(Gθ , G⋆ ),

next,AP

dist(Gθ

, G⋆ ) ≤ 1 − η λmin (Fθ ⊤ Fθ ) dist(Gθ , G⋆ ).

The proof of Proposition 4.1 can be found in Appendix B.4. The first statement exposes the bias of GD under non-isotropic covariance, and the non-convergence of GD and contraction of AP in feature space. Notably, under AP the resulting update (4.2) behaves as if the features were isotropic Σs = I, which is uniquely where GD and AP experience the same feature learning, noting G⋆ Σs P⋆⊥ = G⋆ P⋆⊥ = 0. Notably, we show in Figure 4 that pure GP (e.g. Adam, Muon) do not yield the same improvements, even if the Lval convergence is seemingly accelerated. Conceptually, the pathology illustrated in Proposition 4.1 can be inductively applied for multi-layer networks, where under TTF, poor feature learning at early layers will under TTF destroy signal for all downstream layers; see e.g., Davis and Drusvyatskiy [2025], Shumaylov et al. [2026] for detailed investigations of optimizer dynamics in similar settings.

4.2

Lifting AP to General Architectures: An Invariance Principle

In Proposition 4.1, we see that optimizers suffer from impaired feature learning, thus exacerbating TTF, when the covariances of layerwise activations are ill-conditioned, i.e. highly anisotropic, and observed that AP mitigates this by feature-learning as if the covariances were isotropic. Here, we leverage this observation to yield a general invariance principle, providing a reliable schematic for generalizing DoPr to general network layers, such as convolutions and self-attention. We first consider the following thought experiment: consider arbitrary layerwise full-rank affine transformations: zℓ ≜ Aℓ zℓ , Wℓ ≜ Wℓ A−1 ℓ , ℓ ∈ [L],

(4.3)

and denote the resulting network output fθ̄ (·), by construction, the pre-activations remain the same: Wz = Wz, and thus the network outputs are identical fθ̄ (x) = fθ (x). Remark 4.1 (Coordinate-dependent optimization paths). By definition of (4.3), we have fθ (x) = fθ̄ (x) for all x, x = A1 x. Assume the loss gradient ∇θ L is non-zero. Applying the backpropagation formula 8

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

Figure 5: When an affine transform is applied to the input distribution, with the initial weights transformed accordingly (4.3), the SGD trajectories (left) diverge, while the DoPr-SGD trajectories (right) match exactly, demonstrating the invariance induced by DoPr under affine transforms (Proposition 4.2). See Appendix C.3 for experiment details.

reveals that after a step of (full-batch) gradient descent: ∂L ⊤ z W+ = W − η E ∂h ∂L ⊤ W+ = W − η E z ∂ h ∂L ⊤ ⊤ =W−ηE z A , h≡h ∂h ∂ L −1 = W+ A−1 + η (A − A⊤ ), ∂W

(4.4)

the updated layer outputs are non-identical Wnext zℓ ̸= Wℓ zℓ barring the special case of orthogonal Aℓ ℓ for all ℓ ∈ [L], and thus the updated models no longer have identical outputs: fθ next(x) ̸≡ fθ̄ next(x). In other words, two NNs of the same architecture initialized to have identical output, can take different optimization paths depending on the “coordinate system” of the internal activations, independently of the final output of the network. Consequently, the activation distributions propagated from the training distribution, e.g., Σs = Eπdemo [s t s⊤ t ] versus Σs̄ = A1 Σs A1 , impart a coordinate-dependent bias on the optimizer trajectory, which Proposition 3.1 and Proposition 3.2 demonstrate can exacerbate TTF. Deriving AP for generic layer architectures boils down to the following principle: under an affine transformation on the input to the layer (and appropriate inverse transform to the applied weights) that preserve the layer output, what is the update rule that renders the layer outputs of the updated network invariant to the transform? In the feedforward case, this precisely yields the AP update as described in (DoPr). Proposition 4.2. Consider a feedforward network, and consider the weights and activations {(Wℓ , zℓ )}ℓ∈[L] under non-degenerate layer-wise affine transforms defined in (4.3). Then, given the following layer-wise update rule: ⊤ Wnext = W − η∇W L ( fθ ), Σ−1 z , Σz = E[zz ], next

the updated weights satisfy Wnext zℓ = Wℓ ℓ

(AP)

zℓ for all ℓ ∈ [L], z1 = x, z1 = A1 x.

In other words, an AP update for a given layer architecture is one that induces invariance to affine transforms. Thus, AP updates for convolutional or self-attention layers can be similarly derived; see Appendix C.3. This invariance principle can be viewed as a customization of Martens and Grosse [2015, Section 10], itself a specialization of the natural-gradient method’s parameterization invariance; see Appendix A. We emphasize that the activation-invariance of AP is an exact property of the optimizer trajectory (see Figure 5) and Proposition 4.2 does not imply accelerated loss convergence. 9

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

Loss

6.40

muP | DoPr-AdamW 5.20 4.80

5.60

4.40

4.80

4.00

4.00 2 w=128

1 0 1 LR Multiplier (log2) w=256

w=512

w=1024

2 w=2048

2 w=128

1 0 1 LR Multiplier (log2) w=256

w=512

w=1024

2 w=2048

Figure 6: µP scaling behavior. Left: DoPr-AdamW’s scaling trends under standard (SP) and AdamW’s µP parameterizations on a GPT2 model. We find base AdamW’s µP-scaling also enables hyperparameter transfer for DoPr-AdamW. Right: update-to-weight norm ratio scaling trend, under standard constant weight decay (SP) and AdamW’s weight decay µP scaling. See Appendix C.4 for full details.

Remark 4.2. Isotropic activations have been understood to possess desirable properties for enabling feature learning, and thus are often (approximately) enforced by normalization layers such as BatchNorm [Ioffe and Szegedy, 2015], LayerNorm [Ba et al., 2016] etc. However, misplacement of normalization layers may have unintended consequences on the conditioning and expressivity of the network. The AP gradient provides a complementary approach: the effective model change under (AP) evolves as if the activations were isotropic, where ∇W L ( fθ )Σ−1 z = ∇W L ( f θ ).

4.3

Gradient-Preconditioning and Hyperparameter Scaling for DoPr

Determining favorable hyperparameters for large-scale training runs on a novel optimizer can prove challenging. This is a key motivation behind the “maximum-update parameterization” (µP) [Yang et al., 2022, Dey et al., 2025], where the broad goal is to make hyperparameter choices invariant across network scales (e.g., width and depth); see Appendix C.4 for full discussion. Once made scale-invariant, hyperparameters can be tuned at small-scale and zero-shot transferred to the large-scale run. Recent work has demonstrated that in practice µP is largely determined by coarse statistics such as the magnitude of the update direction [Yang et al., 2023, Hong and Wang, 2025]. Thus, in addition to the practical stabilizing and accelerating properties of GPs, a key benefit of using GPs in DoPr is as follows: whereas deriving scaling rules for a new optimizer can often be laborious, the normalizing property of GP trivializes this process for DoPr. Observation 4.3. Normalization by definition returns outputs of the same norm regardless of input. Thus, substituting the raw gradient with the AP gradient does not affect the magnitude of the update direction. Consequently, DoPr hyperparameter scaling rules can be ported directly from established rules based on the GP of choice, e.g. Adam, Muon, Shampoo etc. [Everett et al., 2024, Qiu et al., 2025]. We exhibit the immediate transfer of hyperparameter scaling rules in Figure 6, where we show learning rate and weight decay scaling transfer. Key Takeaways 1. As shown in Section 3, what is crucial is not necessarily the training objective, but rather the error directions most exacerbated under TTF. Thus, the design of DoPr aims to decouple uniform feature learning and accelerated training/validation loss convergence. 2. Activation-preconditioning (AP) debiases the raw gradient direction from the statistics of the activations, which encourages uniform feature learning in all error directions, including those that are underweighted for the training objective but sensitive under TTF. 3. Gradient-preconditioning (GP), such as in Adam and Muon, stabilizes the AP gradient, accelerating training, while preserving the AP geometry. Furthermore, GP-ed updates allow seamless transfer of existing hyperparameter scaling rules.

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

Figure 7: Humanoid-v5 DoPr performance across AdamW, Muon, Signum, and AdaMuon. DoPr variants attain higher terminal reward which does not consistently correlate with train or validation loss improvements.

Capabilities

In Section 4, we derive DoPr to address the poor feature learning of standard optimizers in TTF settings, with the promise of inducing better downstream behavior. To demonstrate this, we run experiments on diverse tasks and metrics across continuous control, robot policy learning, and LLM training. As noted in Section 4.2, DoPr need not improve the train/val loss convergence compared to the base GP, which we will revisit across our experiments. We also provide in Appendix C.1 general operating guidelines of our optimizer used in the ensuing experiments.

5.1

Setting the Stage: Drop-in for Continuous Control

We first evaluate DoPr for imitation learning in Gymnasium [Towers et al., 2025], which precisely aligns with our formalism of behavior cloning in MDPs. Here, we focus on the Humanoid-v5 task and refer to Appendix C.5 for full details and further experiments. We consider four GP primitives: Adam, Muon, Signum, and AdaMuon, that roughly cover instantaneous versus momentumized entrywise- and matrixbased normalization. We use a residual MLP architecture [He et al., 2015], and conduct a full sweep over the optimizer hyperparameters and report the results over independent evaluation trajectories and seeded training runs. We additionally use an EMA-ed copy of the policy parameters [Block et al., 2024] for evaluation, which has been shown to be crucial for stabilizing BC policy performance in continuous control. We observe from Figure 7 that: 1. different base GPs all have similar terminal rewards, 2. DoPr always improves terminal reward compared to the baseline counterparts, 3. the train/validation losses of DoPr variants are not uniformly better than the base GP.

Figure 8: Tool Hang (PH) and Transport (PH) Best Success Rate for AdamW, Muon, and DoPr variants. Each curve shows the min/median/max over 3 random seeds. DoPr-variants outperform their baselines.

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

Final Train loss

Accuracy (%)

40 30 20 10

AdamW DoPr-AdamW

0 10 5

AdamW DoPr-AdamW

100

10 5

10 3 10 4 Learning rate

Figure 9: GSM8K 3B SFT sweep. Peak GSM8K accuracy across training steps vs. learning rate, comparing AdamW with DoPr-AdamW. The comparison between Muon and DoPr-Muon is provided in Appendix C.7.1.

5.2

Image-Based Robot Policy Learning

We evaluate DoPr on pixel-based imitation learning with generative policies on Robomimic tasks. We focus on challenging tasks Tool-Hang Proficient-Human (PH) and Transport (PH), which are not solvable to ∼ 100% under modern BC recipes [Pan et al., 2025], and test complementary capabilities: dexterous precision in Tool-Hang and long-horizon coordination in Transport. Following Pan et al. [2025], we train a flow-based policy [Lipman et al., 2023] with a U-Net backbone [Chi et al., 2023]. As in Section 5.1, we apply model EMA to stabilize policy performance across all settings. Tuned hyperparameters are listed in Tables 4 and 5, and full experiment details are in Appendix C.6. We report the best checkpoint success rate across training in Figure 8. We observe that DoPr invariably improves the task success rate compared to the base GPs AdamW and Muon, while yielding worse or equal training loss—see Figure 15. Notably, we see neither base GP is better than the other on both tasks simultaneously.

5.3

Language Models

We finetune an LLM for mathematical reasoning, where training optimizes token likelihood but task performance depends on sequence accuracy. Appendix C.7 contains full experimental details. 3B Supervised Fine-tuning. To characterize the behavior of DoPr, we first run a smaller-scale SFT experiment with the Llama-3.2-3B base model [AI@Meta, 2024, Grattafiori et al., 2024], using LoRA [Hu et al., 2022] on a 100K-sample subset of OpenMathInstruct-2 [Toshniwal et al., 2024] for one epoch. This experiment probes learning-rate sensitivity, sample efficiency, and the relationship between downstream accuracy and token-level training loss. We report our results in Figure 9. Across most learning rates, DoPr improves peak GSM8K performance. Notably, these improvements are not explained by lower token-level training loss: the final token-level training loss of DoPr is comparable to, or higher than, that of the base GP. We report similar trends for Muon (Figure 16) in Appendix C.7.1. 8B Supervised Fine-tuning: TTF at Larger Scale. Following the core SFT setup of OpenMathInstruct-2 [Toshniwal et al., 2024], we fine-tune Llama-3.1-8B [AI@Meta, 2024, Grattafiori et al., 2024] with LoRA [Hu et al., 2022] on the OpenMathInstruct-2 train_1M split for two epochs. We sweep learning rates for AdamW and DoPr-AdamW, and evaluate the final checkpoints on GSM8K, GSM8K-CoT, and MATH-500, and estimate validation NLL on a 10K held-out subset of the SFT data. We examine in Figure 10 how validation loss corresponds to downstream performance. For DoPr-AdamW, lower validation NLL generally corresponds to stronger downstream accuracy, especially on GSM8K-CoT and MATH-500. This is desirable, as holding data and architecture equal, one would expect a model with the best held-out NLL also has the best performance. However, this is not true for baseline AdamW. For larger learning rates, AdamW-trained models predictably improve in-distribution validation loss but task performance (e.g., GSM8K) degrades. This illustrates our thesis that, without intervention, 12

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

validation loss can be a poor proxy for full-sequence performance. Finally, DoPr-AdamW also reduces cross-task conflict: optimal learning rates on GSM8K-CoT are also (near-)optimal on MATH-500, whereas no AdamW checkpoints are simultaneously near-optimal for GSM8K and MATH-500 performance. This suggests that, as previewed in Section 4.1, DoPr encourages more robust feature learning. Additional results from the full sweep can be found in Appendix C.7.2. GSM8K

GSM8K-CoT

MATH-500

Final accuracy (%)

80 80

48 60 40

70 40 60

Better

20 24

50 5

28 0.

30 0.

5 31

0 33

5 28

Validation Loss

AdamW

DoPr-AdamW

Figure 10: 8B SFT sweep. We plot final accuracy on GSM8K, GSM8K-CoT, and MATH-500 against validation loss (NLL), computed on a 10K held-out subset of the training data. Each point is a final checkpoint from the learning-rate sweep while label numbers denote index of learning rates sorted by size: 2e-5, 5e-5, 7e-5, 1e-4, 2e-4, 5e-4, 7e-4. Error bars indicate one standard error.

We lastly remark additional experiments and discussion on flow-based generative modeling, another TTF setting, can be found in Appendix C.8. Key Takeaways 1. We identify across the different TTF domains such as robot policy learning and language modeling a mismatch between train/validation loss and downstream performance metrics. 2. We find that wrapping DoPr around existing GPs yields plug-in improvements on downstream performance, despite not necessarily accelerating train/validation loss, indicating an actionable facet of feature learning that is decoupled from loss convergence. 3. DoPr may enable commensurate scaling between training objective and downstream performance. 4. As predicted in Section 4.3, optimal GP hyperparameters in DoPr are inherited from the base GP.

Discussion

Toward the ultimate goal of optimizing for downstream performance, we: 1. introduce Test-Time Feedback as a unified setting for many modern applications, 2. identify the mismatch between features that accelerate training convergence versus those sensitive under TTF, 3. prescribe Double Preconditioning (DoPr) that combines Activation Preconditioning (AP) to equalize feature learning and Gradient Preconditioning (GP) to stabilize training. We provide evidence for TTF and DoPr being useful abstractions via experiments across distinct TTF applications. Importantly, we find that improved downstream performance may not accompany improved train/validation loss, suggesting that there remains a viable design space for improving deep learning optimization that is orthogonal to accelerating loss convergence. Notably, our proposed recipe of combining activation- and gradient-preconditioning is just one candidate of double preconditioning; we posit that decoupling considerations for directionality (e.g., AP) and adaptivity (e.g., GP) may lead to fruitful optimizer design. We conclude with a question: Given the incommensurability between validation loss and downstream performance under TTF, we conclude by asking: is there still an optimizer free-lunch left on the table? 13

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

Acknowledgments TZ gratefully acknowledges a gift from AWS AI to Penn Engineering’s ASSET Center for Trustworthy AI. TZ and NM are supported in part by NSF Award SLES-2331880, NSF CAREER award ECCS-2045834, NSF EECS-2231349, and AFOSR Award FA9550-24-1-0102. MS acknowledges support from a Google Robotics Award, and a Jane Street Fellowship. MS and TZ also thank Jeremy Cohen, Alex Damian, Hamed Hassani, and Behrad Moniri for helpful conversations.

References Kwangjun Ahn, Noah Amsel, and John Langford. Dion2: A simple method to shrink matrix in muon. arXiv preprint arXiv:2512.16928, 2025a. Kwangjun Ahn, Byron Xu, Natalie Abreu, Ying Fan, Gagik Magakyan, Pratyusha Sharma, Zheng Zhan, and John Langford. Dion: Distributed orthonormalized updates, 2025b. AI@Meta. Llama 3 model card, 2024. URL https://github.com/meta-llama/llama3/blob/ main/MODEL_CARD.md. Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251–276, 1998. Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry, volume 191. American Mathematical Soc., 2000. Ehsan Amid, Rohan Anil, and Manfred Warmuth. Locoprop: Enhancing backprop via local loss optimization. In International Conference on Artificial Intelligence and Statistics, 2022. Kang An, Yuxing Liu, Rui Pan, Yi Ren, Shiqian Ma, Donald Goldfarb, and Tong Zhang. Asgo: Adaptive structured gradient optimization. Advances in Neural Information Processing Systems, 38:126775– 126814, 2026. Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness-aware minimization leads to low-rank features. Advances in Neural Information Processing Systems, 36: 47032–47051, 2023. Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. arXiv:1607.06450, 2016.

Layer normalization.

arXiv preprint

Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7360–7371, 2022. Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015. Frederik Benzing. Gradient descent on neurons and its link to approximate second-order optimization. In International Conference on Machine Learning, 2022. 14

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

Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Power lines: Scaling laws for weight decay and batch size in llm pre-training, 2025. URL https://arxiv.org/ abs/2505.13738. Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409.20325, 2024. Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning. In Forty-second International Conference on Machine Learning, 2025. Adam Block, Dylan J Foster, Akshay Krishnamurthy, Max Simchowitz, and Cyril Zhang. Butterfly effects of SGD noise: Error amplification in behavior cloning and autoregression. In The Twelfth International Conference on Learning Representations, 2024. David E Carlson, Edo Collins, Ya-Ping Hsieh, Lawrence Carin, and Volkan Cevher. Preconditioned spectral descent for deep learning. Advances in neural information processing systems, 28, 2015. Roger Creus Castanyer, Johan Obando-Ceron, Lu Li, Pierre-Luc Bacon, Glen Berseth, Aaron Courville, and Pablo Samuel Castro. Stable gradients for stable learning at scale in deep reinforcement learning. arXiv preprint arXiv:2506.15544, 2025. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021. Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36:49205–49233, 2023. Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023. Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International Conference on Machine Learning, 2021. Michael Crawshaw, Chirag Modi, Mingrui Liu, and Robert M Gower. An exploration of non-euclidean gradient descent: Muon and its many variants. arXiv preprint arXiv:2510.09827, 2025. Alex Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, 2022. Felix Dangel, Bálint Mucsányi, Tobias Weber, and Runa Eschenhagen. Kronecker-factored approximate curvature (kfac) from scratch. arXiv preprint arXiv:2507.05127, 2025. Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memoryefficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022. Damek Davis and Dmitriy Drusvyatskiy. When do spectral gradient updates help in deep learning? arXiv preprint arXiv:2512.04299, 2025. 15

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

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition, 2009. Nolan Dey, Bin Claire Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, and Joel Hestness. Don’t be lazy: Completep enables compute-efficient deep transformers. arXiv preprint arXiv:2505.01618, 2025. Zhehang Du and Weijie Su. The newton-muon optimizer. arXiv preprint arXiv:2604.01472, 2026. John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011. Runa Eschenhagen, Alexander Immer, Richard Turner, Frank Schneider, and Philipp Hennig. Kroneckerfactored approximate curvature for modern neural network architectures. Advances in Neural Information Processing Systems, 36:33624–33655, 2023. Katie E Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter J Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, and Jeffrey Pennington. Scaling exponents across parameterizations and optimizers. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 12666–12700. PMLR, 21–27 Jul 2024. Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. Dylan J Foster, Adam Block, and Dipendra Misra. Is behavior cloning all you need? understanding horizon in imitation learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. Kevin Frans, Sergey Levine, and Pieter Abbeel. A stable whitening optimizer for efficient neural network training. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2026. Thomas Frerix, Thomas Möllenhoff, Michael Moeller, and Daniel Cremers. Proximal backpropagation. In International Conference on Learning Representations, 2018. Samir Yitzhak Gadre, Georgios Smyrnis, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, et al. Language models scale reliably with over-training and on downstream tasks. In International Conference on Learning Representations, volume 2025, pages 67661–67682, 2025. Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in neural information processing systems, 31, 2018. Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029–1054, 2021a. Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Journal of Statistical Mechanics: Theory and Experiment, 2021(12), 2021b. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad AlDahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. 16

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

Roger Grosse and James Martens. A kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pages 573–582. PMLR, 2016. Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, 2018. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. URL https://arxiv.org/abs/1512.03385. Naicheng He, Kaicheng Guo, Arjun Prakash, Saket Tiwari, Ruo Yu Tao, Tyrone Serapio, Amy Greenwald, and George Konidaris. Spectral collapse drives loss of plasticity in deep continual learning. arXiv preprint arXiv:2509.22335, 2025. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, pages 30016–30030, 2022. Letong Hong and Zhangyang Wang. On the provable separation of scales in maximal update parameterization. In Forty-second International Conference on Machine Learning, 2025. Stefan Horoi, Guy Wolf, Eugene Belilovsky, and Gintare Karolina Dziugaite. Less is more: Undertraining experts improves model upcycling. arXiv preprint arXiv:2506.14126, 2025. Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022. Zarif Ikram, Arad Firouzkouhi, Stephen Tu, Mahdi Soltanolkotabi, and Paria Rashidinejad. Crispedit: Low-curvature projections for scalable non-destructive llm editing. arXiv preprint arXiv:2602.15823, 2026. Sergey Ioffe and Christian Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37, page 448–456. JMLR.org, 2015. Satoki Ishikawa and Ryo Karakida. On the parameterization of second-order optimization effective towards the infinite width. arXiv preprint arXiv:2312.12226, 2023. Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, 2018. Keller Jordan.

cifar-10 airbench,

2025.

URL https://github.com/KellerJordan/

cifar10-airbench. Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan. github.io/posts/muon/. Michael Kelly, Chelsea Sidrane, Katherine Driggs-Campbell, and Mykel J Kochenderfer. Hg-dagger: Interactive imitation learning with human experts. In 2019 International Conference on Robotics and Automation (ICRA), pages 8077–8083. IEEE, 2019. Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from Adam to SGD. arXiv preprint arXiv:1712.07628, 2017. 17

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

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks, 2024. URL https://arxiv.org/abs/2305.17212. Atli Kosson, Jeremy Welborn, Yang Liu, Martin Jaggi, and Xi Chen. Weight decay may matter more than µp for learning rate transfer in practice. In OPT 2025: Optimization for Machine Learning, 2025. URL https://openreview.net/forum?id=RVxJ9rWbZk. Dmitry Kovalev. Understanding gradient orthogonalization for deep learning via non-euclidean trustregion optimization. arXiv preprint arXiv:2503.12645, 2025. Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced research). Atharva Kulkarni, Jacob Mitchell Springer, Arjun Subramonian, and Swabha Swayamdipta. Disentangling geometry, performance, and training in language models. arXiv preprint arXiv:2602.20433, 2026. Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approximation for natural gradient descent. Advances in neural information processing systems, 32, 2019. Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. Dart: Noise injection for robust imitation learning. In Conference on robot learning, pages 143–156. PMLR, 2017. Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE transactions on neural networks and learning systems, 29(5):1454–1466, 2017. Xi-Lin Li. Preconditioner on matrix lie group for sgd. arXiv preprint arXiv:1809.10232, 2018. Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, and Tuo Zhao. Normuon: Making muon more efficient and scalable. arXiv preprint arXiv:2510.05491, 2025. Wu Lin, Scott C Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger B Grosse. Understanding and improving shampoo and soap via kullback-leibler minimization. arXiv preprint arXiv:2509.03378, 2025. Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. Andy Zeyi Liu, Elliot Paquette, and John Sous. Spectral lens: Activation and gradient spectra as diagnostics of llm optimization. arXiv preprint arXiv:2605.05683, 2026a. Jiashun Liu, Zihao Wu, Johan Obando Ceron, Pablo Samuel Castro, Aaron Courville, and Ling Pan. Measure gradients, not activations! enhancing neuronal activity in deep reinforcement learning. Advances in Neural Information Processing Systems, 38:79546–79571, 2026b. Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023. Yuejiang Liu, Jubayer Ibn Hamid, Annie Xie, Yoonho Lee, Max Du, and Chelsea Finn. Bidirectional decoding: Improving action chunking via closed-loop resampling. In The Thirteenth International Conference on Learning Representations, 2025. Yuxing Liu, Jianyu Wang, and Tong Zhang. Optimizer-model consistency: Full finetuning with the same optimizer as pretraining forgets less. arXiv preprint arXiv:2605.06654, 2026c. 18

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

I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. Kevin Luk and Roger Grosse. A coordinate-free construction of scalable natural gradient. arXiv preprint arXiv:1808.10340, 2018. Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal. Learning dynamics and generalization in reinforcement learning. arXiv preprint arXiv:2206.02126, 2022. Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pages 23–40. Springer, 2024. Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkarni, Li Fei-Fei, Silvio Savarese, Yuke Zhu, and Roberto Martín-Martín. What matters in learning from offline human demonstrations for robot manipulation. In arXiv preprint arXiv:2108.03298, 2021. James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1–76, 2020. James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International Conference on Machine Learning, 2015. James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018. Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4): 667–766, 2022. AI Meta. Introducing meta llama 3: The most capable openly available llm to date. Meta AI, 2(5):6, 2024. MLCommons. Announcing the rithms benchmark competition,

results 2024.

the inaugural AlgoPerf: Training algoURL https://mlcommons.org/2024/08/

mlc-algoperf-benchmark-competition/. Behrad Moniri and Hamed Hassani. Asymptotics of linear regression with linearly dependent data. arXiv preprint arXiv:2412.03702, 2024. Yann Ollivier. Riemannian metrics for neural networks i: feedforward networks. arXiv preprint arXiv:1303.0818, 2013. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022. Matteo Pagliardini, Pierre Ablin, and David Grangier. The adEMAMix optimizer: Better, faster, older. In The Thirteenth International Conference on Learning Representations, 2025. Chaoyi Pan, Giri Anantharaman, Nai-Chieh Huang, Claire Jin, Daniel Pfrommer, Chenyang Yuan, Frank Permenter, Guannan Qu, Nicholas Boffi, Guanya Shi, et al. Much ado about noising: Dispelling the myths of generative robotic control. arXiv preprint arXiv:2512.01809, 2025. Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv:1301.3584, 2013. 19

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

Francisco Patitucci and Aryan Mokhtari. Adaptive optimization via momentum on variance-normalized gradients. arXiv preprint arXiv:2602.10204, 2026. J Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, Ian Foster, and Zhao Zhang. Kaisa: an adaptive second-order optimizer framework for deep neural networks. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–14, 2021. William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195–4205, 2023. Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018. Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constrained LMOs. In Forty-second International Conference on Machine Learning, 2025. Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Advances in neural information processing systems, 1, 1988. Shikai Qiu, Zixi Chen, Hoang Phan, Qi Lei, and Andrew Gordon Wilson. Hyperparameter transfer enables consistent gains of matrix-preconditioned optimizers across scales. arXiv preprint arXiv:2512.05620, 2025. Marc’Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. arXiv preprint arXiv:1511.06732, 2015. Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 661–668. JMLR Workshop and Conference Proceedings, 2010. Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627–635. JMLR Workshop and Conference Proceedings, 2011. Nicolas Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances in neural information processing systems, 20, 2007. Samuel Schapiro and Han Zhao. Towards understanding the role of sharpness-aware minimization algorithms for out-of-distribution generalization. arXiv preprint arXiv:2412.05169, 2024. Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through a crowded valleybenchmarking deep learning optimizers. In International Conference on Machine Learning, 2021. Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing systems, 35: 22955–22968, 2022. Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International conference on machine learning, pages 4596–4604. PMLR, 2018.

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

Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:2309.06497, 2023. Zakhar Shumaylov, Nathaël Da Costa, Peter Zaika, Bálint Mucsányi, Alex Massucco, Yoav Gelberg, CarolaBibiane Schönlieb, Yarin Gal, and Philipp Hennig. Muon is not that special: Random or inverted spectra work just as well. arXiv preprint arXiv:2605.11181, 2026. Chongjie Si, Debing Zhang, and Wei Shen. Adamuon: Adaptive muon optimizer. arXiv preprint arXiv:2507.11005, 2025. Max Simchowitz, Daniel Pfrommer, and Ali Jadbabaie. The pitfalls of imitation learning when actions are continuous. In Nika Haghtalab and Ankur Moitra, editors, Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pages 5248–5351. PMLR, 2025. Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phenomenon in deep reinforcement learning. In International Conference on Machine Learning, pages 32145–32168. PMLR, 2023. Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning, pages 32211–32252. PMLR, 2023. Jacob Mitchell Springer, Sachin Goyal, Kaiyue Wen, Tanishq Kumar, Xiang Yue, Sadhika Malladi, Graham Neubig, and Aditi Raghunathan. Overtrained language models are harder to fine-tune. In Forty-second International Conference on Machine Learning, 2025. Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15 (1):1929–1958, 2014. Weijie Su. Isotropic curvature model for understanding deep learning optimization: Is gradient orthogonalization optimal? arXiv preprint arXiv:2511.00674, 2025. Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, et al. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.20534, 2025. Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. Coursera: Neural Networks for Machine Learning, 4(2):26, 2012. Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. Shubham Toshniwal, Wei Du, Ivan Moshkov, Branislav Kisacanin, Alexan Ayrapetyan, and Igor Gitman. Openmathinstruct-2: Accelerating ai for math with massive open-source instruction data. arXiv preprint arXiv:2410.01560, 2024. Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U. Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, Rodrigo Perez-Vicente, Andrea Pierré, Sander Schulhoff, Jun Jet Tai, Hannah Tan, and Omar G. Younis. Gymnasium: A standard interface for reinforcement learning environments, 2025. URL https://arxiv.org/abs/2407.17032. 21

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

Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M Kakade. Soap: Improving and stabilizing Shampoo using Adam. In OPT 2024: Optimization for Machine Learning, 2024. Xi Wang and Laurence Aitchison. How to set adamw’s weight decay as you scale model and dataset size. In Forty-second International Conference on Machine Learning, 2025. URL https://openreview. net/forum?id=IszVnczhfz. Ishaan Watts, Catherine Li, Sachin Goyal, Jacob Mitchell Springer, and Aditi Raghunathan. Sharpnessaware pretraining mitigates catastrophic forgetting. arXiv preprint arXiv:2605.02105, 2026. Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, 2017. Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling, 2025. URL https://arxiv.org/abs/2409.15156. Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. In International conference on machine learning, pages 5393–5402. PMLR, 2018. Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. Advances in neural information processing systems, 32, 2019. Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv:2006.14548, 2020.

arXiv preprint

Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, 2021. Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit. arXiv preprint arXiv:2308.01814, 2023. Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022. Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813, 2023. Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: Feature learning in infinite depth neural networks. In The Twelfth International Conference on Learning Representations, 2024. Thomas T. Zhang, Daniel Pfrommer, Chaoyi Pan, Nikolai Matni, and Max Simchowitz. Action chunking and data augmentation yield exponential improvements in behavior cloning for continuous spaces. In The Fourteenth International Conference on Learning Representations, 2026. Thomas TCK Zhang, Leonardo Felipe Toso, James Anderson, and Nikolai Matni. Sample-efficient linear representation learning from non-IID non-isotropic data. In International Conference on Learning Representations, 2024. Thomas Tck Zhang, Behrad Moniri, Ansh Nagwekar, Faraz Rahman, Anton Xue, Hamed Hassani, and Nikolai Matni. On the concurrence of layer-wise preconditioning methods and provable feature learning. In International Conference on Machine Learning, pages 75793–75833. PMLR, 2025a. 22

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

Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. In The Thirteenth International Conference on Learning Representations, 2025b. Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient LLM training by gradient low-rank projection. In Forty-first International Conference on Machine Learning, 2024. Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning. In International conference on machine learning, pages 26982–26992. PMLR, 2022. Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 33:18795–18806, 2020. Liu Ziyin, Zhikang T Wang, and Masahito Ueda. Laprop: Separating momentum and adaptivity in adam. arXiv preprint arXiv:2002.04839, 2020.

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

Contents 1 Introduction

2 Learning under Test-Time Feedback

3 Distribution Shift under Test-Time Feedback 3.1 TTF Shift and Imperfect Feature Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

4 5

4 Double Preconditioning (DoPr) 4.1 Theoretical Motivation: AP Mitigates TTF via Isotropic Feature Learning . . . . . . . . . . 4.2 Lifting AP to General Architectures: An Invariance Principle . . . . . . . . . . . . . . . . . . 4.3 Gradient-Preconditioning and Hyperparameter Scaling for DoPr . . . . . . . . . . . . . . .

6 7 8 10

5 Capabilities 5.1 Setting the Stage: Drop-in for Continuous Control . . . . . . . . . . . . . . . . . . . . . . . . 5.2 Image-Based Robot Policy Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5.3 Language Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

11 11 12 12

6 Discussion

A Extended Related Work

B Extended Method Derivation and Mechanisms B.1 Full DoPr Algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2 AP For General Architectures: Derivations via Invariance . . . . . . . . . . . . . . . . . . . . B.3 Approximation Schemes for AP . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.3.1 Rank-1 + Diagonal . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.3.2 Spatially Uncorrelated Activations (SUA) . . . . . . . . . . . . . . . . . . . . . . . . . B.4 Feature Learning Theory . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5 Dynamical Systems Theory . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5.1 Mismatch between Lideal and Lval : a construction . . . . . . . . . . . . . . . . . . . B.5.2 Proof of Proposition 3.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.6 Limitations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

28 28 30 32 32 32 34 35 36 39 41

C Experiment Details and Additional Results C.1 General Operating Guideline . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.2 Feature Learning: Validation Loss versus Subspace Distance . . . . . . . . . . . . . . . . . . C.3 Activation Invariance of DoPr . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.4 Hyperparameter Scaling DoPr . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.5 State-based Imitation Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.6 Image-Based Robot Policy Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.7 Supervised Fine-Tuning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.7.1 3B SFT . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.7.2 8B SFT . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.8 Generative Modeling . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

42 43 43 44 45 48 49 49 49 51 53

D Additional Experiments and Setup D.1 Moment Grafting Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.2 Stable Ranks of Optimizer and Model Quantities . . . . . . . . . . . . . . . . . . . . . . . . .

55 55 57

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

Extended Related Work

Imitation learning and the compounding errors problem. The compounding errors problem has been recognized as a fundamental issue in policy learning from demonstration data [Pomerleau, 1988], where small one-step learning errors may shift the learned model away from the offline data distribution it was trained on, leading to potentially catastrophic failure. Traditional approaches to mitigating compounding error come from modifying data collection, either through iterative online collection of recovery behavior [Ross et al., 2011, Kelly et al., 2019, Foster et al., 2024] or noising action execution [Laskey et al., 2017, Zhang et al., 2026] to better regularize the policy geometry locally around the data distribution. More recently, modifying policy parameterization has been explored as an intervention for enabling consistent longer-horizon sequential deployment, such as the use of iteratively supervised methods (e.g., diffusion and flow-based parameterizations) [Chi et al., 2023, Pan et al., 2025] and action-chunking (i.e., multi-step prediction and execution) [Shafiullah et al., 2022, Chi et al., 2023, Liu et al., 2025]. Therefore, while data-collection and architecture design have been the primary avenues for mitigating compounding error, comparatively less focus has been directed to the role of the deep learning optimizer. An exception of note is Block et al. [2024], which identifies that small loss fluctuations in gradient-based neural network training can cause disproportionately large or even oscillatory performance fluctuations in behavior cloning settings, precisely aligning with our core thesis that small error directions as measured by training loss can be disproportionately harmful when rolled out. Their proposed intervention is a variance-reduction approach via maintaining an EMA-ed copy of the model parameters for evaluation, which is complementary to our preconditioning approach, and as such we adopt in our continuous-space BC experiments. In most other literature, the optimizer is obviated as an oracle yielding (approximate) empirical risk minimizers (see e.g., Ross et al. [2011], Foster et al. [2024], Zhang et al. [2026]). However, the role of specifically the deep learning optimizer has received comparably more attention in deep reinforcement learning, likely due to the compounding difficulties RL presents over supervised learning, such as input- and target- nonstationarity [Lyle et al., 2022, Sokar et al., 2023, Castanyer et al., 2025]. In this paper, we expose the potential suboptimalities of a wide range of deep learning optimizers that arise even in comparatively simpler offline supervised TTF settings, such as behavior cloning. Fisher Information, KFAC, and prior occurrences of Activation Preconditioning. The Fisher Information matrix is a core object in mathematical statistics. When used as a preconditioner for local optimization in maximum likelihood estimation, the resulting method is known as the natural gradient method. From this perspective, the Fisher Information in some cases is related to other curvature matrices such as the Hessian (or the Generalized Gauss-Newton approximation thereof), thus in settings where feasibly implementable, the natural gradient method often greatly accelerates optimization; see e.g., [Amari, 1998, Amari and Nagaoka, 2000, Martens, 2020]. Specified to neural networks, which have too many parameters for full second-order methods to be feasible, factorized versions of the Fisher Information have been proposed [Roux et al., 2007, Ollivier, 2013, Pascanu and Bengio, 2013]. Most notably, under certain statistical independence assumptions, Martens and Grosse [2015] derive a block-wise Kronecker-Factored approximation of the Fisher Information, which renders preconditioning into left- and right- preconditioning the layerwise gradient matrix A−1 ∇W L B−1 , which spawned the Kronecker-Factored Approximate Curvature (KFAC) line of neural network optimizers, seeing consistent development and adaptation [Grosse and Martens, 2016, George et al., 2018, Pauloski et al., 2021, Eschenhagen et al., 2023, Dangel et al., 2025]. However, most literature deriving from KFAC focus on the standard signal of improving train/validation convergence; due to the additional memory overhead of leftand right- preconditioner states per-layer, the numerical nuance of matrix inversion, and relative scarcity rules-of-thumb compared to popular GPs, AdaGrad-based (i.e., GP) methods typically remain favored in mainstream training set-ups, though KFAC finds utility in many scenarios that directly require curvature estimation [He et al., 2025, Ikram et al., 2026]. Whereas KFAC involves a “left-side” derivative-based 25

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

preconditioner, its “right-side” preconditioner is simply the activation covariance, which only requires forward passes to compute. This has motivated “right-side-only” approximations to KFAC (i.e., AP in our parlance) [Benzing, 2022], which have been independently derived from proximal or local layer-wise loss perspectives [Frerix et al., 2018, Amid et al., 2022]. We note AP has also been encountered from shallow neural-network learning theory [Zhang et al., 2024], which we discuss in the dedicated feature learning theory paragraph. Furthermore, a perhaps underutilized perspective of KFAC’s approximation of the natural gradient method (or equivalently an alternate derivation of it) is reparameterization invariance; recall that natural gradient flow traces a path over distributions that is invariant to smooth reparameterizations. Similarly, KFAC, as an approximation thereof, induces a layerwise invariance to affine reparameterizations, see Martens and Grosse [2015, Section 10]. The resulting optimizers can be equivalently recovered by a Euclidean-norm steepest descent formulation under a weighted local inner product: 〈G, ∇θ L 〉Σ(θ ) , s.t.∥θ ∥2 ≤ η, where 〈·, ·〉Σ(θ ) is the classical “Fisher metric” for the natural gradient method, and the self-explanatory “KFAC” metric for KFAC [Luk and Grosse, 2018]. Similarly, we may further reduce KFAC to the one-sided AP preconditioning that preserves the invariance property that we argue is salient for TTF in Proposition 4.2, where 〈·, ·〉Σ(θ ) is simply the activation-covariance reweighted inner product. Departing from the Euclidean-norm steepest descent to other (layerwise) weight norms recovers DoPr for different GPs [Pethick et al., 2025]. We also note concurrent work in [Du and Su, 2026], which independently derives activationpreconditioning from an “isotropic curvature model” [Su, 2025]. In fact, DoPr-Muon in our parlance precisely recovers their proposed algorithm. Crucially, the focus of their work is on improving pretraining loss convergence, while our central thesis studies the benefit of AP independent of loss convergence; notably, we found that DoPr-Muon does not universally accelerate training/validation loss convergence across applications. Overall, we opted to use AP rather than the full KFAC-preconditioned gradient (which would additionally introduce a “left-side” preconditioner) due to subtleties of either requiring an MLE interpretation of the loss function [Martens and Grosse, 2015], contending with possible “incorrectness” of the approximation due to using the Empirical Fisher [Kunstner et al., 2019], and most importantly computing and possibly storing an additional optimizer preconditioner state. Anecdotally, using APinstead of KFAC- preconditioning in DoPr sufficed to provide the boosts in downstream performance documented in this paper, while avoiding compounding numerical instability or memory boundedness; however, we emphatically do not rule out the possibility for an efficient implementation or adaptation to yield further improvements. Ultimately, we view DoPr as a framework for decoupling directionality (e.g., from preconditioning) and self-normalization (e.g., from popular GP mechanisms). Adaptive gradient optimizers in deep learning. Adaptive optimizers have served as the backbone of training deep learning models for more than a decade. In particular, the AdaGrad family of gradient whiteners/normalizers [Duchi et al., 2011] serves as a unifying template for many ensuing optimizers. RMSprop and Adam use a “diagonal” approximation of the AdaGrad gradient covariance matrix for adaptive optimization, while Shampoo [Gupta et al., 2018, Anil et al., 2020, Shi et al., 2023] can be viewed as a layer-wise Kronecker-Factored approximation (with one-sided variants [An et al., 2026]), and Muon can be viewed as an instantaneous version of Shampoo, which amounts to (approximately) orthogonalizing layerwise gradient matrices. From there, many further reductions or augmentations have been introduced. To name just a few, within the “Adam family”, we have the precursor RMSprop [Tieleman and Hinton, 2012], AdaBelief [Zhuang et al., 2020], LAProp [Ziyin et al., 2020], Lion [Chen et al., 2023], MVN-Grad [Patitucci and Mokhtari, 2026], augmented versions AdEMAMix [Pagliardini et al., 2025], as well as reduced versions AdaFactor [Shazeer and Stern, 2018], GaLore [Zhao et al., 2024], Adam-mini [Zhang et al., 2025b], among many more; see e.g., [Schmidt et al., 2021] for a more comprehensive account. On the matrix-shaped adaptive optimizer side, there have been 26

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

recent developments on top of Shampoo, e.g., through adaptive optimization in an eigenbasis-adjusted gradient (SOAP [Vyas et al., 2024], SPlus [Frans et al., 2026], KL-Shampoo [Lin et al., 2025]), or additional adaptivity applied to Muon, e.g., AdaMuon [Si et al., 2025], NorMuon [Li et al., 2025], MuonMax [Crawshaw et al., 2025]. We also note the PSGD line of work [Li, 2017, 2018], which spawns distinct preconditioners using various Lie group factorizations. Lastly, we mention a line of work that models (and creates) optimizers based on interpretations as steepest descent with respect to different choices of layerwise norm [Carlson et al., 2015, Bernstein and Newhouse, 2025, Kovalev, 2025, Pethick et al., 2025]. We note that, with few exceptions, all the aforementioned optimizers are of a whitening or normalizing flavor, with the goal of ensuring the update direction is (approximately) well-conditioned in various senses, e.g., entry-wise, row/column-wise, or spectrum-wise. Though there are far too many optimizers in this list whose efficacy as GP’s in our DoPr paradigm can be verified, our experiments on representative examples and general hypothesis is that any such “normalizing-type” GP is compatible with the AP gradient. Neural network feature learning theory. A by now classical approach to understanding how neural networks learn is via the neural tangent kernel (NTK) [Jacot et al., 2018, Mei and Montanari, 2022]. However, kernel methods rely on fixed, features, which fundamentally limits their expressivity and leads to suboptimal sample complexity for learning nonlinear functions [Ghorbani et al., 2021a,b]. This has motivated a large body of subsequent work studying feature learning in neural networks. In the setting of isotropic covariates, it has been shown that even a single step of SGD on the first layer of a two-layer network can learn sufficiently informative features to achieve improved sample complexity over kernel methods, provided the target function has low-dimensional structure [Damian et al., 2022, Moniri and Hassani, 2024]. In contrast, for general anisotropic covariates, Zhang et al. [2024] and Zhang et al. [2025a] identify that SGD is heavily distorted by anistropy of the inputs, which can cause the subspace distance between the learned and ground-truth representations to grow, even as training loss decreases. They reconcile this issue with a layerwise preconditioning schemes; the proposed algorithm in Zhang et al. [2024] is a two-layer version of pure AP (or equivalently DoPr-GD). Hyperparameter scaling conditions. A complementary line of work studies how optimizer hyperparameters should scale with model width, depth, and parameterization to ensure stable training and hyperparameter transfer. The Tensor Programs framework [Yang, 2019, 2020] provides a unifying theory of infinite-width limits, and maximal update parameterization (µP) [Yang and Hu, 2021, Yang et al., 2022] derives learning-rate and initialization scalings that keep activations Θ(1) while preserving maximal feature learning across widths. Subsequent work extends these principles to second-order and matrix-preconditioned optimizers such as KFAC and Shampoo [Ishikawa and Karakida, 2023], alternative formulations based on spectral norms [Yang et al., 2023], depth-wise scaling laws and residual multipliers [Dey et al., 2025, Yang et al., 2024], and newer optimizers including SOAP and Muon [Qiu et al., 2025]. Recent work further shows that for long training runs, weight decay scaling can dominate stability, with some works advocating constant learning-rate–weight-decay products [Wang and Aitchison, 2025, Kosson et al., 2025] and others proposing width-dependent weight decay scaling [Qiu et al., 2025]. Broadly, these methods derive optimizer-specific scaling rules, often through single-step update analyses, to enable efficient hyperparameter tuning across model sizes. Our work enables painless integration of these findings: by layering the GPs studied in these works on top of the AP update, we are able to make immediate use of these scaling rules. Mismatch between training performance and downstream behavior. Beyond anecdotal or folklore understanding that training/validation loss across optimizers or even hyperparameters need not correlate to better downstream performance, there have been various works trying to understand or ameliorate this 27

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

phenomenon. From the optimization perspective, recent work has identified that increasing pre-training effort (thus improving training performance) may lead to worse downstream finetunability [Springer et al., 2025, Horoi et al., 2025]. On the other hand, works since Chinchilla [Hoffmann et al., 2022] have steadily supported that pretraining token-per-parameter scaling laws can be in practice be much larger [Gadre et al., 2025] while downstream performance improves; notably, this includes prominent open-weight models such as Llama 3 [Meta, 2024]. Beyond coarse optimization interventions such as data quantity and training duration, certain works in the context of LLM fine-tuning have identified suboptimality in downstream performance when using different optimizers during pretraining and finetuning, e.g., Adam v.s. Muon, notably documented in Moonshot’s whitepaper introducing large-scale Muon-trained models, as well as follow-up studies [Liu et al., 2026c]. We note that our fine-tuning experiments use open-weight pre-trained models, and thus using DoPr for finetuning departs from the optimizer-consistency paradigm—it remains to verify: 1. if pretraining with DoPr can further extend its benefits, 2. if our observed benefits using DoPr for finetuning hold at scale and for general architectures. In a related perspective to ours regarding the benefits of “uniform” feature learning, some works have proposed using Sharpness-Aware Minimization (SAM) and extensions thereof as an optimization intervention to yield better zero-shot or few-shot performance of the model on downstream tasks [Foret et al., 2021, Bahri et al., 2022, Zhao et al., 2022], even at the cost of (pre-)training validation loss [Watts et al., 2026]. However, the mechanisms behind when SAM can yield reliable benefits as well as the feature geometry it actually induces do not necessarily reflect the nominal purpose of SAM [Andriushchenko et al., 2023, Schapiro and Zhao, 2024]; in the same light, the precise mechanisms by which DoPr—or any modern optimizer—yield improvements in general remain largely ill-understood. As a parallel line of literature, many works have explored proxy statistics or metrics beyond naive held-out/validation loss on the training objective toward predicting downstream performance, robustness, or adaptability. As a small sample, we highlight certain works that explore jointly accounting for gradientand activation-statistics [Liu et al., 2026b,a]. In Liu et al. [2026b], a gradient-weighted activation statistic is proposed in the context of maintaining plasticity in RL networks, a phenomenon by which much of a network’s capacity is silently wasted; notably, in RL contexts the training objective is nonstationary, and thus instantaneous training losses are often less predictive than in supervised learning. As our work focuses on supervised contexts, we leave investigating DoPr’s effectiveness in RL contexts for future work. Lastly, we note evidence cautioning using internal network statistics for concluding downstream performance; e.g. recent work finds that certain network trends typically associated with better/worse downstream performance may rather reflect current training practices, rather than a causal relationship with performance [Kulkarni et al., 2026, Liu et al., 2026a].

Extended Method Derivation and Mechanisms

B.1

Full DoPr Algorithm

In this section, we list the full DoPr algorithm and draw attention to various important design decisions we made to obtain performant results. Handling Different Layer Types While we define AP for linear layers, we extend it to other common architectural primitives by reducing them to equivalent linear maps. We unfold Conv1d layers and treat them akin to linear layers. For Conv2d layers, see Appendix B.2, and for Attention Layers, see Appendix B.2. Finally, for affine conditioning modules such as FiLM [Perez et al., 2018] and AdaLN [Peebles and Xie, 2023], we recommend first trying DoPr, as the performance of these modules can be sensitive to the relative learning speeds of the network backbone (which is presumably being trained with DoPr). We note that zero-initialized networks (e.g., in AdaLN-Zero) may require slight attention 28

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

Algorithm 2 Double Preconditioning (Feedforward Layer, layer index ℓ suppressed) 1: Input: Base Optimizer GP,learning rate η, weight decay λ, damping γ, EMA rate β 2: for k = 1 to K do

Sample batch B, |B| = n c(θ (k−1) ) G ← ∇W L Pn (k−1) (k−1)⊤ 1 Σ ← n i=1 zi zi

3: 4:

(k)

← βG

(k−1)

← βΣ

(k−1)

(k)

+ (1 − β)G + (1 − β)Σ

(k) −1 M ← G Σ + λ tr(Σ )I 9: D(k) ← GP(M(k) , ∗∗ kwargs) 10: W(k) ← (1 − ηλ)W(k−1) − ηD(k) 11: end for

(k)

(AP) (DoPr)

as initial activations may be zero, potentially destabilizing early training due to ill-conditioned AP. We found that with damping and the self-normalizing property of popular GPs, DoPr can work without modification, but if early training is unstable, we recommend a short warmup period with the base GP before switching to the full DoPr. Decoupled Weight Decay This stabilizes training by controlling parameter norms independently of the gradient update [Loshchilov, 2017]. EMA We provide the option to maintain exponential moving averages of both the gradient and the activation covariance. For a generic estimator G(k) , the EMA update is Ḡ(k) = β Ḡ(k−1) + (1 − β)G(k) . Unrolling the recursion gives Ḡ(k) = (1 − β)

k−1 X

β t G(k−t) + β k Ḡ(0) ,

t=0

i.e., a geometrically weighted average of past iterates. Assuming G(k) are unbiased estimators of some quantity µ with covariance ΣG , we have E[Ḡ(k) ] = (1 − β k )µ + β k Ḡ(0) , so the estimator is biased at early iterations, but converges to µ at rate β k . For the variance, assuming independence for simplicity, Var(Ḡ(k) ) = (1 − β)2

k−1 X

β 2t ΣG −→

t=0 1−β

1−β ΣG 1+β

as k → ∞.

Thus EMA reduces variance by a factor of 1+β while introducing a bias, which becomes negligible for sufficiently large k. We provide the option of applying EMA to both the gradient and the activation covariance. We recommend first trying AP in DoPr without additional EMAs: 1. most of our experiments compute AP with just the batch-wise gradient and activation covariance, 2. many base GPs are equipped with their own gradient EMAs (e.g. Adam and Muon), and thus applying gradient EMA in the AP computation requires accounting for doubly applying EMA, 3. using batch-wise statistics in AP removes the additional persistent memory requirement for the activation covariance buffer, making DoPr’s memory cost equal to its base GP. 29

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

Damping We add a damping term to the activation covariance estimate of the form Σz + λ tr(Σ )

tr(Σz ) dz I to

ensure invertibility. We note that the damping is set relative to the magnitude of Σz via d z such that z the damping hyperparameter is scale-invariant [Ishikawa and Karakida, 2023, Dey et al., 2025]. Though damping the activation covariance has conceptual connections with ridge regression, we posit that damping, and indeed using generic PD linsolvers that requires full-rankness such as Cholesky, is likely crude, as damping causes small activation directions (which AP aims to upweight) to be boosted equally as zero directions (which AP should in principle keep zero). We leave designing flexible, fast linsolvers that better handle rank degeneracy and pseudo-inversion (e.g., extensions of the lauded Newton-Schulz iteration in Muon) to future work. Inversion and Approximations For most layers, we compute the inverse of the activation covariance using PyTorch’s implementation of Cholesky decomposition torch.linalg.cholesky_ex. In our experiments, this was tractable for the considered model sizes, and usually not the compute bottleneck. This is in line with empirical insight that matrix operations that scale independently of the batch size (such as in computing GΣ−1 ) become negligible in large-batch training. For certain settings with inflated activation dimensions, such as in the im2col parameterization of convolutional layers, we also support structured approximations such as rank-1+diagonal (derived in Appendix B.3.1) and rank-k plus diagonal decompositions [Ahn et al., 2025a]. For convolutional layers specifically, in line with KFAC’s development, we employ the Spatially Uncorrelated Approximation (SUA) Appendix B.3.2 to bypass the (mostly redundant) cost of forming and inverting large patch covariance matrices. We also remark that for LLM embedding layers lead to input (i.e., “activation”) dimensions that scale with comparatively huge vocabulary size: however, we note that due to the one-hot nature of the inputs, the activation covariance is diagonal, where each non-zero entry is just the count of that token in the batch, and thus AP-ing the embedding layer gradient can be performed extremely efficiently and does not require forming a full dvocab × dvocab matrix.

B.2

AP For General Architectures: Derivations via Invariance

This section contains proofs for invariance under non-degenerate layer-wise affine transforms for feedforward, convolutional, and attention layers. Feedforward Affine Invariance Proposition 4.2. Consider a feedforward network, and consider the weights and activations {(Wℓ , zℓ )}ℓ∈[L] under non-degenerate layer-wise affine transforms defined in (4.3). Then, given the following layer-wise update rule: ⊤ Wnext = W − η∇W L ( fθ ), Σ−1 z , Σz = E[zz ], next

the updated weights satisfy Wnext zℓ = Wℓ ℓ

zℓ for all ℓ ∈ [L], z1 = x, z1 = A1 x.

Proof. Recall that: ⊤ −1 z z Σ−1 = E ℓ ℓ z ⊤ −1 = E Aℓ zℓ z⊤ ℓ Aℓ −1 −1 −1 = A⊤ Σ ℓ Aℓ . ℓ

(AP)

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

By the chain rule ∇Wℓ L fθ̄

∂L ⊤ z =E ∂h ℓ ∂L ⊤ ⊤ =E z Aℓ . ∂h ℓ

Therefore: next

Wℓ

zℓ = Wℓ − η∇Wℓ L fθ̄ Σ−1 zℓ zℓ ∂ L ⊤ ⊤ ⊤ −1 −1 −1 Aℓ zℓ Σ A = Wℓ A−1 − ηE A A z zℓ ℓ ℓ ℓ ℓ ∂h ℓ ∂ L ⊤ −1 = Wℓ − ηE zℓ Σz zℓ ℓ ∂h next = Wℓ zℓ

Since the argument holds for arbitrary ℓ, the conclusion holds for all layers, and thus the network output.

Convolutional Affine Invariance A convolutional layer can be written as a linear map applied to unfolded input patches via the standard im2colparameterization. Specifically, let Zℓ = im2col(zℓ ) ∈ Rdpatch ×N denote the matrix of flattened patches, where each column zℓ,b,p corresponds to a spatial patch indexed by batch element b and location p, and let Wℓ ∈ RCout ×dpatch denote the flattened convolutional kernel. Then the layer can be written as Hℓ = Wℓ Zℓ . Corollary B.1 (Convolutional Affine invariance). Consider a convolutional layer in im2col form with patch matrix Zℓ and weights Wℓ , and consider the transformed variables Zℓ = Aℓ Zℓ , Wℓ = Wℓ A−1 for invertible ℓ next ⊤ −1 Aℓ . Then under the AP update Wℓ = Wℓ − η∇Wℓ L ( fθ )ΣZ with ΣZℓ = E b,p [zℓ,b,p zℓ,b,p ], the updated ℓ

next weights satisfy Wℓ Zℓ = Wnext Zℓ . ℓ

Proof. Under the im2col representation, the convolutional layer is a linear map applied to patch vectors. The transformation Zℓ = Aℓ Zℓ , Wℓ = Wℓ A−1 preserves pre-activations, i.e. Wℓ Zℓ = Wℓ Zℓ . Thus the ℓ setting reduces exactly to the feedforward case of Proposition 4.3 applied to the patch vectors zℓ,b,p , yielding the claim. Attention Affine Invariance A self-attention layer first applies learned linear projections to an input token matrix Zℓ ∈ Rdmodel ×N : Q Qℓ = Wℓ Zℓ , Kℓ = WℓK Zℓ , Vℓ = WℓV Zℓ . The attention output is then computed as a deterministic function of these projected activations, for example p Attn(Qℓ , Kℓ , Vℓ ) = Vℓ softmax K⊤ Q / d , ℓ ℓ up to convention-dependent transposes. Thus, if an affine change of coordinates is applied to the input token representation and the inverse transformation is absorbed into each projection matrix, the projected queries, keys, and values are unchanged.

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

Corollary B.2 (Attention projection AP invariance). Consider a self-attention layer with input activations Zℓ r and projection weights {Wℓr } r∈{Q,K,V } . Let Aℓ be invertible, and define Zℓ = Aℓ Zℓ and Wℓ = Wℓr A−1 for each ℓ r next r −1 r r ∈ {Q, K, V }. Then under the AP update (Wℓ ) = Wℓ − η∇W L ( fθ )ΣZ , where ΣZℓ = E b,t [zℓ,b,t z⊤ ], ℓ,b,t ℓ

ℓ

the updated projections satisfy (Wℓ )next Zℓ = (Wℓr )next Zℓ for each r ∈ {Q, K, V }. Consequently, the updated attention queries, keys, values, and attention outputs are invariant under the transformed parameterization. Proof. For each projection r ∈ {Q, K, V }, the map Zℓ 7→ Wℓr Zℓ is a feedforward linear layer applied to token activations. Moreover, the transformed parameterization preserves the corresponding projected activations: r Wℓ Zℓ = Wℓr A−1 ℓ A ℓ Zℓ = Wℓr Zℓ . Therefore Proposition 4.3 applies separately to the query, key, and value projections, yielding r

(Wℓ )next Zℓ = (Wℓr )next Zℓ

for each r ∈ {Q, K, V }.

Thus the updated queries, keys, and values are identical in the two parameterizations. Since the attention output is a deterministic function of Qℓ , Kℓ , Vℓ , the updated attention output is also identical.

B.3

Approximation Schemes for AP

For larger models, particularly those with 2d convolutions, storing and computing the full activation covariance used in the preconditioner can be expensive. To subvert this cost, we develop multiple preconditioning approximation schemes. Rank-1 + Diagonal Let µ = E [X] and Var = E XX⊤ − µ ⊙ µ⊤ . We approximately decompose B.3.1

Σ + λI ≈ diag (Var) + λI + µµ⊤ Let D = diag (Var) + λI. Since D is diagonal and µµ⊤ is rank-1, by the Sherman-Morrison formula we may compute (Σ + λI)−1 = D + µµ⊤

−1

= D−1 −

(B.1) ⊤

−1

D µµ D

1 + µ⊤ D−1 µ

(B.2)

Which can be computed in linear time with respect to the hidden dimension D. B.3.2

Spatially Uncorrelated Activations (SUA)

We derive the SUA [George et al., 2018] approximation for a single 2D convolutional layer using an im2col parameterization. Let the (tensor-shaped) convolutional weights be W ∈ RCout ×Cin ×kh ×kw and let θ = vec(W). For each example i ∈ [n] and output spatial index u ∈ [m] (where m = Hout Wout ), let ai,u ∈ Rd denote the vectorized input patch (concatenating channels and kernel offsets) with d = Cin kh kw ,

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

and let gi,u ∈ RCout denote the gradient of the loss w.r.t. the pre-activation at location u. Define the stacked matrices ⊤ ⊤ Ai ≜ ai,1 · · · ai,m ∈ Rm×d , Gi ≜ gi,1 · · · gi,m ∈ Rm×Cout . Then the per-example weight gradient (reshaped to Cout × d) is Cout ×d ∇W ℓ(xi ) = G⊤ . i Ai ∈ R

Using vec(XYZ) = (Z⊤ ⊗ X)vec(Y) gives ⊤ ⊤ ∇θ ℓ(xi ) ≜ vec(∇W ℓ(xi )) = vec(G⊤ i Ai ) = (Ai ⊗ ICout ) vec(Gi ).

The empirical Fisher for this layer is n m m 1 XXX ⊤ b ∇θ ℓ(x) ∇θ ℓ(x)⊤ = F̂ ≜ E ai,u a⊤ ⊗ g g i,u i,v . i,v n i=1 u=1 v=1

The double sum over (u, v) shows that the exact curvature couples pairs of spatial locations, making exact storage/inversion expensive when m is large. KFAC Grosse and Martens [2016] applies the standard approximation that (patch) activations and backpropagated gradients are independent: ⊤ ⊤ E (au a⊤ ) ⊗ (g g ) ≈ E[au a⊤ u v v v ] ⊗ E[gu g v ]. This reduces the Fisher to products of second moments of a and g. SUA further assumes spatial uncorrelatedness of activations: E[au a⊤ v] = 0

for u ̸= v,

(and analogously, in many implementations, E[gu g⊤ ̸ v). Under this approximation, only v ] ≈ 0 for u = the diagonal terms u = v remain, yielding m X ⊤ E[au a⊤ F ≈ u ] ⊗ E[gu gu ]. u=1

If we additionally assume stationarity across spatial locations (i.e. the second moments do not depend on u), then d×d Cout ×Cout Σa ≜ E[au a⊤ , Σg ≜ E[gu g⊤ , u]∈R u]∈R and thus F ≈ m (Σa ⊗ Σg ), where the scalar factor m is absorbed into the step size / damping. We estimate the factors by averaging over both batch and spatial indices: n m n m 1 XX 1 XX ⊤ ⊤ ⊤ b b ai,u ai,u = Ei,u [ai,u ai,u ], gi,u g⊤ Σ̂a = Σ̂g = i,u = Ei,u [gi,u gi,u ]. nm i=1 u=1 nm i=1 u=1 Notably, Σ̂a has size (Cin kh kw ) × (Cin kh kw ) and is independent of the feature-map resolution or batch size (and reduces to Cin × Cin for 1 × 1 convolutions). Let gW ≜ ∇W ℓ be the weight gradient in matrix form RCout ×d . Under the Kronecker-factored approximation F ≈ Σa ⊗ Σg , the (damped) KFAC/SUA preconditioner applies ∆W ∝ (Σg + λI)−1 gW (Σa + λI)−1 , followed by reshaping ∆W back to the tensor shape of W. Importantly, SUA does not drastically reduce the expressivity of the preconditioner as the discarded curvature directions correspond to spatial interaction modes that are unidentifiable under weight sharing. As such, they lie largely in the nullspace of the parameterization. Note that we only use the right-sided SUA. 33

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

B.4

Feature Learning Theory

We adapt parts of Zhang et al. [2024] and Zhang et al. [2025a] toward establishing the feature learning result in Proposition 4.1. We note for presentational convenience, we make some technical assumptions: Assumption B.1. The current iterate Gθ satisfies: 1. Gθ ∈ Rd×n is row-orthonormal. (F ⊤ F ) λ

(Σ )

2. dist(Gθ , G⋆ ) ≤ λ min (F⋆ ⊤ F⋆ ) λmin (Σs ) . max

⋆

max

Proposition B.1 (Proposition 4.1, formal ver.). Let Assumption B.1 hold. Given 0 < η ≤ 0.5 λmax (F⋆ ⊤ F⋆ )−1 . Then, one step of full-batch GD (resp. AP) satisfies: next,GD

P⋆⊥ = (Gθ − ηFθ ⊤ Fθ Gθ Σs )P⋆⊥ + ηFθ ⊤ F⋆ G⋆ Σs P⋆⊥

next,AP

P⋆⊥ = (I − ηFθ ⊤ Fθ )Gθ P⋆⊥ .

Gθ

(B.3) next,AP

Then, the AP update satisfies the following one-step contraction in subspace distance: dist(Gθ , G⋆ ) ≤ ⊤ 1 − η λmin (Fθ Fθ ) dist(Gθ , G⋆ ). On the other hand, one may construct Fθ , F⋆ , G⋆ , Σs such that next,GD dist(Gθ , G⋆ ) > dist(Gθ , G⋆ ). Proof of Proposition 4.1. First, recalling the loss: L (Fθ Gθ ) = Eπdemo ∥Fθ Gθ − s t ∥22 = ∥(Fθ Gθ − 2 F⋆ G⋆ )Σ1/2 s ∥ F , we have the GD and AP update: next,GD

Gθ

= Gθ − η∇Gθ L = Gθ − ηFθ ⊤ Fθ Gθ Σs + ηFθ ⊤ F⋆ G⋆ Σs

next,AP

Gθ

= Gθ − η∇Gθ L · Σ−1 s = Gθ − ηFθ ⊤ Fθ Gθ + ηFθ ⊤ F⋆ G⋆ .

Noting G⋆ P⋆⊥ = 0, applying P⋆⊥ to the above yields (4.2). Now, given that next,AP

Gθ

P⋆⊥ = (I − ηFθ ⊤ Fθ )Gθ P⋆⊥ ,

applying ∥·∥op and submultiplicativity on the right-hand side yields: next,AP

dist(Gθ

, G⋆ ) ≤ I − ηFθ ⊤ Fθ op Gθ P⋆⊥ op = (1 − η λmin (Fθ ⊤ Fθ ))dist(Gθ , G⋆ ),

completing the one-step contraction bound for a step of AP. next,GD To show that conversely Gθ may in general be non-convergent in subspace distance, we construct a numerical example to illustrate this. Take d = 1, n = 2, and let 0 0 ⊥ G⋆ = 1 0 , P⋆ = . 0 1 p 3/2 , we have Choosing Gθ = cos(π/3) sin(π/3) = 1/2 dist(Gθ , G⋆ ) = ∥Gθ P⋆⊥ ∥2 = 34

p 0 3/2 2 =

p 3 . 2

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

Now setting

5 2 Σs = ≻ 0, 2 1

Fθ = 1,

F⋆ = 1,

η = 1,

we have next,GD

Gθ

P⋆⊥ = (Gθ − Gθ Σs )P⋆⊥ + G⋆ Σs P⋆⊥ .

Computing the three terms:

p 3/2 , Gθ P⋆⊥ = 0 p Gθ Σs P⋆⊥ = 0 1 + 3/2 , G⋆ Σs P⋆⊥ = 0 2 ,

this yields next,GD

Gθ

p p 3/2 − 0 1 + 3/2 + 0 2 P⋆⊥ = 0 = 0 1 .

Hence next,GD dist(Gθ , G⋆ ) = 1 >

p 3 = dist(Gθ , G⋆ ), 2

completing the construction and the proof.

B.5

Dynamical Systems Theory

Proof of Proposition 3.1.

We first establish the linear systems preliminaries in Proposition 3.1 P t−1 Proposition 3.1. Define the state-covariance matrix Γ t (K) = s=0 (A + BK)s Σw (A + BK)s⊤ , and Γ̄ t (K) = Pt πdemo π 1 = N(0, Γ t (K⋆ )), Pŝ θ = N(0, Γ t (Kθ )), and t s=1 Γ s (K). Then, Ps t t

Lval (πθ ) = (K⋆ − Kθ )Γ̄ T (K⋆ ) /2 F , 1

(3.2)

−Rtest (πθ ) = Lideal (πθ ) = (K⋆ − Kθ )Γ̄ T (Kθ ) /2 F . 1

(3.3)

If D denotes the Wasserstein-2 distance, we have Dt (πθ ) = ∥Γ t (K⋆ ) /2 − Γ t (Kθ ) /2 ∥2F . 1

Proof. Recall the state transition s t+1 = As t + Ba t + w t , w t ∼ N(0, Σw ). Under a t = Ks t , we have for given s t : s t+1 ∼ N((A + BK)s t , ΣW ). Applying this recursively to s0 = 0 yields: t−1 X s t ∼ N 0, (A + BK)s Σw (A + BK)s⊤ =: N(0, Γ t (K)).

s=0 π

Plugging in K = K⋆ , Kθ completes the proof of the distribution of the state marginals: Ps tdemo = π N(0, Γ t (K⋆ )), Pŝ θ = N(0, Γ t (Kθ )). As for the losses, we have: t

Lval (πθ ) = Eπdemo ∥Kθ s − a∥22 =E

πdemo 1

T X ∥(Kθ − K⋆ )s t ∥22

T t=1

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

1 X πdemo ⊤ E tr (Kθ − K⋆ )s t s⊤ = t (Kθ − K⋆ ) T t=1 T 1 X πdemo ⊤ = tr (Kθ − K⋆ ) E [s t s⊤ t ](Kθ − K⋆ ) T t=1 = tr (Kθ − K⋆ )Γ̄ t (K⋆ )(Kθ − K⋆ )⊤ = ∥(Kθ − K⋆ )Γ̄ t (K⋆ )1/2 ∥2F , Lideal (πθ ) = Eπθ ∥Kθ ŝ − â∥22 T 1 X πθ ⊤ ⊤ = tr (Kθ − K⋆ ) E [ŝ t ŝ t ](Kθ − K⋆ ) T t=1 = ∥(Kθ − K⋆ )Γ̄ t (Kθ )1/2 ∥2F . Finally, when the divergence D is the Wasserstein-2 distance, we may use the formula for Wasserstein-2 distance between two multivariate Gaussians P = N(0, Σ1 ), Q = N(0, Σ2 ): 1/2

1/2

W22 (P, Q) = tr(Σ1 + Σ2 − 2(Σ1 Σ2 Σ1 )1/2 ). Plugging in Σ1 = Γ̄ t (K⋆ ), Σ2 = Γ̄ t (Kθ ) and simplifying yields the result and completes the proof.

B.5.1

Mismatch between Lideal and Lval : a construction

We want to demonstrate that due to feedback between a policy and the dynamics, two policies that may satisfy Lval (π1 ) ≪ Lval (π2 ) may in fact satisfy Lideal (π1 ) ≫ Lideal (π2 ). Consider the dynamical system 0 0 A= , B = I, Σw = I (B.4) 0 1 Next, we set 0 0 K⋆ = , 0 −α

F⋆ = e2 ,

G⋆ = −αe⊤ 2

(B.5)

Finally, let us consider two learned parameters −1

Kθ = Fθ Gθ ,

Fθ = µe2 , Gθ = −αµ α Kψ = Fψ Gψ , Fψ = , Gψ = G⋆ 1

cos θ sin θ

⊤ (B.6) (B.7)

Then, we have that

0 Kθ = K⋆ − α − cos θ 0 1 Kψ = K⋆ − ε 0 0 And the following computation will be useful: 36

0 1 − sin θ

(B.8) (B.9)

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

Lemma B.2. It holds that cos2 θ − cos θ (1 − sin θ ) (Kθ − K⋆ ) (Kθ − K⋆ ) = α − cos θ (1 − sin θ ) (1 − sin θ )2 0 0 (Kψ − K⋆ )⊤ (Kψ − K⋆ ) = 0 ε2 ⊤

(B.10) (B.11)

For simplicity, we work now in the limit where T → ∞. The following is standard: Lemma B.3. Let K be such that ρ(A + BK) < 1, where ρ(·) denotes the spectral radius. Then, Γ ∞ (K) := lim T →∞ Γ T (K) exists and is finite, is equal to lim T →∞ Γ̄ T (K), and is given by the unique solution to the following Riccati equation: Γ ∞ (K) = Σw + (A + BK)Γ ∞ (K)(A + BK)⊤ .

(B.12)

We now compute the stationary covariances under K⋆ , Kθ , Kψ . Lemma B.4. Γ ∞ (K⋆ ) exists, and is given by

1 Γ ∞ (K⋆ ) = 0

(B.13)

1 1−(1−α)2

Proof. A + BK⋆ has eigenvalues 0.5, 1 − α, so Lemma B.3 holds. In this case, the limiting Γ satisfies 1 0 0 0 1 0 (B.14) + Γ Γ= 0 1 0 1−α 0 1−α decoupling across coordinates. So we select Γ to be diagonal with coordinates γ1 , γ2 on the diagonals. These must satisfy γ1 , and γ2 (1 − α)2 + 1 = γ2 , so that γ = 1/(1 − (1 − α)2 ) Lemma B.5. For θ ∈ (0, π/2), Γ ∞ (Kθ ) exists, Γ ∞ (Kθ ) =

1+α2 cos2 θ α sin θ (2−α sin θ )

(B.15)

Proof of Lemma B.5. The spectral radius condition can be checked for θ ∈ (0, π/2) by upper triangularity. Now set ⊤ 0 0 0 1 0 Γ= Γ + (B.16) −α cos θ 1 − α sin θ −α cos θ 1 − α sin θ 0 1 We now solve the above:

0 X := −α cos θ

0 , 1 − α sin θ

1 0 Q= . 0 1

Now write

γ11 γ12 Γ= . γ12 γ22 The Lyapunov equation to solve is Γ = XΓ X⊤ + Q. Computing 0 0 ⊤ XΓ X = . 0 α2 cos2 θ γ11 − 2α cos θ (1 − α sin θ ) γ12 + (1 − α sin θ )2 γ22 37

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

Matching entries gives γ11 = 1, γ12 = 0, γ22 = α2 cos2 θ γ11 − 2α cos θ (1 − α sin θ ) γ12 + (1 − α sin θ )2 γ22 + 1. Using γ11 = 1 and γ12 = 0, the final equation becomes γ22 = α2 cos2 θ + (1 − α sin θ )2 γ22 + 1. Therefore, 1 − (1 − α sin θ )2 γ22 = 1 + α2 cos2 θ , and since 1 − (1 − α sin θ )2 = α sin θ (2 − α sin θ ), we obtain γ22 =

1 + α2 cos2 θ . α sin θ (2 − α sin θ )



 . 1 + α2 cos2 θ α sin θ (2 − α sin θ )

Hence

 Γ =



Lemma B.6. Γ ∞ (Kψ ) exists, and satisfies Γ ∞ (Kψ ) =

ε2 1 + α(2−α) ε(1−α)

− α(2−α)

ε(1−α)

− α(2−α) 1 α(2−α)

Proof of Lemma B.6. Again, we can check the stability conditions from upper tri-angularity. We have ⊤ 0 −ε 1 0 0 −ε Γ= + Γ 0 1−α 0 1 0 1−α | {z }

(B.17)

:=X

0 −ε X := , 0 1−α

1 0 Q= . 0 1

γ11 γ12 Now writing Γ = , the Lyapunov equation to solve is Γ = XΓ X⊤ + Q. First, we compute γ12 γ22  XΓ X⊤ = 

ε2 γ22

−ε(1 − α)γ22

(1 − α)2 γ22

Matching entries gives γ11 = ε2 γ22 + 1, γ12 = −ε(1 − α)γ22 , γ22 = (1 − α)2 γ22 + 1. 38

 .

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

We solve from bottom to top. The third equation gives 1 − (1 − α)2 γ22 = 1. Since 1 − (1 − α)2 = 1 α(2 − α), we obtain γ22 = α(2−α) . Therefore, γ12 = −ε(1 − α)γ22 = −

ε(1 − α) γ11 α(2 − α)

= 1 + ε2 γ22 = 1 +

ε2 . α(2 − α)

Hence   Γ = 

1+ −

ε2 α(2 − α)

ε(1 − α) α(2 − α)

 ε(1 − α) α(2 − α)  .  1 α(2 − α)

−

From here, we can compute the various error terms Lemma B.7. Let θ ∈ (0, π/4). Then, taking the T → ∞ limit, Lval (Kθ ) ∼ α,

α , sin θ

(B.18)

ε2 α

(B.19)

Lideal (Kθ ) ∼

where as Lval (Kψ ), Lideal (Kψ ) ∼

Proof. From Proposition 3.1, the fact that ∥∆Γ 1/2 ∥2F = tr(∆⊤ ∆Γ ), and that ∆ = K⋆ − Kθ , we have that tr(∆⊤ ∆Γ ) = α2 Γ 22 (1 − sin θ )2 by Lemma B.2, we obtain Lval (Kθ ) = α2 cos2 θ + α2 (1 − sin θ )2 ·

1 ∼α 1 − (1 − α)2

(B.20)

for α small and θ ∈ (0, π/4). Moreover, from these same argument Lideal (Kθ ) = α2 cos2 θ + α2 (1 − sin θ )2 ·

1 + α2 cos2 θ α α ∼ α2 + ∼ . α sin θ (2 − α sin θ ) sin θ sin θ

We may compute the quantities on Kψ similarly. By modulating ε, α appropriately, the following proposition follows immediately: Proposition B.8. For ε ≪ α and sin θ ≪ ε2 /α2 , we have we have that Lval (Kθ ) ≪ Lval (Kψ ) ∼ Lideal (Kψ ) ≪ Lideal (Kθ ).

(B.21)

Continuing with these quantities, we may go on to prove Proposition 3.2. B.5.2

Proof of Proposition 3.2

We recall Proposition 3.2: Proposition 3.2 (Poor Feature Learning causes TTF Shift, Informal). There exists (F1 , G1 ), (F2 , G2 ), (F⋆ , G⋆ ) and linear dynamics A, B, Σw such that Lval (F1 G1 ) ≪ Lval (F2 G2 ), but dist(G1 , G⋆ ) > dist(G2 , G⋆ ), and for any possible “last-layer” Fθ , we have Lideal (F2 G2 ) ≪ Lideal (Fθ G1 ). 39

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

Notably, compared to the constructions above, the difference is that Proposition 3.2 introduces the additional complication of allowing the flexibility of re-fitting the last-layer Fθ . Inheriting K⋆ , Kθ , Kψ and relevant quantities therein from earlier, the key result to establish is: 2

α Lemma B.9. For θ ∈ (0, π/8), we have infF Lideal (FGθ ) ≥ 4 sin θ.

By rescaling Gθ , Fθ by µ−1 and µ, and using Proposition 3.1 as in the proof of Lemma B.7, it suffices to prove the following. Lemma B.10. Let θ ≤ π/8. Define the vector

cos θ w = −α sin θ

⊤ ,

0 0 A= , 0 1

1 0 Q= 0 1

(B.22)

For any v, define two terms: Xv = A − vw⊤ , ∆v = αA − vw⊤ . When v is such that a solution Γ v to the Riccati equation Γ v = Xv Γ v X⊤ v + Q.

(B.23)

exists, where Γ v is defined as above. Then, we have inf

v∈R2 :Γ v exists

tr ∆⊤ ∆ · Γ ≥ v v v

α2 . 4 sin θ

(B.24)

Proof.

cos θ w = −α , sin θ

0 0 A= , 0 1

1 0 Q= . 0 1

For any v ∈ R2 , define Xv = A − vw⊤ and ∆v = αA − vw⊤ . Since −vw⊤ = αv cos θ Xv = A + αv cos θ

sin θ ,

∆v = αA + αv cos θ

Let c := cos θ , s := sin θ , and τ := tan θ = cs . Writing v =

pc Xv = qc

sin θ , we have

sin θ .

v1 , define p := αv1 and q := αv2 . Then v2

ps , 1 + qs

pc ∆v = qc

ps . α + qs

Now reparameterize by r := pc = αv1 cos θ and h := −αv2 sin θ θ . Since s, c > 0, this is an equivalent parameterization. In these variables, r rτ r rτ Xv = , ∆v = . − τh 1 − h − τh α − h The characteristic polynomial of Xv is χ(z) = z 2 −(1+r−h)z+r; therefore, by the Schur stability conditions, the Lyapunov/Riccati solution Γ v = Xv Γ v X⊤ v + Q exists if and only if r < 1, h > 0, and 2+ 2r − h > 0. Equivalently, −1 < r < 1 and 0 < h < 2(1 + r). Define the objective J(r, h) := tr ∆⊤ v ∆v Γ v . Solving the Lyapunov equation and substituting gives J(r, h) =

N (r, h, τ, α) hτ2 (1 − r)(2 + 2r − h) 40

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

where N (r, h, τ, α) = α2 (1 + r)h2 + α2 τ2 (1 − r)2 (1 + r) + 2(1 − α)h3 + τ2 2α2 hr − 2αh2 r + 2αhr 2 − 2αh + h2 (1 + r) + 2hr 2 + τ4 r 2 (1 + r). For fixed 0 < α ≤ 1 and θ sufficiently small, the numerator satisfies the lower bound N (r, h, τ, α) ≥

α2 (1 + r) h2 + τ2 (1 − r)2 . 2

Also, since 0 < h < 2(1 + r), we have 2 + 2r − h ≤ 2(1 + r). Therefore, α2 2 2 2 2 (1 + r) h + τ (1 − r) J(r, h) ≥ hτ2 (1 − r) 2(1 + r)

α2 h2 + τ2 (1 − r)2 . 4 hτ2 (1 − r)

By AM-GM, h2 + τ2 (1 − r)2 ≥ 2hτ(1 − r). Hence α2 2hτ(1 − r) · 4 hτ2 (1 − r) α2 = . 2τ

J(r, h) ≥

Thus, inf

v: Γ v exists

tr ∆⊤ v ∆v Γ v ≥

α2 . 2 tan θ

Since tan θ ≤ 2 sin θ for θ sufficiently small, this further implies inf

v: Γ v exists

tr ∆⊤ v ∆v Γ v ≥

α2 . 4 sin θ

Thus, chaining Proposition B.8 and Lemma B.9 together yields the following quantitative form of Proposition 3.2. 2

Proposition B.11. For sin θ ≪ αε 3 and ε ≪ α, we have Lval (Kθ ) ≪ Lval (Kψ ) ∼ Lideal (Kψ ) ≪ inf Lideal (FGθ ) ≤ Lideal (Kθ ). F

B.6

(B.25)

Limitations

DoPr incurs a nontrivial computational and memory overhead. For a linear layer with input dimension din , output dimension dout , and batch size N , the per-step overhead consists of accumulating the input covariance, factoring the damped preconditioner, and applying the inverse preconditioner to the gradient: S = X⊤ X,

X ∈ RB×din , 41

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss 2 which costs approximately Bdin FLOPs (due to symmetry of S. Computing the Cholesky factorization of 3 the damped matrix A = S + λI ∈ Rdin ×din costs approximately 13 din FLOPs. Finally, applying the inverse preconditioner to the weight gradient

G(S + λI)−1 ,

G ∈ Rdout ×din ,

2 via two triangular solves costs 2din dout FLOPs. Thus, the total per-layer per-step idealized FLOP overhead is 1 3 2 2 + 2din dout . N din + din 3

In line with empirical know-how (see e.g., Jordan et al. [2024]), in the large-batch regime, the cost of forming the activation covariance S = X⊤ X dominates. Thus, for layers where din ≈ dout , the overhead is approximately equal to an additional forward pass. Indeed, throughout our experiments, we found the overhead to be ≈ 1.25x in wall-clock time, commensurate with the estimate of a backward pass equal to ∼ 3 forward passes, such that AP’s overhead coarsely comes out to 5/4 = 1.25× regular backprop. DoPr also introduces additional memory overhead from storing the matrix-valued preconditioner statistics. With an exponential moving average (EMA), this requires storing one din × din symmetric 2 matrix per preconditioned layer, i.e., ∼ din /2 additional parameters of optimizer state per layer. Our current implementation of DoPr only supports Distributed Data Parallelism, and not model parameter sharding such as FSDP. However, we do not anticipate any fundamental barriers for such schemes beyond existing related matrix-preconditioning methods such as Shampoo, which have distributed and sharded implementations [Shi et al., 2023, Ahn et al., 2025b]. Finally, DoPr can be sensitive to the choice of matrix-inversion solver and damping parameter. In practice, we found PyTorch’s Cholesky solve together with trace-scaled damping to provide a sufficiently stable and effective default. DoPr may improve downstream performance without improving training or validation loss, and it may not improve training-loss convergence relative to the base GP optimizer. Moreover, exact AP for convolutional or other weight-sharing layers can incur substantial memory overhead due to unfolding. While we introduce approximation schemes to account for the redundancy therein, developing hardwareaccelerated approximations and exploring model–optimizer codesign remain important directions for future work. Finally, we observe that large dropout rates can conflict with AP EMA buffers and accelerate staleness due to zero-ing out random subsets of activations; therefore, in our current practice, we recommend using no or low dropout. In principle, we note more advanced schemes for updating the activation buffer (e.g., analogous to Ahn et al. [2025a]) exist, but may require case-by-case treatment depending on how the architecture is implementated. Improving the compatibility of DoPr with dropout is left for future work.

Experiment Details and Additional Results

Hardware • All imitation learning experiments were performed on 4× NVIDIA RTX A5000 GPUs. • All language modeling experiments were performed on 8× NVIDIA RTX PRO 6000 Blackwell GPUs. • All image generation experiments were performed on 8× NVIDIA A100 GPUs.

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

C.1

General Operating Guideline

In practice, we summarize a small set of operating guidelines that we have found to be broadly useful across settings. • EMA smoothing for very small-batch, continuous-valued inputs. For continuous-valued inputs (in contrast to discrete-token embeddings in LLMs), we suggest applying to the AP an exponential moving average (EMA) to both the activation covariance and gradient to reduce update variance when the batch size is significantly smaller than the model width. We note that appropriate considerations should be made if the GP of choice also contains EMAs, as this may doubly apply smoothing on the gradient. • Use no/low dropout. Dropout [Srivastava et al., 2014] is designed to prevent spurious co-adaptation between features in the network. Firstly, due to masking different subsets of activations per update, large dropout can conflict with accumulation in the activation preconditioner buffers, accelerating staleness. Furthermore, AP can be viewed as similarly de-biasing the gradient from its activation (i.e. feature) correlations. Therefore, we recommend using no dropout for most settings, or reduced amounts in large-scale convolutional architectures. • Handling high-redundancy/weight-sharing layers. AP is most directly derived for linear layers. For layers that have high-redundancy and weight-sharing properties, such as convolutional layers, the exact AP (resulting from Proposition 4.2) can incur needlessly large memory overhead from unfolding. Following prior literature in Kronecker-Factored preconditioners [Grosse and Martens, 2016, George et al., 2018], we may derive an efficient approximation of AP that empirically preserves the performance of the full AP as in Appendix B.3.2. We also note that one of the motivations of heavy weight-sharing is to make the architecture “well-conditioned” [Wilson et al., 2017, Keskar and Socher, 2017, Xiao et al., 2018], and thus may be partially redundant to our optimizer that is designed to handle heavy anisotropy. As we do not modify baseline model architectures in this work, we leave exploring model-optimizer co-design for future work.

C.2

Feature Learning: Validation Loss versus Subspace Distance

This section provides details on the experiments shown in Figure 4. Data generation. We follow the synthetic two-layer setup of Zhang et al. [2025a]. We first generate a base matrix F0 ∈ RdY ×k with i.i.d. N(0, 1) entries. For each task s ∈ {train, test}, we sample a matrix Bs ∈ RdY ×dY and define the task-specific head Fs⋆ = exp 0.005(Bs − B⊤ ) F0 , s which produces small random rotations of a shared base feature matrix. The shared representation G⋆ ∈ Rk×dX is sampled uniformly from the set of row-orthonormal matrices. For each task s, inputs are generated as dX xsi ∼ Σ1/2 x,s · Unif({±1} ), and outputs are generated according to ysi = Fs⋆ G⋆ xsi + ϵ si , Anisotropy.

2 ϵ si ∼ N(0, σϵ,s IdY ).

We sample a random rotation O ∈ RdX ×dX and define Σx,s = ODO⊤ ,

D = diag(logspace(0, 5, dX )),

yielding highly ill-conditioned covariances. This setting amplifies the feature-learning bias described in Proposition 4.1. 43

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

Model and training. We train linear models of the form Fθ Gθ on the training task train using squared loss, i.e., Ltrain (θ ) = E(x,y)∼Dtrain ∥Fθ Gθ x − y∥2 . We compare SGD, Adam, and DoPr-SGD, using batch size n = 1024 and averaging results over 10 random seeds. All methods are initialized identically and use matched learning rates. Metrics. We report (i) the in-distribution validation loss Lval on held-out samples from train, and (ii) the subspace distance dist(Gθ , G⋆ ), which measures alignment between the learned and ground-truth feature subspaces. The latter isolates representation learning independent of downstream fitting. Results. As shown in Figure 4, validation loss suggests successful optimization across all methods. However, the subspace distance reveals a clear discrepancy: standard and gradient-preconditioned methods fail to recover the true feature subspace under high anisotropy, whereas AP significantly improves alignment. This demonstrates that low validation loss can mask poor feature learning, and supports the claim that AP mitigates TTF by promoting more uniform feature recovery.

C.3

Activation Invariance of DoPr

This section provides details on the experiments shown in Figure 5. Models and datasets. We use 3 layer MLPs trained on 2 synthetically generated datasets. For the first dataset, the input data is generated with an anisotropic covariance Σ ∈ R32×32 , where Σ has log space eigenvalues ranging from 1 to 100. Then, the output is generated as y = M2 σ(M1 x), where M1 ∈ R256×32 , σ is the GeLU activation function, and M2 ∈ R10×256 , for each input data x. M1 is 1 initialized with each element drawn from N (0, 32 ) and M2 is initialized with each element drawn from 1 N (0, 256 ). The second dataset, denoted the transformed dataset, generates each sample as x t = Qx + d, where x is a sample from the original. Here, Q ∈ R32×32 is a symmetric matrix with linear space eigenvalues ranging from 1 to 100, and d ∈ R32 is a random vector with elements drawn from the standard normal. Then, y t = M2 σ(M1 x t ), using the same method as in the other dataset. Both datasets are split into a 0.8 : 0.2 into training data and validation data. This induces an affine transform of the input activations. We then initialize 4 models, 2 to be trained on the original dataset (one optimized with SGD, and one optimized with DoPr-SGD), and 2 to be trained on the transformed dataset (one optimized with SGD, and one optimized with DoPr-SGD). The two models trained on the original dataset are initialized identically, and the two trained on the transformed dataset are also initialized identically. All models have an input layer of dimension (32, 256), hidden layer (256, 256), and output layer (256, 10), and use SiLU activations. For all models, the hidden layer and output layers are identical. For the original dataset models, let the input layer have weight W, and bias b. Then, for the transformed dataset models, the input layer is initialized with weight WQ−1 , and bias b − WQ−1 d. We note here, inputting x into the models trained on the original data yield h = Wx + b. Then, x t = Qx + d, and when passed through the transformed models the output will be h t = WQ−1 (Qx + d) + b − WQ−1 d = Wx + b So, for all models, the pre-activations after the first layer will be identical, and thus at initialization, all models will have identical outputs. From here, both models are trained for 50 epochs of 20 steps per epoch, with batch size 2048. At each step, the models are validated on the 0.2 validation set for their respective datasets, and the results 44

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

reported. The models are trained so that at each training step, each batch exactly corresponds (e.g. for a batch X ∈ R32×2048 of the Original data, the transformed batch will exactly be X t = QX + b1 T ). For the models trained using DoPr, damping is set at 10−8 · tr(Σ), where Σ is the preconditioner. The identical trajectories of both DoPr models demonstrates the affine invariance demonstrated in Proposition 4.2 – when learning the same function, and with identical initialization, affine transforms to the input distribution do not affect the loss trajectory when trained with DoPr. However, loss trajectories under vanilla SGD do exhibit a drift, demonstrating the coordinate-dependent trajectory as stated in Section 4.2.

C.4

Hyperparameter Scaling DoPr

This section provides details on the experiments shown in Figure 6, Figure 12, and Figure 11. The goals of these experiments are to provide guidelines for scaling hyperparameters as width increases, as discussed in Yang et al. [2022] and related work (see Appendix A). These scalings allow for zero-shot hyperparameter transfer across models with differing width.

Figure 11: We train a 3-layer MLP on CIFAR, and sweep over learning rates. We compare Standard Parametrization (SP, left) and Maximal Update parametrization (µP, right). The final train loss is reported for both µP and SP. µP achieves learning rate transfer as width increases.

Models and datasets. In Figure 6 and Figure 12, for language modeling, we use a GPT2-style Transformer with 4 layers, head dimension 64, sequence length 1024, and batch size 12, trained on the OpenWebText dataset. Each experiment processes approximately 100M training tokens. For Figure 11 and the weight-decay µP plots in Figure 6, we use 3 layer MLPs trained on CIFAR10 Krizhevsky et al., normalizing input data to be between [−1, 1]. The MLP layers have no bias, and for width w, the layers are of size (3072, w), (w, w), and (w, 10), with cross entropy loss. For the learning rate sweeps (Figure 11), we use batch size 512, and train for 8 epochs. For the weight decay experiment, we use batch size 4096 and train for 200 epochs. Optimization. For the language modeling and CIFAR learning rate experiments we sweep over learning rate. For the language modeling we use base weight decay of 0.1, and for CIFAR we use base weight decay of 0.01. For language modeling, a base learning rate of 6 · 10−3 was used. For the Standard 45

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

Parametrization (SP) sweep for CIFAR, we use base learning rate of 2−10 , and for the µP sweep for CIFAR, we use base learning rate of 2−4 . Then, the learning rate used is: η base = η · 2k where k is the sweep multiplier. In all experiments, all layers are preconditioned, with damping that scales with tr(Σ), where Σ is the preconditioner. The language model used damping 10−4 · tr(Σ), and the CIFAR experiments used damping 10−3 · tr(Σ). Noting that the preconditioner Σ is the same as the preconditioner in KFAC, this damping aligns with the rescaled damping in Ishikawa and Karakida [2023], which naturally achieves the µP scaling with width. We also EMA the preconditioner with β = 0.9, only in the CIFAR experiments. In language modeling, we report the best training loss achieved during training. For the CIFAR learning rate sweep, we report the training loss at the final epoch for the learning rate sweep. ||∆W|| For the CIFAR weight decay experiment, we report ||W|| F , the ratio of the update norm to the weight F norm for the hidden layer W. We report this at each step in training. We use a base learning rate of η base = 2−4 , and a base weight decay of λ base = 0.01. µP vs SP. In all experiments, we follow the µP scaling rules as described in Yang et al. [2022]. As further explained below, we note that for DoPr-AdamW and DoPr-Muon in particular, the same scaling rules apply as for AdamW and Muon. This is because of the normalizing nature of these optimizers - despite DoPr altering the gradient that AdamW and Muon act on, the final object still has the same normalized size. −1 For the µP sweeps, the input layers, are initialized as Wi j ∼ N (0, din ). The hidden layers are −1

initialized as Wi j ∼ N (0, w ), and the output layers are initialized as Wi j ∼ N (0, w−2 ). The input layer learning rate is η base , the hidden layer learning rate is η base · w−1 , and the output layer learning rate is qT k

qT k

η base · w−1 . In µP, during attention instead of computing p , d is used in attention. d Given a base weight decay of λ base = 0.01, the input layer’s weight decay is λ base , the hidden layers’ weight decay is λ base · w, and the output layer’s weight decay is λ base · w. This style of weight decay keeps η · λ width-independent has been proposed by several past work [Kosson et al., 2024, 2025, Wang and Aitchison, 2025, Bergsma et al., 2025]. Other works, such as Qiu et al. [2025], Xiao [2025], propose other methods of scaling weight decay, such as scaling weight decay with w1 . We find that the scaling of weight decay with width allows the relative update size to remain invariant towards the infinite width limit. For SP in the learning rate sweeps, all layers are set at their PyTorch defaults. In particular, the input −1 layers are initialized at Wi j ∼ N (0, din ). The hidden and output layers are initialized as Wi j ∼ N (0, w−1 ). All layers have constant learning rate, and constant weight decay. For the weight decay experiments, for the SP run we keep the µP initialization and learning rate scaling, and use a constant weight decay for all layers. Justification of DoPr-AdamW and DoPr-Muon µP scaling Here we provide explanation for why DoPr does not affect the µP scaling rules for AdamW and Muon. Intuitively, the reason for this is because AdamW and Muon can both be seen as normalizing the gradient object - so regardless of whether it is applied to the raw gradient, or the DoPr-gradient, the result will still have similar magnitude elements. Let the weight in question be W at initialization, with raw gradient G, and preconditioner Σ. For simplicity we consider batch size of 1, let W be square, and let z ∈ Rw be the incoming data that the optimizer updates on. We assume z has Θ(1) entries with respect to width. We examine the one-step update. The preconditioned gradient is thus, for a damping constant ρ M = G(Σ + ρ · tr(Σ)I)−1 46

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

Loss

SP | DoPr-Muon

muP | DoPr-Muon

7.50

6.00

4.50

4.50 2

w=128

0 2 4 LR Multiplier (log2) w=256

w=512

w=1024

w=2048

w=128

0 2 4 LR Multiplier (log2) w=256

w=512

w=1024

6 w=2048

Figure 12: µP scaling behavior. Learning-rate sweeps comparing for DoPr-Muon under non-µP(SP) and µP parameterizations on GPT2 language model.

D = Muon(M) And so, the update is W+ = W − ηD We note that Muon orthonormalizes the gradient matrix (in this case M), up to numerical errors in the Newton-Schulz iterations. Thus, we may treat D as orthonormal. Then, we get: W+ z = Wz − ηDz The preconditioning only affects the Dz term, so we focus on this term. Since D is approximately orthonormal, and z has Θ(1) entries, Dz will also have Θ(1) entries, as D is orthonormal (for square M). Thus, DoPr-Muon achieves learning rate transfer across width with constant learning rate, matching Muon. In the case of AdamW, we consider W+ = W − ηAdamW(M) We consider AdamW without the EMA buffers, since momentum does not affect µP [Qiu et al., 2025]. In this case, AdamW is simply an elementwise sign function. We can also note that G can be written as δz⊤ , where δ is the gradient of the loss with respect to Wz. Similarly, we have Σ = zz⊤ . Thus, we see M = (δz⊤ )(zz⊤ + ρ · tr(zz⊤ )I)−1 Thus, applying an elementwise sign function to this will create a matrix D = AdamW(M) with all elements Θ(1), but correlated with z. Thus, we get W+ = Wz − ηDz Again, the first term is the same as in regular AdamW. In the second term, as D and z are correlated, by the Law of Large Numbers the resulting terms will be Θ(w), and so η should be scaled with w−1 , exactly matching the learning rate scaling for regular AdamW. Notice in this case, w is just the dimension of z. So, in either case, DoPr-AdamW and DoPr-Muon retains the same µP scaling rules as AdamW and Muon themselves.

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

Algorithm 3 Bootstrap Mean Estimation of Success Rate 1: Input: checkpoints {ck }K , rollouts per checkpoint R, bootstrap trials B, sample size m k=1

2: D ← [ ]

3: for k = 1, . . . , K do

Load checkpoint ck for r = 1, . . . , R do 6: Run one rollout episode 7: Append outcome x ∈ {0, 1} to D {1 = success} 8: end for 9: end for 10: for b = 1, . . . , B do 11: Draw subset P S b ⊂ D of size m (without replacement) 1 12: µ̂ b ← m x∈S b x 13: end for 14: Return: mean(µ̂), q0.10 , q0.50 , q0.90 4:

Figure 13: Half-Cheetah-v5 DoPr performance across AdamW, Muon, Signum, and AdaMuon. DoPr variants attain higher terminal reward which does not consistently correlate with train or validation loss improvements.

C.5

State-based Imitation Learning

This section provides implementation details and supporting results for Humanoid-v5 and Half-Cheetah-v5 Gymnasium benchmarks [Towers et al., 2025] Model and architecture. To maintain consistency with prior work, we use the Residual MLP architecture together with model EMA and a cosine learning rate scheduler following Block et al. [2024]. Full architectural hyperparameters are reported in Table 1. Training setup. For each optimizer variant, we conduct grid searches over learning rate, damping, and weight decay. The best-performing hyperparameter configurations are reported in Table 2 for Humanoid-v5 and Table 3 for Half-Cheetah-v5. We train three independent seeds over 1000 training steps. Evaluation Protocol. We evaluate model performance using episodic return on the corresponding control task. For each seed, we use a bootstrapped-mean evaluation procedure described in Appendix C.5 and report the 10/50/90 quantiles. Results. Across both Humanoid-v5 (Figure 7) and Half-Cheetah-v5 (Figure 13), DoPr variants consistently improve reward relative to their baseline optimizer counterparts. These gains do not reliably track training or validation loss trends.

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss Table 1: State-based IL Network Architecture and Training Configuration

C.6

Parameter

Value

Hidden layers Activation Batch size Training steps Eval trajectories LR scheduler Warmup fraction Min LR (ηmin ) EMA decay EMA update freq

[256, 256] ReLU 256 1000 100 Cosine annealing 0.05 0.0 0.99 Every step

Image-Based Robot Policy Learning

This section provides additional implementation details and supporting results for the Tool-Hang and Transport Robomimic tasks [Mandlekar et al., 2021] Model and architecture. We evaluate DoPr performance on state-of-the-art imitation learning architectures with generative control policies (GCPs) using the image-based Tool Hang and Transport tasks at Proficient Human (PH) expert level from Robomimic [Mandlekar et al., 2021]. We use the training pipeline of Pan et al. [2025], inheriting all architectural and environment hyperparameters. Our primary policy architecture is ChiUNet [Chi et al., 2023], with approximately 20M parameters. Training setup. We conduct learning rate sweeps over AdamW, Muon, and their corresponding DoPr variants. We precondition all layers (including FiLM) with full preconditioning (via Cholesky) except for Conv2d layers which use SUA Appendix B.3.2. For each method, we tune learning rates while keeping the remaining training configuration fixed to the reference implementation. Full optimizer hyperparameters are reported in Table 4 and Table 5. Evaluation Protocol. Policies are evaluated using task success rate under the standard Robomimic benchmark protocol. During training, we evaluate checkpoints every 20k steps using 9-step action sampling and report the corresponding rollout success rate using the bootstrapped mean procedure Appendix C.5. Results. Results are summarized in Figure 8. Across both tasks, DoPr variants consistently match or exceed their baseline counterparts, demonstrating that the proposed preconditioning strategy transfers effectively to high-dimensional visuomotor control.

C.7

Supervised Fine-Tuning

This section provides full details for the supervised fine-tuning (SFT) experiments used in Section 5.3. We report two SFT settings. The first is a small-scale Llama-3.2-3B GSM8K sweep used to characterize learning-rate sensitivity, training-time sample efficiency, and the relationship between downstream accuracy and token-level training loss. The second is our primary large-scale math SFT experiment, which fine-tunes Llama-3.1-8B with LoRA on the OpenMathInstruct-2 train_1M split and evaluates on GSM8K, GSM8K-CoT, and MATH-500. C.7.1

3B SFT

Overview. This experiment uses the Llama-3.2-3B base model AI@Meta [2024], Grattafiori et al. [2024] with LoRA Hu et al. [2022] on a 100K-sample subset of OpenMathInstruct-2 Toshniwal et al. 49

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss Table 2: Best optimizer hyperparameters across trajectory counts for Humanoid. Method

10 traj

20 traj

30 traj

40 traj

50 traj

AdamW

η = 3×10−3 λ = 1×10−2

η = 3×10−3 λ=0

η = 1×10−3 λ=0

DoPr-AdamW

η = 1×10−3 λ=0 d = 1×10−4

η = 1×10−3 λ=0 d = 1×10−3

η = 1×10−3 λ=0 d = 1×10−4

η = 1×10−3 λ=0 d = 1×10−3

η = 1×10−3 λ=0 d = 1×10−5

Muon

ηh = 1×10−2 ηh = 2×10−2 ηh = 1×10−2 ηh = 2×10−2 ηh = 1×10−2 λ = 1×10−2 λ = 5×10−3 λ = 1×10−2 λ = 1×10−2 λ = 5×10−3

DoPr-Muon

ηh = 2×10−2 ηh = 2×10−2 ηh = 2×10−2 ηh = 2×10−2 ηh = 2×10−2 λ = 1×10−2 λ = 5×10−3 λ = 1×10−2 λ = 1×10−2 λ = 1×10−2 d = 1×10−4 d = 1×10−4 d = 1×10−4 d = 1×10−5 d = 1×10−5

Signum

η = 3×10−4 λ=0

η = 1×10−3 λ=0

η = 3×10−4 λ=0

DoPr-Signum

η = 1×10−3 λ=0 d = 1×10−3

η = 3×10−4 λ=0 d = 1×10−4

η = 3×10−4 λ=0 d = 1×10−5

η = 3×10−4 λ=0 d = 1×10−4

η = 3×10−4 λ=0 d = 1×10−5

AdaMuon

η = 1×10−3 λ = 1×10−1

η = 3×10−3 λ = 1×10−2

η = 1×10−3 λ = 1×10−2

η = 3×10−3 λ = 1×10−1

DoPr-AdaMuon

η = 3×10−3 λ = 1×10−2 d = 1×10−2

η = 3×10−3 λ = 1×10−1 d = 1×10−2

η = 1×10−3 λ = 1×10−1 d = 1×10−4

η = 1×10−3 λ = 1×10−2 d = 1×10−4

η = 3×10−3 λ = 1×10−2 d = 1×10−3

[2024] for one epoch. Additional results We report the Muon SFT sweep results in Figure 16. Consistent with the DoPr-AdamW comparison, DoPr-Muon also achieves stronger downstream GSM8K performance across the learning rate sweep, while often demonstrating higher final training loss than the base Muon optimizer. This further supports our main observation that training loss alone can be a poor proxy for downstream performance, and that DoPr improves the optimization in a way that is better aligned with downstream generalization. Model and architecture. We fine-tune the Llama-3.2-3B base model AI@Meta [2024], Grattafiori et al. [2024] using the Llama-3.2-3B-Instruct tokenizer. All experiments use bfloat16 precision for weights and activations, and implement attention with FlashAttention-2. Dataset and preprocessing. Training data are drawn from OpenMathInstruct-2 Toshniwal et al. [2024], using a fixed 100K-example subset. We tokenize sequences to a maximum context length of 1024 tokens and enable sequence packing for efficiency. Parameter-efficient fine-tuning. We apply LoRA Hu et al. [2022] to all linear layers with rank r = 16 and scale α = 16. We additionally train the token embedding matrix to accommodate chat-template special tokens, while preserving weight tying between the input embeddings and output projection. 50

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss Table 3: Best optimizer hyperparameters across trajectory counts for HalfCheetah. Method

10 traj

20 traj

30 traj

40 traj

50 traj

AdamW

η = 2×10−3 λ = 3×10−3

η = 2×10−4 λ=0

η = 3×10−3 λ = 1×10−2

η = 2×10−3 λ = 3×10−2

DoPr-AdamW

η = 1×10−4 λ=0 d = 1×10−5

η = 3×10−4 λ=0 d = 1×10−3

η = 2×10−3 λ=0 d = 1×10−3

η = 2×10−3 λ = 1×10−2 d = 1×10−6

η = 2×10−3 λ = 1×10−2 d = 1×10−3

Muon

ηh = 2×10−2 ηh = 2×10−2 ηh = 3×10−2 ηh = 3×10−2 ηh = 3×10−2 λ = 5×10−3 λ = 5×10−3 λ = 5×10−3 λ = 1×10−2 λ = 2×10−2

DoPr-Muon

ηh = 5×10−3 ηh = 3×10−2 ηh = 3×10−2 ηh = 3×10−2 ηh = 3×10−2 λ = 5×10−3 λ = 2×10−2 λ = 2×10−2 λ = 5×10−3 λ = 2×10−2 d = 3×10−4 d = 1×10−3 d = 1×10−3 d = 1×10−3 d = 1×10−3

Signum

η = 1×10−3 λ=0

η = 3×10−4 λ=0

η = 3×10−3 λ=0

DoPr-Signum

η = 1×10−4 λ=0 d = 1×10−4

η = 1×10−3 λ=0 d = 1×10−5

η = 1×10−3 λ=0 d = 1×10−3

AdaMuon

η = 1×10−2 λ = 1×10−2

η = 3×10−3 λ = 1×10−2

DoPr-AdaMuon

η = 1×10−2 λ = 1×10−2 d = 1×10−3

η = 1×10−2 λ = 1×10−2 d = 1×10−2

η = 1×10−3 λ = 1×10−2 d = 1×10−3

η = 1×10−2 λ = 1×10−2 d = 1×10−2

η = 3×10−3 λ = 1×10−2 d = 1×10−3

Optimization and training. We train for one epoch using either AdamW or Muon, each with and without the DoPr wrapper. The base learning rate is 1 × 10−5 , with cosine scheduling and 5% warmup. We use per-device batch size 2 and gradient accumulation over 8 steps. Evaluation protocol. Downstream evaluation uses lm-eval with the vLLM backend in bfloat16 precision. We evaluate on GSM8K with 5-shot prompting, enabling chat-style prompting via the model’s native template and formatting few-shot exemplars as multi-turn conversations. C.7.2

8B SFT

Overview. Our primary SFT experiment uses meta-llama/Llama-3.1-8B checkpoint AI@Meta [2024], Grattafiori et al. [2024] as the frozen base model. We train only LoRA adapters Hu et al. [2022]; the base model weights remain frozen throughout training. The tokenizer and chat template are taken from meta-llama/Llama-3.1-8B-Instruct. Training data are drawn from nvidia/OpenMathInstruct-2 Toshniwal et al. [2024], split train_1M. The primary comparison in this setting is between AdamW and DoPr-AdamW, where DoPr-AdamW applies AP to the trainable LoRA modules before the AdamW update. Additional results. Results are reported in Table 7. Across the full learning-rate grid, DoPr-AdamW improves the average final accuracy over AdamW by +3.9, +10.8, and +0.9 percentage points on GSM8K, 51

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss Table 4: Best optimizer hyperparameters for Tool Hang PH Image using ChiUNet (seeds 42, 43, 44). Hyperparameter

AdamW

DoPr-AdamW

Muon

DoPr-Muon

lr weight_decay muon_momentum grad_clip_norm gradient_steps batch_size precond_ema_rate grad_ema_rate precond_include_film damping damping_mode

1e−4 1e−5 — 10.0 300k 256 — — — — —

1e−4 1e−5 — 10.0 300k 256 0.9 0.9 true 1e−4 trace

3e−4 0.1 0.95 1.0 200k 256 — — — — —

3e−4 0.1 0.95 1.0 200k 256 0.9 0.9 true 1e−4 trace

Table 5: Best optimizer hyperparameters for Transport MH Image using ChiUNet (seeds 42, 43, 44). Hyperparameter

AdamW

DoPr-AdamW

Muon

DoPr-Muon

lr weight_decay muon_momentum grad_clip_norm gradient_steps batch_size precond_ema_rate grad_ema_rate precond_include_film damping damping_mode

1e−4 1e−5 — 10.0 300k 256 — — — — —

1e−4 1e−5 — 10.0 300k 256 0.0 0.0 true 1e−4 trace

1e−4 0.1 0.95 1.0 200k 256 — — — — —

1e−4 0.1 0.95 1.0 200k 256 0.0 0.0 true 1e−4 trace

GSM8K-CoT, and MATH-500, respectively. The gains are especially pronounced in the high-learning-rate regime: at 5 × 10−4 , DoPr-AdamW improves GSM8K-CoT from 51.6% to 81.9%, and at 7 × 10−4 , from 14.5% to 78.4%. Thus, beyond improving the best observed endpoint performance, DoPr substantially broadens the stable and effective learning-rate range for large-scale SFT. Taking the best final checkpoint over the sweep, DoPr-AdamW reaches 80.5% on GSM8K, 81.9% on GSM8K-CoT, and 47.4% on MATH-500, improving over the best AdamW result on all three benchmarks. Training data and prompt construction. Each training example consists of a math problem and a generated solution from OpenMathInstruct-2. The problem field is inserted into the fixed user prompt:

Solve the following math problem. answer) inside \boxed{}.

Make sure to put the answer (and only

The target completion is the corresponding generated_solution. Parameter-efficient fine-tuning. All large-scale SFT runs use LoRA Hu et al. [2022]; full fine-tuning is not used. LoRA adapters are applied to the standard Llama attention and MLP projection matrices: {q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj}. The LoRA rank is r = 64, the LoRA scale is α = 128, LoRA dropout is 0.05, and no bias terms are trained. 52

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

(a) Tool Hang: AdamW vs DoPr-AdamW

(b) Tool Hang: Muon vs DoPr-Muon

(d) Transport: Muon vs DoPr-Muon

Figure 14: Best mean success rate throughout training on Tool-Hang and Transport tasks from robomimic. We compare AdamW, Muon, and their DoPr variants, reporting the best mean success rate over evaluation checkpoints every 20k steps.

Evaluation protocol. Evaluation is performed with lm-evaluation-harness using the vLLM backend. Fine-tuned models are evaluated by loading the frozen Llama-3.1-8B base weights together with the LoRA adapter and the saved adapter tokenizer when present. The evaluation wrapper applies the chat template for fine-tuned adapters. Inference uses bfloat16 precision, tensor parallelism over 8 GPUs, greedy decoding, temperature 0.0, and top-p = 1.0. For the learning-rate sweep evaluation, we use max_model_len=8193, max_gen_toks=2048, gpu_memory_utilization=0.9, batch_size=auto, and max_batch_size=64. The one-token context margin is introduced to avoid vLLM scoring failures at exactly the model length. We report the following evaluation tasks: • gsm8k: 5-shot direct GSM8K generation, exact match with the flexible extraction filter. • gsm8k_cot: 8-shot chain-of-thought GSM8K generation, exact match with the flexible extraction filter. • minerva_math500: 4-shot Minerva-style MATH-500 generation, scored with math_verify. Reported uncertainty values are the standard errors emitted by lm-evaluation-harness.

C.8

Generative Modeling

Generative modeling as TTF. Generative modeling is a natural TTF setting: training is local teacherforced velocity regression, whereas deployment rolls out the learned velocity field under its own induced 53

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

101

100

Train Loss

Tool Hang

101

20 40 60 80 100120140160180200220240260280300

Training Steps (k) AdamW

DoPr-AdamW

Muon

(a) Tool Hang: AdamW vs DoPr-AdamW

Train Loss

Train Loss 60

100 120 140 160 180 200

101

Training Steps (k) AdamW

DoPr-Muon

Transport

101

100 120 140 160 180 200

(b) Tool Hang: Muon vs DoPr-Muon

Transport

Training Steps (k)

100 120 140 160 180 200

Training Steps (k)

DoPr-AdamW

Muon

DoPr-Muon

(d) Transport: Muon vs DoPr-Muon

Figure 15: Training Loss on Tool-Hang and Transport tasks from robomimic. We compare AdamW, Muon, and their DoPr variants

trajectory. As in flow matching, the model minimizes errors on (x t , y t ), where y t = v(x t , t) is the closed-form target velocity. At generation time, however, samples are produced by iteratively integrating d b zt = b v(b x t , t), so velocity errors can accumulate and shift future states away from the learned field, e.g. dt the training distribution. Consequently, per-step training loss can misalign with final sample quality metrics such as FID. We evaluate DoPr on generative modeling using the SiT-S model [Ma et al., 2024] trained at 256×256 resolution on ImageNet [Deng et al., 2009]. As in Ma et al. [2024], the baseline model is trained with AdamW, and generates using a number of function evaluations (NFE) fixed to 250. We follow our general operating guidelines and swap the optimizer to DoPr-AdamW, maintaining the same learning rate and all other architectural, training, and evaluation protocols as the baseline. We then repeat this using baseline Muon and DoPr-Muon. Table 10 summarizes the model, training, and evaluation configuration. We plot the FID-50K scores against training steps as the “downstream metric”; refer to Figure 17. Across both AdamW and Muon, DoPr consistently accelerates convergence and improves final FID. In addition, qualitative inspection of intermediate generations in Figure 18 indicates that DoPr induces a distinct optimization trajectory compared to the base optimizers. Notably, as seen in Figure 17, the relative improvements in FID are somewhat proportional to 54

Muon DoPr-Muon

30 20 10

Final Train loss

Accuracy (%)

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

0 10 4

10 3 Learning rate

Muon DoPr-Muon

100

10 4

10 2

10 3 Learning rate

10 2

Figure 16: GSM8K 3B SFT sweep. Peak GSM8K accuracy across training steps vs. learning rate, comparing between Muon and DoPr-Muon.

improvements in train loss convergence. We hypothesize that generative modeling is a relatively unique TTF setting, in that one has control over the data-generating dynamical system. For example, many works study “straightening” the transport trajectory (cf. standard flow matching versus Rectified Flow [Liu et al., 2023] and Optimal-Transport CFM [Tong et al., 2024]), which can enable low NFE sampling. Therefore, the severity of TTF is tied to the generative algorithm. We leave exploring this space for co-design to future work.

Additional Experiments and Setup

D.1

Moment Grafting Experiments

Datasets. The “original” dataset used is CIFAR10 Krizhevsky et al.. We split this into a 0.8:0.2 split of train/validation data. Then, we generate a “transformed” dataset by applying a random spatial rotation of CIFAR by 90◦ , and after normalizing all coordinates to be between [−1, 1], we apply tanh elementwise. We note this induces a nonlinear transform. Training. Next, we train 3 models, each a 3 layer MLP with layer dimensions (3072, 256), (256, 256), and (256, 10), and SiLU activations. All models are initialized identically. One model is trained on the original CIFAR dataset Krizhevsky et al., called “Original." Another model is trained on the transformed CIFAR data, denoted “Transformed." We also coordinate training so that at each step, the batches the Transformed model sees corresponds exactly to the transformed version of the batches the Original model sees. The last model is called “Grafted." All models are trained simultaneously. The Grafted model is trained as follows: during training, the activations inputted into each layer of the Transformed model are recorded, and their means/covariances calculated. Then, the Grafted model takes as input the Original data batch. Then, before passing the data through each layer’s weights, an affine transform is applied to the incoming activations to match their mean and covariance to the cached means/covariances of the Transformed model’s forward pass. Concretely, let µ t , Σ t be the mean and covariance of transformed batch, X t ∈ R3072×b , where b is the batch size. When the corresponding batch X is passed into the Grafted model, let X have mean µ and covariance Σ. Then, before we pass X through the layer, we apply the following transform to X: 1/2

Xnew = Σ t Σ−1/2 (X − µ1⊤ ) + µ t 1⊤ We can note that as X − µ1⊤ has mean 0, then Xnew will have mean µ t . Similarly, we can find the 55

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss Table 6: Configuration for the small-scale GSM8K SFT sweep. Component

Setting

Base model Tokenizer Precision Attention implementation

Llama-3.2-3B AI@Meta [2024], Grattafiori et al. [2024] Llama-3.2-3B-Instruct tokenizer

Training dataset Data subset Context length Packing

OpenMathInstruct-2 Toshniwal et al. [2024] 100K examples sampled from the 1M split 1024 tokens Enabled

PEFT method LoRA target modules LoRA rank / scale Additional trainables Weight tying

LoRA Hu et al. [2022] All linear layers r = 16, α = 16 Token embedding matrix, used for chat-template special tokens Preserved between input embeddings and output projections

Optimizers compared Training length Base learning rate LR schedule Per-device batch size Gradient accumulation

AdamW and Muon, each with and without the DoPr wrapper

Eval framework Eval precision Benchmark Prompting format

lm-eval with vLLM backend bfloat16 GSM8K, 5-shot Chat template enabled; few-shot examples are formatted as multi-turn conversations

bfloat16 weights and activations FlashAttention-2 Dao et al. [2022], Dao [2024]

1 epoch 1 × 10−5 Cosine decay with 5% warmup 2 8 steps

covariance as: 1/2

1/2

(Xnew − µ t 1⊤ )(Xnew − µ t 1⊤ )⊤ = Σ t Σ−1/2 (X − µ1⊤ )(X − µ1⊤ )⊤ (Σ−1/2 )⊤ (Σ t )⊤ Noting that covariances matrices are symmetric: 1/2

1/2

= Σ t Σ−1/2 Σ(Σ−1/2 )⊤ (Σ t )⊤ = Σ t Σ t

= Σt

Therefore, this transform does give Xnew a mean of µ t , and covariance of Σ t . However, since a nonlinear transform was used to create the transformed data, this affine transform will not be an exact recovery of the transformed data, only matching the mean and covariance. We repeat this process for the hidden layer and output layer, with the respective activations for each. In these experiments, we use a batch size of 2048, and µP scaling. The base learning rate was −3 10 for SGD, 10−5 for AdamW, and 10−4 for Muon. Otherwise, the default PyTorch parameters are used for AdamW and Muon. To compute square roots and inverse square roots of the covariance matrices, symmetric eigenvalue decomposition was used, with a damping constant of 10−6 · tr(A), where A is the matrix to be decomposed. Evaluation. Finally, the models are validated on the Transformed data after each training step, and the validation losses plotted. The results are displayed in Figure 19. Despite not being trained on the Transformed data, the Grafted model’s validation performance on the transformed data matches much more closely with the Transformed model than the Original 56

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss Table 7: 8B SFT sweep. Final accuracy is reported as percent ± one standard error for AdamW and DoPr-AdamW on GSM8K, GSM8K-CoT, and MATH-500. ∆ denotes the improvement of DoPr-AdamW over AdamW in percentage points. Learning rate

2×10−5 5×10−5 7×10−5 1×10−4 2×10−4 5×10−4 7×10−4 1×10−3 Mean ∆ over LR grid

GSM8K

GSM8K-CoT

MATH-500

AdamW

DoPr-AdamW

∆

AdamW

DoPr-AdamW

∆

AdamW

DoPr-AdamW

∆

71.4±1.2 68.1±1.3 69.7±1.3 68.1±1.3 78.0±1.1 66.2±1.3 56.0±1.4 2.9±0.5

71.9±1.2 70.1±1.3 69.5±1.3 66.7±1.3 80.1±1.1 80.5±1.1 70.4±1.3 2.2±0.4

+0.5 +2.0 −0.2 −1.4 +2.1 +14.3 +14.4 −0.7

72.2±1.2 69.4±1.3 69.9±1.3 72.5±1.2 81.0±1.1 51.6±1.4 14.5±1.0 2.8±0.5

71.9±1.2 71.6±1.2 70.1±1.3 68.5±1.3 75.7±1.2 81.9±1.1 78.4±1.1 2.0±0.4

−0.2 +2.1 +0.2 −3.9 −5.2 +30.3 +63.9 −0.8

33.8±2.1 34.4±2.1 29.6±2.0 33.0±2.1 41.6±2.2 46.4±2.2 42.6±2.2 3.6±0.8

33.8±2.1 35.0±2.1 34.4±2.1 32.2±2.1 38.6±2.2 47.4±2.2 45.8±2.2 4.6±0.9

0.0 +0.6 +4.8 −0.8 −3.0 +1.0 +3.2 +1.0

+3.9

+10.8

+0.9

model (trained on the original data) across optimizers, suggesting that when trained on data with similar means and covariances, models often learn similarly, even if the data itself does not perfectly match. While not displayed, we further note that increasing batch size and decreasing learning rate both help to improve how closely the Grafted model fits the Transformed model in validation loss, as they reduce drift arising from mean/covariance estimation noise and higher-order mismatches from the non-linear data transform. Notably, we remark that since there is no corrective mechanism to port the Grafted model’s weights onto the Transformed model’s “ground truth” weights, the drift will monotonically compound over the training horizon. The Grafted model matching the Transformed model in validation demonstrates that even ignoring the higher-order nonlinear aspects of the forward pass, the mean and covariance of a forward pass are sufficient to graft much of the downstream performance of a model that is truly trained on the transformed data onto another model, despite the latter only seeing the original data in its forward/backward pass. This lends credence to the premise of AP, where the working hypothesis is that “error directions” (as captured by the mean and covariance of each layer’s activations) strongly bias the optimizer’s downstream behavior. Grafting demonstrates that, in the hypothetical case where we knew the feature directions of the “downstream” model, we can graft its bias onto the training on the training distribution. However, since we typically cannot know these features ahead of time, AP— which essentially treats all feature directions equally — is the “worst-case” optimal approach.

D.2

Stable Ranks of Optimizer and Model Quantities

Following prior work that hypothesizes the superior training efficiency of Muon (compared to, e.g., SGD) due to ratios of certain model/optimizer quantities, we plot these relevant quantities using the terminal weights of the small-scale 3B SFT run in Section 5.3 and Appendix C.7. In particular, we follow the predictions in Davis and Drusvyatskiy [2025] and compare the ratio of the “nuclear rank” of the layer-wise gradients and the stable rank of the batch activation covariances of a forward pass, see Figure 20a. However, we note that DoPr-Muon does not orthogonalize the raw gradient, rather, it takes the AP gradient as input. Following the hypothesis in Davis and Drusvyatskiy [2025], we also compare the nuclear rank of the AP gradient to the stable rank of the activation covariances (Figure 20b). On the other hand, we plot the standalone stable ranks of the weight matrices themselves (Figure 20c) and the batch activation covariances (Figure 20d). In this case, we do not see a consistent trend across layers, which supports our working hypothesis that our intervention on the optimizer is unrelated to 57

Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss Table 8: Core configuration for the large-scale 8B SFT experiment. Component

Setting

Base model Tokenizer / chat template Precision Attention implementation Trainer

meta-llama/Llama-3.1-8B meta-llama/Llama-3.1-8B-Instruct

Training dataset Input fields Objective Maximum sequence length Packing Training length Processed tokens Optimizer steps

nvidia/OpenMathInstruct-2, split train_1M problem and generated_solution

Effective global batch size World size Per-device micro-batch size Gradient accumulation DeepSpeed configuration Gradient checkpointing

512 packed sequences 8 GPUs 2 32 steps ZeRO-3 Enabled, reentrant checkpointing path

bfloat16 model and compute; float32 preconditioner statistics FlashAttention-2 Dao et al. [2022], Dao [2024] trl.SFTTrainer with PEFT LoRA

Completion-only next-token loss; user-prompt tokens are masked 4096 tokens Enabled 2 epochs 997,968,955 tokens 482

regularizing the optimizer dynamics in weight geometry, but rather targets the error geometry, captured by the correspondence between the (AP) gradient and the activation.

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

Table 9: LoRA, optimizer, learning-rate sweep, and preconditioner configuration for the large-scale 8B SFT experiment. Component

Setting

PEFT method LoRA target modules LoRA rank / scale LoRA dropout LoRA bias Trainable adapter parameters

LoRA Hu et al. [2022]

AdamW betas and epsilon Weight decay LR schedule

q,k,v,o,gate,up,down projection matrices r = 64, α = 128 0.05

none 167,772,160 (0.9, 0.98), ε = 10−8 0.01 Constant, no warmup

Table 10: Experimental setup for SiT-S generative modeling on ImageNet at 256 × 256. Across all runs, we vary only the optimizer choice and whether DoPr is enabled; all other training and evaluation settings match the standard SiT protocol. Component

Setting

Model Dataset Resolution Architecture / parameterization / conditioning

SiT-S Ma et al. [2024] ImageNet Deng et al. [2009] 256 × 256 Exactly follows the original SiT implementation Ma et al. [2024]

Optimizers compared Training steps Learning rate

AdamW; Muon; each with and without the DoPr wrapper AdamW: 400K steps; Muon: 200K steps AdamW: matches SiT defaults Ma et al. [2024]; Muon: 10−3 (per prior recom-

Other training settings

mendations) Held fixed across runs (data preprocessing, batch size, LR schedule, regularization)

Sampler NFE Metric FID reference set

ODE-based sampler from SiT Ma et al. [2024] 250 FID-50K (unless otherwise specified) ImageNet training set Deng et al. [2009]

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

Figure 17: SiT-S on ImageNet256. Top row: perceptual quality (FID-50K) vs. training steps. Bottom row: training loss trajectories. We see that DoPr improves sample quality as well as loss convergence.

AdamW (20K steps)

DoPr-AdamW (20K steps)

Figure 18: Generated ImageNet256 samples at 20K training steps (SiT-S, AdamW). Although both models are trained under identical settings and evaluated using the same ODE sampler, the samples generated by DoPr-AdamW exhibit qualitatively different structure and composition compared to AdamW at the same early training stage. 60

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

Figure 19: We train 3 MLPs on CIFAR, and a nonlinear invertible transform of CIFAR. We plot the validation loss for the 3 models for each of AdamW (left), SGD (middle), and Muon (right). Across all 3, it can be seen that the Grafted model’s trajectory is closer to the Transformed model than the Original model.

3.0

Muon

DoPr-Muon

ratio

2.4

1.8 1.2

0.6

0.0

0.4 0.0 0.4 0.8

0.5 0.0 0.5

k_ pr o o_ j: pr B o o_ j: ga pr B te oj _p :A r up oj _p :B ro k_ j: pr B o q_ j: pr A up oj : _ ga pr A te oj _p :A ro q_ j: pr A o v_ j: do pr B wn oj : _ do pr A wn oj _p :A ro v_ j: pr B oj :B

1.0

k_ pr o o_ j: pr B o o_ j: ga pr B te oj _p :A r up oj _p :B ro k_ j: pr B o q_ j: pr A up oj : _ ga pr A te oj _p :A ro q_ j: pr A o v_ j: do pr B wn oj : _ do pr A wn oj _p :A ro v_ j: pr B oj :B

ratio

Muon

3.0

2.4

(a)

(b)

Muon

DoPr-Muon

100 stable rank

10.0 7.5 5.0 2.5

0.6 0.3 0.0 0.3 0.6

75 50 25

15 0

j: pr B oj o_ : ga pr B te oj _p :A r up oj _p :B ro k_ j: pr B o q_ j: pr A up oj : _ ga pr A te oj : _p A ro q_ j: A pr o v_ j: do pr B wn oj : _ do pr A wn oj : _p A ro v_ j: pr B oj :B

pr o

j: pr B oj : ga pr B te oj _p :A r up oj _p :B ro k_ j: pr B o q_ j: pr A up oj : _ ga pr A te oj _p :A ro q_ j: pr A o v_ j: do pr B wn oj : _ do pr A wn oj _p :A ro v_ j: pr B oj :B

pr o

DoPr-Muon

0.0

Muon

(c)

stable rank

12.5

(d)

Figure 20: Metric values aggregated by layer group for Muon and DoPr-Muon. The bottom panel shows the mean difference (DoPr-Muon minus Muon).

Record · ID 259467 · SHA-256 0fceaf8c80734e31
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.