# Unitary-dilated stochastic router

- ID: 1786
- Canonical URL: https://synthcore.org/idea/1786/unitary-dilated-stochastic-router
- API JSON: https://synthcore.org/api/idea/1786.json
- API Markdown: https://synthcore.org/api/idea/1786.md
- Verification status: unverified
- Source: [arXiv:2608.04354](https://arxiv.org/abs/2608.04354)
- Category: architecture
- Solves: stability, accuracy, generalization
- ML areas: moe-routing, transformer, world-model
- Math tags: linear-algebra, probability, geometry
- Ratings: usefulness 5/10; difficulty 6/10; novelty 5/10

## Idea description

Replace a softmax transition or mixture-of-experts router by probabilities obtained from squared amplitudes of an isometric latent transition. Each input state is mapped to an orthogonal latent subspace, and summing probability over the latent index produces the desired expert or next-state distribution. The latent amplitudes can retain information that would be destroyed by directly averaging expert outputs, while normalization is guaranteed by construction.

## Mathematical statement

A stochastic matrix \(\boldsymbol{\Gamma}\in\mathbb{R}^{M\times N}\) satisfies \(\Gamma_{ij}\geq 0\) and \(\sum_{i=1}^{M}\Gamma_{ij}=1\), with \(\Gamma_{ij}=p(i\mid j)\). The paper's Born representation expresses entries as \(\Gamma_{ij}=\operatorname{Tr}(E_iP_j)\), where \(\{E_i\}_{i=1}^M\) is a POVM satisfying \(E_i\succeq0\) and \(\sum_iE_i=I\), and \(\{P_j\}_{j=1}^N\) is a PVM with orthogonal projectors \(P_jP_{j'}=\delta_{jj'}P_j\) and \(\sum_jP_j=I\). A concrete neural dilation uses latent output labels \((i,j)\): define an isometry \(V:\mathbb{C}^{N}\to\mathbb{C}^{MN}\) by \(V|j\rangle=\sum_{i=1}^{M}\sqrt{\Gamma_{ij}}e^{\mathrm{i}\theta_{ij}}|i,j\rangle\). Its columns are orthonormal because different inputs use disjoint latent subspaces, and \(\sum_i|\langle i,j|V|j\rangle|^2=\sum_i\Gamma_{ij}=1\). Extending \(V\) to a square unitary \(U\) gives a unistochastic representation in the enlarged space; marginalizing the latent label recovers \(\Gamma_{ij}\). The phases \(\theta_{ij}\) do not affect marginal probabilities but can carry trainable information for downstream coherent processing.

## Key formulas

- $$\Gamma_{ij}\geq 0,\qquad \sum_{i=1}^{M}\Gamma_{ij}=1,\qquad \Gamma_{ij}=p(i\mid j).$$
- $$\Gamma_{ij}=\operatorname{Tr}(E_iP_j),\qquad E_i\succeq0,\quad \sum_{i=1}^{M}E_i=I,\qquad P_jP_{j'}=\delta_{jj'}P_j,\quad \sum_{j=1}^{N}P_j=I.$$
- $$V|j\rangle=\sum_{i=1}^{M}\sqrt{\Gamma_{ij}}\,e^{\mathrm{i}\theta_{ij}}|i,j\rangle,\qquad V^{\dagger}V=I_N,$$
- $$\Gamma_{ij}=\sum_{a=1}^{A}\left|U_{(i,a),j}\right|^2,\qquad U^{\dagger}U=I_D,\qquad D\geq N.$$

## Implementation notes

Integrate this at a mixture-of-experts router or a discrete world-model transition layer. Let the router input be a hidden vector \(h\in\mathbb{R}^d\), let \(N\) be the number of source states or tokens and \(M\) the number of experts or destination states, and produce column-stochastic probabilities with \(\Gamma_{:j}=\operatorname{softmax}(g_j(h)/\tau)\). For an explicit dilated module, use amplitudes \(A_{ij}=\sqrt{\Gamma_{ij}}\exp(i\theta_{ij})\), create latent states \(z_{i,j}=A_{ij}x_j\), and aggregate over the latent source label only after an expert-specific transformation: \(y_i=\sum_j A_{ij}F_i(x_j)\) for complex features, or \(y_i=\sum_j|A_{ij}|^2F_i(x_j)\) for an ordinary probabilistic router. Pseudocode is: `logits = router(h); Gamma = softmax(logits / tau, dim=expert); theta = phase_net(h); A = sqrt(Gamma) * exp(1j*theta); z[i,j] = A[i,j] * x[j]; z = expert_blocks(z); y[i] = sum_j(z[i,j])`; apply a real projection after the sum if needed. If a square unitary is required, complete the \(N\) orthonormal columns of \(V\) with QR decomposition, but do not form the full \(MN\times MN\) matrix initially. The paper supplies the exact stochastic/Born normalization; the engineer estimates \(\Gamma\) from router logits and learns phases \(\theta\), while orthogonality is checked through \(\|V^{\dagger}V-I\|_F\). Start with a 4-expert, 8-layer Transformer on WikiText-2, comparing standard top-1 softmax routing against the dilated router at equal parameter count and FLOPs. Measure validation perplexity, expert-load variance, gradient-norm spikes, and performance after noisy-token corruption. Success means lower routing collapse or load imbalance and improved perplexity or corruption robustness at matched compute; failure is increased cost without gains.

## Disclaimer

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