Don’t Just “Throw Adam at It”: Misunderstanding Adam Will Cost You

Editor
21 Min Read


on a particularly hard reinforcement learning problem. The research demonstrated similar agents learning comparable tasks. Yet our model was dead in the water.

I simplified the architecture. I added layers. I removed layers. I swapped LSTMs for Transformers, added attention, removed it. I rebuilt the input features at least 20 times. I even tried exotic memory architectures. I spent thousands on GPU hours in pathological hyperparameter search mode and I couldn’t get the results.

Nada.

The fix, when I finally found it, was humiliating: I changed β₂ from 0.999 to 0.95 and reduced β₁ from 0.9 to 0.5.

# The "humiliating" fix that actually worked:
optimizer = optim.AdamW(
    model.parameters(), 
    lr=3e-4, 
    betas=(0.5, 0.95), # The magic numbers
    eps=1e-4           # Preventing division by zero in RL
)

Most deep learning engineers, myself included, just do the following:

Just useAdamW with the default parameters, and set the learning rate to 3e-4.

I don’t vibe-code deep learning models out of personal preference. If you do, you will most certainly see this as the optimizer choice and learning rate.

In fact, let’s test GPT 5.6 with the prompt: "create a training script for a deep learning model"

Unsurprisingly, GPT happily spits out:

CONFIG = {
    "lr": 3e-4,
    "weight_decay": 1e-5,
     ...
}
....
optimizer = optim.AdamW(
   model.parameters(),
   lr=CONFIG["lr"],
   weight_decay=CONFIG["weight_decay"]
)

In Karpathy we trust.

Nope.

The problem with this dangerous assumption, is that engineers glaze over this very important aspect of training, treating this part of the equation as solved.

You may be able to get away with it, as it provides a sensible default, but in the long run, especially when encountering a non-stationary or difficult optimization landscape, you will fail. Hard.

In this article, I’ll cover:

  1. How Adam actually works
  2. Where Adam fails spectacularly
  3. The intuition behind modifying the defaults

This isn’t another generic review of parameter optimization.

This is why your incorrect assumption of “just use Adam” is wrong with the mathematical intuition as to why.

Who’s this for: You use Adam as your optimizer of choice with a surface level (or no) understanding of how it works and where it fails.

How Adam actually works

Vanilla stochastic gradient descent (SGD) applies the same learning rate to every parameter:

θt=θt1αgt\theta_t = \theta_{t-1} – \alpha \cdot g_t

This is the equation most practitioners are familiar with. If you have a CS degree, you may have implemented it from scratch.

Multiply α\alpha (the learning rate) by the gtg_t (gradient) and voila, you’ve got an update.

This has some fundamental problems, as heavily documented in existing research, including fundamental instability caused by noisy gradients. Everybody knows this, which is why SGD is mostly treated like a second rate citizen*.

Variants soon were invented:

  • SGD with Momentum
  • RMSProp
  • etc.

* I say mostly because SGD is sometimes used as the first choice in large scale image problems

In comes Adam

Adam (Kingma & Ba, 2015) was designed to solve both problems simultaneously: smooth noisy gradients and adapt step sizes per parameter based on observed gradient statistics.

Adam maintains two running statistics for each parameter in your model, which can be mathematically elegant but means triple the memory requirements with gigantic state_dicts.

You may have seen this 1000 times. But do you actually understand it? Read this thoroughly until it is understood.

First moment estimate (the momentum):

mt=β1mt1+(1β1)gtm_t = \beta_1 \cdot m_{t-1} + (1 – \beta_1) \cdot g_t

We need to understand each variable in this equation.

  • mtm_t: The exponentially weighted average of gradients up to the current timestep tt. ->This is the actual direction applied to the model params
  • β1\beta_1: the decay of the momentum, set to 0.9 as a default. This is what you can control.
  • gtg_t: the calculated gradient for the timestamp. -> This is not applied to the model

So the gradient applied for each timestep has accumulated directional information. It holds onto the average direction, instead of using the direction calculated for each timestep.

You must understand the actual gradient is never applied. It only directionally affects the final gradient, which may be quite different.

Your actual gradient only perturbs the accumulated momentum, which is the direction in which the gradient actually moves. Image by Author

Second moment estimate (the variance):

vt=β2vt1+(1β2)gt2v_t = \beta_2 \cdot v_{t-1} + (1 – \beta_2) \cdot g_t^2

  • vtv_t: The exponentially weighted average of the squared gradient up to the current timestep tt. -> This determines how the update is normalized
  • β2\beta_2: the decay of the the squared gradient, set to 0.999 as a default. This is what you can control.
  • gt2g^2_t: the calculated squared gradient for the timestamp.

This second moment controls the actual step size.

It’s not like the learning rate doesn’t matter, but Adam effectively bumps up the applied gradient for parameters which typically receive a small actual gradient and dampens the applied gradient for those which receive large gradients.

V_t controls how much the gradient is normalized per step. Large gradients dampen the step size. Small gradients can increase the step size. Image by Author.

Note that the default here: 0.999 averages over the last ~1000 steps.

Bias correction

Because m₀ = 0 and v₀ = 0, the estimates are biased toward zero in early training. Adam corrects for this:

m^t=mt1β1tv^t=vt1β2t\hat{m}_t = \frac{m_t}{1 – \beta_1^t} \qquad \hat{v}_t = \frac{v_t}{1 – \beta_2^t}

In early steps (small t), the denominator (1 – β^t) is significantly less than 1, which scales up the estimates to compensate for the zero initialization. As t grows, the correction becomes negligible.

Putting it all together: the actual parameter update

The update formula then becomes:

θt=θt1αm^tv^t+ϵ\theta_t = \theta_{t-1} – \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

  • θt\theta_t: current parameter value
  • α\alpha learning rate
  • m^t\hat{m}_t​: bias-corrected first moment
  • v^t\hat{v}_t​: bias-corrected second moment
  • v^t\sqrt{\hat{v}_t}​​: estimated gradient scale
  • ϵ\epsilon: stability floor (a small constant to prevent huge updates from a small scale)

The intuition is that Adam gives each parameter its own adaptive learning rate. Parameters in volatile, high-gradient regions get small, cautious steps. Parameters in quiet, low-gradient regions get larger, more confident steps.

Even if the gradients are random, or noisy, they get washed out. In theory, this should help the optimizer find the true local minima of the loss, but is obviously not guaranteed.

This is elegant. It mostly works, which is why it’s become a meme in the community. It’s also the source of nearly every failure mode.

A visual example of one of Adam’s epic fails.

In this example, an adaptive step size causes Adam to dramatically overshoot the true minima, Image by Author.

Where Adam fails spectacularly

If the visual explanations aren’t making you fall out of love with Adam, these fails might.

Adam can fail to converge on trivial convex optimization problems

This isn’t hypothetical.

Reddi, Kale, and Kumar’s 2018 ICLR paper, On the Convergence of Adam and Beyond, constructed an explicit simple convex optimization problem where Adam does not converge to the optimal solution, and pinpointed the flaw: the proof of convergence given in the original Adam paper was flat out wrong.

The short-term memory of the exponential moving average kills off informative gradients too fast, and the authors even show a case where Adam converges to the worst possible solution.

Adam converges to a suboptimal solution, but faster!

It looks like your model is training in the right direction. Those loss curves look pretty decent, directionally accurate.

Nah.

Wilson et al.’s 2017 NeurIPS paper The Marginal Value of Adaptive Gradient Methods in Machine Learning is an interesting example.

In the paper, they build a simple binary classification task. SGD achieves zero test error, while the adaptive methods, including Adam have much larger errors. These are trivial problems which are easily solved by non adaptive learning.

Model architectures matter little. In this research, each solution (different architectures) produce the same result:

Adaptive models perform worse, not better, than vanilla SGD on the test set.

The bigger the model, the more fail from Adam!

Researchers at Meta (Molybog et al.) put out A Theory on Adam Instability in Large-Scale Machine Learning which surfaces unexplained loss spikes during large language model training.

The researchers show that Adam is to blame here, with adaptive gradient rescaling. Their analysis is supported by experiments across models ranging from 7 billion to 546 billion parameters

The intuition is this:

Some parameters enter states where the gradient update is quite small, shrinking v_t.

As explained above, Adam, uses v_t to normalize the gradient update, which can be a major problem when it’s small.

E.g, dividing by:

vt+ϵ\sqrt{v_t} + \epsilon​ ​

Is a problem when v_t has shrunk down over 1000 steps with very small calculated gradients. If large gradient values come along, increasing m_t at a rate which is misaligned with v_t and 💥: gigantic update! Loss spike! Weird unexplainable dynamics where you can’t explain what happened since no one ever inspects the second order moments of their optimizer state_dict with over 100M+ parameters.

This is what we hypothesize was happening with our training. We were getting strange peaks in our loss which were unexplainable. I don’t have the patience to sift though the moments of a very large network to find out which parameters were causing the exploding update, but it’s my best guess.

RL (can) destroy Adam

If you work in RL, this is especially important, as it’s rarely discussed in the actual research and Adam pathologies are only accounted for in hard to find places (e.g. you must read the source code to find it).

One of the core assumptions which makes Adam such a reliable default is that the gradient distribution is stationary: the volatility of a parameter update now will inform the volatility of the update in the future, making adaptive step sizes a reasonable choice in supervised learning.

In RL, this is simply not the case.

The policy which generated the data 10 steps ago may be completely different than the policy which is generating the data now. The optimization landscape is constantly shifting. As a result, v_t is always chasing a moving target.

Two very common things are incorporated into RL research which account for Adam’s quirks:

  1. Increasing ϵ\epsilon, the default constant used in the adaptive normalization (see below)
  2. Clipping the gradient to a much smaller value7

Aside on increasing ϵ\epsilon

Adam’s default ε is 1e-8: a stability floor meant to be negligible.

In Dopamine4, a deep RL framework put out by Google research, the following ϵ\epsilon values are used as defaults

Agent Adam ϵ\epsilon vs. Adam default ϵ\epsilon
DQN 1.5e-4 ~15,000× larger
Rainbow 1.5e-4 ~15,000× larger
IQN / M-IQN 3.125e-4 ~31,250× larger

Why? Mysteriously never discussed in the research and is only an artifact of smart people massaging Adam to make RL optimization work.

The reason why it’s up to 32,000 times larger is to avoid a potential division by near zero problem. The gradient magnitudes cause v_t to shrink to very small values, which increases the likelihood of large destabilizing updates. Big no-no in RL.

Nonstationary compounds with the moving target problem

In actor-critic methods for RL, the target itself (the value estimate which you’re regressing toward) changes quite significantly as the network updates.

This means the true gradient direction for a parameter is constantly shifting and can completely flip scale or sign in less than 100 steps. v_t‘s exponential memory, which averages over ~1000 steps given the default setting of 0.999 is a recipe for poor performance, given that this target, and advantages calculated from them, are in constant flux.

🤔Should you throw Adam away in favor of another optimizer, like RMSProp?

Not necessarily.

In Revisiting Rainbow, Obando-Ceron & Castro, 2021, researchers directly A/B test Adam vs. RMSProp across 60 Atari games and conclude that Adam paired with MSE loss outperforms the “standard practice” of RMSProp + Huber loss, which was typical in AI Labs between 2015-2018.

This is why I didn’t drop Adam immediately. It can work, but needs to be coaxed into it.

The key takeaway isn’t “Adam works in RL” or “Adam doesn’t work in RL”, it’s that adaptive optimization works differently in RL than it does in classical supervised learning, and “just throwing Adam at it” will break down.

Don’t do it. Think about the defaults. If you see some of the pathologies I’ve mentioned, like unexplained spikes in your loss, look straight at your Adam settings.

The intuition behind modifying the defaults

Adam (and AdamW) became the default for a reason. It converges quickly on many problems, typically doesn’t require tuning, and is (mostly) forgiving compared to SGD.

Your mistake isn’t using Adam. The mistake is using Adam with the assumption that optimization is a solved problem that doesn’t require thought.

Adam’s defaults:

  • β1\beta_1 : 0.9
  • β2\beta_2 : 0.999
  • ϵ\epsilon : 1e-8

Are used because they work reasonably well across a wide spectrum of optimization problems. As I’ve expressed, this can lead to unexplainable catastrophic failure.

Some of the symptoms of “pathological optimization” that you need to be aware of (and associated intuitions):

  • Large unexplained loss spikes?
    • Consider reducing β₂ (first to 0.99, then to 0.95 and then possibly lower). This shortens the second moment memory and allows the variance estimate to react more quickly.
  • Highly non-stationary problems (especially RL)?
    • Try a larger ϵ\epsilon and a smaller β2\beta_2. The optimizer should forget stale statistics faster.
  • Noisy gradients?
    • Increase gradient clipping before reducing the learning rate. Sometimes the instability isn’t the learning rate, it’s Adam’s adaptive normalization.
  • Poor generalization despite excellent training loss?
    • Try swapping the optimizer. SGD and SGD+Momentum (Vanilla or Nesterov), can sometimes outperform Adam.
  • Training enormous models?
    • Look for sudden gradient norm explosions and investigate optimizer state if they appear. Likely that v_t is to blame.

The PyTorch trap you probably miss

One, also strange note that you may (or may not be) aware of.

If you use PyTorch, Adam has a default weight_decay=0, but, if you swap it out for AdamW, the default is weight_decay=0.01. This was also quite problematic for me, as you would assume that the defaults would be the same across optimizer implementations in the same library.

Wrapping up

The optimizer is not just a brief blip in your training script. It is the primary decision, handling every weight update in your billion parameter model.

The next time you encounter inexplicable loss spikes, unstable reinforcement learning, or a model that won’t converge despite weeks of architectural changes, don’t immediately redesign the network.

  • Open the optimizer.
  • Inspect the hyperparameters
  • Question defaults
  • Think critically (now armed with better intuition)

You might discover, as I did after burning thousands of GPU hours, that the problem wasn’t your model at all.

It was three little numbers.

References

[1] Kingma, D. P., & Ba, J. (2015). Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR).

[2] Molybog, I., Denoyelle, N., Bakhtin, A., Sankararaman, K. A., Sinha, A., & Goyal, G. (2023). A theory on Adam instability in large-scale machine learning. arXiv preprint arXiv:2304.09871.

[3] Reddi, S. J., Kale, S., & Kumar, S. (2018). On the convergence of Adam and beyond. International Conference on Learning Representations (ICLR).

[4] Castro, P. S., Moitra, S., Gelada, C., Kumar, S., & Bellemare, M. G. (2018). Dopamine: A research framework for deep reinforcement learning. arXiv preprint arXiv:1812.06110.

[5] Tieleman, T., & Hinton, G. (2012). Lecture 6.5-RMSProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 4(2), 26–31.

[6] Wilson, A. C., Roelofs, R., Stern, M., Srebro, N., & Recht, B. (2017). The marginal value of adaptive gradient methods in machine learning. Advances in Neural Information Processing Systems (NeurIPS), 30.

[7] Zhang, J., He, T., Sra, S., & Jadbabaie, A. (2020). Why gradient clipping accelerates training: A theoretical justification for adaptivity. International Conference on Learning Representations (ICLR).

Share this Article
Please enter CoinGecko Free Api Key to get this plugin works.