HyperAIHyperAI

Command Palette

Search for a command to run...

ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION

Diederik P. Kingma Jimmy Lei Ba

Abstract

We introduce Adam, an algorithm for first-order gradient-based optimization of stochastic objective functions, based on adaptive estimates of lower-order moments. The method is straightforward to implement, is computationally efficient, has little memory requirements, is invariant to diagonal rescaling of the gradients, and is well suited for problems that are large in terms of data and/or parameters. The method is also appropriate for non-stationary objectives and problems with very noisy and/or sparse gradients. The hyper-parameters have intuitive interpretations and typically require little tuning. Some connections to related algorithms, on which Adam was inspired, are discussed. We also analyze the theoretical convergence properties of the algorithm and provide a regret bound on the convergence rate that is comparable to the best known results under the online convex optimization framework. Empirical results demonstrate that Adam works well in practice and compares favorably to other stochastic optimization methods. Finally, we discuss AdaMax, a variant of Adam based on the infinity norm.

One-sentence Summary

Researchers at the University of Amsterdam, OpenAI, and the University of Toronto propose Adam, a first-order stochastic optimization method based on adaptive estimates of lower-order moments that is computationally efficient, memory-light, requires minimal hyperparameter tuning, and performs well on large-scale problems with noisy or sparse gradients.

Key Contributions

  • The paper introduces Adam, a computationally efficient first-order optimizer for stochastic objective functions that uses adaptive estimates of first and second gradient moments and is suited to large datasets, high-dimensional parameters, non-stationary objectives, noisy gradients, and sparse gradients.
  • The paper analyzes convergence and derives a regret bound under the online convex optimization framework, and explains Adam's relation to RMSProp and AdaGrad, including how bias correction prevents excessively large parameter updates when beta_2 is near 1.
  • Experimental results show Adam compares favorably with other stochastic optimization methods on non-convex machine-learning optimization problems, and the paper introduces AdaMax, a variant based on the infinity norm.

Introduction

Stochastic first-order optimization is central to machine learning and scientific computing because many objectives are defined over large, noisy datasets or regularizers such as dropout. In high-dimensional parameter spaces, second-order methods are often impractical, while earlier first-order optimizers have tradeoffs: AdaGrad handles sparse gradients well, and RMSProp suits online or non-stationary settings, but RMSProp without bias correction can produce unstable updates when sparse-gradient behavior is needed. The authors introduce Adam, which uses running estimates of the first and second moments of gradients, with bias correction, to compute per-parameter adaptive learning rates. Adam combines sparse-gradient robustness with non-stationary adaptivity, requires little memory, and provides convergence analysis and strong empirical results on large-scale machine learning problems.

Method

The authors design the method as an adaptive stochastic optimization algorithm for minimizing an expected objective function. Let f(θ)f(\theta)f(θ) be a noisy scalar objective that is differentiable with respect to parameters θ\thetaθ, and suppose the algorithm observes a sequence of stochastic realizations f1(θ),…,fT(θ)f_1(\theta), \ldots, f_T(\theta)f1​(θ),…,fT​(θ), for example through minibatches or through inherently noisy gradient estimates. At timestep ttt, the gradient is

gt=∇θft(θ).g_t = \nabla_\theta f_t(\theta).gt​=∇θ​ft​(θ).

The central idea is to maintain running estimates of both the first and second moments of the gradient. Instead of using the raw gradient alone, the optimizer normalizes a smoothed gradient direction by an estimate of gradient variance. This produces parameter updates that are adaptive in scale and robust to noisy or sparse gradients.

The core update can be written as a repeated module. The optimizer begins with θ0\theta_0θ0​, m0=0m_0 = 0m0​=0, and v0=0v_0 = 0v0​=0, where mtm_tmt​ tracks the gradient and vtv_tvt​ tracks the squared gradient. At each step,

gt=∇θft(θt−1),mt=β1mt−1+(1−β1)gt,vt=β2vt−1+(1−β2)gt2,m^t=mt1−β1t,v^t=vt1−β2t,θt=θt−1−αm^tv^t+ϵ^.\begin{aligned} g_t &= \nabla_\theta f_t(\theta_{t-1}), \\ m_t &= \beta_1 m_{t-1} + (1-\beta_1)g_t, \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2)g_t^2, \\ \widehat{m}_t &= \frac{m_t}{1-\beta_1^t}, \\ \widehat{v}_t &= \frac{v_t}{1-\beta_2^t}, \\ \theta_t &= \theta_{t-1} - \alpha \frac{\widehat{m}_t}{\sqrt{\widehat{v}_t} + \widehat{\epsilon}}. \end{aligned}gt​mt​vt​mt​vt​θt​​=∇θ​ft​(θt−1​),=β1​mt−1​+(1−β1​)gt​,=β2​vt−1​+(1−β2​)gt2​,=1−β1t​mt​​,=1−β2t​vt​​,=θt−1​−αvt​​+ϵmt​​.​

Here, β1\beta_1β1​ and β2\beta_2β2​ control the exponential decay rates of the first and second moment estimates. The hyperparameter α\alphaα controls the base stepsize, and ϵ^\widehat{\epsilon}ϵ provides numerical stability in the denominator.

The term mtm_tmt​ is an exponential moving average of the gradient and serves as an estimate of the first moment. The term vtv_tvt​ is an exponential moving average of the elementwise squared gradient and serves as an estimate of the second raw moment. Because both averages are initialized at zero, their early values are biased toward zero. This bias is especially pronounced when β1\beta_1β1​ or β2\beta_2β2​ is close to 111, since old values are heavily discounted and the initial zeros persist for many steps.

To counteract this, the method divides the moving averages by the terms 1−β1t1-\beta_1^t1−β1t​ and 1−β2t1-\beta_2^t1−β2t​. The bias correction for the second moment can be understood by writing the recursive average as

vt=(1−β2)∑i=1tβ2t−igi2.v_t = (1-\beta_2)\sum_{i=1}^{t}\beta_2^{t-i}g_i^2.vt​=(1−β2​)i=1∑t​β2t−i​gi2​.

Taking expectations gives

E[vt]≈E[gt2](1−β2t),\mathbb{E}[v_t] \approx \mathbb{E}[g_t^2](1-\beta_2^t),E[vt​]≈E[gt2​](1−β2t​),

with additional discrepancy that can remain small when the decay rate gives little weight to very old gradients. Dividing by 1−β2t1-\beta_2^t1−β2t​ removes the initialization bias and yields a scale correct estimate v^t\widehat{v}_tvt​. The same logic applies to the first moment.

The resulting update uses the bias corrected first moment and the square root of the bias corrected second moment. The effective parameter change can be viewed as

Δt=αm^tv^t.\Delta_t = \alpha \frac{\widehat{m}_t}{\sqrt{\widehat{v}_t}}.Δt​=αvt​​mt​​.

The ratio

m^tv^t\frac{\widehat{m}_t}{\sqrt{\widehat{v}_t}}vt​​mt​​

acts like a signal to noise ratio. When the smoothed gradient estimate is consistent across recent steps, the magnitude of this ratio can be close to one, allowing larger effective moves. When the gradient direction is noisy or uncertain, the ratio shrinks, reducing the effective step. Near an optimum, the signal to noise ratio often becomes small, which naturally reduces the update magnitude. This behavior gives the optimizer a form of automatic annealing without requiring an external stepsize schedule in many practical settings.

The update also has a bounded effect. For the common case where (1−β1)≤1−β2(1-\beta_1) \leq \sqrt{1-\beta_2}(1−β1​)≤1−β2​​, the effective step is bounded by approximately α\alphaα. In more severe sparse cases, a looser bound is

∣Δt∣≤α1−β11−β2.|\Delta_t| \leq \alpha \frac{1-\beta_1}{\sqrt{1-\beta_2}}.∣Δt​∣≤α1−β2​​1−β1​​.

This bounded step behavior creates a trust region around the current parameters: the optimizer does not allow a single noisy gradient estimate to produce arbitrarily large movements. Because α\alphaα controls the upper scale of parameter changes, it is often possible to choose its magnitude based on prior knowledge of the parameter geometry, such as expected distance from initialization to good solutions.

Another important property is scale invariance. If gradients are rescaled by a constant factor ccc, then mtm_tmt​ scales by ccc and vtv_tvt​ scales by c2c^2c2. The ratio

cm^tc2v^t=m^tv^t\frac{c\widehat{m}_t}{\sqrt{c^2\widehat{v}_t}} = \frac{\widehat{m}_t}{\sqrt{\widehat{v}_t}}c2vt​​cmt​​=vt​​mt​​

remains unchanged. Therefore, multiplying all gradients by a constant does not change the optimizer direction or step magnitude, except for the numerical stabilizer.

The computational flow can be simplified by folding the bias correction terms into the stepsize. Instead of computing m^t\widehat{m}_tmt​ and v^t\widehat{v}_tvt​ separately, one can use

αt=α1−β2t1−β1t\alpha_t = \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t}αt​=α1−β1t​1−β2t​​​

and then update

θt=θt−1−αtmtvt+ϵ^.\theta_t = \theta_{t-1} - \alpha_t \frac{m_t}{\sqrt{v_t} + \widehat{\epsilon}}.θt​=θt−1​−αt​vt​​+ϵmt​​.

This formulation is mathematically equivalent to the bias corrected update while avoiding an extra division in the main loop.

The method is also analyzed through an online learning framework. Given an arbitrary sequence of convex cost functions f1(θ),…,fT(θ)f_1(\theta), \ldots, f_T(\theta)f1​(θ),…,fT​(θ), the algorithm chooses parameters θt\theta_tθt​ before observing the current cost. The performance is measured by regret relative to the best fixed parameter θ∗\theta^*θ∗:

R(T)=∑t=1T[ft(θt)−ft(θ∗)].R(T) = \sum_{t=1}^{T}\left[f_t(\theta_t) - f_t(\theta^*)\right].R(T)=t=1∑T​[ft​(θt​)−ft​(θ∗)].

Under bounded gradient and bounded parameter movement assumptions, the update rule achieves an O(T)O(\sqrt{T})O(T​) regret bound. Consequently, the average regret satisfies

R(T)T=O(1T),\frac{R(T)}{T} = O\left(\frac{1}{\sqrt{T}}\right),TR(T)​=O(T​1​),

which implies that the average regret converges to zero as TTT grows. For sparse data, the adaptive normalization can improve over non adaptive methods because dimensions that are rarely updated should not receive the same scaling treatment as frequently updated dimensions.

An important variant arises by generalizing the second moment normalization from an L2L^2L2 based quantity to an LpL^pLp based quantity. If the stepsize is made inversely proportional to an LpL^pLp norm surrogate, the second moment estimator can be written as

vt=β2pvt−1+(1−β2p)∣gt∣p.v_t = \beta_2^p v_{t-1} + (1-\beta_2^p)|g_t|^p.vt​=β2p​vt−1​+(1−β2p​)∣gt​∣p.

Taking the limit as ppp approaches infinity yields a simpler recursive quantity:

ut=max⁡(β2ut−1,∣gt∣).u_t = \max(\beta_2 u_{t-1}, |g_t|).ut​=max(β2​ut−1​,∣gt​∣).

This quantity behaves like a decayed running maximum of gradient magnitudes. Using this in place of v^t\sqrt{\widehat{v}_t}vt​​ leads to a more stable variant for settings where large norm generalizations become numerically difficult. In this formulation, initialization bias correction is not required, and the effective parameter update has the simpler bound

∣Δt∣≤α.|\Delta_t| \leq \alpha.∣Δt​∣≤α.

The core training module can also be augmented with temporal averaging. Since stochastic optimization produces noisy iterates, averaging recent parameters often improves generalization. A simple extension is to maintain an exponential moving average of the parameter vector:

θˉt=β2θˉt−1+(1−β2)θt,\bar{\theta}_t = \beta_2 \bar{\theta}_{t-1} + (1-\beta_2)\theta_t,θˉt​=β2​θˉt−1​+(1−β2​)θt​,

with θˉ0=0\bar{\theta}_0 = 0θˉ0​=0. Bias correction can again be applied using

θˉ^t=θˉt1−β2t.\widehat{\bar{\theta}}_t = \frac{\bar{\theta}_t}{1-\beta_2^t}.θˉt​=1−β2t​θˉt​​.

Alternatively, a uniform average over iterates can be used, but the exponential form places more weight on more recent parameters and fits naturally into the streaming gradient optimization loop.

Experiment

The experiments evaluate Adam across logistic regression, multilayer neural networks, convolutional neural networks, and a variational autoencoder, comparing it with momentum SGD, Adagrad, SFO, and RMSProp under consistently tuned settings. Logistic regression experiments confirm that Adam converges comparably to momentum SGD and handles sparse features effectively, matching Adagrad. On non-convex deep models, Adam generally trains efficiently, often outperforming or matching competing optimizers, while in convolutional networks it benefits from adaptive layer-wise scaling and variance reduction through momentum. The bias-correction study shows that this term improves stability when slow-decay estimates are used, making Adam at least as strong as RMSProp across the tested configurations.


Build AI with AI

From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.

AI Co-coding
Ready-to-use GPUs
Best Pricing

HyperAI Newsletters

Subscribe to our latest updates
We will deliver the latest updates of the week to your inbox at nine o'clock every Monday morning
Powered by MailChimp