Maths Catch-Up · Summation Is What Reuse Costs: sum(0) or sum(1), Just Compare the Indices

The next stretch of work is to derive the backward pass of a two-layer network entirely by hand — twenty-odd steps, each one written out and then checked against loss.backward(). Before starting, the tools need to be in order.

There are two of them so far: part 3 gave the summation of the multivariable chain rule, and part 13 used it to derive the linear layer’s backward formula. But that was one special case. What actually snags is a different sort of line:

dbngain = (bnraw * dhpreact).sum(0, keepdim=True)
dcounts_sum_inv = (counts * dprobs).sum(1, keepdim=True)

Why is there a sum here at all? Why is one axis 0 and the other 1? And why do some places need a broadcast instead? Neither part 3 nor part 13 answers those — this one does.

A few things from part 13 are used here without being re-explained: the gradient always has the shape of the thing it is a gradient of (denominator layout), the linear-layer formula itself, and what a VJP is.

Where the addition comes from

The multivariable chain rule, written out in both part 3 and part 13:

∂L/∂xᵢ = Σⱼ (∂L/∂yⱼ) · (∂yⱼ/∂xᵢ)

Both parts state it as a fact: x reaches L along several paths, so add up the contributions of those paths. This part derives that plus sign first — because all four shapes below are what it turns into, and whether the summation survives or not is the one and only test for sum versus broadcast.

The smallest example

y₁ = 2a, y₂ = 3a, L = y₁² + y₂

Substituting directly: L = 4a² + 3a, so dL/da = 8a + 3.

Going path by path and adding:

via y₁：(∂L/∂y₁)·(∂y₁/∂a) = 2y₁ × 2 = 4a × 2 = 8a
via y₂：(∂L/∂y₂)·(∂y₂/∂a) = 1   × 3 =          3
                                       sum  = 8a + 3

The two agree. At a = 1.7 the numerical derivative is 16.600000, and 8×1.7+3 = 16.6 ✅

Why addition and not something else

When a changes by da, every yⱼ changes with it:

dy₁ = (∂y₁/∂a)·da　　dy₂ = (∂y₂/∂a)·da

And the change in L is caused by all of the yⱼ together:

dL = (∂L/∂y₁)·dy₁ + (∂L/∂y₂)·dy₂

Substitute the first two into the third, and the coefficient of da is dL/da — the chain rule.

The addition comes from “the total change in L is the sum of the changes each path causes”. That is the definition of the total differential, not a convention anyone chose.

The only thing that varies across the four shapes below is how many entries of ∂yⱼ/∂xᵢ are zero.

Compare the indices once and the four shapes sort themselves out

That table of partials has a name, and part 13 already drew one — the Jacobian. The one it drew was for C[X]: 36 cells with six 1s in them. What this part adds is that the other faces of that table come to four in total.

And working out which one applies takes no arithmetic. There is one test: compare the variable’s index set with the output’s.

The next four sections take one shape each. Section five is the important one — it answers the question in the title.

A, one-to-one: the Jacobian is diagonal

The mark of it is that the indices match on both sides of the equals sign:

y[b,c] = f( x[b,c] )

x[b,c] appears only in the definition of y[b,c]; no other cell involves it. Substituting into the general formula, 863 of the 864 terms under the summation are zero and one survives:

∂L/∂x[b,c] = (∂L/∂y[b,c]) · f′(x[b,c])

dx = f_prime(x) * dy       # shape unchanged

What the Jacobian looks like: non-zero on the diagonal only. “Elementwise” is the name for that diagonal matrix — multiplying a vector by a diagonal matrix degenerates into an elementwise product by definition.

Part 3 put this well when it was still about single-output networks: “a single-output network does not lack the summation; its summation simply has one term.” Shape A is the tensor version of the same sentence — the summation is still there, 863 terms of it eaten by zeros.

Instances

The last two take an output the forward pass already computed and use it as the local derivative rather than recomputing anything. It is a standard way of saving work in a backward pass, and it is the flip side of part 13’s “tanh’s derivative depends on the current value” — it depends on the current value, so the current value can simply be reused.

The reciprocal line is easy to write as a tautology

When differentiating counts_sum_inv = counts_sum ** -1, do not write S⁻¹ = S⁻¹ — that says nothing. The left-hand side is a different variable:

T = S⁻¹　⟹　dT/dS = −S⁻²

The power rule d(xⁿ)/dx = n·xⁿ⁻¹ with n = −1 gives −x⁻². Checking numerically: at S = 10.000, T = 0.100000; at S = 10.001, T = 0.099990; the ratio is −0.00001/0.001 = −0.01, and −1/10² = −0.01 ✅

The minus sign says something plain: a larger denominator makes a smaller reciprocal, so the direction flips.

B, one-to-many: sum that axis away

This section is the heart of the part. The mark of it is that the variable has one index fewer than the output:

y[b,c] = x[b,c] · a[b]

Fix attention on a[b]. Hold b fixed and write c out:

y[b,0] = x[b,0]·a[b]
y[b,1] = x[b,1]·a[b]
   …                     ← 27 equations, every one of them containing a[b]
y[b,26] = x[b,26]·a[b]

This is exactly the “one a, several y” structure from section two, with the number of paths gone from 2 to 27. The local derivative is the coefficient of a linear function:

∂y[b,c]/∂a[b] = x[b,c]

so:

∂L/∂a[b] = Σ_c (∂L/∂y[b,c]) · x[b,c]

da = (x * dy).sum(1, keepdim=True)     # (32,27) → (32,1)

Why the summation runs over c only, and not over b

Written strictly it should be a double summation over every component of y:

∂L/∂a[b] = Σ_b′ Σ_c  (∂L/∂y[b′,c]) · (∂y[b′,c]/∂a[b])

But when b′ ≠ b, the equation y[b′,c] = x[b′,c]·a[b′] contains no a[b] at all, so that partial is zero. Of the 32 × 27 = 864 terms only 27 are non-zero, and the double summation collapses into a single one.

Which summation survives depends on which partials are non-zero. The b axis collapsed entirely; the c axis survived entirely.

🔑 A “missing index” is not a differentiation rule — it is a piece of notation

Two separate questions are easy to fuse into one here:

The accurate way to put it is a comparison of two index sets:

A “missing index” is not itself a rule of differentiation; it is notation recording how many times this quantity was reused. Reused → it affects several things downstream → the chain rule has to add those paths up. Summation is what reuse costs.

Fixing the axis

In plain terms it is more telling still: γ[j] is shared by all 32 samples in the batch, each of which has its own opinion, and sum(0) is what pools those 32 opinions into one number. The indices only write down on paper who shares what.

The general form

The position of the summation never changes; only the local derivative does:

keepdim=True cannot be dropped. Part 8 covered what it does (whether the summed axis stays, with length 1); in a backward pass it is a hard requirement — without it (1,64) collapses to (64,) and the gradient no longer has the shape of the thing it is a gradient of.

⚠️ Worse, getting the axis wrong does not raise an error. Part 8 covered this as the section’s biggest trap: the shapes often broadcast just fine, and PyTorch assumes that is what you meant.

C, many-to-one: broadcast it back

B’s mirror image. The mark of it is that the output has one index fewer than the variable — which is to say the forward line carries a summation:

y[b] = Σₖ x[b,k]

Fix attention on x[b,c]: it contributes to the single number y[b].

Why every coefficient is 1

Split the equation into the part containing x[b,2] and the part that does not:

y[b] = x[b,2] + ( x[b,0] + x[b,1] + x[b,3] + … )
       ↑has it   ↑does not — all constant as far as x[b,2] is concerned; call it C

That is the simplest possible linear function, y = x + C, so ∂y[b]/∂x[b,2] = 1 + 0 = 1.

Checking numerically: [1.0, 2.0, 6.0, 1.0] sums to 10.000; add 0.001 to the 6.0 and the sum is 10.001. The input rose by 0.001, the output by 0.001, a ratio of 1 ✅

Differentiating by x[b,0] instead splits the same way — all 27 terms come apart into y = x + C, so all 27 partials are 1. The underlying reason: summation is linear, and every term carries a weight of exactly 1, with no coefficient in front.

Which immediately gives the version that does have a coefficient:

μ[j] = (1/n)·Σ_b x[b,j]　⟹　∂μ[j]/∂x[b,j] = 1/n

dhprebn += (1.0/n) * torch.ones_like(hprebn) * dbnmeani     # ← no longer 1

And the result is a broadcast

∂L/∂x[b,c] = 1 · (∂L/∂y[b])

The right-hand side has no index c — all 27 positions in that row receive the same gradient, which is a broadcast:

dx = torch.ones_like(x) * dy           # (32,1) → (32,27)
# with all coefficients equal to 1 this is just dx = dy; broadcasting expands it
# writing ones_like out only makes the local derivative visible

B and C are strict inverses

Know one and the other comes free.

Note that the same index c going missing produces opposite operations depending on where it goes missing: missing from the variable means sum, missing from the output means broadcast.

D, many-to-many: part 13’s formula, this time derived

Part 13 gave the conclusion:

z = xW   ⟹   ∂L/∂x = ∂L/∂z · Wᵀ        ∂L/∂W = xᵀ · ∂L/∂z

and a genuinely useful way to get there — match the shapes. With dy at (32,27) and W at (64,27), producing (32,64) leaves exactly one of four arrangements with matching inner dimensions.

dy @ W.T    (32,27) @ (27,64) → (32,64)   ✅ same shape as x
W @ dy      (64,27) @ (32,27)   inner 27 ≠ 32   ✗
dy.T @ W    (27,32) @ (64,27)   inner 32 ≠ 64   ✗
W.T @ dy    (27,64) @ (32,27)   inner 64 ≠ 32   ✗

What this section adds is the why — and the thing matching shapes cannot tell you.

The forward pass is y[b,j] = Σᵢ x[b,i]·W[i,j]. Two index sets: the variable x is [b,i] and the output y is [b,j]. Each has an axis the other lacks — which is shape D. The test says: contract the output’s own axis, keep the variable’s own.

Differentiating by x contracts j:

∂L/∂x[b,i] = Σⱼ (∂L/∂y[b,j]) · W[i,j]

Differentiating by W, the entry W[i,j] appears in the j-th output of all 32 samples, so it contracts b:

∂L/∂W[i,j] = Σ_b (∂L/∂y[b,j]) · x[b,i]

dx = dy @ W.T          # (32,27) @ (27,64) → (32,64)
dW = x.T @ dy          # (64,32) @ (32,27) → (64,27)

And what the transpose is doing is now clear: it moves the index being summed over into the position matrix multiplication expects. Σⱼ W[i,j]·dy[b,j] is “row i of W” dotted with “row b of dy”, while matrix multiplication wants a row times a column — transpose, and row i becomes column i. The transpose is not a trick; it is index alignment.

🔑 This is the whole difference between matching shapes and deriving from indices. Matching shapes tells you what to write; deriving tells you whose opinions this gradient is pooling — ∂L/∂x runs over j and pools the 27 outputs, while ∂L/∂W runs over b and pools the 32 samples, which is the same thing γ[j] was doing in the previous section.

It also explains why it is x that gets transposed in dW and not dy: the axis that has to be laid alongside for the inner product is the sample axis.

The same Σ, read a second way

One thing remains, and it is the easiest to miss: a variable that appears in more than one equation.

counts_sum = counts.sum(1, keepdims=True)    # counts used once
probs = counts * counts_sum_inv              # and again

Part 13 already settled the principle here: “this summation is the one and only source of every later ‘why does this get added’.” So what follows is not a fifth shape and not a new rule — it is the same Σ read a second way:

Two paths

Writing it out makes it plainest:

p[b,c] = e[b,c] · S[b]⁻¹　　S[b] = Σₖ e[b,k]

e[b,2] appears on the right of both equals signs. Nudge it and two paths fire at once:

e[b,2] moves by ε
  ├─→ p[b,2] moves with it directly              (short path, shape A)
  └─→ S[b] moves → S[b]⁻¹ moves → all 27 p in that row move   (long path, shape C)

∂L/∂e[b,c] = S[b]⁻¹·(∂L/∂p[b,c])　+　1·(∂L/∂S[b])

dcounts = counts_sum_inv * dprobs                   # first path
dcounts += torch.ones_like(counts) * dcounts_sum    # second path — the += is mandatory

❗ Why dropping the second path is so hard to catch

Walk sample 0’s four numbers through: e₀ = [1.0, 2.0, 6.0, 1.0], S₀ = 10.0, S₀⁻¹ = 0.10, and an upstream dp₀ = [0, 0, −0.556, 0]. First work out dS₀:

dS⁻¹ = Σ_c e·dp = 6.0 × (−0.556)        = −3.3333     ← shape B
dS   = (−S⁻²)·dS⁻¹ = (−0.01)×(−3.3333) = +0.0333     ← shape A

Then take the two paths cell by cell:

Look at the second row: the first path is zero and the second is not. Write = instead of += and those cells’ gradients go to zero outright — while the first path on its own is perfectly correct. The formula is right, nothing raises, and the result is a well-formed tensor. Part 13’s “writing = means losing the summation” is losing exactly this.

It also explains something in passing: why dcounts stops being sparse. The second path adds a non-zero value to every single cell, filling in the 832 zeros that were there before.

Finding every fork

Search the variable’s name through the whole set of forward equations. However many right-hand sides it appears on is how many paths it has.

Scan the forward pass before deriving anything and circle the forks; then no += gets left out while writing the code.

The test, end to end

Take one forward line and fix attention on the variable being differentiated
        ↓
① How many times does it appear across the whole set of equations?
   more than once → one pass per path, added together (+=)
        ↓
② For each path, compare the variable's index set with the output's
        ↓
   identical         → A: multiply by the local derivative, shape unchanged
   variable lacks one → B: sum that axis away, the shape collapses
   output lacks one   → C: broadcast along that axis, the shape expands
   each lacks one     → D: contract the output's own, keep the variable's own — a matmul
        ↓
③ Check the shapes: dX.shape must equal X.shape

Step three is part 13’s iron rule, and the cheapest check available — if the shape is wrong, one of the first two steps must be too.

This part introduced no new mechanism. It only pushed part 3’s summation all the way down: the addition comes from the total differential, the four shapes come from which partials are zero, and that in turn comes from how the two index sets differ.

The one-line version: summation is not a rule to memorise, it is what reuse costs. However many times a quantity was copied in the forward pass, that many gradients have to be added back in the backward one; and whichever axis it was copied along is the axis to sum.

🔮 “Let the indices decide the summation” has a formal name: the Einstein summation convention. It is what torch.einsum runs on, and its rule is every index absent from the output subscript gets summed over — exactly isomorphic to what this part derived. Multimodal work later on leans heavily on einsum, and the intuition built here transfers directly.

Next comes the twenty-odd steps, by hand. The tools are in order.