# Poisson–Kingman expert-capacity prior

- ID: 262
- Canonical URL: https://synthcore.org/idea/262/poisson-kingman-expert-capacity-prior
- API JSON: https://synthcore.org/api/idea/262.json
- API Markdown: https://synthcore.org/api/idea/262.md
- Verification status: unverified
- Source: [arXiv:2607.04264](https://arxiv.org/abs/2607.04264)
- Category: regularization
- Solves: generalization, stability, scalability
- ML areas: moe-routing, regularization, training
- Math tags: probability, stochastic-processes, combinatorics
- Ratings: usefulness 5/10; difficulty 5/10; novelty 7/10

## Idea description

Replace the usual uniform expert-load target in sparse MoE training with a random, heavy-tailed capacity allocation generated by a conditioned Poisson point process. The constant profile reproduces a Poisson–Dirichlet-like allocation, while a profile such as \(\phi_\gamma(x)=1+e^{-\beta\gamma x}\) deliberately changes the frequency of large versus small expert allocations.

## Mathematical statement

The paper constructs a marked Poisson point process \(\mathcal{N}=\sum_i\delta_{(J_i,M_i)}\) on \((0,\infty)\times\mathsf{M}\) with intensity \(x^{-1}\eta_x(dm)\,dx\), where \(J_i>0\) is an atom mass, \(M_i\in\mathsf{M}\) is an optional mark, and \(\eta_x\) is a finite measure describing marks at mass scale \(x\). Its unmarked mass intensity is \(x^{-1}\phi(x)\,dx\), with \(\phi(x)=\eta_x(\mathsf{M})\). The Poisson–Kingman bridge is this process conditioned on total mass \(\sum_iJ_i=S\), where \(S>0\) is the available mass; normalize by \(p_i=J_i/S\), sort decreasingly, and use the resulting \(p_i\) as expert-load proportions. If \(\phi(x)\equiv1\), the normalized ranked masses have \(\mathrm{PD}(0,1)\) law. A finite visible-spectrum profile is \(\phi(x)=\sum_{r=1}^Q\theta_r e^{-\beta\lambda_rx}\), with \(\theta_r>0\), \(\lambda_r\ge0\), and \(\beta>0\). The double-well example is \(\phi_\gamma(x)=1+e^{-\beta\gamma x}\). The transferable property is the explicit scale-dependent intensity: it controls the distribution of large capacity blocks while conditioning enforces \(\sum_i p_i=1\).

## Key formulas

- $$\mathcal{N}=\sum_i\delta_{(J_i,M_i)},\qquad \mathbb{E}\,\mathcal{N}(dx,dm)=x^{-1}\eta_x(dm)\,dx$$
- $$\phi(x)=\eta_x(\mathsf{M}),\qquad \mathbb{E}\,\mathcal{N}(dx)=x^{-1}\phi(x)\,dx$$
- $$\mathcal{N}\ \bigg|\ \sum_iJ_i=S,\qquad p_i=\frac{J_i}{S},\qquad \sum_i p_i=1$$
- $$\phi_\gamma(x)=1+e^{-\beta\gamma x},\qquad \phi(x)=\sum_{r=1}^{Q}\theta_r e^{-\beta\lambda_rx}$$

## Implementation notes

(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.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.