Hand-Written Gradients: Prefix-Product Backward Pass and Stable Cross-Entropy

This machine learning coding round asks you to write gradients by hand, the way an automatic differentiation (autograd) engine would, without using any autograd library. ### Clarifying Questions - Can the input arrays contain zeros, or negative values? - Should the functions work on plain Python lists, NumPy arrays or batches? - For cross-entropy, are the inputs raw scores (logits) or probabilities, and is the loss averaged over the batch or summed? - What precision is expected, and how will correctness be checked? ### Part 1 — The prefix-product operation The prefix-product operation maps an array `a` of length `n` to an array `y` with `y[i] = a[0] × a[1] × … × a[i]`. Implement its forward pass, and its backward pass: given the upstream gradient `g[i] = ∂L/∂y[i]` for some scalar loss `L`, return `∂L/∂a[k]` for every `k`, applying the chain rule. Start with the special case where the loss depends only on the total product `y[n − 1]`: then the gradient for `a[k]` is the product of all the other elements. Your backward pass must run in `O(n)` time and must not divide by elements of `a`, since the input can contain zeros. ```hint Why division fails, and what replaces it Dividing the total product by `a[k]` breaks as soon as one element is zero; think about what you can accumulate from the left and from the right instead. ``` ```hint Reuse the sum from the right In the general case, the gradient for `a[k]` sums contributions from every output at or after position `k`; look for a recurrence that computes these sums from right to left. ``` #### What This Part Should Cover - A correct forward pass and the chain-rule expression for each input's gradient - The special case of the total product, solved without division using products from both sides - A linear-time backward pass for general upstream gradients - A numerical gradient check, including inputs with one zero and with two zeros ### Part 2 — Cross-entropy from scratch Given logits of shape `N × C` (a batch of `N` examples and `C` classes) and integer labels, implement the mean softmax cross-entropy loss and its gradient with respect to the logits, in NumPy, without calling any library loss function. ```hint Stay in log space Consider what happens to exponentials of large logits, and how subtracting a per-row constant changes the softmax. ``` #### What This Part Should Cover - A numerically stable forward pass that does not overflow for large logits - The gradient with respect to the logits, and its derivation - Correct scaling by the batch size - A finite-difference gradient check ### What a Strong Answer Covers - Correct chain-rule reasoning, with clear notation for upstream and local gradients - Division-free, linear-time gradients that remain correct when inputs contain zeros - Numerical stability in the cross-entropy implementation - Verification of every gradient against finite differences ### Follow-up Questions - How would you implement the backward pass of the prefix sum? Of the running maximum? - How would you combine these pieces into a tiny autograd engine that builds a graph and runs backward automatically? - How does label smoothing change the cross-entropy gradient? - Why does the gradient of softmax combined with cross-entropy have such a simple form, and why do frameworks fuse the two?