General Back-Propagation · Example MLP Training · Complications

§6.5.6 General back-propagation — one bprop rule per operation

The MLP algorithms were hand-specialized. Real libraries implement back-propagation generically: every operation carries its own local bprop method that knows how to convert a gradient on its output into a gradient on each of its inputs — a Jacobian-vector product. The back-propagation algorithm itself then knows no differentiation rules at all; it only calls each operation’s bprop with the right arguments and sums the results where paths converge.

Term by term

\texttt{op.bprop(inputs, } \boldsymbol{X} \texttt{, } \boldsymbol{G}) = \sum_i \big(\nabla_{\boldsymbol{X}}\, \texttt{op.f(inputs)}_i\big)\, G_i

$\boldsymbol{G}$	the gradient already computed on this operation’s OUTPUT	downstream grad
$\nabla_{\boldsymbol{X}} \texttt{op.f}$	the operation’s local Jacobian w.r.t. the input X we want — this is the only math the OPERATION must know	local Jacobian
$\texttt{bprop}$	returns the gradient on X. Example — matrix multiply C = AB: bprop for A returns GBᵀ, for B returns AᵀG. The algorithm just wires these together	the local rule

Aside — bprop must pretend inputs are distinct

A subtle rule: op.bprop should treat all its inputs as distinct even when they’re the same variable. If mul is handed two copies of $x$ to compute $x^2$ , its bprop returns $x$ as the derivative w.r.t. both inputs; the outer algorithm then sums them to the correct total $2x$ . This keeps each operation’s rule purely local.

The full driver (algorithm 6.5) prunes the graph to the relevant nodes and calls the recursive build_grad (algorithm 6.6) that fills each node’s gradient by summing over its consumers:

Algorithms 6.5 + 6.6 — general back-propagation (click lines)

1prune $G$ to $G'$ : nodes that are both ancestors of z and descendants of the targets $\mathbb{T}$ ⓘ
2grad_table[z] $\leftarrow 1$ ; for V in $\mathbb{T}$ : build_grad(V, …)ⓘ
3build_grad(V): if V in grad_table, return it (memoize)ⓘ
4for C in consumers(V): $G^{(i)} \leftarrow$ op(C).bprop(inputs(C), V, build_grad(C))ⓘ
5 $G \leftarrow \sum_i G^{(i)}$ ; grad_table[V] $\leftarrow G$ ; return Gⓘ

Cost

Term by term

O(\#\,\text{edges}) \;\;\approx\;\; O(n) \text{ for chain-structured nets}

$O(n^2)$	worst case: back-prop adds one Jacobian-vector product (O(1) nodes) per edge, and a DAG has ≤ O(n²) edges	general bound
$O(n)$	most neural nets are roughly chain-structured → O(n) — vastly better than the exponentially-many paths a naive expansion would traverse	in practice
$\approx \text{forward cost}$	backprop’s work is the SAME ORDER as the forward pass — the reason training is affordable at all	the punchline

§6.5.7 Example: back-propagation for MLP training

Put it together on a one-hidden-layer classifier: $\boldsymbol{H} = \max\{0, \boldsymbol{X}\boldsymbol{W}^{(1)}\}$ , unnormalized log-probs $\boldsymbol{H}\boldsymbol{W}^{(2)}$ , cross-entropy cost $J_{\mathrm{MLE}}$ , plus weight decay $\lambda(\sum W^{(1)2} + \sum W^{(2)2})$ . The cost’s computational graph (fig 6.11) has two backward paths into each weight matrix:

Fig 6.11 — the MLP-training cost graph; two paths reach each weight (cross-entropy + weight decay)

J = J_MLE + λ(‖W¹‖² + ‖W²‖²). Weight decay adds a short second path to each Wⁱ contributing 2λWⁱ.

Tracing the two paths

The weight-decay path is trivial: it always adds $2\lambda\boldsymbol{W}^{(i)}$ to weight $i$ ‘s gradient. The cross-entropy path is the real work: let $\boldsymbol{G}$ be the gradient on the log-probs $\boldsymbol{U}^{(2)}$ ; the shorter branch adds $\boldsymbol{H}^\top\boldsymbol{G}$ to $\nabla_{\boldsymbol{W}^{(2)}}$ ; the longer branch computes $\nabla_{\boldsymbol{H}} = \boldsymbol{G}\boldsymbol{W}^{(2)\top}$ , has the ReLU zero out gradient where its input was negative, then adds $\boldsymbol{X}^\top(\cdot)$ to $\nabla_{\boldsymbol{W}^{(1)}}$ . Cost is dominated by the matrix multiplies — $O(w)$ forward and backward alike; memory is $O(m\, n_h)$ to hold the hidden-layer pre-nonlinearity input.

§6.5.8 Complications

The textbook version is simpler than production code, which must also handle: operations with multiple outputs (e.g. returning both a max and its argmax in one pass); memory management (accumulate sums into a single buffer rather than materializing every tensor first); many data types (32-/64-bit float, integers); and genuinely undefined gradients (track and report them). None is insurmountable — this chapter gives the intellectual tools — but the subtleties are real.

Next: how the rest of the world computes derivatives (forward vs reverse mode), the Hessian, and the history — differentiation outside deep learning, and historical notes.

Check yourself — general backprop and MLP training

1.In general back-propagation, where does the differentiation knowledge live?

2.For matrix multiply C = AB with output gradient G, what does bprop return for A and for B?

3.Why must op.bprop pretend its inputs are distinct even when they're the same variable?

4.What is back-propagation's cost, and why does that make training feasible?

5.In the fig 6.11 MLP-training graph, why does each weight matrix have two backward paths?

5 questions