Reverse Sweep Mechanics

## The graph, and what "reuse" means A training step builds a directed acyclic graph: nodes are tensors, edges are operations, and the single scalar loss `L` sits at the sink. **Reuse** means one node has more than one outgoing edge — the same activation feeds two layers, the same parameter matrix is used by two operations, the same embedding table is read twice. That is *fan-out*. Reverse-mode autodiff computes, for every node `v`, its **adjoint** `v_bar = dL/dv`. It walks the graph backwards from the loss, and at each operation it turns the adjoint of the output into adjoints of the inputs (a vector-Jacobian product). The question is what happens at a node with several outgoing edges. ## The chain rule says sum Suppose `x` feeds two consumers, `y1 = f(x)` and `y2 = g(x)`, and `L` depends on `x` only through them. Perturb `x` by a small `dx`. Both `y1` and `y2` move, and both movements change `L`. To first order the two effects add: ``` dL/dx = (dL/dy1)(dy1/dx) + (dL/dy2)(dy2/dx) ``` In vector form each term is `J_i^T * y_i_bar`, the transposed Jacobian of that operation applied to that consumer's adjoint. With `k` consumers you get `k` terms. This is the total derivative — it is not an approximation, and it is not a convention that could have been chosen differently. So the rule is the neat duality: **copy in the forward pass, add in the backward pass**. A broadcast, an explicit copy, using a variable twice in an expression, and a parameter shared between two layers are all the same node with fan-out, and all of them accumulate. ## Why a buffer, and why reverse topological order An implementation cannot know a node's full gradient the moment it processes the first consumer. So each node holds an accumulator initialised to zero; each consumer's backward step *adds into* it. The node's own backward step — pushing gradient further upstream — may only run once every consumer has contributed, otherwise it would propagate a partial sum. Processing nodes in reverse topological order is exactly the scheduling discipline that guarantees this. If you ever hand-write a backward pass, the two bugs are (a) assigning instead of adding, so the last consumer overwrites the others, and (b) propagating before the last consumer has reported. ## Sum, not average Averaging over consumers would be a different, wrong derivative. If you want one path to count for less, scale it in the loss — for example by weighting a head's loss term — and let summation do the rest. A consequence worth internalising: a shared trunk feeding three heads receives the sum of three gradients, so its parameters see a systematically larger update than a single-head model at the same learning rate. The head loss weights and the trunk's learning rate are coupled. ## The silent bug Because accumulation is invisible when it works, dropping a path is one of the quietest failures in deep learning. The canonical case is **weight tying**: one matrix serves as the input embedding (a row lookup for the tokens in the batch — a sparse gradient) and as the output projection (a dense matmul against every row — a dense gradient). The correct gradient is the sum of the sparse and dense contributions. An implementation that keeps two copies of the matrix and only updates one, or that severs one path, produces a model that trains, converges and looks fine on the loss curve — it is simply optimising a different objective from the one you wrote down. The same shape appears in multi-head models: three heads over one trunk, and all three upstream gradients must arrive before the trunk's parameters are updated *once*. Updating the trunk three times, once per head, is not equivalent — three sequential steps of size `lr` from three different points is a different trajectory from one step along the summed direction, and the difference grows with the learning rate. ## How to check Craft a loss that depends on only one path, run one backward pass, and inspect which parameters received a nonzero gradient. Then compare the shared parameter's gradient against the sum of the two single-path gradients. For small models, a finite-difference check on a shared parameter catches a missing path immediately, because the numerical derivative naturally includes every route from that parameter to the loss.