Unverified 2026

Bures-Shaped Latent State Space

Implementation & benchmark of arXiv:2608.19754 — A Controllability Gramain Shaping with LMI Constraints under Bures--Wasserstein Distance

Usefulness6/10

Difficulty6/10

Novelty7/10

Architecture Stability Accuracy Generalization Ssm Rnn Regularization Training dynamics Dynamical systems Optimal transport Convex analysis Linear algebra Control theory

Source paper: A Controllability Gramain Shaping with LMI Constraints under Bures--Wasserstein Distance arXiv:2608.19754 ⓘ · analyzed Sep 1, 2026

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

Idea description

Add a stable linear latent state-space block whose controllability Gramian is trained toward a chosen positive-definite target using squared Bures–Wasserstein distance. Direction-specific semidefinite constraints can suppress disturbance amplification in nuisance coordinates while preserving controllability in coordinates needed for prediction.

Formulas

$$\dot z=Az+B\eta,\qquad AP+PA^{\top}+BB^{\top}=0,\qquad P=\int_{0}^{\infty}e^{At}BB^{\top}e^{A^{\top}t}\,dt.$$

$$d_{BW}^{2}(P,P_d)=\operatorname{tr}(P)+\operatorname{tr}(P_d)-2\operatorname{tr}\!\left[(P^{1/2}P_dP^{1/2})^{1/2}\right].$$

$$\operatorname{tr}\!\left[(P^{1/2}P_dP^{1/2})^{1/2}\right]=\max_{X}\ \operatorname{tr}(X)\quad\text{s.t.}\quad\begin{bmatrix}P&X\\X^{\top}&P_d\end{bmatrix}\succeq0.$$

$$\mathcal L=\mathcal L_{task}+\lambda_{BW}d_{BW}^{2}(P,P_d)+\lambda_{dir}\sum_{i\in\mathcal N}\operatorname{softplus}\!\left(\frac{P_{ii}-\varepsilon_i}{\tau}\right).$$

Mathematical statement

For a stable continuous-time latent system $\dot z=Az+B\eta$, with latent state $z\in\mathbb R^d$, disturbance $\eta$ having identity covariance, Hurwitz matrix $A$, and input matrix $B$, the controllability Gramian $P\succ0$ satisfies $AP+PA^{\top}+BB^{\top}=0$, equivalently $P=\int_0^{\infty}e^{At}BB^{\top}e^{A^{\top}t}\,dt$. For target Gramian $P_d\succ0$, use the squared Bures–Wasserstein distance $d_{BW}^2(P,P_d)=\operatorname{tr}(P)+\operatorname{tr}(P_d)-2\operatorname{tr}[(P^{1/2}P_dP^{1/2})^{1/2}]$. The cross term has the semidefinite representation $\operatorname{tr}[(P^{1/2}P_dP^{1/2})^{1/2}]=\max_X\operatorname{tr}(X)$ subject to $\begin{bmatrix}P&X\\X^{\top}&P_d\end{bmatrix}\succeq0$, where $X\in\mathbb R^{d\times d}$. A directional constraint such as $P_{ii}\leq\varepsilon_i$ limits disturbance energy in latent coordinate $i$. The neural adaptation computes $P$ from a discretized state-space recurrence or a differentiable Lyapunov solve.

Implementation notes

Integrate this into an S4-like or gated recurrent block, replacing its unconstrained latent transition with a stable linear recurrence. Use discrete time for efficient training: $z_{t+1}=A_dz_t+B_d\eta_t+K_du_t$, where $A_d=\exp(\Delta A)$, $A$ is parameterized as $A=-LL^{\top}-\alpha I$ with $\alpha>0$, and $B_d$ is learned. At each training step, compute the controllability covariance by solving the discrete Lyapunov equation $P=A_dPA_d^{\top}+B_dB_d^{\top}$, or use the truncated stable series $P_K=\sum_{k=0}^{K-1}A_d^kB_dB_d^{\top}(A_d^{\top})^k$. Add the Bures loss between $P_K+\epsilon I$ and a fixed diagonal target $P_d$, using eigendecomposition for the matrix square roots. Choose large target entries for task-relevant channels and small entries for nuisance channels. Initially implement directional restrictions with the softplus penalties in the displayed loss; a second version can alternate network updates with a small semidefinite projection enforcing $P_{ii}\leq\varepsilon_i$. The math-derived quantities are the Gramian recursion and Bures objective; $K$, $\epsilon$, $\lambda_{BW}$, $\lambda_{dir}$, $\tau$, and target entries are hyperparameters. The first experiment should use a 16- or 32-dimensional SSM in a noisy sequential-MNIST or permuted-MNIST classifier, comparing an unconstrained SSM, Frobenius covariance matching, and Bures shaping at equal parameter count. Track validation accuracy under injected latent noise, the largest eigenvalue of $P$, directional variances $P_{ii}$, gradient norms, and loss stability. Success means lower accuracy degradation under nuisance noise without clean-accuracy loss, bounded Gramian spectra, and fewer exploding-gradient runs.

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.