Optimizers in Deep Learning

Qiang Liu
UT Austin
Optimization is the Key of AI Training

□ Training = Optimization: Training an AI model is about solving:

min L(θ).
θ

◦ L(θ): loss function over model parameters θ.

◦ Large AI training has billions of parameters and data, leading to massive
optimization problems.

□ Everything depends on optimization:

◦ Accuracy
◦ Speed
◦ Cost
◦ Scalability
Idea: (Stochastic) Gradient Descent

□ Gradient descent (GD):

θt+1 = θt − ϵ · ∇L(θt ),

◦ ϵ: a small step size

◦ converges to local minima of loss.
□ In practice, the loss function is a sum over a large dataset:
n
1X
L(θ) = ℓ(θ; x (i) ).
n
i=1

◦ Each ℓ(θ; x (i) ) is the loss contributed by data point x (i) .

◦ The data size n can be millions or more in modern tasks.
□ Stochastic gradient descent (SGD):

θt+1 = θt − ϵ · ∇ℓ(θt ; x (it ) ).

◦ Uses gradient from a data point x (it ) , or a mini batch.

◦ Cheap, noisy, and widely used in practice.
Idea: Momentum

□ Gradient can vary dramatically across iterations:

◦ Due to data noise in SGD.
◦ Due to poorly conditioned loss landscape.
□ Momentum helps smooth out noise and accelerate progress:

mt = βmt−1 + (1 − β)gt
θt+1 = θt − ϵ · mt

◦ mt is a moving average of batch gradients gt = ∇L(θt ; x (ii ) ).

Unrolling the recurrence:

mt = (1 − β) gt + βgt−1 + β 2 gt−2 + · · · + β t−1 g1 + β t m0

– Momentum mt is an exponentially weighted average of past gradients.

– β ∈ [0, 1) controls smoothing magnitude (e.g., β = 0.9).

□ Benefits:
◦ Smooths out gradient noise.
◦ Accelerates motion in consistent directions.
◦ Damps oscillations in narrow valleys.
Idea: Normalizing Gradients

□ Gradients can have highly imbalanced magnitudes across coordinates:

◦ Some components are large, others nearly zero.
◦ This leads to unstable or slow updates.
◦ Can we balance the updates across coordinates?

□ Signed Gradient: normalize the update magnitude of coordinates:

Signed GD: θt+1 = θt − η · sign(gt )

Signed Momentum: θt+1 = θt − η · sign(mt ).

◦ Updates depend only on the direction of the gradient, not its magnitude.
◦ To provide better trade off, we can use soft normalization:

SoftSign GD: θt+1 = θt − η · softsign(gt )

SoftSign Momentum: θt+1 = θt − η · softsign(mt )
x
where softsign(x) = |x|+ϵ with ϵ ≥ 0.
Adam

□ Adam is an “adaptive” variant of softign momentum.

mt = β1 mt−1 + (1 − β1 )gt

vt = β2 vt−1 + (1 − β2 )gt2
mt
θt+1 = θt − η √
vt + ϵ

◦ Widely popular due to good performance.

mt mt
◦ Replaces softsign with √ .
|mt | + ϵ vt + ϵ
◦ Can be viewed as adding an adapative learning rate:

|mt | + ϵ
θt+1 = θt − ηt · softsign(mt ), where ηt = η √ .
vt + ϵ

– Effect: slows down when gradients are conflicting across iterations.

Adam

□ Bias correction: In standard implements of Adam, mt and vt are scaled

as follows

mt = β1 mt−1 + (1 − β1 )gt

vt = β2 vt−1 + (1 − β2 )gt2
mt
m̂t =
1 − β1t
vt
v̂t =
1 − β2t
m̂t
θt+1 = θt − η √ .
v̂t + ϵ
The purpose of this is to ensure that the exponential moving averages are
properly normalized.

Recall that (assume m0 = 0)

mt = (1 − β1 ) gt + β1 gt−1 + β12 gt−2 + · · · + β1t−1 g1 .

Hence
(1 − β1 ) gt + β1 gt−1 + β12 gt−2 + · · · + β1t−1 g1


m̂t =
1 − β1t
gt + β1 gt−1 + β12 gt−2 + · · · + β1t−1 g1
= .
1 + β1 + β12 + · · · + β1t−1

□ This ensures that m̂t = gt and v̂t = gt if all gt are equal.

□ But because 1 − β1t ≈ 1 and 1 − β2t ≈ 1 when t is large. The impact can
be small.
Weight Decay: Regularization and Constraint

□ To prevent overfitting, it is common to use L2 regularization:

λ
Lreg (θ) = L(θ) + ∥θ∥2 .
2

◦ Applying gradient descent yields a “weight decay” term:

θt+1 = θt − η (∇L(θt ) + λθt ) .

□ Ingeneral, it has been found that adding such weight decay term on
different algorithms helps stabilize training and improve generalization.

SignGD+WD: θt+1 = θt − η (sign(∇L(θt )) + λθt .) .

 
m̂
AdamW: θt+1 = θt − η √ t + λθt .
v̂t + ϵ

□ However, the effect of weight decay are different when applied on

different algorithms.