Branching-Fanout Sparse MoE

Integrate this at the sparse-MoE router, after the router has produced a token difficulty score but before expert dispatch. Use a small base fanout $j$ and generate a random requested fanout independently for each token, or use the same random seed across ablations. Pseudocode: (1) input token representations $h_b$ and router scores $r_{b,e}$; (2) compute a difficulty score $d_b=\operatorname{softplus}(w^\top h_b)$; (3) choose $p_b=\sigma(c_0+c_1d_b)$, clipped to $[p_{\min},p_{\max}]$; (4) draw $N_b$ by repeatedly continuing with probability $b_b=1/(1+2p_b)$, stopping with probability $a_b=2p_b/(1+2p_b)$; (5) initialize a population of $j$ dispatch slots and for each of $N_b$ generations replace every slot by one child with probability $1-p_b$ or two children with probability $p_b$; (6) set requested fanout to $K_b=\min(X_b,K_{\max})$; (7) select the top $K_b$ experts according to $r_{b,e}$ and dispatch the token, while applying capacity overflow handling. The paper’s formula determines the stopping probabilities and offspring law. Estimate empirically the realized mean, tail exponent, and periodic residual by logging $(K_b)$; do not assume the capped mean remains exactly $2j$. Add a compute regularizer $\lambda_c\max(0,\frac1B\sum_bK_b-C)^2$ and optionally a router loss encouraging larger $K_b$ only when validation loss falls. Start with a 12-layer, 8-expert Switch-style MLP on CIFAR-100 or WikiText-2, comparing fixed top-1/top-2 routing and learned stochastic fanout at equal average dispatched experts. The first success signal is lower validation loss or perplexity at the same average expert FLOPs; secondary signals are improved performance on high-loss tokens and a predictable heavy-tailed but bounded dispatch histogram. Stress-test $K_{\max}$, router noise, and expert-capacity overflow, since the uncapped mathematical tail can create severe batch imbalance.