Backpropagation as a Computation Graph: A Two-Layer MLP by Hand
AI Learning Project
Backpropagation as a Computation Graph: A Two-Layer MLP by Hand
Trace a two-layer MLP through local gradients, ReLU, softmax cross-entropy, and a reproducible NumPy backward pass.
AI Learning Project

Backpropagation as a Computation Graph: A Two-Layer MLP by Hand

Reading info

Level: Intermediate Reading time: 14 min
  • Backpropagation
  • Computation Graph
  • Softmax
Open knowledge map

Backpropagation is often introduced as the mysterious engine powering deep learning, but at its core, it is simply reverse-mode automatic differentiation applied to a computation graph. It is not a neural-network-specific trick; rather, it is a disciplined, highly optimized way to move local gradients backward through program operations using the chain rule of calculus.

While forward propagation computes the model’s prediction by traversing the graph from inputs to the loss, backpropagation traverses the graph in reverse. In this article, we rigorously work through a two-layer Multi-Layer Perceptron (MLP): x -> W1x+b1 -> ReLU -> W2h+b2 -> softmax cross-entropy, expanding on both the mathematical theory and the engineering realities of implementing it from scratch.

1. The Computation Graph Unveiled

To compute derivatives systematically, we decompose the complex neural network into a Directed Acyclic Graph (DAG) of primitive operations. Each node represents a simple mathematical operation (like matrix multiplication or ReLU), and edges represent the flow of tensors. Crucially, each node must be capable of doing two things: calculating its forward output, and calculating its local Vector-Jacobian Product (VJP) during the backward pass.

graph TD
    x["Input x"] --> z1["z1 = x @ W1 + b1"]
    W1["Weights W1"] --> z1
    b1["Bias b1"] --> z1
    z1 --> h["h = ReLU(z1)"]
    h --> logits["logits = h @ W2 + b2"]
    W2["Weights W2"] --> logits
    b2["Bias b2"] --> logits
    logits --> p["p = Softmax(logits)"]
    p --> L["Loss = CrossEntropy(p, target)"]
    target["Target y"] --> L

    classDef fwd fill:#e1f5fe,stroke:#039be5,stroke-width:2px;
    classDef param fill:#fce4ec,stroke:#d81b60,stroke-width:2px;
    class z1,h,logits,p,L fwd;
    class W1,b1,W2,b2 param;

The forward pass must cache intermediate values (like x, z1, and h) because they are required to compute the local gradients during the backward pass. This is why training neural networks is heavily memory-bound compared to inference.

Two-layer MLP computation graph and backward path
A computation graph turns a network into small differentiable steps, calculating exact gradients without the numerical approximation errors of finite differences.

2. The Softmax Cross-Entropy Shortcut

In theory, you could calculate the Jacobian of the Cross-Entropy loss with respect to the Softmax probabilities, and then multiply that by the Jacobian of the Softmax with respect to the logits. In practice, doing this explicitly is a recipe for numerical disaster and wasted compute.

When you combine Softmax and Cross-Entropy, the mathematical terms elegantly cancel out, resulting in a beautifully simple gradient with respect to the logits:

dL/dlogits = p - one_hot(target)

For example, if the target is class 1 and the model predicts probabilities [0.1, 0.7, 0.2], the gradient is simply [0.1, 0.7 - 1.0, 0.2] = [0.1, -0.3, 0.2]. The negative sign correctly pushes the correct logit higher, while the positive signs penalize the incorrect logits. This algebraic simplification is why production frameworks like PyTorch fuse these operations into CrossEntropyLoss.

3. Mathematical Derivations of the MLP

Once we have the gradient of the loss with respect to the logits (dlogits), we propagate it backward using the chain rule. Notice that we never instantiate full Jacobian matrices; instead, we compute Vector-Jacobian Products efficiently using matrix transposes.

# Layer 2 gradients
dW2 = h^T dlogits
db2 = sum(dlogits, axis=0)
dh  = dlogits W2^T

# Layer 1 gradients
dz1 = dh * ReLU'(z1)  # Element-wise multiplication
dW1 = x^T dz1
db1 = sum(dz1, axis=0)

By checking the norm of these gradients (e.g., norm_dW1=0.999823, norm_dW2=0.993682), we can verify that the network isn’t suffering from vanishing or exploding gradients.

4. Real-World Numpy Implementation

Translating the math into executable code reveals how frameworks actually operate under the hood. Here is a batched, robust implementation using NumPy:

import numpy as np

def relu(x): 
    return np.maximum(0, x)
def relu_backward(dout, cache_x): 
    return dout * (cache_x > 0).astype(float)

def softmax(x):
    # Subtract max for numerical stability
    exps = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return exps / np.sum(exps, axis=-1, keepdims=True)

# 1. Forward Pass
x = np.random.randn(32, 10)  # Batch size 32, features 10
W1 = np.random.randn(10, 64) * 0.1
b1 = np.zeros((1, 64))
W2 = np.random.randn(64, 5) * 0.1
b2 = np.zeros((1, 5))
targets = np.random.randint(0, 5, size=(32,))

z1 = x @ W1 + b1
h = relu(z1)
logits = h @ W2 + b2
probs = softmax(logits)

# 2. Backward Pass
# Softmax-CE gradient
batch_size = x.shape[0]
dlogits = probs.copy()
dlogits[np.arange(batch_size), targets] -= 1
dlogits /= batch_size  # Average over batch

# Layer 2
dW2 = h.T @ dlogits
db2 = np.sum(dlogits, axis=0, keepdims=True)
dh = dlogits @ W2.T

# Layer 1
dz1 = relu_backward(dh, z1)
dW1 = x.T @ dz1
db1 = np.sum(dz1, axis=0, keepdims=True)

print(f"Gradient norms: dW1={np.linalg.norm(dW1):.4f}, dW2={np.linalg.norm(dW2):.4f}")

5. Personal Experience / Engineer’s Perspective

From years of writing and debugging custom CUDA kernels and deep learning architectures, here are my takeaways on backpropagation in the wild:

The Memory Wall: Beginners often think training is slow because of math, but it’s actually constrained by memory bandwidth. During the forward pass, we have to stash the activations (like z1 and h) in HBM because the backward pass needs them. This is why techniques like Gradient Checkpointing (recomputing activations on the fly) exist—they trade FLOPs to save VRAM.

  • Broadcasting Bugs: In NumPy and PyTorch, silent broadcasting is a silent killer. If you compute db1 = dlogits without summing over the batch axis, the tensor shapes might accidentally broadcast later, producing garbage gradients without throwing an error. Always use keepdims=True.
  • Gradient Checking: When writing a custom C++ or CUDA backward pass, your first step should always be writing a finite-difference gradient checker. Compare your analytical gradient against (f(x + h) - f(x - h)) / 2h. If they don’t match up to 1e-4, your backprop is wrong.
  • Numerical Stability: Never compute np.exp(logits) directly. Always subtract the maximum logit first. A logit of 1000 will overflow float32 instantly, resulting in NaN gradients that poison the entire network.

6. Visualizing the Flow

The animation starts at the loss and lights up the gradient path through logits, hidden activations, and the first-layer weights.

When watching the animation, do not only watch arrow direction. Observe which forward values each node must store and reuse during the backward pass. The next article studies how these gradients move parameters and why optimizers take different paths to the minima.

7. Backpropagation Verification Matrix

When reproducing this article, split the forward and backward pass into auditable stages. The table below turns “is the formula correct?” into visible evidence, so a decreasing loss is not mistaken for proof that every gradient is correct.

Stage Values to cache or inspect Common symptom when it fails
Forward cache x, z1, h, logits, and probs. The ReLU mask cannot be reconstructed, or gradient shapes are only “fixed” by broadcasting.
Softmax-CE dlogits = probs - one_hot(target), averaged over the batch. The loss moves but gradients are too large and training is overly sensitive to batch size.
Matrix gradients dW2 = h.T @ dlogits and dW1 = x.T @ dz1. The gradient shape does not match the parameter, or a transpose error silently prevents learning.
Numerical stability Subtract max before softmax and inspect NaN, Inf, and gradient norms. Loss becomes NaN early, or a layer’s gradient norm collapses to zero.

Internal links

Related tutorials

Review foundation Matrix Calculus for Neural Networks

Derive dL/dW for y = Wx + b and verify it with finite differences.

 Continue next Gradient Descent and Optimizer Geometry

Compare gradient descent, momentum, and Adam on a visible quadratic loss surface.

 Learning path AI Learning Path

Read AI foundations, machine learning, CNN work, traffic defense, and AI security in order.

 Resources Resource Library

Download companion code, datasets, diagrams, model files, and notes.

 Topic hub Student AI Projects

Connect handwritten digits, browser demos, and CNN projects into a portfolio path.

Search questions

FAQ

Who is this article for?

This article is for readers who want an intermediate-level guide to Backpropagation as a Computation Graph. It takes about 14 min and focuses on Backpropagation, Computation Graph, Softmax.

What should I read next?

The recommended next step is Gradient Descent and Optimizer Geometry, so the article connects into a longer learning route instead of ending as an isolated note.

Does this article include runnable code or companion resources?

Yes. Use the run notes, resource cards, and download links on the page to reproduce the example or inspect the companion files.

How does this article fit into the larger site?

It is connected to the article context block, learning routes, resources, and project timeline so readers can move from concept to implementation.

Article context

AI Learning Project

A practical route from AI concepts to machine learning workflow, evaluation, neural networks, Python practice, handwritten digits, a CIFAR-10 CNN, adversarial traffic-defense notes, and AI security.

Level: Intermediate Reading time: 14 min
  • Backpropagation
  • Computation Graph
  • Softmax

Your next step

Continue: Gradient Descent and Optimizer Geometry
Review the foundation Open resource
Other language version 反向传播计算图：两层 MLP 的前向、局部梯度和反向传播
Share summary Backpropagation as a Computation Graph

Trace local gradients through ReLU and softmax cross-entropy in a two-layer MLP.

Download share card Open share center

Companion resources

Leave a Reply

Project timeline

Published posts

  1. AI Basics Learning Roadmap Separate AI, machine learning, and deep learning before going into implementation details.
  2. Machine Learning Workflow Follow the practical path from data and features to training, prediction, and evaluation.
  3. Model Training and Evaluation Understand loss, overfitting, train/test splits, accuracy, recall, and F1.
  4. Neural Network Basics Move from perceptrons to activation, forward propagation, backpropagation, and training loops.
  5. Matrix Calculus for Neural Networks Derive dL/dW for y = Wx + b and verify it with finite differences.
  6. Backpropagation as a Computation Graph Trace local gradients through ReLU and softmax cross-entropy in a two-layer MLP.
  7. Gradient Descent and Optimizer Geometry Compare gradient descent, momentum, and Adam on a visible quadratic loss surface.
  8. Convolution and Receptive Field Math Compute convolution output size, receptive fields, channel mixing, and im2col layout.
  9. Transformer Attention Math Hand-calculate Q/K/V scores, softmax weights, masks, multi-head structure, and KV cache.
  10. Python AI Mini Practice Run a small scikit-learn classification task and read the experiment output.
  11. Handwritten Digit Dataset Basics Read train.csv, test.csv, labels, and the flattened 28 by 28 pixel layout before training the classifier.
  12. Handwritten Digit Softmax in C Follow the C implementation from logits and softmax probabilities to confusion matrices and submission export.
  13. Handwritten Digit Playground Notes See how the offline classifier was adapted into a browser demo with drawing input and probability output.
  14. CIFAR-10 Tiny CNN Tutorial in C Build and train a small convolutional neural network for CIFAR-10 image classification, then read its loss and accuracy output.
  15. High-Entropy Traffic Defense Notes Study encrypted metadata leaks, entropy, traffic classifiers, and a defensive Python chaffing prototype.
  16. AI Security Threat Modeling Build a defense map with NIST adversarial ML, MITRE ATLAS, and OWASP LLM risks.
  17. Adversarial Examples and Robust Evaluation Evaluate clean and perturbed accuracy with an FGSM-style digits experiment.
  18. Data Poisoning and Backdoor Defense Study poison rate, trigger behavior, attack success rate, and training pipeline controls.
  19. Model Privacy and Extraction Defense Measure membership inference signal and surrogate fidelity against a local toy model.
  20. LLM, RAG, and Agent Security Separate instructions from data and enforce tool permissions against indirect prompt injection.

Published resources

  1. Python AI practice code guide The article includes a runnable scikit-learn classification script.
  2. digit_softmax_classifier.c The C source for the handwritten digit softmax classifier.
  3. train.csv.zip Compressed handwritten digit training set with 42000 labeled samples.
  4. test.csv.zip Compressed handwritten digit test set with 28000 unlabeled samples.
  5. sample_submission.csv The official submission format example for checking the final output columns.
  6. submission.csv The prediction file generated by the current C project.
  7. digit-playground-model.json The compact softmax demo model and sample set used by the browser playground.
  8. digit-sample-grid.svg A small handwritten digit preview grid extracted from the training set.
  9. Handwritten digit project bundle Contains the source file, compressed datasets, submission files, browser model, and preview grid.
  10. cifar10_tiny_cnn.c source Single-file C tiny CNN with CIFAR-10 loading, convolution, pooling, softmax, and backpropagation.
  11. model_weights.bin sample weights Model weights generated by one local small-sample run.
  12. test_predictions.csv sample predictions Sample test prediction output from the CIFAR-10 tiny CNN.
  13. CNN project explanation PDF Companion explanation material for the CNN project.
  14. Virtual Mirror redacted code skeleton A redacted mld_chaffing_v2.py control-flow skeleton with secrets, node topology, and target lists removed.
  15. Virtual Mirror stress-test template A redacted CSV template for CPU, memory, peak threads, pulse rate, latency, and error measurements.
  16. Virtual Mirror classifier-evaluation template A CSV template for TP, FN, FP, TN, accuracy, precision, recall, F1, ROC-AUC, entropy, and JS divergence.
  17. Virtual Mirror resource notes Notes explaining why the public resources include only redacted code, test templates, and architecture context.
  18. AI Security Lab README Setup, safety boundaries, and quick-run commands for the AI Security series.
  19. AI Security Lab full bundle Includes safe toy scripts, result CSVs, risk register, attack-defense matrix, and architecture diagram.
  20. AI security risk register CSV risk register template for AI threat modeling and release review.
  21. AI attack-defense matrix Maps attack surface, toy demo, metric, and defensive control into one CSV table.
  22. AI Security Lab architecture diagram Shows threat modeling, robustness, data integrity, model privacy, and RAG guardrails.
  23. FGSM digits robustness script FGSM-style perturbation and accuracy-drop experiment for a local digits classifier.
  24. Data poisoning and backdoor toy script Demonstrates poison rate, trigger behavior, and attack success rate on digits.
  25. Model privacy and extraction toy script Outputs membership AUC, target accuracy, surrogate fidelity, and surrogate accuracy.
  26. RAG prompt injection guard toy script Uses a deterministic toy agent to demonstrate external-data demotion and tool-policy blocking.
  27. Deep Learning Math Lab README Setup commands, script entry points, generated outputs, and figure notes for the math series.
  28. Deep learning math full lab bundle Bundles NumPy scripts, CSV outputs, formula diagrams, loss contours, convolution figures, and attention heatmaps.
  29. Gradient check results CSV Stores MSE analytic gradients, finite-difference gradients, and error norms.
  30. Optimizer path CSV Step-by-step coordinates and loss for gradient descent, momentum, and Adam on a 2D quadratic.
  31. Attention weights CSV Scores, softmax weights, and context vectors for a three-token scaled dot-product attention example.
  32. Deep learning math figure set Includes matrix shapes, computation graphs, loss contours, convolution scans, and attention heatmaps.
  33. Deep learning math interactive visualizer Browser modules for gradient checking, optimizer paths, convolution output size, and attention heatmaps.
  34. Deep Learning topic share card A 1200x630 SVG card for sharing the Deep Learning / CNN topic hub.
  35. Machine Learning From Scratch share card A 1200x630 SVG card for the K-means, Iris, and ML workflow topic hub.
  36. Student AI Projects share card A 1200x630 SVG card for handwritten digits, C classifiers, and browser demos.
  37. CNN convolution scan animation An 8-second Remotion animation showing how a 3x3 convolution kernel scans an input and builds a feature map.

Current route

  1. AI Basics Learning Roadmap Learning path step
  2. Machine Learning Workflow Learning path step
  3. Model Training and Evaluation Learning path step
  4. Neural Network Basics Learning path step
  5. Matrix Calculus for Neural Networks Learning path step
  6. Backpropagation as a Computation Graph Learning path step
  7. Gradient Descent and Optimizer Geometry Learning path step
  8. Convolution and Receptive Field Math Learning path step
  9. Transformer Attention Math Learning path step
  10. LLM Visualizer Learning path step
  11. Python AI Mini Practice Learning path step
  12. Handwritten Digit Dataset Basics Learning path step
  13. Handwritten Digit Softmax in C Learning path step
  14. Handwritten Digit Playground Notes Learning path step
  15. CIFAR-10 Tiny CNN Tutorial in C Learning path step
  16. High-Entropy Traffic Defense Notes Learning path step
  17. AI Security Threat Modeling Learning path step
  18. Adversarial Examples and Robust Evaluation Learning path step
  19. Data Poisoning and Backdoor Defense Learning path step
  20. Model Privacy and Extraction Defense Learning path step
  21. LLM, RAG, and Agent Security Learning path step

Next notes

  1. Add more image-classification and error-analysis cases
  2. Turn common metrics into a quick reference
  3. Add more AI security defense experiment notes
Scroll down