Research map / Optimization / Optimizers
Careful with that Scalpel: Improving Gradient Surgery with an EMA
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Bloop, a gradient surgery method that follows the main-loss gradient plus the auxiliary gradient projected orthogonally to an EMA of main-loss gradients, approximating simple bilevel optimization in stochastic settings.
Problem
Optimizing a main training loss while controlling an auxiliary objective (e.g., robustness, performance on another dataset, weight norm); naive regularized (mixed) training conflicts with gradient directions and the simple bilevel formulation is hard to solve stochastically for neural networks.
Approach
Project the mini-batch auxiliary gradient orthogonally onto an exponential moving average of main-loss gradients (rather than the noisy current batch gradient), so the update preserves first-order main-loss descent and approximates the bilevel solution; the EMA is what makes the orthogonality hold stochastically.
Key results
- On transformer pre-training with c4 and RCV-1, Bloop achieves a significantly better Pareto front than mixed, PCGrad and DB, which perform similarly to mixed.
- On the Cifar10Mnist multi-task problem, Bloop gives a slightly improved Pareto front over the other methods.
- Theory: with an EMA rate rho near the step size, stochastic Bloop converges at O(T^{-1/4}) for non-convex L_main and linearly to a radius O(eta^{1/3}) under PL.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use an EMA of the main-loss gradient in the projection | Bloop · MNIST smooth-network training | training Pareto front | Helped |
| Use an EMA of the main-loss gradient in the projection | Bloop · ImageNet with L2 regularization | test Pareto front | Helped |
| Use an EMA of the main-loss gradient in the projection | Bloop · Cifar10Mnist | Pareto front | Helped |
| Use an EMA of the main-loss gradient in the projection | Bloop · c4 and RCV-1 language modeling | evaluation loss Pareto front | Helped |
| Use an EMA of the main-loss gradient in the projection | Bloop · Paracrawl and WMT translation | evaluation loss Pareto front | Helped |
| Remove the EMA from the projection (ablation) | Bloop · Transformer pre-training | main and auxiliary losses | Hurt |
| Set the EMA parameter to ρ=0.001 (ablation) | Bloop · Transformer pre-training | main and auxiliary losses | Hurt |
| Set the EMA parameter to ρ=0.9 (ablation) | Bloop · Transformer pre-training | Pareto-front behavior | Hurt |
| Replace Bloop with PCGrad or Dynamic Barrier without EMA | Bloop · Experiments across the paper | Pareto front | Hurt |
| Use Bloop instead of the other gradient-surgery methods | ImageNet and CIFAR-10 classification with separate heads · ImageNet/CIFAR-10 | Pareto front | No clear effect |
Surprising findings
- PCGrad and Dynamic Barrier behave very similarly to the mixed method in nearly all experiments, attributed to high main-loss gradient variance.
- On the ImageNet + CIFAR-10 joint classification experiment, all methods including Bloop perform similarly because the two tasks barely conflict.
Limitations
- Does not prove convergence to KKT points of the simple bilevel problem; guarantees are weaker than convex-case results.
- Theory requires assumptions such as bounded directions, bounded variance and a PL condition for linear rates.
- Gains are mostly on training losses; test-loss improvements are problem dependent (e.g., only slight on some experiments).
Experimental setup
Datasets: MNIST ImageNet CIFAR-10 Cifar10Mnist (constructed) c4 RCV-1 Paracrawl WMT 2009-2019 WMT 2020
Benchmarks: MNIST ImageNet Cifar10Mnist c4 and RCV-1 language modeling Paracrawl/WMT translation
Metrics: training loss test loss Pareto front of main vs auxiliary loss
Baselines: Mixed (regularized) training Dynamic Barrier PCGrad
Builds on: Adam PCGrad Dynamic Barrier A-GEM
Closest papers by meaning
- Beyond the Mean: Fisher-Orthogonal Projection for Natural Gradient Descent in Large Batch Training (AAAI 2026 · similarity 0.83)
- The AdEMAMix Optimizer: Better, Faster, Older (ICLR 2025 · similarity 0.82)
- BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization (NeurIPS 2026 · similarity 0.82)
- The Iterative Optimal Brain Surgeon: Faster Sparse Recovery by Leveraging Second-Order Information (NeurIPS 2024 · similarity 0.82)
- Implicit Regularization for Tubal Tensor Factorizations via Gradient Descent (ICML 2025 · similarity 0.81)
- Stacey: Promoting Stochastic Steepest Descent via Accelerated $\ell_p$-Smooth Nonconvex Optimization (ICML 2025 · similarity 0.81)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.