AdamW › Implementing AdamW from Scratch
Language AI HandbookAll chapters AdamW › Implementing AdamW from ScratchOne section · Sections in this chapter7 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement both Adam with L2 regularization and AdamW from scratch to see the difference in practice. We'll train a small regression network on noisy data and compare how the two approaches handle regularization.
import numpy as np
# Reproducibility
np.random.seed(42)
# Generate noisy regression data
n_samples = 200
X = np.linspace(-3, 3, n_samples).reshape(-1, 1)
y_true = (
0.5 * X.squeeze() ** 3 - 2 * X.squeeze() + np.random.randn(n_samples) * 1.5
)
# Normalize
X_mean, X_std = X.mean(), X.std()
y_mean, y_std = y_true.mean(), y_true.std()
X_norm = (X - X_mean) / X_std
y_norm = (y_true - y_mean) / y_std# Simple 2-layer MLP with NumPy
def relu(x):
return np.maximum(0, x)
def relu_grad(x):
return (x > 0).astype(float)
class MLP:
def __init__(self, input_dim, hidden_dim, output_dim, seed=0):
rng = np.random.RandomState(seed)
scale = 0.1
self.W1 = rng.randn(input_dim, hidden_dim) * scale
self.b1 = np.zeros(hidden_dim)
self.W2 = rng.randn(hidden_dim, output_dim) * scale
self.b2 = np.zeros(output_dim)
self.params = [self.W1, self.b1, self.W2, self.b2]
self.grads = [None] * 4
def forward(self, X):
self.h1_pre = X @ self.W1 + self.b1
self.h1 = relu(self.h1_pre)
self.out = self.h1 @ self.W2 + self.b2
return self.out.squeeze()
def backward(self, X, y_pred, y_true):
n = len(y_true)
dout = 2 * (y_pred - y_true) / n
dW2 = self.h1.T @ dout.reshape(-1, 1)
db2 = dout.sum(keepdims=True)
dh1 = dout.reshape(-1, 1) @ self.W2.T
dh1_pre = dh1 * relu_grad(self.h1_pre)
dW1 = X.T @ dh1_pre
db1 = dh1_pre.sum(axis=0)
self.grads = [dW1, db1, dW2, db2.squeeze()]def adam_l2_update(
params,
grads,
moments,
step,
lr=1e-3,
beta1=0.9,
beta2=0.999,
eps=1e-8,
l2_lambda=0.01,
):
"""Adam with L2 regularization: adds lambda*param to gradient before moment updates."""
ms, vs = moments
updated = []
for i, (p, g, m, v) in enumerate(zip(params, grads, ms, vs)):
# L2 regularization: add to gradient BEFORE updating moments
g_reg = g + l2_lambda * p
m_new = beta1 * m + (1 - beta1) * g_reg
v_new = beta2 * v + (1 - beta2) * g_reg**2
m_hat = m_new / (1 - beta1**step)
v_hat = v_new / (1 - beta2**step)
p_new = p - lr * m_hat / (np.sqrt(v_hat) + eps)
ms[i] = m_new
vs[i] = v_new
updated.append(p_new)
return updated
def adamw_update(
params,
grads,
moments,
step,
lr=1e-3,
beta1=0.9,
beta2=0.999,
eps=1e-8,
weight_decay=0.01,
):
"""AdamW: weight decay applied AFTER moment-based update, not inside gradient."""
ms, vs = moments
updated = []
for i, (p, g, m, v) in enumerate(zip(params, grads, ms, vs)):
# Moments computed from original gradient only
m_new = beta1 * m + (1 - beta1) * g
v_new = beta2 * v + (1 - beta2) * g**2
m_hat = m_new / (1 - beta1**step)
v_hat = v_new / (1 - beta2**step)
# Weight decay applied separately, outside adaptive scaling
p_new = p - lr * (m_hat / (np.sqrt(v_hat) + eps) + weight_decay * p)
ms[i] = m_new
vs[i] = v_new
updated.append(p_new)
return updatedThe key difference is in a single line. In adam_l2_update, the regularization is added to the gradient before it enters the moment estimates (g_reg = g + l2_lambda * p). In adamw_update, the moments are computed from the original gradient, and the weight decay is added as a completely separate term outside the adaptive scaling. The regularization term never sees the adaptive denominator .
def train_with_norms(
optimizer_fn,
n_epochs=500,
lr=1e-2,
reg=0.01,
seed=0,
use_weight_decay=False,
):
"""Train the MLP and track weight norms at each epoch."""
model = MLP(1, 32, 1, seed=seed)
ms = [np.zeros_like(p) for p in model.params]
vs = [np.zeros_like(p) for p in model.params]
moments = [ms, vs]
losses = []
weight_norms = []
for epoch in range(1, n_epochs + 1):
y_pred = model.forward(X_norm)
loss = np.mean((y_pred - y_norm) ** 2)
losses.append(loss)
model.backward(X_norm, y_pred, y_norm)
if use_weight_decay:
new_params = optimizer_fn(
model.params,
model.grads,
moments,
epoch,
lr=lr,
weight_decay=reg,
)
else:
new_params = optimizer_fn(
model.params, model.grads, moments, epoch, lr=lr, l2_lambda=reg
)
model.W1, model.b1, model.W2, model.b2 = new_params
model.params = new_params
total_norm = sum(np.sum(p**2) for p in model.params)
weight_norms.append(np.sqrt(total_norm))
return losses, weight_norms, model
losses_adam_l2, norms_adam_l2, model_adam = train_with_norms(
adam_l2_update, lr=1e-2, reg=0.05, use_weight_decay=False
)
losses_adamw, norms_adamw, model_adamw = train_with_norms(
adamw_update, lr=1e-2, reg=0.05, use_weight_decay=True
)Adam + L2 final loss: 0.3916 AdamW final loss: 0.2395 Adam + L2 final weight norm: 3.1271 AdamW final weight norm: 5.3782
The final loss values tell only part of the story. The more interesting observation is in the weight norms: AdamW typically achieves a lower and more stable final weight norm than Adam+L2 with the same regularization coefficient, because the weight decay is applied uniformly rather than being mediated by gradient history.
The weight norm comparison reveals a key behavioral difference. AdamW achieves more controlled weight norms because the decay is applied uniformly without interference from the adaptive learning rate scaling. Adam+L2's weight norms are less predictable because the effective regularization strength varies by parameter, and the optimizer is not actually minimizing the same objective you think it is.