Optimization

Training models requires finding parameters $\theta$ that minimize a cost function $J(\theta)$ by updating them iteratively.

Stochastic Gradient Descent (SGD)

SGD updates parameters using the gradient calculated from a single random instance (or a mini-batch) at each step:

$\theta \leftarrow \theta - \eta \nabla_{\theta} J(\theta; x^{(i)}, y^{(i)})$

While computationally efficient, SGD’s updates can oscillate, slowing convergence.

Momentum and Nesterov Accelerated Gradient

Momentum accelerates SGD by adding a fraction $\beta$ of the previous update vector $m$, acting like a physical ball rolling down a hill:

$m \leftarrow \beta m + \eta \nabla_{\theta} J(\theta)$ $\theta \leftarrow \theta - m$

NAG calculates the gradient ahead of the current position ($J(\theta - \beta m)$) to stabilize convergence.

Adaptive Learning Rates

Adaptive algorithms adjust the learning rate per parameter based on historical gradients:

• RMSProp: Decays past squared gradients to focus on recent updates. $s \leftarrow \beta s + (1 - \beta) (\nabla_{\theta} J(\theta))^2$ $\theta \leftarrow \theta - \frac{\eta}{\sqrt{s + \epsilon}} \nabla_{\theta} J(\theta)$
• Adam (Adaptive Moment Estimation): Tracks first (mean) and second (variance) moments of the gradients: $m \leftarrow \beta_1 m + (1 - \beta_1) \nabla_{\theta} J(\theta)$ $s \leftarrow \beta_2 s + (1 - \beta_2) (\nabla_{\theta} J(\theta))^2$

Bias correction is applied: $\hat{m} = m/(1-\beta_1^t)$ and $\hat{s} = s/(1-\beta_2^t)$.

AdamW (Weight Decay)

AdamW improves on Adam by applying weight decay directly to the parameters rather than incorporating it into the gradient update step.

Example: Optimizing Parameters

The following example demonstrates parameter optimization using Adam in PyTorch:

import torch

# Start parameter with gradient tracking
w = torch.tensor([5.0], requires_grad=True)

# Adam optimizer with learning rate of 0.1
optimizer = torch.optim.Adam([w], lr=0.1)

# Minimize w^2 over 5 iterations
for _ in range(5):
  optimizer.zero_grad() # Reset gradients
  loss = w ** 2         # Compute loss
  loss.backward()       # Backprop
  optimizer.step()      # Update w
  print(f"w: {w.item():.4f}")

w: 4.9000
w: 4.8001
w: 4.7003
w: 4.6009
w: 4.5019

Exercise

Test your understanding of adaptive optimization updates:

References & Further Reading

• Géron, A. (2022). Hands-On Machine Learning with Scikit-Learn, Keras, and TensorFlow (3rd ed.). O’Reilly Media.
• Smola, A., & Vishwanathan, S. V. N. (2008). An Introduction to Machine Learning. Cambridge University Press.