Unverified 2026

Unitary-dilated stochastic router

Implementation & benchmark of arXiv:2608.04354 — The Born Representation Theorem and the Unistochastic Theorem

Usefulness5/10

Difficulty6/10

Novelty5/10

Architecture Stability Accuracy Generalization Moe routing Transformer World model Linear algebra Probability Geometry

Source paper: The Born Representation Theorem and the Unistochastic Theorem arXiv:2608.04354 ⓘ · analyzed Aug 31, 2026

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

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.

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

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.

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.

Verification

This idea has not been verified yet.

Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.

Artifacts

Artifacts unavailable.