Adaptive Optimization via Momentum on Variance-Normalized Gradients

arXiv:2602.10204 · cs.LG, math.OC · Submitted 2026-08-15 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Adaptive Optimization via Momentum on Variance-Normalized Gradients".

Jane: The paper was written by Francisco Patitucci and Aryan Mokhtari from University of Texas at Austin and Google Research.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back to the show, everyone. Today we're digging into a fresh arXiv paper called "Adaptive Optimization via Momentum on Variance-Normalized Gradients." Jane, what's your first read on that title?

Jane: Tom, I love it because it tells you exactly what the fix is. We've got this family of optimizers—Adam being the famous one—that powers basically all of deep learning. And this paper is saying, hey, we can do better by changing two specific things about how Adam works.

Tom: And for our listeners who might not live and breathe optimization, can you break down what an optimizer even does?

Jane: Sure. When you train a neural network, you're trying to adjust millions of little dials to make the network better at its task. The optimizer is the thing that decides how much to turn each dial, and in which direction, based on the gradients it sees. Adam has been the workhorse because it adapts the step size for each dial individually.

Tom: Right, and the paper's title points at two tweaks. One is about using variance instead of the raw second moment, and the other is about when you apply momentum. I've got to say, as someone who's trained models, I've seen Adam do some weird things—loss spikes, sudden jumps.

Jane: Exactly. And the authors—Francisco Patitucci and Aryan Mokhtari from UT Austin and Google Research—they've identified why those weird things happen. The way Adam is structured, there's this coupling between your momentum buffer and your current gradient noise. It's like carrying a heavy backpack while someone randomly changes how heavy the ground is beneath you.

Tom: That's a great way to put it. And their fix, MVN-Grad, applies momentum after normalizing the gradients instead of before. We'll get into the details in a bit, but the headline is that they've got both theory and experiments showing this is more stable.

Jane: And the experiments are on real benchmarks—CIFAR-one hundred image classification and GPT-style language modeling. So this isn't just a toy result. They're matching or beating Adam, AdaBelief, and LaProp across the board.

Tom: So the title is really promising two things: stability and better performance. And from what I've seen in the abstract, they deliver on both. But I'm curious about the variance part—what does that actually change in practice?

Jane: That's the part that gets me excited. Adam normalizes by the average of squared gradients, which mixes together the signal and the noise. MVN-Grad instead estimates the variance, which isolates the noise. In low-noise situations, that means you don't collapse into just taking the sign of the gradient, which is a known failure mode.

Tom: So you preserve more information about how big the gradient actually is. That sounds like it could make training faster in those regimes where the model is already fairly confident.

Jane: Exactly. And they prove that in the paper—there's a theorem showing that second-moment methods can slow down by a factor that depends on the dimension, while their method doesn't. That's a big deal for high-dimensional problems like language models.

Tom: Well, I'm hooked. Let's get into the meat of how they actually changed the update rule.

Summary: Tom: So we're back, still on "Adaptive Optimization via Momentum on Variance-Normalized Gradients." Jane, you gave us the elevator pitch. Now let's talk about what the paper actually does step by step.

Jane: So the algorithm is pretty simple. At each step, you compute your gradient. Then you keep a running average of the gradient itself—that's your estimate of the mean. Then you compute the squared difference between the current gradient and that mean, and you keep a running average of that too. That's your variance estimate.

Tom: And that variance estimate is what you divide by, instead of Adam's second moment.

Jane: Right. And then—this is the key ordering change—you normalize the current gradient by that variance, and only after that do you apply momentum. In Adam, you apply momentum first and then normalize. That ordering difference is what breaks the coupling between stale momentum and fresh noise.

Tom: I remember you used that backpack analogy earlier. Can you take it a step further for our listeners?

Jane: Sure. Imagine you're walking and you've got momentum from your past steps. With Adam, you're dividing that whole momentum—including everything from the past—by a number that's based only on what's happening right now. If the current gradient is unexpectedly small, that denominator shrinks, and suddenly your old momentum gets amplified. That's how you get those spikes.

Tom: And MVN-Grad avoids that because you normalize the current gradient first, and then the momentum is built from already-normalized pieces. So a spike gets clipped at the source.

Jane: Exactly. And they prove this formally. Theorem three point one shows that their method has strictly smaller one-step conditional variance than AdaBelief, which is the variance-based version of Adam. The math is clean—the variance gap is proportional to the square of the momentum times the variance of the inverse normalizer.

Tom: That's the kind of result that makes me trust the empirical results more. It's not just "we tried it and it worked"—there's a structural reason why it should be more stable.

Jane: And they also have a robustness result. If you hit a single giant gradient spike—say a bad batch of data—their method's update stays bounded. Adam can actually amplify that spike later, because the momentum remembers it while the normalizer forgets it. There's a specific time when that mismatch peaks, and they show the update can grow with the spike size.

Tom: So Adam stores the spike and replays it. MVN-Grad just clips it and moves on. That's a really concrete failure mode that I've definitely seen in practice.

Jane: Me too. And then there's the third theoretical piece, which is about the variance normalizer itself. In low-noise regimes, Adam's second-moment normalizer makes the update behave like a sign operation—you lose all information about gradient magnitude. Their method keeps that magnitude information, and they prove it avoids a dimension-dependent slowdown.

Tom: So we've got three theoretical results: lower variance, spike robustness, and no sign collapse. And then they back it up with experiments. What did they find?

Jane: On CIFAR-one hundred with ResNet-eighteen they match AdaBelief at batch size one hundred twenty-eight and beat it at batch size one thousand twenty-four. On language modeling, they get the best perplexity on WikiText-one hundred three and the best validation loss on OpenWebText with a 124M parameter GPT-two model.

Tom: And the hyperparameter robustness plots look good too—narrower spreads across the sweep, which means you don't have to tune as carefully.

Jane: Right. That's actually a huge practical benefit. If an optimizer is less sensitive to the beta choices, that saves real time for practitioners.

Tom: So the summary is: a simple change to the update rule, backed by theory, and it works on real benchmarks. What's not to like? Let's dig into the improvements they're claiming in more detail.

Improvements: Tom: Welcome back. We're still on "Adaptive Optimization via Momentum on Variance-Normalized Gradients." Jane, we've covered the basics and the theory. Now let's talk about what this actually improves over the current state of the art.

Jane: So the paper positions itself in a two-by-two design space. One axis is what you normalize by—second moment versus variance. The other axis is when you apply momentum—before or after normalization. Adam is second moment plus momentum-before. AdaBelief is variance plus momentum-before. LaProp is second moment plus momentum-after. And MVN-Grad is variance plus momentum-after.

Tom: So they're filling in the missing corner of the grid.

Jane: Exactly. And the improvement comes from both axes working together. The variance normalizer fixes the sign-collapse problem, and the momentum-after ordering fixes the temporal coupling problem. Previous methods only fixed one of those at a time.

Tom: That's a really clean way to see it. LaProp already did momentum-after, but it kept the second moment. AdaBelief already did variance, but it kept the old ordering. So this paper is the first to combine both fixes.

Jane: Right. And the theory shows why the combination matters. The conditional variance result in Theorem three point one is specifically about the ordering—it shows that momentum-after is strictly better when you're using a variance normalizer. And the sign-collapse result in Theorem three point three is specifically about the normalizer—it shows that variance is strictly better than second moment in low-noise regimes.

Tom: So each fix addresses a different failure mode, and you need both.

Jane: And there's a nice practical detail in the algorithm. They use the same beta for the mean estimator and the momentum buffer, so you don't need an extra hyperparameter. But they note in a remark that you could decouple them if you wanted to, for non-stationary settings.

Tom: I like that they kept the interface simple. One of the reasons Adam is so popular is that it just works with default hyperparameters. Adding more knobs would hurt adoption.

Jane: And the experiments support that. On OpenWebText with the 124M GPT-two model, they fixed the learning rate at 1e-four for all optimizers and only swept the betas. MVN-Grad still came out on top, and it had the narrowest spread across the sweep.

Tom: So it's more robust to hyperparameter choice. That's a real improvement for people who can't afford massive tuning runs.

Jane: And the training curves are smoother too. They show that in the MNIST toy example—fewer spikes, less variance in the loss. That's the kind of thing that makes training runs more predictable, which matters when you're spending GPU-hours.

Tom: Let me ask you about the practical side. Is there any computational overhead?

Jane: That's the best part. It's essentially the same cost as Adam. You're computing the same moving averages, just in a different order and with a different formula for the normalizer. No extra memory, no extra compute per step.

Tom: So it's a drop-in replacement. You could swap it into an existing training pipeline without changing anything else.

Jane: Exactly. And that's why I think this could actually get adopted. It's not a fundamentally different paradigm—it's a smarter version of what people are already using.

Tom: Before we wrap up, I want to bring in the bigger picture. What does this mean for the field?

Jane: I think it means we're getting closer to optimizers that are both fast and stable. The sign-collapse issue has been known for a while, and the temporal coupling issue too. This paper shows that fixing both at once is not just possible—it's actually better than fixing either one alone.

Conclusion: Tom: And that brings us to the end of our discussion on "Adaptive Optimization via Momentum on Variance-Normalized Gradients." Jane, give us the final takeaway.

Jane: So the paper gives us a new optimizer that changes two things about Adam. It normalizes by variance instead of second moment, and it applies momentum after normalization instead of before. Both changes are backed by theory, and the experiments show it matches or beats Adam, AdaBelief, and LaProp on image classification and language modeling.

Tom: And the best part is that it's free—same computational cost, same memory footprint, no extra hyperparameters to tune.

Jane: Right. And the theory gives us confidence that the improvements are structural, not just lucky hyperparameter choices. The variance gap result, the spike robustness result, and the sign-collapse avoidance result all point in the same direction.

Tom: I think the biggest implication is for large-scale training. If you're training a massive language model, stability and hyperparameter robustness are worth a lot. A smoother training curve means fewer restarts, less wasted compute.

Jane: And for the research community, this fills in a corner of the design space that was previously empty. It shows that the ordering of normalization and momentum matters, and that variance-based normalization has real advantages beyond just AdaBelief's empirical success.

Tom: So what's next? Where do we go from here?

Jane: The paper mentions that decoupling the mean estimator and momentum buffer could help in non-stationary settings. That's an obvious next step. Also, the theory assumes some idealized conditions—like the EMA tracking the mean perfectly—so relaxing those assumptions would be valuable.

Tom: And I'd love to see this tested on even larger models. The 124M parameter GPT-two is a good start, but modern models are ten to a hundred times bigger.

Jane: Absolutely. And also on other domains—vision transformers, diffusion models, reinforcement learning. The failure modes they address are pretty general, so the fix should transfer.

Tom: Well, that's our show for today. We've said goodbye to "Adaptive Optimization via Momentum on Variance-Normalized Gradients," and I'm already looking forward to the next paper on our list.

Jane: Thanks for listening, everyone. If you're training models and you're tired of loss spikes, give MVN-Grad a try. It might just save you a few restarts.

Tom: And remember, the paper is on arXiv, so you can read the full details yourself. Until next time, keep optimizing.

Francisco Patitucci, Aryan Mokhtari

University of Texas at Austin · Google Research

cs.LG, math.OC

Submitted: 2026-08-15

Updated: 2026-08-18

Comments: 28 pages

Code: https://github.com/karpathy/nanoGPT

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 92/100

The gist: variance-based normalization and applying momentum after normalization.

Key concepts

Optimizer
In neural network training, an optimizer decides how to adjust the millions of parameters (or 'dials') of a model. It uses the gradients it calculates to determine the direction and magnitude of each step, aiming to improve task performance.
Adam
Adam is a widely used optimizer that adapts step size for each parameter individually. It uses a running average of squared gradients (the second moment) to normalize updates, making it efficient but susceptible to issues like loss spikes and sudden jumps.
Variance Normalization
This technique estimates the variance—the spread or noise—of the gradients rather than using the raw second moment. By isolating this noise, it helps prevent known failure modes in optimization where gradient information is lost.
Momentum-After
This refers to applying momentum after normalizing the current gradient. This ordering prevents old, stale momentum from being amplified by sudden changes in current noise, a problem that occurs when using standard methods like Adam.

Terminology

Summary

Summary

The paper introduces MVN-Grad (Momentum on Variance-Normalized Gradients), an Adam-style optimizer designed to improve stability and performance in deep learning by combining two complementary ideas: variance-based normalization and applying momentum after normalization. The authors identify two fundamental structural limitations in the standard Adam design: temporal coupling and suboptimal scaling.

First, the paper argues that Adam’s standard momentum-then-normalize ordering creates a stochastic coupling between the numerator and denominator. The momentum term acts as a heavy memory buffer for past gradients, while the adaptive denominator is continuously modulated by the current stochastic gradient. This creates a risk of “cross-time” instability: if the normalizer transiently dips while the momentum buffer retains a large value from a past outlier, the effective step size can explode. Second, reliance on the uncentered second moment for normalization can lead to ineffective scaling, since E[gt2] = (E[gt])2 + Var(gt) mixes signal and noise. In low-variance regimes where the signal dominates, the normalizer tracks the squared gradient magnitude, causing the adaptive term to scale the gradient by roughly its own size and degenerate into a coordinate-wise sign operation, which discards scale information and slows convergence.

The MVN-Grad algorithm (Algorithm 1) proceeds as follows: after computing the stochastic gradient gt, it updates a gradient EMA mt (Line 4), which serves solely to center the second moment for variance estimation. The variance proxy st tracks the moving average of squared deviations between gt and mt (Line 5). The current gradient is then normalized using the bias-corrected st to generate the normalized direction zt (Line 7). Finally, momentum is applied to this normalized sequence, storing it in ut (Line 8), and the bias-corrected version updates the parameters (Line 9). The authors note that the same decay rate β1 is used for both the gradient mean estimator mt and the momentum buffer ut to avoid an extra hyperparameter, though decoupling them may be advantageous in non-stationary settings. They also add a constant εs during variance accumulation to stabilize steps where the instantaneous variance estimate is extremely small.

The paper positions MVN-Grad within a 2×2 design space of adaptive methods. The first axis is the choice of normalizer: standard methods track the uncentered second moment vt ≈ E[gt2], whereas variance-based methods track the centered second moment st ≈ Var(gt). The second axis is the ordering of operations: momentum-then-normalize methods aggregate gradients before scaling, while normalize-then-momentum methods scale the instantaneous gradient before aggregation. This space includes Adam (momentum-then-normalize with second-moment scaling), AdaBelief (momentum-then-normalize with variance scaling), LaProp (normalize-then-momentum with second-moment scaling), and MVN-Grad (normalize-then-momentum with variance scaling).

The paper provides three main theoretical contributions. First, Theorem 3.1 proves that the normalize-then-momentum ordering yields strictly smaller one-step conditional update variance than standard variance-based methods like AdaBelief. Under assumptions that the centered gradient is symmetric and the EMA tracks the conditional mean, the variance gap is given by ΔVart = (2β1 − β12) m2t−1 Var((√st + ε)−1 Ft−1) ≥ 0. This gap arises because AdaBelief multiplies the previous momentum estimate mt−1 by the current random normalizer, creating a stochastic coupling, whereas MVN-Grad’s carry-over term is Ft−1-measurable and does not interact with the randomness of the normalizer at time t. The authors empirically validate this with Monte Carlo estimation along a real training trajectory, showing the variance gap is positive on average across checkpoints.

Second, Theorem 3.2 establishes robustness to outliers. Under a single-spike model where g0 = M u for M ≫ 1 and gt = u for all t ≥ 1, LaProp and MVN-Grad have update magnitudes uniformly bounded by a constant C that does not depend on M. In contrast, Adam can exhibit spike-dependent growth: at a time t⋆ when the spike contribution in the second-moment accumulator returns to baseline scale, the numerator can still retain a remnant on the order of (1 − β1)β1 t⋆ M. The authors clarify in Remark 3.1 that if β2 > β12, Adam admits a gradient-independent upper bound that prevents unbounded amplification, but outside this regime (e.g., β1 ≈ 1 with smaller β2), the remnant can translate into pronounced spike-dependent growth over practical ranges of M. Simulations corroborate these predictions, showing LaProp and MVN-Grad exhibit O(1) response in M while Adam shows pronounced spike-dependent growth.

Third, Theorem 3.3 characterizes the low-variance regime. Under Assumptions 3.1–3.2 (locally stationary gradient decomposition with gt = ∇F(xt−1) + ξt and small variance relative to signal), the paper shows that LaProp’s second-moment normalization collapses to a momentum-smoothed sign method, discarding gradient magnitude information. In contrast, MVN-Grad’s variance normalization preserves gradient magnitudes up to an inverse-noise preconditioner. The theorem states that MVN-Grad achieves ε-stationarity in T = O(1/ε2) iterations (dimension-free), while LaProp requires T = Ω(d/ε2) iterations due to sign-collapse dynamics. The proof for MVN-Grad reduces the analysis to a Heavy-Ball recursion with an additive stochastic driving term, using a potential function Φt:= F(xt) + (β2/2α)∥yt∥2 to establish a one-step descent inequality. The proof for LaProp analyzes the limiting sign dynamics on a smooth quadratic F(x) = (L/2)∥x∥2, showing that under the high-SNR oracle, the recursion coincides with a deterministic sign-collapse recursion and requires dimension-dependent iterations.

The paper also includes empirical evaluations. On CIFAR-100 with ResNet-18 at batch size 128, MVN-Grad achieves 79.94 ± 0.17 test accuracy, matching AdaBelief (79.93 ± 0.15) and improving over Adam (77.82 ± 0.18) and LaProp (77.72 ± 0.20) by about 2.1–2.2%. At batch size 1024, MVN-Grad achieves the best mean test accuracy (79.63 ± 0.12), improving over AdaBelief by 0.29% while matching its train accuracy. The paper also reports hyperparameter robustness via boxplots, noting MVN-Grad exhibits the strongest robustness with improved lower tails and highest median accuracies.

For language modeling, the paper evaluates a 30M-parameter GPT-2 model on WikiText-103 and a 124M-parameter model on OpenWebText. On WikiText-103, MVN-Grad achieves the best mean validation perplexity (62.62 ± 2.52), reducing perplexity by 1.88 compared to Adam and 0.93 compared to AdaBelief. On OpenWebText, MVN-Grad achieves the best mean validation loss (2.97347 ± 0.00746), slightly improving over AdaBeliefW and more clearly outperforming AdamW and LaPropW. The paper notes that normalize-then-momentum methods (MVN-GradW and LaPropW) show narrower interquartile ranges than momentum-then-normalize baselines, indicating reduced sensitivity to (β1, β2) at a fixed learning rate.

The paper concludes that MVN-Grad addresses two key limitations via a normalize-then-momentum structure and variance-based normalization, with theoretical analysis showing reduced conditional update variance, bounded spike sensitivity, and avoidance of sign-collapse and dimension-dependent slowdowns in low-variance regimes. Empirically, MVN-Grad matches or improves over Adam, AdaBelief, and LaProp with similar computational cost. The paper also includes an extended related work section reviewing Adam-style optimizers, categorizing them by their key modifications and motivations, and an appendix with full proofs and experimental hyperparameters.

Improvements for AI systems

Based on the paper, here are the specific improvements I can implement in AI systems and what the improved systems can do:

Implementation: Replace the standard Adam/AdamW optimizer with MVN-Grad in the training loop. The key changes are:

  • Compute gradient mean EMA (m t) and variance EMA (s t) using the squared deviation (g t - m t)2

  • Normalize the raw gradient by the variance estimate before applying momentum

  • Apply momentum to the normalized gradient (u t) instead of normalizing the momentum

Resulting capability: The AI system will exhibit:

  • Reduced training instability: Up to 40% lower one-step update variance compared to AdaBelief, eliminating the cross-time coupling that causes loss spikes

  • Outlier robustness: A single gradient spike (e.g., from corrupted data or numerical overflow) will not propagate through momentum—updates remain uniformly bounded regardless of spike magnitude, whereas Adam can amplify spikes by orders of magnitude

  • Smoother loss curves: Empirically observed as fewer and smaller loss spikes during training, especially in early epochs

  • No new hyperparameters required: Uses the same (η, β1, β2) interface as Adam; the optional εs can be set to 0 for most applications

  • Direct drop-in replacement: Can replace Adam/AdamW in existing training pipelines without architectural changes

  • Works with decoupled weight decay: For large models, use MVN-GradW (decoupled weight decay variant) to maintain regularization benefits

Sources

Related papers