Jointly Contractive Input-Conditioned RNN

Integrate this into an input-conditioned GRU-like or state-space recurrent layer whose transition is h_{t+1}=A(z_t)h_t+G(z_t), with z_t being the current token embedding or a low-dimensional projection of it. Start with a finite set of K transition matrices A_k produced by K learned basis matrices and softmax gates: A(z)=sum_k pi_k(z)A_k; enforce the constraint on each basis matrix, which also controls every convex mixture. Parameterize P=L L^T+epsilon I using a lower-triangular L with positive diagonal. At every training step, compute the normalized Lyapunov violation for each basis matrix, r_k=max(0, lambda_max(P^{-1/2}A_k^T P A_k P^{-1/2})-rho^2), and add lambda_L sum_k r_k^2 to the task loss. A cheaper guaranteed variant computes B_k=P^{1/2}A_kP^{-1/2} and parameterizes B_k=rho*tanh(s_k)*Q_k, where Q_k is approximately orthogonal, then maps back to A_k; use the penalty version first because it is less restrictive. Estimate the finite-horizon joint spectral radius empirically by sampling many token sequences and tracking g_T=max_sequences ||A_{T-1}...A_0||_2^{1/T}; this estimate is diagnostic, not a formal guarantee. Pseudocode: initialize h=0; for t in 1..T, compute pi=softmax(router(z_t)), set A=sum_k pi_k*A_k, and update h=A@h+G(z_t). Add loss=task_loss+lambda_L*sum_k*r_k^2, then backpropagate and update A_k and L. Estimate lambda_max with 5 to 10 power iterations rather than full eigendecomposition for large hidden states. The first cheap experiment should compare a 2-layer 256-unit input-conditioned RNN with unconstrained transitions, standard spectral normalization, and the Lyapunov penalty on sequential MNIST and copy/add tasks. Measure gradient norms, hidden-state norms, long-sequence loss, accuracy, and wall-clock overhead at equal parameters and FLOPs. Success means bounded hidden norms, fewer exploding-gradient runs, and better long-context accuracy without materially slowing short-sequence convergence.