Don’t Simply “Throw Adam at It”: Misunderstanding Adam Will Value You

on a very onerous reinforcement studying drawback. The analysis demonstrated comparable brokers studying comparable duties. But our mannequin was lifeless within the water.

I simplified the structure. I added layers. I eliminated layers. I swapped LSTMs for Transformers, added consideration, eliminated it. I rebuilt the enter options not less than 20 occasions. I even tried unique reminiscence architectures. I spent 1000’s on GPU hours in pathological hyperparameter search mode and I couldn’t get the outcomes.

Nada.

The repair, after I lastly discovered it, was humiliating: I modified β₂ from 0.999 to 0.95 and lowered β₁ from 0.9 to 0.5.

# The "humiliating" repair that truly labored:
optimizer = optim.AdamW(
    mannequin.parameters(), 
    lr=3e-4, 
    betas=(0.5, 0.95), # The magic numbers
    eps=1e-4           # Stopping division by zero in RL
)

Most deep studying engineers, myself included, simply do the next:

Simply useAdamW with the default parameters, and set the training charge to 3e-4.

I don’t vibe-code deep studying fashions out of private choice. If you happen to do, you’ll most actually see this because the optimizer alternative and studying charge.

The truth is, let’s take a look at GPT 5.6 with the immediate: "create a coaching script for a deep studying mannequin"

Unsurprisingly, GPT fortunately spits out:

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

In Karpathy we belief.

Nope.

The issue with this harmful assumption, is that engineers glaze over this essential side of coaching, treating this a part of the equation as solved.

You might be able to get away with it, because it supplies a wise default, however in the long term, particularly when encountering a non-stationary or tough optimization panorama, you’ll fail. Exhausting.

On this article, I’ll cowl:

1. How Adam really works
2. The place Adam fails spectacularly
3. The instinct behind modifying the defaults

This isn’t one other generic evaluation of parameter optimization.

Because of this your incorrect assumption of “simply use Adam” is mistaken with the mathematical instinct as to why.

Who’s this for: You employ Adam as your optimizer of alternative with a floor degree (or no) understanding of the way it works and the place it fails.

How Adam really works

Vanilla stochastic gradient descent (SGD) applies the identical studying charge to each parameter:

$theta_t = theta_{t-1} – alpha cdot g_t$

That is the equation most practitioners are accustomed to. If in case you have a CS diploma, you’ll have carried out it from scratch.

Multiply $alpha$ (the training charge) by the $g_t$ (gradient) and voila, you’ve obtained an replace.

This has some elementary issues, as closely documented in current analysis, together with elementary instability attributable to noisy gradients. Everyone is aware of this, which is why SGD is generally handled like a second charge citizen*.

Variants quickly have been invented:

• SGD with Momentum
• RMSProp
• and so forth.

_{* I say principally as a result of SGD is usually used as the primary alternative in giant scale picture issues}

In comes Adam

Adam (Kingma & Ba, 2015) was designed to resolve each issues concurrently: easy noisy gradients and adapt step sizes per parameter primarily based on noticed gradient statistics.

Adam maintains two working statistics for every parameter in your mannequin, which may be mathematically elegant however means triple the reminiscence necessities with gigantic state_dicts.

You’ll have seen this 1000 occasions. However do you really perceive it? Learn this completely till it’s understood.

First second estimate (the momentum):

$m_t = beta_1 cdot m_{t-1} + (1 – beta_1) cdot g_t$

We have to perceive every variable on this equation.

• $m_t$: The exponentially weighted common of gradients as much as the present timestep $t$. ->That is the precise course utilized to the mannequin params
• $beta_1$: the decay of the momentum, set to 0.9 as a default. That is what you’ll be able to management.
• $g_t$: the calculated gradient for the timestamp. -> This isn’t utilized to the mannequin

So the gradient utilized for every timestep has accrued directional data. It holds onto the common course, as a substitute of utilizing the course calculated for every timestep.

You will need to perceive the precise gradient is rarely utilized. It solely directionally impacts the ultimate gradient, which can be fairly totally different.

Second second estimate (the variance):

$v_t = beta_2 cdot v_{t-1} + (1 – beta_2) cdot g_t^2$

• $v_t$: The exponentially weighted common of the squared gradient as much as the present timestep $t$. -> This determines how the replace is normalized
• $beta_2$: the decay of the the squared gradient, set to 0.999 as a default. That is what you’ll be able to management.
• $g^2_t$: the calculated squared gradient for the timestamp.

This second second controls the precise step measurement.

It’s not like the training charge doesn’t matter, however Adam successfully bumps up the utilized gradient for parameters which generally obtain a small precise gradient and dampens the utilized gradient for these which obtain giant gradients.

Word that the default right here: 0.999 averages during the last ~1000 steps.

Bias correction

As a result of m₀ = 0 and v₀ = 0, the estimates are biased towards zero in early coaching. Adam corrects for this:

$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 considerably lower than 1, which scales up the estimates to compensate for the zero initialization. As t grows, the correction turns into negligible.

Placing all of it collectively: the precise parameter replace

The replace components then turns into:

$theta_t = theta_{t-1} – alpha cdot frac{hat{m}_t}{sqrt{hat{v}_t} + epsilon}$

• $theta_t$: present parameter worth
• $alpha$ studying charge
• $hat{m}_t$​: bias-corrected first second
• $hat{v}_t$​: bias-corrected second second
• $sqrt{hat{v}_t}$​​: estimated gradient scale
• $epsilon$: stability ground (a small fixed to forestall enormous updates from a small scale)

The instinct is that Adam offers every parameter its personal adaptive studying charge. Parameters in unstable, high-gradient areas get small, cautious steps. Parameters in quiet, low-gradient areas get bigger, extra assured steps.

Even when the gradients are random, or noisy, they get washed out. In idea, this could assist the optimizer discover the true native minima of the loss, however is clearly not assured.

That is elegant. It principally works, which is why it’s turn out to be a meme in the neighborhood. It’s additionally the supply of practically each failure mode.

A visible instance of certainly one of Adam’s epic fails.

The place Adam fails spectacularly

If the visible explanations aren’t making you fall out of affection with Adam, these fails may.

Adam can fail to converge on trivial convex optimization issues

This isn’t hypothetical.

Reddi, Kale, and Kumar’s 2018 ICLR paper, On the Convergence of Adam and Past, constructed an express easy convex optimization drawback the place Adam doesn’t converge to the optimum resolution, and pinpointed the flaw: the proof of convergence given within the authentic Adam paper was flat out mistaken.

The short-term reminiscence of the exponential transferring common kills off informative gradients too quick, and the authors even present a case the place Adam converges to the worst attainable resolution.

Adam converges to a suboptimal resolution, however quicker!

It appears like your mannequin is coaching in the suitable course. These loss curves look fairly first rate, directionally correct.

Nah.

Wilson et al.’s 2017 NeurIPS paper The Marginal Worth of Adaptive Gradient Strategies in Machine Studying is an fascinating instance.

Within the paper, they construct a easy binary classification activity. SGD achieves zero take a look at error, whereas the adaptive strategies, together with Adam have a lot bigger errors. These are trivial issues that are simply solved by non adaptive studying.

Mannequin architectures matter little. On this analysis, every resolution (totally different architectures) produce the identical consequence:

Adaptive fashions carry out worse, not higher, than vanilla SGD on the take a look at set.

The larger the mannequin, the extra fail from Adam!

Researchers at Meta (Molybog et al.) put out A Concept on Adam Instability in Massive-Scale Machine Studying which surfaces unexplained loss spikes throughout giant language mannequin coaching.

The researchers present that Adam is in charge right here, with adaptive gradient rescaling. Their evaluation is supported by experiments throughout fashions starting from 7 billion to 546 billion parameters

The instinct is that this:

Some parameters enter states the place the gradient replace is sort of small, shrinking v_t.

As defined above, Adam, makes use of v_t to normalize the gradient replace, which is usually a main drawback when it’s small.

E.g, dividing by:

$sqrt{v_t} + epsilon​ ​$

Is an issue when v_t has shrunk down over 1000 steps with very small calculated gradients. If giant gradient values come alongside, rising m_t at a charge which is misaligned with v_t and 💥: gigantic replace! Loss spike! Bizarre unexplainable dynamics the place you’ll be able to’t clarify what occurred since nobody ever inspects the second order moments of their optimizer state_dict with over 100M+ parameters.

That is what we hypothesize was taking place with our coaching. We have been getting unusual peaks in our loss which have been unexplainable. I don’t have the endurance to sift although the moments of a really giant community to seek out out which parameters have been inflicting the exploding replace, however it’s my greatest guess.

RL (can) destroy Adam

If you happen to work in RL, that is particularly vital, because it’s not often mentioned within the precise analysis and Adam pathologies are solely accounted for in onerous to seek out locations (e.g. you need to learn the supply code to seek out it).

One of many core assumptions which makes Adam such a dependable default is that the gradient distribution is stationary: the volatility of a parameter replace now will inform the volatility of the replace sooner or later, making adaptive step sizes an affordable alternative in supervised studying.

In RL, that is merely not the case.

The coverage which generated the information 10 steps in the past could also be utterly totally different than the coverage which is producing the information now. The optimization panorama is continually shifting. Consequently, v_t is all the time chasing a transferring goal.

Two quite common issues are integrated into RL analysis which account for Adam’s quirks:

1. Rising $epsilon$, the default fixed used within the adaptive normalization (see under)
2. Clipping the gradient to a a lot smaller worth⁷

Apart on rising $epsilon$

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

In Dopamine⁴, a deep RL framework put out by Google analysis, the next $epsilon$ values are used as defaults

Why? Mysteriously by no means mentioned within the analysis and is barely an artifact of sensible individuals massaging Adam to make RL optimization work.

The explanation why it’s as much as 32,000 occasions bigger is to keep away from a possible division by close to zero drawback. The gradient magnitudes trigger v_t to shrink to very small values, which will increase the chance of huge destabilizing updates. Large no-no in RL.

Nonstationary compounds with the transferring goal drawback

In actor-critic strategies for RL, the goal itself (the worth estimate which you’re regressing towards) modifications fairly considerably because the community updates.

This implies the true gradient course for a parameter is continually shifting and may utterly flip scale or check in lower than 100 steps. v_t‘s exponential reminiscence, which averages over ~1000 steps given the default setting of 0.999 is a recipe for poor efficiency, provided that this goal, and benefits calculated from them, are in fixed flux.

🤔Do you have to throw Adam away in favor of one other optimizer, like RMSProp?

Not essentially.

In Revisiting Rainbow, Obando-Ceron & Castro, 2021, researchers immediately A/B take a look at Adam vs. RMSProp throughout 60 Atari video games and conclude that Adam paired with MSE loss outperforms the “normal apply” of RMSProp + Huber loss, which was typical in AI Labs between 2015-2018.

Because of this I didn’t drop Adam instantly. It may work, however must be coaxed into it.

The important thing takeaway isn’t “Adam works in RL” or “Adam doesn’t work in RL”, it’s that adaptive optimization works otherwise in RL than it does in classical supervised studying, and “simply throwing Adam at it” will break down.

Don’t do it. Take into consideration the defaults. If you happen to see a few of the pathologies I’ve talked about, like unexplained spikes in your loss, look straight at your Adam settings.

The instinct behind modifying the defaults

Adam (and AdamW) turned the default for a motive. It converges shortly on many issues, usually doesn’t require tuning, and is (principally) forgiving in comparison with SGD.

Your mistake isn’t utilizing Adam. The error is utilizing Adam with the idea that optimization is a solved drawback that doesn’t require thought.

Adam’s defaults:

• $beta_1$ : 0.9
• $beta_2$ : 0.999
• $epsilon$ : 1e-8

Are used as a result of they work fairly effectively throughout a large spectrum of optimization issues. As I’ve expressed, this may result in unexplainable catastrophic failure.

Among the signs of “pathological optimization” that you just want to concentrate on (and related intuitions):

• Massive unexplained loss spikes?
  • Take into account lowering β₂ (first to 0.99, then to 0.95 after which presumably decrease). This shortens the second second reminiscence and permits the variance estimate to react extra shortly.
• Extremely non-stationary issues (particularly RL)?
  • Strive a bigger $epsilon$ and a smaller $beta_2$. The optimizer ought to neglect stale statistics quicker.
• Noisy gradients?
  • Improve gradient clipping earlier than lowering the training charge. Typically the instability isn’t the training charge, it’s Adam’s adaptive normalization.
• Poor generalization regardless of wonderful coaching loss?
  • Strive swapping the optimizer. SGD and SGD+Momentum (Vanilla or Nesterov), can generally outperform Adam.
• Coaching huge fashions?
  • Search for sudden gradient norm explosions and examine optimizer state if they seem. Seemingly that v_t is in charge.

The PyTorch lure you in all probability miss

One, additionally unusual notice that you could be (or might not be) conscious of.

If you happen to use PyTorch, Adam has a default weight_decay=0, however, in case you swap it out for AdamW, the default is weight_decay=0.01. This was additionally fairly problematic for me, as you’ll assume that the defaults could be the identical throughout optimizer implementations in the identical library.

Wrapping up

The optimizer isn’t just a quick blip in your coaching script. It’s the main resolution, dealing with each weight replace in your billion parameter mannequin.

The following time you encounter inexplicable loss spikes, unstable reinforcement studying, or a mannequin that received’t converge regardless of weeks of architectural modifications, don’t instantly redesign the community.

• Open the optimizer.
• Examine the hyperparameters
• Query defaults
• Suppose critically (now armed with higher instinct)

You may uncover, as I did after burning 1000’s of GPU hours, that the issue wasn’t your mannequin in any respect.

It was three little numbers.

References

[1] Kingma, D. P., & Ba, J. (2015). Adam: A technique for stochastic optimization. Worldwide Convention on Studying Representations (ICLR).

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

[3] Reddi, S. J., Kale, S., & Kumar, S. (2018). On the convergence of Adam and past. Worldwide Convention on Studying Representations (ICLR).

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

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

[6] Wilson, A. C., Roelofs, R., Stern, M., Srebro, N., & Recht, B. (2017). The marginal worth of adaptive gradient strategies in machine studying. Advances in Neural Data Processing Techniques (NeurIPS), 30.

[7] Zhang, J., He, T., Sra, S., & Jadbabaie, A. (2020). Why gradient clipping accelerates coaching: A theoretical justification for adaptivity. Worldwide Convention on Studying Representations (ICLR).