PyTorch Optimizer: AdamW and Adam with Weight Decay

PyTorch, one of the leading deep learning frameworks, provides a plethora of optimizers to facilitate training neural networks. Among them, Adam and its variant AdamW are particularly popular due to their adaptive learning capabilities. Understanding the distinction between AdamW and Adam with weight decay is crucial for researchers and practitioners in deep learning.

Understanding Adam

Adam, derived from "Adaptive Moment Estimation," is a robust optimization algorithm that combines the advantages of two other popular extensions of stochastic gradient descent: AdaGrad and RMSProp. It computes adaptive learning rates for each parameter, making it particularly suitable for problems with sparse gradients or noisy data.

Key Features of Adam:

• Adaptive Learning Rate: It adjusts the learning rate based on first and second moments of past gradients.
• Bias Correction: It includes bias correction terms, especially crucial in the early stages of training for low-magnitude initial gradients.
• Momentum: Incorporates a momentum term with first-order derivatives, enhancing parameter update speed and stability.

The basic equations for updating parameters in Adam are:

• First moment estimate (mean of gradients): $$m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t$$
• Second moment estimate (uncentered variance of gradients): $$v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2$$
• Bias-corrected estimates: $$\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t}$$
• Parameter update rule: $$\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \cdot \hat{m}_t$$ where:
• $g_t$ is the gradient of the objective function w.r.t. the parameter at time step $t$.
• $\beta_1$ and $\beta_2$ are hyperparameters that control the decay rates of the moving averages.
• $\epsilon$ is a small constant to avoid division by zero.
• $\eta$ is the learning rate.

Weight Decay in Adam

Weight decay is a regularization technique that adds a penalty to the loss function to prevent overfitting by discouraging complex models. In the context of optimizers, this involves adding an additional term to the gradient update.

Adam with Weight Decay

Using weight decay directly in Adam, the parameter update rule becomes: $$\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \cdot \hat{m}_t - \eta \cdot \lambda \cdot \theta_{t-1}$$ where $\lambda$ is the weight decay coefficient. This formulation incorporates weight decay by manually modifying the gradient, potentially leading to undesired effects due to Adam's internal adaptive learning rate adjustment mechanism.

AdamW: A Corrected Version

AdamW, introduced to address the shortcomings of implementing weight decay directly in Adam, changes how weight decay is applied. The goal of AdamW is to decouple weight decay from the optimization step and intuition behind Adam, enhancing generalization performance.

AdamW Update Rule

In AdamW, weight decay is decoupled from the gradient-based update:

• Compute the decoupled weight decay update: $$\theta_t = (1 - \eta \cdot \lambda) \cdot \theta_{t-1}$$
• Then proceed with the corrected Adam update as usual: $$\theta_t = \theta_t - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \cdot \hat{m}_t$$ This separation ensures that the weight decay does not interfere with the adaptive learning rate mechanism provided by Adam.

Practical Implications

When to Use Adam vs. AdamW:

• Adam: Suitable for quick prototyping or tasks where weight decay is not critical.
• AdamW: Preferable in scenarios where regularization and generalization performance are prioritized, i.e., larger neural networks prone to overfitting.

Example Code Usage in PyTorch:

Here's how you can implement both using PyTorch:

1import torch
2import torch.nn as nn
3import torch.optim as optim
4
5# Sample model and data
6model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 2))
7criterion = nn.CrossEntropyLoss()
8data, target = torch.rand((16, 10)), torch.randint(0, 2, (16,))
9
10# Adam with weight decay
11optimizer_adam = optim.Adam(model.parameters(), lr=0.01, weight_decay=0.01)
12
13# AdamW
14optimizer_adamw = optim.AdamW(model.parameters(), lr=0.01, weight_decay=0.01)
15
16# Sample training loop
17def train(optimizer):
18    optimizer.zero_grad()
19    output = model(data)
20    loss = criterion(output, target)
21    loss.backward()
22    optimizer.step()
23
24# Train with Adam
25train(optimizer_adam)
26
27# Train with AdamW
28train(optimizer_adamw)

Summary Table

Understanding and using the appropriate optimizer variant can significantly impact the success of training and the final performance of the model. AdamW provides a more theoretically sound and empirically effective way to leverage weight decay, making it a preferred choice in many deep learning applications.