Exact finite-support reverse AD

(1) Integrate at a categorical latent or small discrete router. Implement a PyTorch module receiving an input tensor $x$, logits $a\in\mathbb{R}^{B\times K}$, and a branch function $f(i,x)$. For each batch item, evaluate all $K$ branch outputs and return the exact expectation $z=\sum_i p_i f(i,x)$; use this only for small $K$, such as 4--32. A custom autograd Function is optional because ordinary tensor operations already realize the mathematical reverse pass, but a fused implementation can reduce memory.

(2) Pseudocode:

p = softmax(a, dim=-1)
for i in 0..K-1:
    y[i] = branch_fn(i, x)
    L_i = loss_fn(y[i], target)
L = sum_i p[i] * L_i
backward(L)
# explicit equivalent logit gradient:
mean_L = sum_i p[i] * L_i
grad_a[i] = p[i] * (L_i - mean_L)
# for expected feature z instead of scalar loss:
z = sum_i p[i] * y[i]
grad_y[i] = p[i] * grad_z
grad_a[i] = p[i] * dot(grad_z, y[i] - z)

The addition fibre rule $v\mapsto\langle v,v\rangle$ means the upstream cotangent is sent to every branch contribution. The probability derivative and value derivative are both retained; do not use stop-gradient on $p$.

(3) Compute exactly from the paper's mathematics: enumerate branches, normalize probabilities, aggregate the weighted result, and apply the product rule. No gradient estimator is needed. If enumeration is too expensive, use top-$k$ support $S$ and renormalize $\tilde p_i=p_i/\sum_{j\in S}p_j$; record discarded mass $1-\sum_{i\in S}p_i$ as an approximation diagnostic.

(4) First experiment: train a categorical VAE on MNIST with a 10-way latent variable. Compare exact enumeration with one-sample REINFORCE plus a learned baseline and straight-through Gumbel-Softmax, matching batch size and optimizer. Measure ELBO, reconstruction error, gradient variance over repeated minibatches, wall-clock time, and time to reach a fixed ELBO. The expected signal is lower gradient variance and faster early ELBO improvement for exact AD at small $K$; also sweep $K\in\{4,8,16,32,64\}$ to locate the compute crossover. As a second test, use a 4- or 8-expert MLP router and compare validation loss and routing overhead at equal expert FLOPs.