Skip to content

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

Aug 30, 20261 min
TL;DRHand-write NumPy MLP + backprop → verify with PyTorch Autograd, fully reproducing OCW HW1 core concepts

🌏 中文版

Video status: Related supplementary video included; the original lecture recording has not been verified. Source details

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):

  1. Problem 1: Softmax & Cross-Entropy — Derive vectorized formulas, numerically stable log-sum-exp
  2. Problem 2: Two-Layer MLP Forward — Matrix multiply, ReLU, Softmax chaining
  3. Problem 3: Backpropagation by Hand — Chain rule layer-by-layer derivation of ∂L/∂W, ∂L/∂b
  4. 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 SymptomLikely CauseFix
Loss is NaNSoftmax missing max subtraction, log(0)Use z - max(z), add 1e-12
Gradient check fails (rel error > 1e-3)Backprop derivation wrong, dimension mismatchPrint shapes per layer, verify against chain rule
Training doesn't convergeLR too large, init wrongReduce LR 10x, check He/Xavier init
dW1 shape wrongMatrix multiply order reverseddz1.T @ x not x.T @ dz1
ReLU gradient all zeroz1 <= 0 all negative, dead neuronsCheck init, consider LeakyReLU
  • Lec 01: Introduction (YouTube: 6FkRvTtUc-o, taught by Sara Beery) — roughly 40%–80% introduces the basic building blocks of neural networks and MLPs (loss, optimization, nonlinearities such as ReLU, stacking layers); the first part covers course rules and the history of deep learning.
  • Lec 02: How to Train a Neural Net (YouTube: vidCX_dMCu0, taught by Sara Beery) — roughly 45%–70% is backpropagation and the chain rule (including that the backward pass of a linear layer is the forward operation with transposed weights).
  • PyTorch Tutorial (YouTube: o5gPABcGZwc, taught by Jamie Meindl, about 29 min) — creating and manipulating tensors, autograd (.backward() and torch.no_grad()), SGD and Adam optimizers, nn.Linear and nn.Sequential, a simple training loop.

None of the three videos walks through HW1 problem by problem; positions above are approximate, estimated from where topics fall in each transcript, not exact timecodes.

Extension Exercises: Add These Yourself

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

Course video sources

Rechecked against the live MIT OCW Fall 2024 gallery on 2026-10-10: Homework 1 has no dedicated recording. The two videos below are related public OCW videos: Lec 02 How to Train a Neural Net (SGD, backpropagation, automatic differentiation) and the PyTorch Tutorial. They are background supplements, not a walkthrough of HW1.

YouTube:MIT 6.7960 Fall 2024 — Lec 02. How to Train a Neural Net
YouTube:PyTorch Tutorial

Original videos: MIT 6.7960 Fall 2024 — Lec 02. How to Train a Neural Net, PyTorch Tutorial

Course and recording entries:

Checked: 2026-10-10.

Content check: verified against the video transcript (2026-10-10): read the Lec 01 and Lec 02 transcripts and the full PyTorch Tutorial transcript (read directly, not saved to a file); all three are related background rather than a HW1 walkthrough; the misleading "video timestamps (relevant lectures)" heading was renamed and each video's actual content added by approximate relative position, with no exact timecodes.

Update Log

  • 2026-10-10: Added explicit video status and checked recording sources and access notes.
  • 2026-10-10: Rechecked video status. HW1 has no dedicated recording, so OCW Lec 02 and the PyTorch Tutorial are embedded as related supplementary videos and the status is now “Related supplementary video included; the original lecture recording has not been verified.”
  • 2026-10-10: Checked the video content against its transcript. Checked the three related videos against their transcripts; renamed the heading to "Related Lecture Videos" and added each video's actual content (no exact timecodes).

References