Katkat
Back

Hand-Rolled Backprop: Running Product and Stable Cross-Entropy

OpenAIMachine Learning Engineer · Onsite
HardMachine Learning EngineerOnsiteMachine Learning

Last updated

Derive and implement, without an autodiff library, the backward pass of a running (cumulative) product in linear time without division, then a numerically stable softmax cross-entropy and its gradient. Tests chain-rule fluency, prefix/suffix tricks, and gradient checking.

You are writing a tiny autodiff library for a compounding-returns model, where period multipliers are combined with a cumulative product. No PyTorch, JAX, or autograd: you write each backward pass by hand. Plain Python lists are fine (NumPy is allowed if you prefer, but explain the vectorized form).

Part 1: Running product

Forward: given a of length n, return y with y[i] = a[0] · a[1] · … · a[i].

1a. Warm-up: if the loss depends only on the final value y[n-1], the gradient for a[k] is the product of every other entry. Compute all n of these in O(n) time without division. Entries may be zero, even several of them.

1b. General backward: given an upstream gradient g[i] = ∂L/∂y[i], return ∂L/∂a[k] for all k in O(n) time, again without dividing by any a[k].

running_product([2, 3, 0, 5])            -> [2, 6, 0, 0]
all_but_one([2, 3, 0, 5])                -> [0, 0, 30, 0]
all_but_one([3, 0, 5, 0])                -> [0, 0, 0, 0]
backward(a=[2, 3, 4], g=[1, 1, 1])       -> [16, 10, 6]

(For the last line: L = a0 + a0·a1 + a0·a1·a2, so ∂L/∂a0 = 1 + 3 + 12.)

1c. Write a finite-difference gradient checker and use it on random inputs with zero, one, and two zeros.

Part 2: Softmax cross-entropy

Given logits (N rows × C classes) and integer labels, return the mean cross-entropy loss and its gradient with respect to the logits. It must not overflow or produce nan for logits like [1000, 1001, 999]. Do not call a library loss function.

logits = [[2.0, 1.0, 0.1], [1000.0, 1001.0, 999.0]], labels = [0, 2]
loss  ≈ 1.4123

Verify the gradient with your checker.

Constraints

  • n up to 10^6 for Part 1, so O(n) time and O(n) extra memory.
  • Inputs may be negative.

Hints

Hint 1

Product of all but k = (product of everything left of k) × (product of everything right of k).

Hint 2

In 1b, a[k] appears in every output y[i] with i ≥ k. Factor out the left prefix and look for a right-to-left recurrence for the rest.

Hint 3

S[k] = g[k] + a[k+1]·S[k+1] collects the suffix contributions without division.

Hint 4

log Σ exp(z) = m + log Σ exp(z − m) with m = max(z). The gradient of mean CE is (softmax − onehot) / N.

Follow-up questions

  • How would you implement the running-product backward on a GPU for n = 10^8?
  • Your logits come in float16. What extra precautions do you take in the loss?
  • Add label smoothing to the cross-entropy. What does the gradient become?
  • How would you compute the running product in log space, and what breaks when entries can be zero or negative?