Example number
68
Example slug
example_68_linear_maps_forward_backward_as_matrix_products
Background

Matrix products are the computational substrate of deep learning: GPUs and TPUs are engineered around GEMM (general matrix multiply) kernels optimized for throughput and memory bandwidth. Forward passes in neural networks chain affine maps $Y = XW + b$, and backpropagation systematically applies transposes to flow gradients backward. The transpose patterns ($X^\top$ for weight gradients, $W^\top$ for input gradients) emerge from the chain rule and ensure shapes compose correctly across layers. Understanding these patterns eliminates memorization and enables confident implementation of custom layers, gradient checks, and mixed-precision training.

Purpose

Build intuition for forward and backward passes in neural networks as matrix products with systematic transpose patterns. You should come away comfortable that $Y = XW + b$ is a batch affine map, that gradients flow backward via transposes ($\frac{\partial L}{\partial W} = X^\top \frac{\partial L}{\partial Y}$, $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} W^\top$), and that these shapes and operations generalize to all linear layers in deep learning. Shape discipline and transpose reasoning make debugging and implementing layers systematic rather than error-prone.

Problem

Compute forward + backward passes for a linear layer Y=XW+b. Verify: dL/dW = X^T dL/dY and dL/dX = dL/dY W^T for loss L = sum(Y).

Solution (Math)

With $Y = XW + b$ (where $X \in \mathbb{R}^{n\times d_{\text{in}}}$, $W \in \mathbb{R}^{d_{\text{in}}\times d_{\text{out}}}$, $b \in \mathbb{R}^{d_{\text{out}}}$, $Y \in \mathbb{R}^{n\times d_{\text{out}}}$) and scalar loss $L(Y)$, reverse-mode (backprop) differentiation gives

\[ \frac{\partial L}{\partial W} = X^\top \frac{\partial L}{\partial Y},\qquad \frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y}\, W^\top,\qquad \frac{\partial L}{\partial b} = \mathbf{1}^\top \frac{\partial L}{\partial Y}, \]

where $\frac{\partial L}{\partial Y} \in \mathbb{R}^{n\times d_{\text{out}}}$ is the upstream gradient.

We use:

  • Data matrix $X\in\mathbb{R}^{n\times d}$ (rows are examples).
  • Vectors are column vectors by default.
  • $\|x\|_2$ is Euclidean norm; $\langle x,y\rangle=x^Ty$.
Solution (Python)

import numpy as np

X = np.array([[1.,0.],
              [0.,1.],
              [1.,1.]])
W = np.array([[2.,1.],
              [-1.,3.]])
b = np.array([0.5,-2.0])

Y = X@W + b
dL_dY = np.ones_like(Y)

dL_dW = X.T @ dL_dY
dL_dX = dL_dY @ W.T
dL_db = dL_dY.sum(axis=0)

print("Y:
", Y)
print("dL/dW:
", dL_dW)
print("dL/dX:
", dL_dX)
print("dL/db:", dL_db)
Code Introduction

This snippet implements a batch affine map and its reverse-mode (backprop) gradients. With $X \in \mathbb{R}^{3\times 2}$, $W \in \mathbb{R}^{2\times 2}$, and $b \in \mathbb{R}^2$, the forward pass computes $Y = XW + b$, where $b$ is broadcast across rows so $Y \in \mathbb{R}^{3\times 2}$. Setting dL_dY = np.ones_like(Y) corresponds to the upstream gradient of a loss $L$ that sums all outputs, i.e., $L = \sum_{i,j} Y_{ij}$.

The gradients follow directly from the chain rule applied to the affine layer. Writing $G = \frac{\partial L}{\partial Y}$, the standard formulas are \[ \frac{\partial L}{\partial W} = X^\top G,\quad \frac{\partial L}{\partial X} = G\, W^\top,\quad \frac{\partial L}{\partial b} = \sum_{i=1}^{n} G_{i,:}, \] so each parameter receives a sum of per-example contributions. Intuitively, for each row $y_i = x_i W + b$, the contribution to $\frac{\partial L}{\partial W}$ is $x_i^\top G_{i,:}$ (outer product), to $\frac{\partial L}{\partial b}$ is $G_{i,:}$, and to $\frac{\partial L}{\partial X}$ is $G_{i,:} W^\top$. With $G$ set to all ones, these reduce to simple row/column sums.

Shape notes and gotchas: $b$ should be shape $(2,)$ or $(1,2)$ to broadcast across $n=3$ rows; mismatched shapes produce silent broadcasting errors. Ensure $G$ has the same shape as $Y$; dL_dY.sum(axis=0) sums over the batch dimension to produce $\frac{\partial L}{\partial b} \in \mathbb{R}^2$. Verify dimensions: $X^\top G$ yields $(2\times 3)(3\times 2)\to(2\times 2)$, $G W^\top$ yields $(3\times 2)(2\times 2)\to(3\times 2)$.

This matches the affine (linear) layer backward pass used in deep learning frameworks: parameters accumulate gradients across examples, while input gradients propagate via the transpose of the weight matrix.

Numerical Implementation Details

Numerical and Shape Notes

  • Shapes first: Declare shapes (e.g., $X \in \mathbb{R}^{n imes d}$, $w \in \mathbb{R}^{d}$, $b \in \mathbb{R}^{n}$). Vectors are column by convention; keep row/column usage consistent.
  • Axis discipline: Be explicit with axis in reductions and normalizations. For attention-like ops, softmax over keys (row-wise) so rows sum to ≈1.
  • Broadcasting: Check that broadcasts are intended (e.g., (n,1) with (n,d)). Prefer reshape/expand-dims to make semantics clear.
  • Stability eps: Add $arepsilon$ for divisions/logs and $arepsilon I$ (jitter) for SPD solves; use log-sum-exp for softmax.
  • Masking preserves shape: Masks should broadcast to the score/activation tensor; verify masked outputs keep the same shape and zero out excluded entries.
  • Dtype choices: Use float64 for clarity in scripts; with mixed precision, keep reductions/factorizations in float32/float64 to avoid under/overflow.
  • Sanity checks: Print shapes and residuals (e.g., ||Ax-b||, reconstruction error, row-sum ≈ 1). Assert finiteness and expected monotonicity where applicable.

Numerical and Implementation Notes

  • Dtype & precision: Prefer float64 for clarity; if using mixed precision, keep reductions (norms, softmax sums, factorizations) in float32/float64. Avoid explicit inverses; use solve, lstsq, Cholesky/QR/SVD.
  • Shapes & broadcasting: Annotate shapes (e.g., $X \in \mathbb{R}^{n imes d}$); vectors are column by default. Verify axes for reductions (axis) and ensure broadcasts are intended.
  • Stability: Use log-sum-exp for softmax; add small diagonal $arepsilon I$ (jitter) for SPD solves; prefer QR/SVD for ill-conditioned least squares.
  • Conditioning: Inspect np.linalg.cond(A) when solutions look unstable; regularize (ridge) or rescale features to improve conditioning.
  • Reproducibility: Set NumPy seed for random data; print shapes and residuals (e.g., ||Ax-b||, reconstruction errors) and assert finiteness.
  • Complexity & memory: Matmul ~ $O(n^3)$ for factorizations, $O(n^2)$ for triangular solves/products. Prefer vectorization over Python loops; avoid materializing large intermediates.
  • Masking & indexing: Use boolean masks that broadcast to target shapes; for attention-like ops, add $-\infty$ before softmax or zero-out after, then verify rows sum to ~1.
  • Sanity checks: Compare against references (e.g., lstsq vs. solve), check orthogonality (U.T @ U ≈ I), PSD (x.T @ A @ x > 0), and residual norms within tolerance (~1e-12 for float64).
  • Construct small $X \in \mathbb{R}^{n\times d_{\text{in}}}$, $W \in \mathbb{R}^{d_{\text{in}}\times d_{\text{out}}}$, $b \in \mathbb{R}^{d_{\text{out}}}$ with integer or simple float values for manual verification.
  • Compute forward pass $Y = X @ W + b$; ensure $b$ broadcasts correctly across $n$ rows.
  • Set upstream gradient $G = \frac{\partial L}{\partial Y}$ to np.ones_like(Y) (corresponding to $L = \sum_{i,j} Y_{ij}$).
  • Compute weight gradient $\frac{\partial L}{\partial W} = X^\top @ G$ via X.T @ dL_dY; shape $(d_{\text{in}}\times d_{\text{out}})$.
  • Compute input gradient $\frac{\partial L}{\partial X} = G @ W^\top$ via dL_dY @ W.T; shape $(n\times d_{\text{in}})$.
  • Compute bias gradient $\frac{\partial L}{\partial b} = \sum_{i=1}^n G_{i,:}$ via dL_dY.sum(axis=0); shape $(d_{\text{out}},)$.
  • Verify shapes match expectations; inspect numerical values to ensure correctness.
What This Example Demonstrates

Pedagogical Significance

  • Learning goals: Build intuition for when and why this tool is used in ML, not just how to compute it.
  • ML-first framing: Tie the concept to a concrete task pattern (fit / project / decompose / solve / measure) to anchor understanding.
  • Shape discipline: Habitually annotating dimensions prevents silent bugs and reinforces linear map thinking.
  • Numerical habits: Prefer stable factorizations over inverses; check residuals and condition numbers to separate bugs from ill-conditioning.
  • Transfer: Reuse the same pattern across models (e.g., projection in PCA, orthogonalization in regressions, attention as weighted sums).
  • Assessment ideas: Quick checks: predict sensitivity from $\kappa(A)$, verify projection properties, or compare solver outputs within tolerance.

ML Examples and Patterns

  • Fit: Linear/logistic regression via least squares or softmax; regularization (ridge) improves conditioning and generalization.
  • Project: PCA/SVD for dimensionality reduction; orthogonal projections to subspaces for denoising and feature extraction.
  • Decompose: Eigen/SVD factorizations to expose structure (low rank, PSD) used in recommender systems, LSA, and spectral clustering.
  • Solve: Stable solves without inversion (Cholesky/QR/SVD; CG for SPD) for optimization steps and kernel methods.
  • Measure: Norms, angles, and condition number $\kappa(A)$ to diagnose sensitivity, stability, and training difficulty.
  • Forward pass: $Y = XW + b$ is a batched affine map with broadcast bias addition.
  • Backward pass: gradients follow transpose patterns systematically via the chain rule.
  • Shape discipline: $\frac{\partial L}{\partial W} = X^\top G$ has shape $(d_{\text{in}}\times d_{\text{out}})$, $\frac{\partial L}{\partial X} = G W^\top$ has shape $(n\times d_{\text{in}})$, where $G = \frac{\partial L}{\partial Y}$.
  • Bias gradients sum over examples: $\frac{\partial L}{\partial b} = \sum_i G_{i,:}$.
  • How transposes naturally arise from dimension matching in the chain rule.
  • Verification: set $G$ to all ones to check gradient computation with simple expectations.
Notes
  • Part 1: Forward pass. $Y = XW + b$ applies the same affine map to each row of $X$; bias $b$ is broadcast across the batch.
  • Part 2: Backward pass gradients. Weight gradient $\frac{\partial L}{\partial W} = X^\top G$ accumulates outer products $x_i^\top G_{i,:}$ over examples. Input gradient $\frac{\partial L}{\partial X} = G W^\top$ propagates errors backward via the transpose. Bias gradient sums $G$ rows.
  • Part 3: Numerical checks. With $G$ set to all ones, gradients reduce to sums of rows/columns; verify shapes and inspect values manually on small examples.
  • Why This Matters for ML. Every feedforward layer, every linear projection, every attention block uses this forward/backward pattern. Mastering transposes and shapes is non-negotiable for implementing custom layers, debugging gradients, and reasoning about backprop.
  • ML Examples and Patterns. Fit: linear layers in neural networks; Project: linear projections and embeddings; Decompose: factorized weight matrices (low-rank, sparse); Solve: gradient-based optimization flows gradients backward via transposes.
  • Connection to Linear Algebra Theory. The transpose in $\frac{\partial L}{\partial X} = G W^\top$ is the adjoint operator; backprop is dual to forward propagation. The chain rule for compositions $(f \circ g)'(x) = f'(g(x)) g'(x)$ translates to matrix products with transposes.
  • Numerical and Implementation Notes. Always verify shapes before and after matrix multiplies; mismatched dimensions produce silent broadcasting errors or crashes. Use .T for transposes; ensure $b$ has compatible shape for broadcasting.
  • Numerical and Shape Notes. Annotate $X \in \mathbb{R}^{n\times d_{\text{in}}}$, $W \in \mathbb{R}^{d_{\text{in}}\times d_{\text{out}}}$, $Y \in \mathbb{R}^{n\times d_{\text{out}}}$, $G \in \mathbb{R}^{n\times d_{\text{out}}}$. Check that $X^\top G$ yields $(d_{\text{in}}\times n)(n\times d_{\text{out}}) \to (d_{\text{in}}\times d_{\text{out}})$.
  • Pedagogical Significance. This is the canonical minimal example for forward/backward in deep learning; mastering it generalizes to convolutions, attention, and all differentiable layers.
History and Applications

Backpropagation via systematic transpose patterns traces to Rumelhart, Hinton, and Williams (1986), though the chain rule and reverse-mode automatic differentiation have deeper roots in optimal control (adjoint methods, 1960s–70s). Modern deep learning frameworks (PyTorch, TensorFlow, JAX) automate gradient computation but rely on the same $X^\top G$ and $G W^\top$ formulas for linear layers. Understanding forward/backward passes as matrix products with transposes generalizes to all differentiable layers—convolutions, attention, recurrent cells—and underpins custom operator implementation, mixed-precision training, and gradient checkpointing for memory efficiency.

Connection to Broader Examples
  • Links to Chapter 4 (linear maps): affine layers implement linear transformations plus translation.
  • Connects to Chapter 16 (matrix products): forward/backward passes are chained matrix multiplies with systematic transposes.
  • Relates to Chapter 6 (projections): gradients project errors backward through weight matrices.
  • Bridges to Chapter 14 (conditioning): ill-conditioned $W$ amplifies gradient noise; initialization and normalization mitigate this.
  • Complements Chapter 12 (least squares): linear layers and least-squares solvers both rely on $X^\top X$ and $X^\top y$ patterns.
  • Reinforces shape discipline from early chapters: annotate dimensions rigorously to avoid silent broadcasting bugs.

Comments