Analytic covariance initialization for switched state-space models

Integrate this at the hidden-state normalization point of a Markov-switching RNN, SSM, or recurrent mixture-of-experts layer. Let $A_j$ be the current mode transition, let $Q_j$ be the covariance of the injected input/noise in mode $j$, and let $p_{ij}$ be the mode transition matrix. Estimate the stationary mode probabilities $\pi$ by solving $\pi^\top=\pi^\top p$ or by averaging router assignments. Construct the implicit block operator $T$ using $T_{ji}=p_{ij}(A_j\otimes A_j)$. Solve $(I-T)m=q$ with conjugate gradients or GMRES; each matrix-vector product is implemented without Kronecker storage by reshaping each $d^2$ vector into $d\times d$ and applying $A_jX A_j^\top$. Reshape the solution blocks into $M_j$, form $C_j=M_j/\pi_j+\delta I$, and compute $C_j^{-1/2}$ using an eigendecomposition for small $d$ or 5--10 Newton-Schulz iterations for larger $d$. During training, recompute this covariance every 100--1000 optimizer steps, detach it from the gradient initially, and use an exponential moving average to avoid noisy mode estimates. Apply $C_{s_t}^{-1/2}$ to the recurrent state before the next transition, or use it only to initialize a learned affine normalization layer. The paper supplies the exact fixed-point equation and the finiteness condition; empirical estimates are $p$, $Q_j$, and $\pi$. Start with a 4-mode diagonal SSM on sequential CIFAR or the copy task, comparing zero-state, LayerNorm, and analytic covariance initialization/normalization. Measure burn-in length, per-mode activation variance, gradient variance, long-sequence loss, and validation accuracy. The expected signal is mode variances close to one from the first hundred steps, less dependence on sequence burn-in, and faster loss descent than LayerNorm at equal FLOPs.