Back to Backprop

Table of Contents

Back to Backprop

In the rapidly advancing field of deep learning, it’s easy to get swept up in the excitement of state-of-the-art architectures and massive datasets. Yet, beneath the surface of every breakthrough is the workhorse of gradient-based optimization: backpropagation. For seasoned practitioners, revisiting this foundational algorithm isn’t just a nostalgic exercise—it’s an opportunity to sharpen our intuition and improve our practice.

This post explores backpropagation from a practitioner’s lens. We’ll begin with a concise refresher, implementing the forward and backward passes for foundational layer types. We’ll demonstrate a full training loop for a simple neural network, using our own custom implementation and mini-autograd engine. Along the way, we’ll dive into real-world challenges, explore modern tweaks. Finally, we’ll culminate in building a simplified PyTorch-like tensor and autograd engine.

Here are some quick links to some awesome resources on backpropagation:

• Derivatives, Backpropagation, and Vectorization - Justin Johnson
• cs224n Lecture 3: Neural net learning: Gradients by hand (matrix calculus) and algorithmically (the backpropagation algorithm)
• cs231n BackProp Notes- Andrej Karpathy
• Computing Neural Network Gradients - Kevin Clark

From Derivatives to Gradients and Jacobians

For completelness sake, I think it is worth just restating some of the basic definitions and concepts that underpin backpropagation.

Derivatives

Derivatives are a way to measure change in the scalar case:

The Chain Rule

The chain rule tells us how to compute the derivative of a composition of functions. Suppose we have functions $f,g: \mathbb{R} \rightarrow \mathbb{R}$ and $h = f \circ g$ In other words $y = f(x)$ and $z = g(y)$. Then, the chain rule states:

In modern neural networks, it turns out that backpropagation is essentially the chain rule applied to a composition of vector/tensor-valued functions.

Gradients: Vector in -> Scalar out

Gradients are the generalization of derivatives to the multivariate case, where we have a scalar-valued function of multiple variables. The gradient of a function $f: \mathbb{R}^N \rightarrow \mathbb{R}$ is a vector of partial derivatives, one for each input variable.

When we consider the relationship

it is important to note that $x, \Delta x, \nabla f(x)$ are all vectors, now. $\nabla f(x)_x \dot \Delta x$ is computed by taking the dot product, producing a scalar, which matches the output signature of the gradient.

Jacobians: Vector in -> Vector out

Now we consider a function $\mathbf{f}: \mathbb{R}^N \rightarrow \mathbb{R}^M$. The Jacobian of $\mathbf{f}$ at a point $\mathbf{x} \in \mathbb{R}^N$ is an $M \times N$ matrix of partial derivatives defined such that the $(i, j)-th$ entry of the Jacobian is $\boxed{\frac{\partial \mathbf{f(x)}_i}{\partial x_j}}$. Written out in other forms:

Ways to remember the formula:

• Outputs are rows (M) (outputs vary down a column)
• Inputs are columns (N) (inputs vary across a row)
• Each row is the gradient of the corresponding output with respect to all the inputs.
• The Jacobian is the transpose of the gradient of the function (i.e., when $m=1$ so $f: \mathbb{R}^N \rightarrow \mathbb{R}$, the Jacobian is the transpose of the gradient, $\mathbf{J}_f = \nabla^Tf)$.

Generalized Jacobians: Tensor in -> Tensor out

From scalars to vectors (1-D array) to matrices(2-D array), we can generalize the concept of Jacobians to higher-order tensors, which we can view as $D$-dimensional arrays.

Commonly, in neural networks, it is very common to think of tensor operations in terms of their “shapes”, which are the dimensions of the tensor, i.e., the tuple returned by tensor.shape in numpy/PyTorch. You will frequently see code where comments are tracking the shape of the tensor, e.g., # [B, N, D] x [B, D, M] -> [B, N, M]. This is a shorthand and has even been formalized in the Einstein Notation used in PyTorch, with einsum and the python package einops. ` to indicate a 2-D tensor with $N$ rows and $D$ columns. . or example, a tensor of shape $(N, D)$ is a 2-D tensor with $N$ rows and $D$ columns.

Now, suppose we had a function $f: \mathbb{R}^{N_1 \times N_2 \times \cdots \times N_{D_x}} \rightarrow \mathbb{R}^{M_1 \times M_2 \times \cdots \times M_{D_y}}$. The Jacobian of $f$ at a point $\mathbf{X} \in \mathbb{R}^{N_1 \times N_2 \times \cdots \times N_{D_x}}$ is a tensor of shape $(M_1 \times M_2 \times \cdots \times M_{D_y} \times N_1 \times N_2 \times \cdots \times N_{D_x})$. The generalized Jacobian is another tensor with shape:

The Mechanics of Backpropagation

At its core, backpropagation efficiently computes gradients for training neural networks by leveraging the chain rule. Here’s how it works:

1. Forward Pass: Inputs flow through the network, producing intermediate activations and final outputs.
2. Backward Pass: Gradients are calculated layer by layer in reverse, propagating partial derivatives back to earlier layers.

Suppose we have a function $z = xW$ as a single node in a computational graph. The output $z$ is passed through additional functions (nodes in the graph) before producing a final scalar loss $L$. The chain rule allows us to compute the gradient of $L$ with respect to $x$ and $W$ by multiplying the gradients of each intermediate function:

where $\dots \frac{\partial s}{\partial z}$ represents the chain of gradients from the loss $L$ to the output $z$, known as the upstream gradient. The goal is to compute the downstream gradients:

[downstream gradient] = [upstream gradient] * [local gradient]

Backprop by hand

Let’s walk through a simple example to illustrate backpropagation. We’ll start use a 1-layer neural network. Let’s assume vectorized/batched computation starting with input $x \in \mathbb{R}^{N \times D}$, where $N$ is the number of samples and $D$ is the input dimension.

Linear Layer:

A linear layer consists of weights $W \in \mathbb{R}^{D \times H}$, bias $b \in \mathbb{R}^{H}$.

• Forward Pass: $h = xW + b$ with $h \in \mathbb{R}^{N \times H}$.
• Gradient Identities:

• Backward Pass: $\frac{\partial L}{\partial W} = X^T \cdot \frac{\partial L}{\partial h}$, $\frac{\partial L}{\partial b} = \frac{\partial L}{\partial h} I$, $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial h} \cdot W^T$.

In code:

class Linear:
    def __init__(self, in_features: int, out_features: int):
        self.weight = torch.empty(in_features, out_features)
        self.bias = torch.zeros(out_features)
        self.x: torch.Tensor | None = None
        init_uniform_(self.weight)
        init_uniform_(self.bias)

    def parameters(self):
        return [self.weight, self.bias]

    def forward(self, x: torch.Tensor):
        """Computes the forward pass for a linear layer.

        Args:
            x (Tensor): Input data, of shape (N, in_features) where N is the batch size and
                in_features is the number of input features.
        """
        self.x = x
        return x @ self.weight + self.bias  # @ is torch.matmul

    def backward(self, dout: torch.Tensor):
        """Computes the backward pass for a linear layer.

        Args:
            dout (Tensor): Upstream derivative, of shape (N, out_features).

        Notes:
        ------------
        z = xW + b
        dz/dx = W.T
        dz/dW = x.T
        dz/db = I_{out}

        chain rule
        dL/dx = dout * dz/dx
              = dout * W.T  [N, out] x [out, in]
        dL/dW = dout * dz/dW (need to swap to make shapes work out)
              = x.T  * dout [in, N] x [N, out]
        dL/db = dout * dz/db
              = dout * I_{out} [N, out] x [out, out]
        """
        if self.x is None:
            raise ValueError("No cache found. Run forward pass first.")
        dx = dout.mm(self.weight.T)  # [N, in_features]
        self.weight.grad = self.cache.T.mm(dout)  # [in_features, out_featues]
        self.bias.grad = dout.sum(dim=0)
        return dx

ReLU Activation:

To start with, recall that ReLU(x) = max(x, 0). This means that

which is also $\text{sgn}(\text{ReLU}(x))$.

def relu(x: torch.Tensor) -> torch.Tensor:
	"""Computes max(0, x) element-wise."""
	return x.clamp(min=0)

class ReLU:
    def parameters(self):
        return []

    def forward(self, x: torch.Tensor):
        """Computes max(0, x) element-wise."""
        self.x = x
        return relu(x)

    def backward(self, dout: torch.Tensor):
        """Computes the backward pass for a layer of rectified linear units (ReLUs).

        ReLU(x) = max(0, x)
        ReLU'(x) = 1 if x > 0 else 0
        out = ReLU(x)
        dL/dx = dout * ReLU'(x)
              = dout * 1[x > 0]
              = dout[x < 0] = 0
              = dout * 1[x > 0]

        Args:
            dout (Tensor): Upstream derivative, of shape (N, M).

        """
        return dout * (self.x > 0).float()

    __call__ = forward

Tanh Activation:

To start with, recall that $\text{Tanh}(x) = \frac{e^{2x} - 1}{e^{2x} + 1}$. This means that the derivative is:

class TanhLayer:
    def forward(self, x):
        self.output = torch.tanh(x)
        return self.output

    def backward(self, dout: torch.Tensor):
        return dout * (1 - self.output**2)

    def parameters(self):
        return []

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

Softmax, LogSoftmax and LogSumExp

These functions commonly arise in the context of neural networks, especially in the output layer. They occur when dealing with categorical and multinomial probability distributions, converting raw logits into a probably distrubution over classes, as well as in attention mechanisms in transformers These functions are all closely related, involving sums of exponentials, and have nice relationships with each other.

Definitions

Given a vector $\mathbf{x} \in \mathbb{R}^n$, we define the following functions:

So we see that the softmax and logsoftmax are both functions that take in a vector and return a vector, while the logsumexp is a scalar valued function. The logsoftmax is just the log of the softmax, which also happens to be the input minus its logsumexp.

And in code:

def softmax(z):
    """Compuze softmax values for each sets of scores in z.

    Args:
        z (Tensor): input tensor

    Notes:
        The definition of the softmax:
            softmax(z) = exp(z) / sum(exp(z))
        can be numerical unstable, so we use the following definition:
            softmax(z) = exp(z) / sum(exp(z - max(z)))
        which multiplies the numerator and denominator by a constant factor.
        It is usually chosen to be -max(z) producing negative to zero range
        to avoid numerical overflow.
    """
    e_z = torch.exp(z - torch.max(z))
    return e_z / e_z.sum()

def logsumexp(s: torch.Tensor, dim=-1):
	"""Numerically stable log(sum(exp(s)))"""
	max_val = s.max(dim=dim, keepdim=True)[0]
	return torch.log(torch.sum(torch.exp(s - max_val), dim=dim, keepdim=True)) + max_val

def log_softmax(z: torch.Tensor, dim=-1):
	"""Return logprobs"""
	return z - logsumexp(z, dim=dim)

class SoftmaxLayer:
    def forward(self, x):
        # Save input for backward pass
        self.input = x

        # For numerical stability, subtract max value before exponential
        # Shape: (batch_size, num_classes)
        shifted_x = x - torch.max(x, dim=1, keepdim=True)[0]

        # Compute exponentials
        # Shape: (batch_size, num_classes)
        exp_x = torch.exp(shifted_x)

        # Compute sum of exponentials for normalization
        # Shape: (batch_size, 1)
        sum_exp = torch.sum(exp_x, dim=1, keepdim=True)

        # Compute softmax output
        # Shape: (batch_size, num_classes)
        self.output = exp_x / sum_exp

        return self.output

    def backward(self, grad_output):
        # grad_output shape: (batch_size, num_classes)
        # output shape: (batch_size, num_classes)

        # The Jacobian of softmax for each sample is:
        # J[i,j] = output[i] * (δ[i,j] - output[j])
        # where δ[i,j] is 1 if i=j and 0 otherwise

        # Compute using the compact form:
        # grad = output * (grad_output - (output · grad_output))

        # Shape: (batch_size, 1)
        sum_grad = torch.sum(grad_output * self.output, dim=1, keepdim=True) # s · ∇_s L

        # Shape: (batch_size, num_classes)
        return self.output * (grad_output - sum_grad) # s ⊙ (∇_s L - (s · ∇_s L))

    def __call__(self, x):
        return self.forward(x)

Sigmoid:

While the softmax function is used for multi-class classification, the sigmoid function is used for binary classification. The primary difference is that the softmax enforces that elements of particular dimension sum to 1 (producing a probability distribution), while the sigmoid function squashes the output to the range $(0, 1)$, making it suitable for binary classification.

Definition

The sigmoid function, also known as the logistic function, maps any real number to a value between 0 and 1. It’s particularly useful for binary classification tasks since its output can be interpreted as a probability. The sigmoid function is defined as:

Its derivative has a particularly nice form:

Which makes it efficient to compute during backpropagation. And in code:

def sigmoid(x: torch.Tensor) -> torch.Tensor:
	return 1 / (1 + torch.exp(-x))

class SigmoidLayer:

	def forward(self, x: torch.Tensor):
		self.out = sigmoid(x)
		return self.out

	def backward(self, dout: torch.Tensor):
		"""
		sigmoid'(x) = sigmoid(x) * (1 - sigmoid(x))
		"""
		return dout * self.out * (1 - self.out)

Typically, in frameworks like PyTorch, models output logits (unnormalized scores) and it usually the role of the crtiterion (e.g. CrossEntropyLoss) to implicitly convert the logits to probabilities using the softmax function.

Losses:

class CrossEntropyLoss:
	"""Computes the forward pass for the cross-entropy loss.

    CE(s, y) = -SUM_i y_i * log(softmax(s)_i)
    torch.mean(-torch.sum(targets * torch.log(torch.softmax(inputs, dim=1)), dim=1)))

    Args:
        x (Tensor): Input data, of shape (N, C) where N is the batch size and
            C is the number of classes.
        y (Tensor): Ground truth labels, of shape (N,).
    """

	def forward(self, input: torch.Tensor, target: torch.Tensor):
		self.input = input
		self.target = target
		indices = torch.arange(target.shape[0])
		return -log_softmax(input)[indices, target].mean()

	def backward(self):
		"""Compute the backward pass for the cross-entropy loss.

		dL(x, y)/dx = softmax(x) - y
		where x is the unnormalized score, y is the one-hot encoded target label.

        dL(p, y)/dx = p - y  (p is predicted prob)

        Returns:
            dx (Tensor): Gradient of the loss with respect to the input x.
        """
		N = self.target.shape[0]
		dx = F.softmax(self.inputs)
		dev = self.input.device
		indices = torch.arange(N, dtype=self.input.dtype, device=device)
		dx[indices, self.target] =- 1
		dx /= N
		return dx

	__call__ = forward

import torch


class BCELoss:
    """
    Need to apply sigmoid activation to inputs before passing it here
    """

    def forward(self, probs, target):
        """

        BCELoss(x, y) = -[y * log(x) + (1 - y) * log(1 - x)]

        Args:
            probs (Tensor): Predictions from the model
            target (Tensor): Ground truth labels

        Returns:
            Binary cross-entropy loss averaged over batch
        """
        self.probs = probs
        self.target = target
        return -torch.mean(
            target * torch.log(probs) + (1 - target) * torch.log(1 - probs)
        )

    def backward(self):
        """Compute the gradient of binary cross-entropy loss.

        dBCE(p, y)/dp = (p - y) / (p * (1 - p))

        """
        return (self.probs - self.target) / (self.probs * (1 - self.probs))

import torch


class BCEWithLogitsLoss:
    def forward(self, input: torch.Tensor, target: torch.Tensor):
        """
        input: (*) shape
        target: (*) shape
        """
        # Use logsumexp trick
        # BCE(x, y) = (1 - y)*x + log(1 + e^{-x})
        #           = (1 - y)*x + log (e^0 + e^{-x})

        zeros = torch.zeros_like(input)
        m = torch.maximum(zeros, -input)

        def lse(a, b):
            return m + torch.log(torch.exp(a - m) + torch.exp(b - m))

        return (1 - target) * input + lse(zeros, -input)

import torch


class MSELoss:
    def forward(self, input: torch.Tensor, target: torch.Tensor):
        """Compute mean squared error loss.

        Args:
            y_pred: Model predictions
            y_true: Ground truth values

        Returns:
            Mean squared error averaged over batch
        """
        self.diff = input - target
        return torch.mean(self.diff**2)

    def backward(self, dout: torch.Tensor):
        """Compute gradient of MSE loss.

        Returns:
            Gradient with respect to predictions: 2/N * (pred - true)
        """
        N = self.diff.shape[0]
        return 2 * self.diff / N

import torch


class L1Loss:
    def forward(self, input: torch.Tensor, target: torch.Tensor):
        """Compute L1 loss."""

        self.diff = input - target
        return torch.mean(torch.abs(self.diff))

    def backward(self):
        """Compute gradient of L1 loss.

        Returns:
            Gradient with respect to predictions: 1/N * sign(pred - true)
        """
        return torch.sign(self.diff) / self.diff.shape[0]

import torch
import torch.nn.functional as F


class LogSoftmaxLayer:
    def forward(self, x):
        # x: (N, C)
        # store input and output for backward
        self.x = x
        self.z = F.log_softmax(x, dim=1)
        return self.z

    def backward(self, dz):
        # dz: gradient w.r.t. output z = log_softmax(x)
        # We know dL/dx_i = dz_i - softmax(x)_i * sum_j dz_j
        # But since softmax(x) = exp(x)/sum exp(x), we can quickly compute it:
        # s = F.softmax(self.x, dim=1)
        # or
        s = torch.exp(self.z)

        # Sum over classes: sum_j dz_j
        sum_dz = torch.sum(dz, dim=1, keepdim=True)

        return dz - s * sum_dz  #  dx

import torch


class NLLLoss:
    def forward(self, log_probs: torch.Tensor, target: torch.Tensor):
        """Compute negative log likelihood loss.

        NLL = -1/N * sum_i log_probs[i, target[i]]

        Args:
            log_probs: Log probabilities, shape (N, C)
            target: Ground truth values, shape (N,)

        Returns:
            Negative log likelihood loss
        """
        self.log_probs = log_probs
        self.target = target
        N = target.shape[0]
        return -log_probs[torch.arange(N), target].mean()

    def backward(self):
        """Compute gradient of negative log likelihood loss.

        dL/dlogp = -1/N * one_hot(target)

        Returns:
            Gradient with respect to log probabilities: -1
        """
        N = self.target.shape[0]
        grad = torch.zeros_like(self.log_probs)
        grad[torch.arange(N), self.target] = -1.0
        grad /= N
        return grad

Evaluating Gradients: Numerical Gradient Checking

While we’ve derived the gradients for common operations, it’s crucial to validate our implementations. One way to do this is by comparing the analytical gradients with numerical approximations. Numerical gradient checking is a technique to verify the correctness of our gradients by approximating them using finite differences.

The central difference formula approximates the gradient of a function $f$ at a point $x$ by evaluating the function at two nearby points $x + \epsilon$ and $x - \epsilon$:

By comparing the analytical gradients with numerical approximations, we can catch errors in our implementations and ensure that our models are learning correctly.

Here’s a simple implementation of gradient checking for any layer implementing the Layer protocol (which you’ll see in the next section):

import torch
from activations import ReLU, SigmoidLayer, SoftmaxLayer, TanhLayer
from layers import Layer, Linear


def eval_numerical_gradients(
    layer: Layer,
    x: torch.Tensor,
    epsilon: float = 1e-6,
    tolerance: float = 1e-5,
):
    print(f"Evaluating grads for: {type(layer).__name__}")
    # Store original parameters
    original_params = [p.data.clone() for p in layer.parameters()]

    # Compute the analytical gradients
    output = layer.forward(x)
    dout = torch.rand_like(output)
    dx = layer.backward(dout)

    analytical_grads = [
        p.grad.data.clone() for p in layer.parameters() if p.grad is not None
    ]
    analytical_input_grad = dx.clone()

    # Compute numerical gradients
    numerical_grads = []
    for i, params in enumerate(layer.parameters()):
        param_grad = torch.zeros_like(params)

        flat_params = params.flatten()
        flat_grads = param_grad.flatten()
        for i in range(params.numel()):
            orig_val = flat_params.data[i].clone()

            # Compute f(x + epsilon)
            flat_params.data[i] = orig_val + epsilon
            output_plus = layer.forward(x)
            f_plus = (output_plus * dout).sum()

            # Compute f(x - epsilon)
            flat_params.data[i] = orig_val - epsilon
            output_minus = layer.forward(x)
            f_minus = (output_minus * dout).sum()

            # Compute numerical gradient
            flat_grads.data[i] = (f_plus - f_minus) / (2 * epsilon)

            # Restore original value
            flat_params.data[i] = orig_val

        numerical_grads.append(param_grad)

    numerical_input_grad = torch.zeros_like(dx)
    flat_x = x.flatten()
    flat_input_grads = numerical_input_grad.flatten()
    for i in range(x.numel()):
        orig_val = flat_x.data[i].clone()

        # Compute f(x + epsilon)
        flat_x.data[i] = orig_val + epsilon
        output_plus = layer.forward(x)
        f_plus = (output_plus * dout).sum()

        # Compute f(x - epsilon)
        flat_x.data[i] = orig_val - epsilon
        output_minus = layer.forward(x)
        f_minus = (output_minus * dout).sum()

        # Compute numerical gradient
        flat_input_grads.data[i] = (f_plus - f_minus) / (2 * epsilon)
        flat_x.data[i] = orig_val

    # Compare analytical and numerical gradients
    for i, (analytical, numerical) in enumerate(zip(analytical_grads, numerical_grads)):
        diff = torch.abs(analytical - numerical).max().item()
        print(f"Max difference for parameter {i}: {diff}")
        print(f"Relative error: {relative_error(analytical, numerical)}")
        assert (
            diff < tolerance
        ), f"Gradient check failed for parameter {i}. Max difference: {diff}"
    # Compare analytical and numerical gradients for input

    input_diff = torch.abs(analytical_input_grad - numerical_input_grad).max().item()
    print(f"Max difference for input: {input_diff}")
    print(
        f"Relative error for input: {relative_error(analytical_input_grad, numerical_input_grad)}"
    )
    assert (
        input_diff < tolerance
    ), f"Gradient check failed for input. Max difference: {input_diff}"

    # Restore original parameters
    for param, original in zip(layer.parameters(), original_params):
        param.data.copy_(original)

    print("Gradient check passed!")


def relative_error(a: torch.Tensor, b: torch.Tensor, eps: float = 1e-11):
    return (torch.abs(a - b) / (torch.abs(a) + torch.abs(b) + eps)).mean()


if __name__ == "__main__":
    # For the Linear layer
    dtype = torch.float64
    linear = Linear(10, 5).to(dtype=dtype)
    x_linear = torch.randn(32, 10, dtype=dtype)
    eval_numerical_gradients(linear, x_linear)

    relu = ReLU()
    x_relu = torch.randn(32, 10, dtype=dtype)
    eval_numerical_gradients(relu, x_relu)

    x = torch.randn(32, 10, dtype=dtype)
    eval_numerical_gradients(SigmoidLayer(), x)

    x = torch.randn(32, 10, dtype=dtype)
    eval_numerical_gradients(SoftmaxLayer(), x)

    x = torch.randn(32, 10, dtype=dtype)
    eval_numerical_gradients(TanhLayer(), x)

Putting it all together: Training a small MLP on MNIST

Now that we have all the building blocks in place, let’s put them together to train a simple multi-layer perceptron (MLP) on the MNIST dataset. This classic computer vision task involves classifying 28x28 grayscale images of handwritten digits (0-9).

Our MLP will consist of multiple linear layers with ReLU activations between them, followed by a final linear layer that outputs logits for each of the 10 digit classes. We’ll use cross-entropy loss and stochastic gradient descent (SGD) for optimization.

The training loop follows these key steps:

1. Forward pass through the network to get predictions
2. Calculate loss using cross-entropy
3. Backward pass to compute gradients
4. Update parameters using SGD
5. Repeat!

This end-to-end example demonstrates how the individual components we’ve built - layers, activations, loss functions, optimizers - work together during training. It also showcases common training practices like:

• Using mini-batches for efficiency
• Tracking metrics like accuracy and loss
• Validating on a test set
• Logging progress during training

While modern deep learning frameworks abstract away much of this complexity, understanding these fundamentals is invaluable for debugging, optimization, and developing new techniques.

Let’s examine the code and dissect what’s happening at each step:

First, let’s examine how we’ve structured our custom layers to mirror PyTorch’s Module interface. Our layers accept input tensors in their forward methods, maintain their own parameters and gradients, and properly chain the backward pass. Just like PyTorch modules, they:

• Have forward() and backward() methods for the computation graph
• Track parameters that need gradients
• Support function-like calling through call
• Maintain state between forward and backward passes

This design pattern makes it easy to compose layers into networks and perform automatic differentiation in a familiar way. The Layer protocol we defined above formalizes this interface.

class Layer(Protocol):
    def forward(self, x: torch.Tensor) -> torch.Tensor: ...
    def backward(self, dout: torch.Tensor) -> torch.Tensor: ...
    def parameters(self) -> list[torch.Tensor]: ...
    def __call__(self, x: torch.Tensor) -> torch.Tensor: ...

class Optimizer(Protocol):
    params: list[torch.Tensor]

    def __init__(self, params: list[torch.Tensor], lr: float):
        self.params = params

    def step(self): ...

    def zero_grad(self): ...

Network Architecture

In our simple network architecture, we have a feedforward neural network (also called a multilayer perceptron or MLP) that takes 784-dimensional input (flattened 28x28 MNIST images) and outputs 10-dimensional logits (one for each digit class). The network consists of:

• Input layer: Flattens 28x28 images into 784-dimensional vectors
• Three hidden layers with ReLU activations between them, each with size hidden_size (default 1024)
• Output layer: Linear layer producing 10 logits
• Cross entropy loss for training

The network uses standard linear layers with learned weight matrices and bias vectors, with ReLU nonlinearities between each layer to allow learning of non-linear functions.

Each linear layer performs an affine transformation: output = input @ weight + bias The ReLU activation applies an elementwise nonlinearity: output = max(input, 0)

This relatively simple architecture is still capable of achieving ~98% accuracy on MNIST with proper training. The multiple hidden layers allow the network to learn hierarchical features from the input digits.

import torch.nn as nn
from activations import ReLU
from layers import Layer, Linear
from losses import CrossEntropyLoss


class Net:
    def __init__(self, hidden_size=1024, num_classes=10):
        self.layers: list[Layer] = [
            Linear(784, hidden_size),
            ReLU(),
            Linear(hidden_size, hidden_size),
            ReLU(),
            Linear(hidden_size, hidden_size),
            ReLU(),
            Linear(hidden_size, num_classes),
        ]
        self.loss = CrossEntropyLoss()

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        for layer in self.layers:
            x = layer(x)
        return x

    def backward(self, dout):
        for layer in reversed(self.layers):
            dout = layer.backward(dout)
        return dout

    def parameters(self):
        return [param for layer in self.layers for param in layer.parameters()]

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)


class PytorchNet(nn.Module):
    def __init__(self, hidden_size=1024, num_classes=10):
        super().__init__()
        self.layers = nn.Sequential(
            *[
                nn.Linear(784, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, hidden_size),
                nn.ReLU(),
                nn.Linear(hidden_size, num_classes),
            ]
        )

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        return self.layers(x)

Optimizers

Let’s look at different optimizers and their implementations:

class SGD:

	params: list[torch.Tensor]

	def __init__(self, params: list[torch.Tensor], lr: float = 0.01):
		self.params = params
		self.lr = lr

	def step(self):
		if self.lr <= 0:
            raise ValueError("Learning rate must be positive")
        if self.params and self.params[0].grad is None:
            raise ValueError("Parameters have no gradients. Call backward first.")
		for param in self.params:
			param -= self.lr * param.grad

	def zero_grad(self):
		for param in self.params:
			param.grad = torch.zeros_like(param)

class SGDMomentum:
	def __init__(self, params: list[torch.Tensor], lr: float = 0.01, momentum: float = 0.9):
		self.params = params
		self.lr = lr
		self.momentum = momentum
		self.velocities = [torch.zeros_like(param) for param in self.params]

	def step(self):
		"""
		dW = momentum * dW - lr * dL/dW (perform grad update on velocity weight)
        W = W + dW
        """
		if self.lr <= 0:
            raise ValueError("Learning rate must be positive")
        if self.params and self.params[0].grad is None:
            raise ValueError("Parameters have no gradients. Call backward first.")
		for i, param in enumerate(self.params):
			self.velocity[i] = self.momentum * self.velocity[i] - self.lr * param.grad
			param.grad += self.velocity[i]

class RMSprop:
    def __init__(
        self,
        params: list[torch.Tensor],
        lr=0.01,
        alpha=0.99,
        momentum: float = 0,
        eps=1e-8,
    ):
        self.params = params
        self.lr = lr
        self.alpha = alpha
        self.eps = eps
        self.square_avg = [torch.zeros_like(p) for p in params]
        self.momentum = momentum
        if momentum > 0:
            self.buffers = [torch.zeros_like(p) for p in params]

    def step(self):
        """
        v = alpha * v + (1 - alpha) * (dL/dW)^2
        W = W - lr * dL/dW / sqrt(v + eps)
        """
        if self.lr <= 0:
            raise ValueError("Learning rate must be positive")
        if self.params and self.params[0].grad is None:
            raise ValueError("Parameters have no gradients. Call backward first.")
        for i, param in enumerate(self.params):
            grad: torch.Tensor = param.grad  # type: ignore
            self.square_avg[i] = (
                self.alpha * self.square_avg[i] + (1 - self.alpha) * grad**2
            )
            if self.momentum > 0:
                self.buffers[i] = self.momentum * self.buffers[i] + self.lr * grad / (
                    torch.sqrt(self.square_avg[i]) + self.eps
                )
                self.params[i] -= self.lr * self.buffers[i]
            else:
                self.params[i] -= (
                    self.lr * grad / (torch.sqrt(self.square_avg[i]) + self.eps)
                )

class Adam:
    def __init__(
        self,
        params: list[torch.Tensor],
        lr=0.001,
        betas=(0.9, 0.999),
        eps=1e-8,
        weight_decay=0,
    ):
        self.params = params
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.m = [torch.zeros_like(p) for p in params]  # First moment
        self.v = [torch.zeros_like(p) for p in params]  # Second moment
        self.t = 0
        self.weight_decay = weight_decay

    def step(self):
        self.t += 1
        for i, (param, m, v) in enumerate(zip(self.params, self.m, self.v)):
            grad: torch.Tensor = param.grad  # type: ignore
            if self.weight_decay > 0:
                grad += self.weight_decay * param

            # Update biased first moment estimate
            self.m[i] = self.beta1 * m[i] + (1 - self.beta1) * grad
            # Update biased second raw moment estimate
            self.v[i] = self.beta2 * v[i] + (1 - self.beta2) * grad**2

            # Bias correction
            m_hat = self.m[i] / (1 - self.beta1**self.t)
            v_hat = self.v[i] / (1 - self.beta2**self.t)

            # Update parameters
            param -= self.lr * m_hat / (torch.sqrt(v_hat) + self.eps)

    def zero_grad(self):
        for param in self.params:
            param.grad = torch.zeros_like(param)