L2 Weight Decay - TorchCode

Problem Description

Implement SGD with L2 regularization (weight decay) from scratch. Adding a penalty λΣw² to the loss contributes 2λw to each weight's gradient, so every step multiplies the weight by a factor just below 1 — it decays toward (but never exactly to) zero.

Signature

class MyL2SGD:
    def __init__(self, params, lr=0.1, lam=0.01): ...
    def step(self): ...
    def zero_grad(self): ...

Algorithm (per parameter)

grad_total = grad + 2 * lam * w # d/dw of λ·w² is 2λw w -= lr * grad_total # ⇒ w ← (1 − 2·lr·λ)·w − lr·grad

Template

Implement the class below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

class MyL2SGD:
    def __init__(self, params, lr=0.1, lam=0.01):
        pass  # store params, lr, lam

    def step(self):
        pass  # grad = p.grad + 2*lam*p ; p -= lr*grad

    def zero_grad(self):
        pass  # zero all gradients

Test Your Implementation

Use this code to debug before submitting. With grad = 0 each step should scale w by (1 − 2·lr·lam).

# 🧪 Debug
w = torch.tensor([0.8], requires_grad=True)
opt = MyL2SGD([w], lr=0.1, lam=0.5)
loss = (w ** 2).sum() * 0   # data gradient = 0 to isolate decay
loss.backward()
opt.step()
print(w.item())              # 0.8 * (1 - 2*0.1*0.5) = 0.72

Reference Solution

Try solving it yourself first! Click below to reveal the solution.

# ✅ SOLUTION

class MyL2SGD:
    def __init__(self, params, lr=0.1, lam=0.01):
        self.params = list(params)
        self.lr = lr
        self.lam = lam

    def step(self):
        with torch.no_grad():
            for p in self.params:
                if p.grad is None:
                    continue
                grad = p.grad + 2 * self.lam * p   # penalty λ·Σw² → +2λw
                p -= self.lr * grad

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

Tips

Run Locally

For interactive practice with auto-grading, run TorchCode locally:
pip install torch-judge then use check("weight_decay")

Key Concepts

L2 penalty → multiplicative weight decay: every step shrinks w by (1 − 2·lr·λ). This equals PyTorch's weight_decay=2λ (which adds wd·w to the gradient); L1 instead subtracts a constant and drives weights to exactly 0. Covered in DL Module 7 §12.