Poisson–Kingman expert-capacity prior

(1) Integrate at the MoE load-balancing loss. Let a batch contain \(T\) tokens and \(E\) experts. The router produces probabilities \(q_{te}\), dispatches by the usual top-1 or top-2 rule, and gives an empirical load vector \(\ell_e=T^{-1}\sum_t\mathbf{1}[e\text{ receives token }t]\). Once per training step, generate a detached target vector \(p\in\Delta^{E-1}\) from the truncated Poisson–Kingman approximation below. Add \(L_{PK}=\alpha\,\mathrm{KL}(\ell+\epsilon\,\|\,p+\epsilon)\), or apply the same loss to expected loads \(\ell_e=T^{-1}\sum_tq_{te}\). Keep the task loss and capacity overflow mechanism unchanged. (2) Pseudocode: choose logarithmic mass bins \([a_k,b_k]\), \(k=1,\ldots,K\), with \(a_1=\delta\) and \(b_K=B\); compute \(\Lambda_k=\int_{a_k}^{b_k}x^{-1}\phi(x)dx\); sample \(C_k\sim\mathrm{Poisson}(\Lambda_k)\); for each of the \(C_k\) atoms sample \(J\) in bin \(k\) from density proportional to \(x^{-1}\phi(x)\); concatenate all atoms, retain the largest \(E\), set \(p_i=J_{(i)}/\sum_jJ_j\), and pad with zeros if fewer than \(E\). To approximate conditioning on total mass \(S=1\), normalize the vector directly and record the approximation error, or reject samples whose total mass lies outside \([1-\varepsilon,1+\varepsilon]\). A practical alternative is to use the exponential tilt \(\phi_\tau(x)=e^{-\tau x}\phi(x)\) and tune \(\tau\) by one-dimensional bisection until the sampled mean total mass is near 1. (3) The paper-derived computations are the \(x^{-1}\phi(x)\) intensity, Poisson bin counts, normalization, and profile choices. The truncation \(\delta\), maximum mass \(B\), number of bins \(K\), rejection tolerance, and \(\alpha\) are empirical parameters. If using a data-dependent profile, estimate the top \(Q\) eigenvalues of router-usage covariance by a few power iterations, map them to nonnegative \(\lambda_r\), and normalize \(\theta_r\) to sum to one. (4) First experiment: train a 16-expert top-1 Switch Transformer, or a small 8-expert Transformer-MoE, on C4 or WikiText-103. Compare the uniform Switch auxiliary loss, Dirichlet/random-load targets, the constant-profile target, and the \(\phi_\gamma\) target at equal batch size, FLOPs, and capacity factor. Measure validation perplexity, coefficient of variation of expert loads, dropped-token fraction, inactive experts, and loss variance across seeds. Success is lower perplexity at the same compute with no increase in overflow, or equal perplexity with a lower capacity factor and fewer routing-collapse events. Test whether constant \(\phi\) is too uniform and whether increasing \(\gamma\) produces useful specialization rather than merely worse imbalance.