RMSProp - TorchCode

Problem Description

Implement RMSProp from scratch. RMSProp keeps a decaying average of squared gradients so each weight gets its own effective learning rate — large-gradient weights take smaller steps, and the step size stays healthy instead of shrinking to zero like AdaGrad.

Signature

class MyRMSProp:
    def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8): ...
    def step(self): ...
    def zero_grad(self): ...

Algorithm (per parameter)

cache = alpha * cache + (1 - alpha) * grad² p -= lr * grad / (sqrt(cache) + eps)

Template

Implement the class below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

class MyRMSProp:
    def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8):
        pass  # store params; init cache buffers to zeros

    def step(self):
        pass  # cache = alpha*cache + (1-alpha)*grad**2 ; p -= lr*grad/(sqrt(cache)+eps)

    def zero_grad(self):
        pass  # zero all gradients

Test Your Implementation

Use this code to debug before submitting.

# 🧪 Debug
torch.manual_seed(0)
w = torch.randn(4, 3, requires_grad=True)
opt = MyRMSProp([w], lr=0.01)
for i in range(5):
    loss = (w ** 2).sum()
    loss.backward()
    opt.step()
    opt.zero_grad()
    print(f'Step {i}: loss={loss.item():.4f}')

Reference Solution

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

# ✅ SOLUTION

class MyRMSProp:
    def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8):
        self.params = list(params)
        self.lr = lr
        self.alpha = alpha
        self.eps = eps
        self.cache = [torch.zeros_like(p) for p in self.params]

    def step(self):
        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                self.cache[i] = self.alpha * self.cache[i] + (1 - self.alpha) * p.grad ** 2
                p -= self.lr * p.grad / (torch.sqrt(self.cache[i]) + self.eps)

    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("rmsprop")

Key Concepts

Per-parameter adaptive learning rate via a decaying (leaky) average of squared gradients. The alpha decay (≈0.99) forgets old gradients, so unlike AdaGrad the effective step does not collapse. Covered in DL Module 7 §5.