double-preconditioning-test-time-optimization

## Context From arXiv:2606.06418 (June 2026) - "Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss" by Thomas T. Zhang, Alok Shah, Yifei Zhang, Vincent Zhang, Nikolai Matni, Max Simchowitz. Addresses train-test feedback (TTF) phenomenon: mismatch between training/validation loss and downstream metrics (task success rate, generation quality) that grows with task length. Proposes double preconditioning as new optimization axis to combat error accumulation in rollouts. ## Core Methodology ### 1. Test-Time Feedback (TTF) Problem **Definition**: Settings where network is: - **Trained on**: One-step prediction loss (regression, cross-entropy) - **Deployed via**: Rolling out along its own predictions **Examples**: - Autoregressive language modeling (next-token prediction → multi-token generation) - Flow-based generative modeling (single-step density → sequential generation) - Robot policy learning (one-step action → trajectory rollout) **Phenomenon**: ``` Training: L_train = ||prediction_single - target_single||² Testing: L_test = task_success_rate, generation_quality Mismatch: L_train improvement ≠ L_test improvement Error accumulation: Mismatch grows with rollout length ``` **Prior solutions** (incomplete): - Data curation (limited scope) - Architecture design (modifies model) - Objective design (changes loss) - **This paper**: Optimization as new design axis ### 2. Double Preconditioning Architecture **Two preconditioning layers**: 1. **Gradient-wise preconditioning**: Adam, Muon (modify gradient updates) 2. **Activation-wise preconditioning (AP)**: KFAC (modify activations before gradient computation) **Combined effect**: ``` Standard SGD: θ ← θ - η * ∇L Adam: θ ← θ - η * M^{-1} * ∇L (gradient preconditioning) DoPr: θ ← θ - η * M^{-1} * K^{-1} * ∇L (double preconditioning) where: M = gradient preconditioner (Adam/Muon) K = activation preconditioner (KFAC-based) ``` ### 3. Implementation Architecture ```python class DoPrOptimizer: def __init__(self, model, lr=1e-3, beta1=0.9, beta2=0.999): self.model = model self.lr = lr # Gradient preconditioning (Adam/Muon) self.gradient_precond = AdamPreconditioner(beta1, beta2) # Activation preconditioning (KFAC-based) self.activation_precond = KFACPreconditioner(model) def step(self, batch): # Step 1: Activation preconditioning activations = self.activation_precond.precond_forward(batch) # Step 2: Forward pass with preconditioned activations output = self.model(activations) loss = compute_loss(output, target) # Step 3: Backward pass gradients = torch.autograd.grad(loss, self.model.parameters()) # Step 4: Gradient preconditioning gradients = self.gradient_precond.precond_gradient(gradients) # Step 5: Parameter update with torch.no_grad(): for param, grad in zip(self.model.parameters(), gradients): param -= self.lr * grad ``` ## Key Applications 1. **Autoregressive language models**: Generation quality ≠ next-token accuracy 2. **Flow-based generative models**: Normalizing flows, diffusion models 3. **Robot policy learning**: Policy rollouts with error accumulation 4. **Time series forecasting**: Multi-step prediction beyond single-step accuracy 5. **Trajectory optimization**: Sequential predictions compound errors ## Pitfalls 1. **KFAC computation cost**: Expensive to compute Kronecker factors. Update periodically. 2. **Memory overhead**: Storing A and G matrices requires memory. Use approximation. 3. **Inversion stability**: Matrix powers need eigenvalue clipping. 4. **Hyperparameter tuning**: Two preconditioner hyperparameter sets. 5. **Validation loss confusion**: Don't expect val loss to predict test-time gains. 6. **TTF setting identification**: DoPr helps only when train-test mismatch exists. 7. **Implementation complexity**: Two preconditioners require careful integration. ## Verification 1. **Test-time metrics**: Evaluate downstream task success rate, generation quality 2. **Validation loss comparison**: Compare DoPr vs Adam/Muon 3. **Rollout analysis**: Measure error accumulation over rollout lengths 4. **Preconditioner diagnostics**: Check A and G matrix statistics 5. **Cross-domain testing**: Language modeling, flow models, robotics 6. **Ablation studies**: Test gradient-only vs activation-only vs double preconditioning ## Key Innovation **Optimization as TTF mitigation axis**: Prior work focused on data/architecture/objective. This introduces optimization preconditioning. **Drop-in intervention**: Add KFAC to existing Adam/Muon. Minimal code change, maximal downstream benefit. **Separation of concerns**: - Gradient preconditioning → Fast training convergence - Activation preconditioning → Test-time stability - Both needed for TTF settings ## Mathematical Foundation **Standard update**: θ ← θ - η * ∇L **DoPr-modified**: θ ← θ - η * M^{-1} * K^{-1} * ∇L **Curvature interpretation**: - M shapes parameter space (gradient direction) - K shapes activation space (loss landscape geometry) - Combined: Richer optimization geometry for TTF **Rollout error intuition**: - Single-step error: ε - N-step error: ε * (1 + ||J|| + ||J||² + ...) where J is Jacobian - DoPr reduces ||J|| → Slower error accumulation ## Experimental Validation (Paper) **Test domains**: Autoregressive language modeling, flow-based generative modeling, robot policy learning **Findings**: - DoPr improves test-time performance - Gains NOT reflected in validation loss improvements - Benefit grows with rollout length - Drop-in intervention: Easy to add **Key observation**: "Gains in test-time performance do not consistently accompany improvements in validation loss" ## Practical Deployment 1. **Model training**: Standard supervised loop 2. **Add KFAC hooks**: Register forward/backward hooks 3. **Compute factors**: Update Kronecker matrices periodically 4. **Apply preconditioning**: Transform activations and gradients 5. **Update parameters**: Standard update with preconditioned gradients 6. **Evaluate both metrics**: Track validation loss AND test-time performance **Code integration**: ```python # Replace Adam with DoPr optimizer = DoPrOptimizer(model, lr=1e-3) # Wraps Adam + KFAC loss.backward() optimizer.step() ``` **Activation**: DoPr optimizer, test-time feedback mitigation, KFAC Adam, autoregressive error accumulation, generation quality optimization, rollout stability, TTF optimization, train-test shift correction