MIT 6.7960 PS1 Walkthrough: From NumPy MLP to PyTorch Autograd Backprop

TL;DR Hand-write NumPy MLP + backprop → verify with PyTorch Autograd, fully reproducing OCW HW1 core concepts

#mit-67960 #deep-learning #pytorch #backpropagation #autograd #mlp #homework #numpy

Table of Contents

HW1 Problem Structure Overview
Core Math: Two-Layer MLP Backpropagation
NumPy Reference Implementation: From-Scratch MLP
PyTorch Autograd Verification: Ensure Numerical Consistency
Common Errors & Debugging Checklist
Video Timestamps (Relevant Lectures)
Extension Exercises: Add These Yourself
References

Series: MIT 6.7960 Fall 2024 OCW Guide (6 / 4)

MIT 6.7960 L04: Regularization in Practice — Weight Decay, Dropout, Batch Norm & Label Smoothing MIT 6.7960 L05: CNN Architectures — From Convolution Kernels to Translation Equivariance

Series: Global AI/CS Course Map (13 / 29)

MIT 6.7960 L04: Regularization in Practice — Weight Decay, Dropout, Batch Norm & Label Smoothing MIT 6.7960 L05: CNN Architectures — From Convolution Kernels to Translation Equivariance

🌏 中文版

MIT 6.7960 Fall 2024 OCW Homework 1 is the course's first programming assignment, with a core goal: implement MLP forward and backward passes from scratch to understand the mechanics of automatic differentiation. This article walks through every HW1 checkpoint, providing runnable NumPy reference implementations and PyTorch verification scripts.

HW1 Problem Structure Overview

HW1 contains four main parts (per the PDF):

Problem 1: Softmax & Cross-Entropy — Derive vectorized formulas, numerically stable log-sum-exp
Problem 2: Two-Layer MLP Forward — Matrix multiply, ReLU, Softmax chaining
Problem 3: Backpropagation by Hand — Chain rule layer-by-layer derivation of ∂L/∂W, ∂L/∂b
Problem 4: Gradient Check — Finite-difference numerical verification of analytical gradients

Grading focus: Not just getting it to run, but explaining where each gradient comes from and why it's written that way.

Core Math: Two-Layer MLP Backpropagation

Architecture: x → W1,b1 → ReLU → W2,b2 → Softmax → CE Loss

Forward:
z1 = x @ W1.T + b1          # [B, H]
a1 = ReLU(z1)               # [B, H]
z2 = a1 @ W2.T + b2         # [B, C]
probs = Softmax(z2)         # [B, C]
L = CE(probs, y)            # scalar

Backward (chain rule):
∂L/∂z2 = probs - y_onehot   # [B, C]  ← Classic result: Softmax+CE combined derivative
∂L/∂W2 = (∂L/∂z2).T @ a1    # [C, H]
∂L/∂b2 = sum(∂L/∂z2, dim=0) # [C]
∂L/∂a1 = ∂L/∂z2 @ W2        # [B, H]
∂L/∂z1 = ∂L/∂a1 * (z1 > 0)  # [B, H]  ← ReLU derivative
∂L/∂W1 = (∂L/∂z1).T @ x     # [H, D]
∂L/∂b1 = sum(∂L/∂z1, dim=0) # [H]

Key insight: Softmax + CrossEntropy derivative merges to probs - y_onehot because:

∂CE/∂z2_i = p_i - 1{y=i}

This avoids numerically unstable explicit log(softmax) computation.

NumPy Reference Implementation: From-Scratch MLP

"""HW1 Reference: NumPy Two-Layer MLP + Hand-written Backprop + Gradient Check"""
import numpy as np

def softmax(z):
    """Numerically stable Softmax"""
    z_max = np.max(z, axis=1, keepdims=True)
    exp_z = np.exp(z - z_max)
    return exp_z / np.sum(exp_z, axis=1, keepdims=True)

def cross_entropy_loss(probs, y):
    """Vectorized CE Loss, y: [B] class indices"""
    B = probs.shape[0]
    log_probs = -np.log(probs[np.arange(B), y] + 1e-12)
    return np.mean(log_probs)

def relu(x):
    return np.maximum(0, x)

def relu_grad(z):
    return (z > 0).astype(float)

class TwoLayerMLP:
    def __init__(self, input_dim, hidden_dim, output_dim, seed=42):
        np.random.seed(seed)
        # He initialization
        self.W1 = np.random.randn(hidden_dim, input_dim) * np.sqrt(2.0 / input_dim)
        self.b1 = np.zeros(hidden_dim)
        self.W2 = np.random.randn(output_dim, hidden_dim) * np.sqrt(2.0 / hidden_dim)
        self.b2 = np.zeros(output_dim)
    
    def forward(self, x):
        """x: [B, D]"""
        self.x = x
        self.z1 = x @ self.W1.T + self.b1      # [B, H]
        self.a1 = relu(self.z1)                 # [B, H]
        self.z2 = self.a1 @ self.W2.T + self.b2 # [B, C]
        self.probs = softmax(self.z2)           # [B, C]
        return self.probs
    
    def loss(self, y):
        return cross_entropy_loss(self.probs, y)
    
    def backward(self, y):
        B = self.x.shape[0]
        y_onehot = np.zeros_like(self.probs)
        y_onehot[np.arange(B), y] = 1.0
        
        # Output layer gradients
        dz2 = self.probs - y_onehot             # [B, C]
        self.dW2 = dz2.T @ self.a1 / B          # [C, H]
        self.db2 = np.mean(dz2, axis=0)         # [C]
        
        # Hidden layer gradients
        da1 = dz2 @ self.W2                     # [B, H]
        dz1 = da1 * relu_grad(self.z1)          # [B, H]
        self.dW1 = dz1.T @ self.x / B           # [H, D]
        self.db1 = np.mean(dz1, axis=0)         # [H]
        
        return dict(dW1=self.dW1, db1=self.db1, dW2=self.dW2, db2=self.db2)
    
    def params_and_grads(self):
        return [
            (self.W1, self.dW1), (self.b1, self.db1),
            (self.W2, self.dW2), (self.b2, self.db2)
        ]

def finite_difference_check(model, x, y, eps=1e-5, tol=1e-4):
    """Numerical gradient check: compare analytical vs finite-difference"""
    model.forward(x)
    model.backward(y)
    analytical = {name: grad.copy() for name, grad in 
                  [('W1', model.dW1), ('b1', model.db1), 
                   ('W2', model.dW2), ('b2', model.db2)]}
    
    max_rel_error = 0.0
    for name, param in [('W1', model.W1), ('b1', model.b1), 
                         ('W2', model.W2), ('b2', model.b2)]:
        grad_num = np.zeros_like(param)
        it = np.nditer(param, flags=['multi_index'], op_flags=['readwrite'])
        while not it.finished:
            idx = it.multi_index
            old_val = param[idx]
            
            param[idx] = old_val + eps
            loss_plus = model.loss(model.forward(x))
            
            param[idx] = old_val - eps
            loss_minus = model.loss(model.forward(x))
            
            param[idx] = old_val
            grad_num[idx] = (loss_plus - loss_minus) / (2 * eps)
            it.iternext()
        
        grad_ana = analytical[name]
        rel_error = np.max(np.abs(grad_num - grad_ana) / (np.abs(grad_num) + np.abs(grad_ana) + 1e-10))
        max_rel_error = max(max_rel_error, rel_error)
        print(f"{name}: max rel error = {rel_error:.2e} {'✓' if rel_error < tol else '✗ FAIL'}")
    
    print(f"Overall max rel error: {max_rel_error:.2e}")
    return max_rel_error < tol

# Run verification
if __name__ == "__main__":
    # Synthetic data
    np.random.seed(123)
    B, D, H, C = 64, 20, 32, 5
    x = np.random.randn(B, D)
    y = np.random.randint(0, C, B)
    
    model = TwoLayerMLP(D, H, C)
    probs = model.forward(x)
    loss = model.loss(y)
    print(f"Initial loss: {loss:.4f}")
    
    # Gradient check
    print("\n=== Gradient Check ===")
    finite_difference_check(model, x, y)
    
    # Train a few steps to verify convergence
    print("\n=== Training Steps ===")
    lr = 0.1
    for step in range(10):
        model.forward(x)
        loss = model.loss(y)
        model.backward(y)
        for param, grad in model.params_and_grads():
            param -= lr * grad
        if step % 2 == 0:
            print(f"Step {step}: loss = {loss:.4f}")

PyTorch Autograd Verification: Ensure Numerical Consistency

"""PyTorch version: same architecture, Autograd auto-backprop, compare values"""
import torch
import torch.nn as nn
import torch.nn.functional as F

class TorchMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim, bias=True)
        self.fc2 = nn.Linear(hidden_dim, output_dim, bias=True)
        # Copy NumPy's He initialization
        nn.init.kaiming_normal_(self.fc1.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc1.bias)
        nn.init.kaiming_normal_(self.fc2.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc2.bias)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)  # logits, no softmax (CrossEntropyLoss includes it)

def compare_numpy_torch():
    torch.manual_seed(123)
    np.random.seed(123)
    
    B, D, H, C = 64, 20, 32, 5
    x_np = np.random.randn(B, D).astype(np.float32)
    y_np = np.random.randint(0, C, B)
    
    x_torch = torch.from_numpy(x_np).requires_grad_(False)
    y_torch = torch.from_numpy(y_np).long()
    
    # NumPy model
    np_model = TwoLayerMLP(D, H, C, seed=123)
    
    # PyTorch model (copy weights)
    torch_model = TorchMLP(D, H, C)
    torch_model.fc1.weight.data = torch.from_numpy(np_model.W1.copy())
    torch_model.fc1.bias.data = torch.from_numpy(np_model.b1.copy())
    torch_model.fc2.weight.data = torch.from_numpy(np_model.W2.copy())
    torch_model.fc2.bias.data = torch.from_numpy(np_model.b2.copy())
    
    # Forward comparison
    np_probs = np_model.forward(x_np)
    torch_logits = torch_model(x_torch)
    torch_probs = F.softmax(torch_logits, dim=1).detach().numpy()
    
    print(f"Forward max diff: {np.max(np.abs(np_probs - torch_probs)):.2e}")
    
    # Loss comparison
    np_loss = np_model.loss(y_np)
    torch_loss = F.cross_entropy(torch_logits, y_torch).item()
    print(f"Loss diff: {abs(np_loss - torch_loss):.2e}")
    
    # Backward comparison
    np_model.backward(y_np)
    torch_loss.backward()
    
    for name, (np_param, torch_param) in [
        ('W1', (np_model.W1, torch_model.fc1.weight)),
        ('b1', (np_model.b1, torch_model.fc1.bias)),
        ('W2', (np_model.W2, torch_model.fc2.weight)),
        ('b2', (np_model.b2, torch_model.fc2.bias))
    ]:
        np_grad = getattr(np_model, f'd{name}')
        torch_grad = torch_param.grad.numpy()
        rel_err = np.max(np.abs(np_grad - torch_grad) / (np.abs(np_grad) + np.abs(torch_grad) + 1e-10))
        print(f"{name} grad rel error: {rel_err:.2e} {'✓' if rel_err < 1e-5 else '✗'}")

if __name__ == "__main__":
    compare_numpy_torch()

Common Errors & Debugging Checklist

Error Symptom	Likely Cause	Fix
Loss is NaN	Softmax missing max subtraction, log(0)	Use `z - max(z)`, add `1e-12`
Gradient check fails (rel error > 1e-3)	Backprop derivation wrong, dimension mismatch	Print shapes per layer, verify against chain rule
Training doesn't converge	LR too large, init wrong	Reduce LR 10x, check He/Xavier init
`dW1` shape wrong	Matrix multiply order reversed	`dz1.T @ x` not `x.T @ dz1`
ReLU gradient all zero	`z1 <= 0` all negative, dead neurons	Check init, consider LeakyReLU

Video Timestamps (Relevant Lectures)

Lec 01: Introduction (YouTube: 6FkRvTtUc-o) — MLP architecture overview
Lec 02: How to Train a Neural Net (YouTube: vidCX_dMCu0) — Backprop derivation, chain rule
PyTorch Tutorial (YouTube: o5gPABcGZwc) — Autograd basics

Extension Exercises: Add These Yourself

Mini-batch SGD: Replace full-batch with DataLoader iteration
Momentum / Adam: Hand-write optimizer classes
Learning rate schedule: Cosine decay implementation
Weight decay: Add L2 to loss or decoupled implementation
Dropout / BatchNorm: Handle correctly in both forward and backward

References

MIT 6.7960 Fall 2024 Homework 1 (PDF) — Official problems
Lec 01: Introduction to Deep Learning — MLP fundamentals
Lec 02: How to Train a Neural Net — Backprop deep dive
PyTorch Tutorial (OCW) — Autograd intro
CS231n: Backpropagation Notes — Classic backprop tutorial (Karpathy)
The Matrix Calculus You Need For Deep Learning (Parr & Howard, 2018) — Matrix calculus cheatsheet
PyTorch Autograd Official Docs — Mechanics & hooks
On this site: MIT 6.7960 L03: Optimization Overview — Optimizer setup
On this site: MIT 6.7960 L04: Regularization in Practice — Regularization deep dive