Implementing models from scratch

Core Motivation

When you import torch.nn.Linear, you're using a computational abstraction that hides:

How gradients actually propagate (backprop mechanics)
Memory layout and computation order
Numerical stability trade-offs
Why certain design choices exist

Building from scratch forces you to answer: What is this layer actually computing, and why does this update rule work?

The Three Pillars of From-Scratch Implementation

1. Forward Pass: Computing Outputs

Derivation for a Simple Dense Layer:

Start with the fundamental question: How do we transform an input vector into an output space?

Linear transformation (the simplest learnable operation): $\mathbf{z} = \mathbf{W}\mathbf{x} + \mathbf{b}$

Why this form?

$\mathbf{W}$ (weight matrix): Rotates and scales input space
$\mathbf{b}$ (bias vector): Translates the decision boundary
This is the most general affine transformation

For input $\mathbf{x} \in \mathbb{R}^{n}$ , output $\mathbf{z} \in \mathbb{R}^{m}$ :

$\mathbf{W} \in \mathbb{R}^{m \times n}$ (each row is a separate linear function)
$\mathbf{b} \in \mathbb{R}^{m}$

Then apply activation function $\sigma$ : $\mathbf{a} = \sigma(\mathbf{z})$

Why activation? Without it, stacking layers just gives another linear function: $\mathbf{W}_2(\mathbf{W}_1\mathbf{x} + \mathbf{b}_1) + \mathbf{b}_2 = \mathbf{W}'\mathbf{x} + \mathbf{b}'$ . Activation injects non-linearity, enabling universal function approximation.

Worked example Forward Pass Implementation

import numpy as np
class DenseLayer:
   def __init__(self, input_dim, output_dim):
       # Xavier initialization: scale by 1/sqrt(n)
       # WHY? Keeps variance stable across layers
       self.W = np.random.randn(output_dim, input_dim) / np.sqrt(input_dim)
       self.b = np.zeros((output_dim, 1))
       
   def forward(self, X):
       # X shape: (input_dim, batch_size)
       self.X = X  # Cache for backward pass
       self.Z = self.W @ X + self.b  # @ is matrix multiply
       return self.Z
 
# Why cache X? We'll need it for gradients!

Step-by-step for batch of3 samples:

Input: $\mathbf{X} \in \mathbb{R}^{4 \times 3}$ (4 features, 3 samples)
Layer: 4 → 5 dimensions
$\mathbf{W} \in \mathbb{R}^{5 \times 4}$ , $\mathbf{b} \in \mathbb{R}^{5 \times 1}$
Output: $\mathbf{Z} = \mathbf{W}\mathbf{X} + \mathbf{b} \in \mathbb{R}^{5 \times 3}$
Broadcasting adds $\mathbf{b}$ to each column

2. Backward Pass: Computing Gradients

Derivation from Chain Rule:

Given loss $L$ , we want $\frac{\partial L}{\partial \mathbf{W}}$ and $\frac{\partial L}{\partial \mathbf{b}}$ .

Assume we receive $\frac{\partial L}{\partial \mathbf{Z}}$ from the next layer (call it $\delta$ ).

For weights, using the chain rule: $\frac{\partial L}{\partial \mathbf{W}} = \frac{\partial L}{\partial \mathbf{Z}} \frac{\partial \mathbf{Z}}{\partial \mathbf{W}}$

Computing $\frac{\partial \mathbf{Z}}{\partial \mathbf{W}}$ :

From $\mathbf{Z} = \mathbf{W}\mathbf{X} + \mathbf{b}$ , for element $Z_i = \sum_j W_{ij}X_j + b_i$ :

$\frac{\partial Z_i}{\partial W_{ij}} = X_j$

In matrix form: $\frac{\partial L}{\partial \mathbf{W}} = \delta \mathbf{X}^T$

Why transpose? Dimensionality check:

$\delta$ : $(m \times n_{batch})$
$\mathbf{X}^T$ : $(n_{batch} \times n)$
Result: $(m \times n)$ ✓ (same shape as $\mathbf{W}$ )

For bias: $\frac{\partial L}{\partial \mathbf{b}} = \delta \mathbf{1}$ where $\mathbf{1}$ is a vector of ones (suming across batch dimension).

Why sum? Bias is shared across all samples—we accumulate their individual gradient contributions.

To pass gradient to previous layer: $\frac{\partial L}{\partial \mathbf{X}} = \mathbf{W}^T \delta$

Why transpose here? We're "undoing" the forward multiplication— $\mathbf{W}$ transformed $\mathbf{X} \to \mathbf{Z}$ , so $\mathbf{W}^T$ transforms gradients backwards $\mathbf{Z} \to \mathbf{X}$ .

Worked example Backward Pass Implementation

class DenseLayer:
   def backward(self, dL_dZ):
       # dL_dZ shape: (output_dim, batch_size)
       batch_size = dL_dZ.shape[1]
       
       # Gradient w.r.t. weights
       # WHY self.XT? Matches dimensions and chain rule
       self.dW = (dL_dZ @ self.X.T) / batch_size
       # Gradient w.r.t. bias
       # WHY sum and keepdims? Bias is per-neuron, shared across batch
       self.db = np.sum(dL_dZ, axis=1, keepdims=True) / batch_size
       
       # Gradient w.r.t. input (pass to previous layer)
       # WHY W.T? Reverse the forward transformation
       dL_dX = self.W.T @ dL_dZ
       
       return dL_dX

Worked Example: ReLU Activation Gradient Forward: $a = \text{ReLU}(z) = \max(0, z)$ Derivative: $\frac{\partial a}{\partial z} = \begin{cases} 1 & \text{if } z > 0 \\ 0 & \text{if } z \leq 0 \end{cases}$

class ReLU:
   def forward(self, Z):
       self.Z = Z  # Cache for backward
       return np.maximum(0, Z)
   def backward(self, dL_dA):
       # Element-wise: gradient flows only where input was positive
       # WHY? ReLU "gates" the gradient—negative inputs had no effect
       return dL_dA * (self.Z > 0)

If $z = -2.5$ , output is $0$ , gradient is $0$ (dead neuron). If $z = 3.7$ , output is $3.7$ , gradient passes through unchanged.

3. Parameter Update: Learning

Why does this work? Taylor expansion around $\boldsymbol{\theta}_t$ :

$L(\boldsymbol{\theta}_t - \alpha \nabla L) \approx L(\boldsymbol{\theta}_t) - \alpha \|\nabla L\|^2$

The second term is negative, so loss decreases (for small enough $\alpha$ ).

Stochastic Gradient Descent (SGD) Derivation:

True gradient over dataset $\mathcal{D}$ : $\nabla_{\boldsymbol{\theta}} L = \frac{1}{|\mathcal{D}|} \sum_{(\mathbf{x}, y) \in \mathcal{D}} \nabla_{\boldsymbol{\theta}} L(\mathbf{x}, y)$

Problem: Computing this requires full pass over dataset (expensive!).

Solution: Estimate using a mini-batch $\mathcal{B} \subset \mathcal{D}$ : $\nabla_{\boldsymbol{\theta}} L \approx \frac{1}{|\mathcal{B}|} \sum_{(\mathbf{x}, y) \in \mathcal{B}} \nabla_{\boldsymbol{\theta}} L(\mathbf{x}, y)$

Why this works: By the law of large numbers, the sample average converges to the true expectation. Even small batches (32-256) give good estimates.

Worked example SGD with Momentum

Plain SGD is noisy. Momentum adds an exponential moving average of past gradients:

class SGDMomentum:
   def __init__(self, learning_rate=0.01, momentum=0.9):
       self.lr = learning_rate
       self.beta = momentum
       self.velocity = {}  # Per-parameter velocity
   def update(self, layer):
       # Initialize velocity for this layer if first update
       if id(layer) not in self.velocity:
           self.velocity[id(layer)] = {
               'W': np.zeros_like(layer.W),
               'b': np.zeros_like(layer.b)
           }
       
       v = self.velocity[id(layer)]
       
       # Momentum update: accumulate past direction
       # WHY beta? Smooths noisy gradients, accelerates in consistent directions
       v['W'] = self.beta * v['W'] + (1 - self.beta) * layer.dW
       v['b'] = self.beta * v['b'] + (1 - self.beta) * layer.db
       
       # Update parameters
       layer.W -= self.lr * v['W']
       layer.b -= self.lr * v['b']

Why does momentum help?

Consistent gradient direction → velocity accumulates → faster progress
Oscillating gradients → velocity cancels out → smother path
Mathematically: acts like a ball rolling downhill with inertia

Common Pitfalls

Common mistake Mistake 2: Shape Mismatches in Backprop

Wrong approach: Just transpose until shapes align, without understanding why.

Why it feels right: "Matrix dimensions need to match, so I'll try different combinations."

The problem: You might get a shape that compiles but computes the wrong gradient. Example: dL_dW = X @ dL_dZ.T has correct shape but wrong semantics.

The rule: Gradient shape must match parameter shape. Work from the chain rule:

If forward was $Y = WX$ , backward is $\frac{\partial L}{\partial W} = \frac{\partial L}{\partial Y} X^T$
Derive it symbolically first, then implement Fix: Draw the computation graph with dimensions. Verify:

assert dW.shape == W.shape, f"Gradient shape {dW.shape} != weight shape {W.shape}"

Common mistake Mistake 3: Forgetting Activation Gradients

Wrong approach:

delta = prev_layer_delta  # Pass gradient through as-is

Why it feels right: "The gradient flows backwards through the network."

The problem: Activations are part of the computation graph! Their derivatives matter.

For $a = \sigma(z)$ , the gradient is: $\frac{\partial L}{\partial z} = \frac{\partial L}{\partial a} \cdot \frac{\partial a}{\partial z} = \frac{\partial L}{\partial a} \cdot \sigma'(z)$

Example where this fails: Sigmoid saturation. If $z = 10$ , then $\sigma(10) \approx 1$ and $\sigma'(10) \approx 0$ . The gradient vanishes—you must account for this!

Fix: Always multiply by activation derivative:

delta = prev_delta * activation.derivative(self.Z)

Complete End-to-End Example

Worked example Training a 2-Layer Network from Scratch

Problem: Binary classification, 2D input, learn XOR function.

import numpy as np
# Data: XOR problem (not linearly separable!)
X_train = np.array([[0, 0], [0, 1], [1, 0], [1, 1]).T  # Shape: (2, 4)
y_train = np.array([[0, 1, 0]])  # Shape: (1, 4)
 
# Architecture: 2 -> 4 -> 1 with ReLU and Sigmoid
layer1 = DenseLayer(2, 4)
relu = ReLU()
layer2 = DenseLayer(4, 1)
def sigmoid(z):
   return 1 / (1 + np.exp(-z))
 
def sigmoid_derivative(a):
   return a * (1 - a)
# Binary cross-entropy loss
def compute_loss(ypred, y_true):
   # WHY clip? Avoid log(0) which is -inf
   y_pred = np.clip(y_pred, 1e-7, 1 - 1e-7)
   return -np.mean(y_true * np.log(y_pred) + (1 - y_true) * np.log(1 - y_pred))
optimizer = SGDMomentum(learning_rate=0.1, momentum=0.9)
 
# Training loop
for epoch in range(1000):
   # === FORWARD PASS ===
   z1 = layer1.forward(X_train)
   a1 = relu.forward(z1)
   z2 = layer2.forward(a1)
   y_pred = sigmoid(z2)
   # Compute loss
   loss = compute_loss(y_pred, y_train)
   
   # === BACKWARD PASS ===
   # Start with loss gradient
   # WHY this form? Derivative of binary cross-entropy w.r.t. output
   dL_dy = (y_pred - y_train) / (y_pred * (1 - y_pred) + 1e-7)
   
   # Backprop through sigmoid
   dL_dz2 = dL_dy * sigmoid_derivative(y_pred)
   # Backprop through layer 2
   dL_da1 = layer2.backward(dL_dz2)
   
   # Backprop through ReLU
   dL_dz1 = relu.backward(dL_da1)
   
   # Backprop through layer 1
   dL_dx = layer1.backward(dL_dz1)
   # === UPDATE ===
   optimizer.update(layer1)
   optimizer.update(layer2)
   
   if epoch % 1000 == 0:
       print(f"Epoch {epoch}, Loss: {loss:.4f}")
# Test
predictions = (y_pred > 0.5).astype(int)
print(f"Predictions: {predictions}")
print(f"True labels: {y_train}")

Why this example?

XOR isn't linearly separable—requires hidden layer (demonstrates need for depth)
Small enough to compute by hand for verification
Shows all components: forward, backward, update, loss computation Step-by-step for one sample $[1, 0]$ :
Forward: $\mathbf{z}^{(1)} = \mathbf{W}^{(1)}[1, 0]^T + \mathbf{b}^{(1)}$ (4 values)
ReLU: $\mathbf{a}^{(1)} = \max(0, \mathbf{z}^{(1)})$ (some neurons activate)
Forward: $z^{(2)} = \mathbf{W}^{(2)}\mathbf{a}^{(1)} + b^{(2)}$ (single value)
Sigmoid: $\hat{y} = \sigma(z^{(2)})$ (probability)
Loss: $L = -[1 \cdot \log(\hat{y}) + 0 \cdot \log(1-\hat{y})] = -log(\hat{y})$
Backward: Compute $\frac{\partial L}{\partial \mathbf{W}^{(2)}}$ , $\frac{\partial L}{\partial \mathbf{W}^{(1)}}$ , etc.
Update: $\mathbf{W} \leftarrow \mathbf{W} - \alpha \nabla_{\mathbf{W}} L$

Key Implementation Patterns

Initialization Strategy

Why it matters: Bad initialization → vanishing/exploding gradients → network won't train.

Xavier/Glorot Initialization: $W_{ij} \sim \mathcal{N}\left(0, \frac{1}{n_{in}}\right)$

Derivation: For linear layer with $n$ inputs, if have variance1: $\text{Var}(z) = \text{Var}\left(\sum_i Wi x_i\right) = \sum_i \text{Var}(W_i) \text{Var}(x_i) = n \cdot \text{Var}(W)$

To keep $\text{Var}(z) = 1$ , we need $\text{Var}(W) = 1/n$ .

He Initialization (for ReLU): $W_{ij} \sim \mathcal{N}\left(0, \frac{2}{n_{in}}\right)$

Why different? ReLU kills half the neurons, so we need 2× the variance to compensate.

Numerical Stability Tricks

# BAD: Direct softmax (can overflow)
def softmax_bad(z):
    return np.exp(z) / np.sum(np.exp(z))
 
# GOOD: Subtract max for stability
def softmax_stable(z):
    # WHY subtract max? exp(z - max) is bounded by exp(0) = 1
    # This prevents overflow while keeping ratios identical
    exp_z = np.exp(z - np.max(z, axis=0, keepdims=True))
    return exp_z / np.sum(exp_z, axis=0, keepdims=True)

Why this works: $\frac{e^{z_i}}{\sum_j e^{z_j}} = \frac{e^{z_i - c}}{\sum_j e^{z_j - c}}$ for any constant $c$ . Choosing $c = \max(z)$ prevents overflow.

Recall Explain to a 12-Year-Old

Imagine you're building a robot that learns to recognize cats. When someone shows you how to build a fancy robot kit from the store, you can use it easily—but you don't understand why it works.

Building a neural network from scratch is like building the robot from individual parts: motors, sensors, wires. You have to:

Wire it correctly (forward pass): Connect sensors → brain → wheels so information flows
Learn from mistakes (backward pass): When the robot messes up, figure out which wires sent bad signals and adjust them
Get better gradually (update): Make tiny adjustments each time so you don't overcorrect

The "from scratch" part means you're using basic tools (like simple math operations) instead of pre-built kits. It's harder, but now you understand exactly how information flows through the robot's brain, how it learns from errors, and how to fix it when something breaks.

The key insight: Learning is just finding patterns by trying things, measuring how wrong you are, and adjusting the "knobs" (weights) a tiny bit in the direction that reduces wrongness. Do this thousands of times, and magically, it learns!

Connections

6.1.01-Neural-Network-Fundamentals - Foundation for what we're implementing
6.2.03-Backpropagation-Algorithm - Mathematical details of gradient computation
6.3.04-Optimization-Algorithms - Advanced update rules beyond SGD
4.2.02-Gradient-Descent - Core optimization concept
6.4.01-PyTorch-Basics - How frameworks abstract these details
6.5.01-Reading-Research-Papers - Understanding novel architectures to implement
2.3.01-Matrix-Operations - Linear algebra foundation for all operations
3.1.04-Chain-Rule - Core calculus for backpropagation

#flashcards/ai-ml

What is the primary purpose of implementing models from scratch? :: To understand the computational machinery—gradient flow, memory layout, numerical stability—that frameworks abstract away, enabling better debugging, optimization, and innovation.

What are the three core components of from-scratch implementation?

1) Forward pass (computing outputs from inputs), 2) Backward pass (computing gradients via backpropagation), 3) Parameter update (applying gradients to learn).

Why must we cache forward pass values?

Because backward pass requires intermediate activations and pre-activation values to compute gradients via the chain rule; recomputing them is inefficient and can be numerically unstable.

What is the formula for a dense layer's forward pass?

\mathbf{z} = \mathbf{W}\mathbf{x} + \mathbf{b}

, where

\mathbf{W}

rotates/scales input space and

\mathbf{b}

translates the decision boundary.

Why do we need activation functions?

Without non-linear activations, stacking layers only produces another linear function:

\mathbf{W}_2(\mathbf{W}_1\mathbf{x}) = \mathbf{W}'\mathbf{x}

. Non-linearity enables universal function approximation.

Derive the gradient of a dense layer with respect to its weights.

From

\mathbf{Z} = \mathbf{W}\mathbf{X} + \mathbf{b}

and chain rule:

\frac{\partial L}{\partial \mathbf{W}} = \frac{\partial L}{\partial \mathbf{Z}} \frac{\partial \mathbf{Z}}{\partial \mathbf{W}} = \delta \mathbf{X}^T

, where

\delta = \frac{\partial L}{\partial \mathbf{Z}}

Why do we transpose X when computing weight gradients?

Dimensionality matching:

\delta

(m \times n_{batch})

\mathbf{X}^T

(n_{batch} \times n)

, producing

(m \times n)

which matches

\mathbf{W}

's shape.

What is the gradient of a dense layer with respect to its input?

\frac{\partial L}{\partial \mathbf{X}} = \mathbf{W}^T \delta

, where

\mathbf{W}^T

"reverses" the forward transformation to flow gradients backward.

Why do we sum gradients across the batch for bias?

Bias is shared across all samples in a batch, so we accumulate gradient contributions from each sample to update the single shared bias parameter.

What is the derivative of ReLU activation?

\frac{\partial}{\partial z}\text{ReLU}(z) = \begin{cases} 1 & \text{if } z > 0 \\ 0 & \text{if } z \leq 0 \end{cases}

— gradient flows through only where input was positive.

Derive the gradient descent update rule.

From Taylor expansion

L(\theta - \alpha \nabla L) \approx L(\theta) - \alpha \|\nabla L\|^2

, the second term is negative, so moving in direction

-\nabla L

decreases loss. Update:

\theta_{t+1} = \theta_t - \alpha \nabla_\theta L

What is mini-batch SGD and why does it work?

Approximating the full gradient

\nabla L = \frac{1}{|\mathcal{D}|}\sum \nabla L_i

with a small batch

\nabla L \approx \frac{1}{|\mathcal{B}|}\sum_{i \in \mathcal{B}} \nabla L_i

. Works by law of large numbers: sample average converges to true expectation.

What is momentum in optimization and how does it help?

Momentum accumulates exponential moving average of gradients:

v_t = \beta v_{t-1} + (1-\beta)\nabla L

. Smooths noisy gradients and accelerates in consistent directions by adding inertia.

What is Xavier/Glorot initialization and why?

W \sim \mathcal{N}(0, 1/n_{in})

. Keeps variance stable:

\text{Var}(\sum W_i x_i) = n \cdot \text{Var}(W)

, so setting

\text{Var}(W) = 1/n

maintains

\text{Var}(z) = 1

Why is He initialization different from Xavier?

He initialization uses

\mathcal{N}(0, 2/n_{in})

because ReLU kills half the neurons (sets them to 0), so we need 2× the variance to compensate for the reduced activations.

Why subtract max in softmax computation?

Numerical stability:

\frac{e^{z_i}}{\sum e^{z_j}} = \frac{e^{z_i-c}}{\sum e^{z_j-c}}

for any

c

. Choosing

c=\max(z)

bounds

e^{z-c} \leq 1

, preventing overflow while keeping ratios identical.

What is the complete backpropagation equation through a layer with activation?

\delta_{\ell-1} = (\mathbf{W}^{(\ell)})^T \delta_{\ell} \odot \sigma'(\mathbf{Z}^{(\ell-1)})

, where

\odot

is element-wise multiplication and

\sigma'

is activation derivative.

Why can't we just transpose randomly until shapes match in backprop?

Correct shape doesn't guarantee correct gradient. Must derive from chain rule symbolically first: wrong transposes compute mathematically invalid gradients that seem to work dimensionally.

What happens if we forget to multiply by activation derivatives?

Gradients don't account for activation's effect on information flow. Example: saturated sigmoid (

z=10

\sigma'(10) \approx 0

) causes vanishing gradients—ignoring

\sigma'

mises this critical signal death.

Why does XOR require a hidden layer?

XOR is not linearly separable—no single line can separate the classes. A hidden layer with non-linear activation creates a new feature space where the classes become separable.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Chalo isko simple tareeke se samajhte hain. Jab aap PyTorch ya TensorFlow mein torch.nn.Linear import karke model banate ho, toh sab kuch ek black box ki tarah kaam karta hai — andar kya ho raha hai, gradients kaise flow kar rahe hain, ye sab hidden rehta hai. From-scratch implementation ka core intuition yahi hai ki aap khud se, numpy jaise basic tools se, ek layer ko banao taaki aap samajh sako ki wo layer actually kya compute kar rahi hai aur kyun kaam karti hai. Jaise cooking sikhna — sirf recipe follow karna alag baat hai, par ye samajhna ki heat se protein kaise change hote hain, wahi asli mastery hai.

Ismein do main cheezein hain: forward pass aur backward pass. Forward pass mein aap input $\mathbf{x}$ ko ek linear transformation $\mathbf{z} = \mathbf{W}\mathbf{x} + \mathbf{b}$ se guzarte ho — jahan $\mathbf{W}$ input space ko rotate aur scale karta hai, aur $\mathbf{b}$ decision boundary ko shift karta hai. Phir ek activation function $\sigma$ lagate ho jo non-linearity add karta hai. Ye non-linearity bahut zaroori hai, kyunki iske bina chahe aap kitni bhi layers stack kar lo, wo sab milke bas ek single linear function ban jaayengi — matlab network kuch complex seekh hi nahi paayega. Backward pass mein aap chain rule use karke nikaalte ho ki har parameter loss ko kitna affect karta hai, aur ye output se input ki taraf ulta chalta hai.

Ye cheez matter kyun karti hai? Kyunki jab aapka model fail hota hai ya galat predictions deta hai, toh debugging tabhi possible hai jab aapko andar ki machinery samajh aati ho. Real research aur practice mein aksar aapko naye ideas try karne padte hain jo ready-made libraries mein hote hi nahi — tab ye foundational knowledge aapko innovate karne ki power deti hai. Ek chhoti si detail bhi note karo: implementation mein hum input $\mathbf{X}$ ko cache karte hain forward pass mein, kyunki backward pass mein gradients compute karne ke liye wahi input chahiye hota hai. Ye chhoti-chhoti practical insights aapko sirf tabhi milti hain jab aap khud haath ganda karke code likhte ho.

Core Motivation

The Three Pillars of From-Scratch Implementation

1. Forward Pass: Computing Outputs

2. Backward Pass: Computing Gradients

3. Parameter Update: Learning

Common Pitfalls

Complete End-to-End Example

Key Implementation Patterns

Initialization Strategy

Numerical Stability Tricks

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Research Frontiers & Practice

Connections