AdamW: Decoupled Weight Decay Optimization

AdamW (Loshchilov & Hutter, 2019) is the state-of-the-art optimizer for modern Transformers, Large Language Models (LLMs), and deep Convolutional Neural Networks. Unlike classical Adam with L2 regularization, AdamW cleanly decouples weight decay from the gradient moment estimation.

[!TIP] Theory & Curriculum Link: For the first-principles mathematical derivation of gradient descent, momentum ravines, and adaptive moments, see Curriculum: Optimization Algorithms.

1. Quick Example: Training with AdamW

import numpy as np
import netcl.autograd as ag
import netcl.nn as nn
import netcl.optim as opt
from netcl.core.device import manager
from netcl.core.tensor import Tensor

q = manager.default("auto").queue

# 1. Initialize model and optimizer
model = nn.Linear(q, in_features=128, out_features=10)
optimizer = opt.AdamW(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,
)

# 2. Allocate training batch on GPU
x = Tensor.from_host(q, np.random.randn(32, 128).astype(np.float32))
target = Tensor.from_host(q, np.random.randn(32, 10).astype(np.float32))

# 3. Complete training iteration
with ag.Tape() as tape:
    pred = model(ag.tensor(x))
    diff = pred - ag.tensor(target)
    loss = ag.mean(diff * diff)

tape.backward(loss)
optimizer.step(max_norm=1.0)  # Fused gradient clipping and parameter update
optimizer.zero_grad()

print(f"Loss after step: {loss.value.to_host()[0]:.4f}")

2. Mathematical Formulation

At every optimization step $t$, the optimizer receives the stochastic gradient vector $\mathbf{g}_t = \nabla_\theta \mathcal{L}(\theta_{t-1})$.

Step 1: Exponential Moving Average of Gradients (First Moment)

$$\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1) \mathbf{g}_t$$

Step 2: Uncentered Exponential Variance of Gradients (Second Moment)

$$\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2$$

Step 3: Bias Correction (Compensating for Zero Initialization)

$$\hat{\mathbf{m}}_t = \frac{\mathbf{m}_t}{1 - \beta_1^t}, \quad \hat{\mathbf{v}}_t = \frac{\mathbf{v}_t}{1 - \beta_2^t}$$

Step 4: Parameter Update with Decoupled Weight Decay

$$\theta_t = \theta_{t-1} - \gamma_t \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon} - \gamma_t \lambda \theta_{t-1}$$

Where $\gamma_t$ is the current learning rate and $\lambda$ is the weight_decay coefficient.

3. Core Intuition: Why L2 Fails in Adam

In classical Adam, practitioners attempted to add L2 regularization by modifying the gradient before the optimizer step:

$$\mathbf{g}_t \leftarrow \mathbf{g}_t + \lambda \theta_{t-1}$$

When plugged into Adam's update rule, the regularization term gets divided by the second moment $\sqrt{\hat{\mathbf{v}}_t}$:

$$\Delta \theta \approx - \frac{\gamma_t}{\sqrt{\hat{\mathbf{v}}_t}} (\dots + \lambda \theta_{t-1})$$

This creates a serious distortion: - Parameters with frequent, large gradients (high $\mathbf{v}_t$) experience almost zero weight decay. - Parameters with sparse, small gradients (low $\mathbf{v}_t$) experience excessive shrinkage.

AdamW solves this pathology by applying weight decay directly to $\theta_{t-1}$, completely independent of gradient variance $\mathbf{v}_t$. All weights decay strictly proportional to the learning rate $\gamma_t$.

4. Hyperparameter Recommendations

5. Fused OpenCL GPU Implementation

In netcl, AdamW.step() launches a single OpenCL kernel where each GPU thread updates one parameter element in local memory registers without writing intermediate moments back to host RAM:

__kernel void adamw_step(
    __global float* restrict param,
    __global const float* restrict grad,
    __global float* restrict m,
    __global float* restrict v,
    const float lr,
    const float beta1,
    const float beta2,
    const float eps,
    const float weight_decay,
    const float bias_correction1,
    const float bias_correction2,
    const int n
) {
    int i = get_global_id(0);
    if (i >= n) return;

    float g = grad[i];
    float m_val = beta1 * m[i] + (1.0f - beta1) * g;
    float v_val = beta2 * v[i] + (1.0f - beta2) * g * g;

    m[i] = m_val;
    v[i] = v_val;

    float m_hat = m_val / bias_correction1;
    float v_hat = v_val / bias_correction2;

    float step = lr * (m_hat / (sqrt(v_hat) + eps) + weight_decay * param[i]);
    param[i] -= step;
}

Memory buffers are sub-allocated and reused via netcl's BufferPool, achieving maximum GPU memory bandwidth.

Related Documentation

• Curriculum: Optimization Algorithms: Full theoretical breakdown of first- and second-order optimization.
• Optimizer API Reference: Complete API docs for AdamW, SGD, RMSprop, and learning rate schedulers.
• Transformer Architecture: Using AdamW with Rotary Attention and RMSNorm.