# WKB-Stokes Sparse Mixer

- ID: 1802
- Canonical URL: https://synthcore.org/idea/1802/wkb-stokes-sparse-mixer
- API JSON: https://synthcore.org/api/idea/1802.json
- API Markdown: https://synthcore.org/api/idea/1802.md
- Verification status: unverified
- Source: [arXiv:2608.04729](https://arxiv.org/abs/2608.04729)
- Category: architecture
- Solves: stability, inference-speedup, memory
- ML areas: ssm, transformer, mlp, inference-speedup
- Math tags: geometry, linear-algebra, dynamical-systems
- Ratings: usefulness 5/10; difficulty 4/10; novelty 7/10

## Idea description

Replace a dense channel-mixing matrix in a sequence layer with alternating diagonal propagation and sparse unipotent Stokes jumps. The diagonal part carries independently controlled exponential phases, while the unipotent factors implement cheap residual-like mode conversion without changing determinant or requiring a dense matrix multiply. Constrain the phase magnitudes and jump coefficients during training to obtain a reversible, norm-monitorable mixer.

## Mathematical statement

The paper uses the four local fourth roots of a quartic differential, $\phi_k=(\sqrt{-1})^{1-k}\,dw$ for $k=1,2,3,4$, and assigns each path segment the phase $\mu_{i,k}^{\eta}=-2\operatorname{Re}\int_{\alpha_i^\eta}\phi_k$ and $\lambda_{i,k}^{\eta}=-2\operatorname{Re}\int_{\delta_i^\eta}\phi_k$. The dominant phase on a segment is the maximum, $\mu_i^\eta=\max_k\mu_{i,k}^\eta$ and $\lambda_i^\eta=\max_k\lambda_{i,k}^\eta$, expressing max-plus selection among asymptotic modes. At Stokes transitions, the paper obtains sparse unipotent jump matrices, including $U_2=\begin{pmatrix}1&0&-s_2&0\\0&1&0&0\\0&0&1&0\\0&0&0&1\end{pmatrix}$ and $U_3=\begin{pmatrix}1&0&0&s_1\\0&1&s_1&0\\0&0&1&0\\0&0&0&1\end{pmatrix}$, with $s_1=\sin(4\theta)/\sin\theta$ and $s_2=\sin(4\theta)\sin(3\theta)/(\sin\theta\sin(2\theta))$, where $\theta=\pi/(d+4)$ and $d$ is the relevant zero order. Both matrices are identity plus nilpotent sparse updates, hence have determinant one and cheap triangular inverses. The neural adaptation uses four learned channel phases $p_k$ and bounded jump coefficients $a_1,a_2$, replacing the paper's geometric coefficients while preserving the same diagonal-plus-unipotent factorization.

## Key formulas

- $$\phi_k=(\sqrt{-1})^{1-k}\,\mathrm{d}w,\qquad \mu_{i,k}^{\eta}=-2\operatorname{Re}\int_{\alpha_i^\eta}\phi_k,\qquad \lambda_{i,k}^{\eta}=-2\operatorname{Re}\int_{\delta_i^\eta}\phi_k,$$
- $$\mu_i^\eta=\max_{1\le k\le4}\mu_{i,k}^\eta,\qquad \lambda_i^\eta=\max_{1\le k\le4}\lambda_{i,k}^\eta,$$
- $$U_2=I-a_2E_{1,3},\qquad U_3=I+a_1(E_{1,4}+E_{2,3}),$$
- $$y=U_3\,U_2\,\operatorname{diag}(e^{p_1},e^{p_2},e^{p_3},e^{p_4})x,\qquad a_j=a_{\max}\tanh(\widehat a_j),\quad p_k=p_{\max}\tanh(\widehat p_k).$$

## Implementation notes

(1) Integration point: in a Transformer or SSM block, take each token vector $x\in\mathbb{R}^{4m}$, reshape it into $m$ groups of four channels, and replace the usual dense channel projection with the proposed mixer. Apply it either before the MLP or as the state-transition map in an SSM; keep the token-mixing operation unchanged initially so the experiment isolates channel mixing. For each group, use trainable unconstrained parameters $\widehat p_1,\ldots,\widehat p_4,\widehat a_1,\widehat a_2$. (2) Pseudocode: `p = p_max*tanh(p_hat); a1=a_max*tanh(a1_hat); a2=a_max*tanh(a2_hat); z=exp(p)*x; z1=z1-a2*z3; z1=z1+a1*z4; z2=z2+a1*z3; y=z`; optionally apply RMSNorm and a learned scalar gate after the mixer. This is exactly the adapted diagonal propagation followed by $U_2$ and $U_3$. For a deeper WKB variant, compute four phase scores from a local feature pair $(u,v)$ as $p_k=\beta\,\mathrm{Re}(i^{1-k}(u+iv))$, then use a temperature-controlled softmax or hard maximum across the four scores. (3) Computed from the paper: the root structure, max-over-modes rule, and sparse unipotent factorization. Estimated empirically: $p_{\max}$, $a_{\max}$, number of mixer repetitions, and whether learned or input-dependent phases work best. Monitor the singular values of the mixer and its inverse; initialize all $p_k$ and jumps at zero, making the initial map the identity. (4) First experiment: train a 4-layer, 256-dimensional causal language model on TinyStories or enwik8, comparing this mixer against a dense linear projection and a standard gated residual channel mixer at equal hidden size and matched parameter/FLOP budgets. Test sequence lengths 1K and 4K. Success means lower validation loss at equal training steps, fewer exploding-gradient events, and stable inverse reconstruction error below $10^{-5}$ after 100 stacked applications. A secondary inference test should measure whether sparse mixing gives lower latency and memory than the dense baseline at equal perplexity.

## Disclaimer

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