# Ultrametric Gibbs MoE Router

- ID: 2506
- Canonical URL: https://synthcore.org/idea/2506/ultrametric-gibbs-moe-router
- API JSON: https://synthcore.org/api/idea/2506.json
- API Markdown: https://synthcore.org/api/idea/2506.md
- Verification status: unverified
- Source: [arXiv:2608.18276](https://arxiv.org/abs/2608.18276)
- Category: architecture
- Solves: scalability, generalization, stability
- ML areas: moe-routing, attention, initialization
- Math tags: stochastic-processes, probability, information-theory
- Ratings: usefulness 6/10; difficulty 5/10; novelty 7/10

## Idea description

Replace flat expert logits with scores generated by Gaussian increments on a binary routing tree. A leaf receives the sum of increments on its root-to-leaf path, so sibling experts have correlated logits and the router can learn nested coarse-to-fine specialization; an inverse-temperature schedule controls the transition from exploratory diffuse routing to sparse routing.

## Mathematical statement

The construction assigns independent standard Gaussian increments $U(v)$ to every non-root vertex $v$ of a binary tree. For a leaf $\sigma$ at depth $L$, its random score is $X_L(\sigma)=\sum_{i=1}^{L}U(\sigma|_i)$, where $\sigma|_i$ is its ancestor at depth $i$. The key property is the ultrametric covariance $\operatorname{Cov}(X_L(\sigma),X_L(\sigma'))=|\sigma\wedge\sigma'|$, where $|\sigma\wedge\sigma'|$ is the depth of the deepest common ancestor. For learned input-dependent scores $s_\sigma(x)$, define router probabilities $p_\sigma(x)=\exp(\beta[s_\sigma(x)+\alpha X_L(\sigma)]) / \sum_\tau\exp(\beta[s_\tau(x)+\alpha X_L(\tau)])$, where $x$ is the token representation, $\beta>0$ is inverse temperature, and $\alpha$ is noise scale. The paper defines $\psi_h(t)=t^2/2+\log 2+\log\cosh(ht)$ and the critical value through $\beta_c\psi_h'(\beta_c)-\psi_h(\beta_c)=0$. For the no-field router, $h=0$, so $\psi_0(t)=t^2/2+\log 2$ and $\beta_c=\sqrt{2\log 2}$. Below $\beta_c$, routing is expected to be diffuse; above it, Gibbs weights become sparse and hierarchical. The paper also proves a free-energy concentration form $\mathbb P(F_L\le f-\delta)\le\exp[-L\delta^2/(8\beta^2)]$, with $F_L=L^{-1}\log Z_L$, $Z_L=\sum_\sigma e^{\beta X_L(\sigma)}$, and $f$ the limiting free energy.

## Key formulas

- $$X_L(\sigma)=\sum_{i=1}^{L}U(\sigma|_i),\qquad U(v)\stackrel{\mathrm{i.i.d.}}{\sim}\mathcal N(0,1).$$
- $$\operatorname{Cov}\mathopen{}\left(X_L(\sigma),X_L(\sigma')\right)=|\sigma\wedge\sigma'|.$$
- $$p_\sigma(x)=\frac{\exp\!\left(\beta\,[s_\sigma(x)+\alpha X_L(\sigma)]\right)}{\sum_{\tau\in\{0,1\}^L}\exp\!\left(\beta\,[s_\tau(x)+\alpha X_L(\tau)]\right)}.$$
- $$\beta_c\psi_0'(\beta_c)-\psi_0(\beta_c)=0,\qquad \psi_0(t)=\frac{t^2}{2}+\log 2\quad\Longrightarrow\quad\beta_c=\sqrt{2\log 2}.$$

## Implementation notes

(1) Integrate the method into the MoE router immediately before expert selection. Use a complete binary tree of depth $L$ with $2^L$ leaves, one expert per leaf, and let $s_\sigma(x)$ be the learned score for leaf $\sigma$. To reduce parameters, implement $s_\sigma(x)$ as the sum of learned node scores along the path, or first use a standard linear projection producing all leaf scores. (2) Pseudocode: sample fixed Gaussian values $U[v]$ for every non-root node; precompute $X[leaf]=\sum_{v\in path(leaf)}U[v]$; for each token representation $x$, compute $s[leaf]$, set $z[leaf]=\beta(t)(s[leaf]+\alpha X[leaf])$, compute $p=softmax(z)$, select top-$k$ leaves, dispatch the token to those experts, and combine outputs using renormalized selected probabilities. Initialize $\beta(t)=0.5\beta_c$ and cosine-ramp it to $1.5\beta_c$ during the first 30% of training. Start with $\alpha=0.1$ and tune it in $\{0.03,0.1,0.3\}$. An alternative scalable implementation selects a branch recursively at each level, retaining the top $b$ children rather than materializing all leaves. (3) Compute the tree covariance analytically from shared path lengths; do not estimate it. Empirically measure free energy $F=B^{-1}\sum_iL^{-1}\log\sum_\sigma\exp(\beta(s_{i\sigma}+\alpha X_\sigma))$, routing entropy, active-leaf count, expert load coefficient of variation, and common-ancestor depth for pairs of selected leaves. Test whether across random seeds the standard deviation of $F$ decreases approximately with $L^{-1/2}$, consistent with the paper's concentration scale. (4) First experiment: compare flat top-2 gating with a depth-4, 16-expert ultrametric router in an equal-parameter Switch-Transformer on WikiText-103, or on CIFAR-100 with a small vision MoE. Keep optimizer, batch size, capacity factor, and FLOPs fixed. Success means lower validation loss or equal loss with better load balance, fewer dropped tokens, a controlled reduction in routing entropy near $\beta_c$, and no increase in seed-to-seed loss variance or dispatch latency.

## Disclaimer

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