Dirichlet Replica-Overlap Routing

Integrate this at the router of a small sparse MoE transformer or an MLP-MoE classifier. The router computes logits \(h(x)\in\mathbb R^K\), converts them to concentrations \(a_i=\alpha K\operatorname{softmax}(h)_i+\epsilon\), and samples two Dirichlet vectors using independent Gamma variables followed by normalization. Use the sampled vectors as soft routing weights, or select top-1 experts independently from them. For each minibatch, record the expert indices \(e_t^{(1)},e_t^{(2)}\) for the same tokens, compute \(L_T/T\), and add \(\lambda(L_T/T-\rho)^2\) to the task loss. A cheaper differentiable version uses the analytic collision statistic \(\sum_i p_i^2\) instead of sampled indices. Pseudocode is: compute router logits; form concentrations; sample two routing vectors; dispatch both replicas; average their outputs; compute task loss; compute collision loss; backpropagate the sum; optionally update \(\alpha\) with a slow controller. The paper supplies the Dirichlet moment identities and the interpretation of second moments through replica collisions; the implementation estimates finite-batch collision local time, output variance, and gradient variance empirically. First test a 4-expert, 2-layer MLP-MoE on CIFAR-10 or synthetic Gaussian classification, comparing ordinary softmax routing with fixed \(\alpha\in\{0.1,1,10,100\}\) and adaptive overlap control. The quantitative prediction is that the isolated-router coincidence ratio should match \(\chi(\alpha,K)=K(\alpha+1)/(K\alpha+1)\) within about 10 percent for large batches. As \(\alpha\) decreases, routing and replica-output variance should increase monotonically. Sweep routing depth and sequence length to test whether output variance is better predicted by cumulative collision local time than by raw depth; instability should occur near a reproducible threshold in collision rate or gradient variance.