Wasserstein Tangent-Space Stability Monitor
Source paper: Pushforward dynamics on Wasserstein spaces and measure rigidity arXiv:2609.00451 ⓘ · analyzed Sep 2, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Treat the empirical hidden-state distribution of a recurrent or state-space model as a Wasserstein-space state and estimate the linearized pushforward operator on perturbation vector fields. Penalize tangent modes whose estimated transfer gains exceed one, while retaining near-unit fixed modes that represent robust invariant distributional structure.
Formulas
Mathematical statement
Let $\mu$ be the hidden-state distribution and $F:M\to M$ one recurrent transition satisfying $F_\ast\mu=\mu$. A Wasserstein tangent perturbation is a gradient vector field $v\in T_\mu\mathcal P_2(M)=\overline{\{\nabla\psi\}}^{L^2(\mu)}$. Define the vector-field transfer operator by $\mathcal L_Fv(y)=\mathbb E[D F(X)v(X)\mid F(X)=y]$ for $X\sim\mu$. The paper's linearization is $D(F_\ast)_\mu[v]=P_\mu\mathcal L_Fv$, where $P_\mu$ is the $L^2(\mu)$ orthogonal projection onto gradient vector fields. The distributional perturbation is linearly stable when the spectral radius $\rho(P_\mu\mathcal L_F)<1$; fixed tangent modes satisfy $P_\mu\mathcal L_Fv=v$ and are first-order invariant deformations. In a neural implementation, estimate leading gains $\widehat\lambda_j$ on probe fields $q_j$ and regularize gains above one.
Implementation notes
Integrate the method at the recurrent transition of a small GRU, SSM, or world model. Let hidden states $h_i$ be sampled from a minibatch and let $F_\theta(h_i,u_i)$ produce the next state; initially use fixed inputs or identical input sequences so the transition is approximately autonomous. Create $r$ probe fields $q_j(h)$ as gradients of random Fourier features or small probe potentials. For every probe and state, use automatic differentiation to compute $a_{ij}=J_F(h_i)q_j(h_i)$, where $J_F$ is the Jacobian of the transition with respect to hidden state. Estimate conditional averaging using destination-state kernel regression: $\widehat{\mathcal L_Fq_j}(y)=\sum_iK_\tau(y,F(h_i))a_{ij}/\sum_iK_\tau(y,F(h_i))$. Estimate $P_\mu$ by least-squares projection of these vectors onto a basis of gradient probes evaluated at sampled states. Compute leading singular values or Rayleigh gains, add $\gamma\mathcal R_{\mathrm{tan}}$ to the task loss, and optionally reduce the learning rate when the largest gain exceeds one. The paper supplies the transfer derivative and fixed-space interpretation; Jacobians, kernels, projections, and spectral estimates are empirical approximations. First test on a GRU predicting a noisy linear dynamical system and on sequential MNIST, against an unregularized GRU and ordinary Jacobian spectral normalization. Multiply the recurrent matrix by a scalar gain $g$ to sweep stability. The prediction is a measurable boundary near $\widehat\rho=1$: below it, tangent perturbation norms decay geometrically; above it, they grow approximately like $\widehat\rho^k$. The rollout-divergence threshold should agree with the estimated spectral boundary within 20 percent.
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.