Memory Determines Learning Direction: A Theory of Gradient-Based Optimization in State Space Models¶

Conference: NeurIPS 2026
arXiv: 2510.00563
Code: https://github.com/state-spaces/s4/tree/simple
Area: Time Series
Keywords: State Space Models, Memory Capacity, Learning Dynamics, Vandermonde Matrix, Reservoir Computing

TL;DR¶

This paper formalizes the gradient optimization dynamics of linear state space models (SSMs) through the memory function framework, proving that the Vandermonde projection operator of recurrent eigenvalues strictly dictates the backpropagation update direction, theoretically establishing the equivalence between S4 and diagonal S4D, and demonstrating that fixing recurrent weights in a reservoir computing scheme prevents overfitting while accelerating convergence.

Background & Motivation¶

Linear recurrent neural networks, particularly structured state space sequence models (such as S4, S4D, and Mamba), have emerged as competitive architectures capable of challenging Transformers in long-range sequence modeling. To effectively retain historical dependencies across thousands of steps, foundational frameworks like HiPPO adopted orthogonal polynomial projection under the continuous-time Lebesgue measure to construct optimal long-term memory matrices, which were subsequently simplified into diagonal recurrent structures via S4D. However, a rigorous theoretical explanation for the learning dynamics of SSMs has been notably absent: why does specific eigenvalue initialization (e.g., S4Dinv and S4Dlin) exert such a decisive impact on downstream convergence? Why must models prioritize long-range memory span over reconstruction precision, despite the fundamental theoretical tradeoff between memory capacity and length? Furthermore, whether the diagonalized S4D is mathematically equivalent to the dense S4 in information processing capability remained unproven.

This lack of theoretical grounding limits both architecture design and optimization efficiency. Existing analyses of linear RNN dynamics typically impose artificial random orthogonal constraints on recurrence matrices or outright neglect the eigenvector structure, failing to trace how past temporal information truly attenuates and drives gradients during backpropagation. In multilayer sequence models, when gradients backpropagate through an individual SSM layer whose recurrence fails to preserve distant error signals, the gradient vector loses essential correlations between historical inputs and supervision targets, causing optimization plateaus or severe overfitting.

To resolve these fundamental questions, this paper leverages Information Processing Capacity (IPC) and its linear counterpart, the Memory Function (MF), to analytically derive the parameter gradient equations of discrete linear SSMs. The analytical findings reveal that the parameter gradient of an SSM readout layer is intrinsically modulated by a normalized Vandermonde projection matrix spanned by recurrent eigenvalues. Core idea: The memory function (MF) of an SSM strictly dictates its gradient learning direction; if initial memory attenuates to zero at long delays, teacher supervision is irrecoverably discarded during backpropagation, necessitating an initialization with maximal memory span and rendering a reservoir computing (RC) training strategy—freezing structured recurrent weights—superior in convergence and generalization.

Method¶

Overall Architecture¶

The analytical and methodological framework progresses through three stages: first, reformulating the discrete linear SSM through the memory function (MF) and information processing capacity (IPC) lens to prove that system memory capacity depends solely on recurrent eigenvalues, thereby establishing the mathematical equivalence between S4 and diagonal S4D; second, analytically expanding the parameter gradients of the linear readout layer under mean squared error (MSE) objectives, uncovering how the normalized Vandermonde projection operator filters the teacher supervision signal; and third, proposing a reservoir computing (RC) training paradigm that completely fixes structured recurrent eigenvalues during optimization, updating only the feedforward projections and readout layers.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input sequence u_t"] --> B["Spectral Decomposition & Memory Function Modeling<br/>Prove theoretical equivalence between S4 and S4D"]
    B --> C["Vandermonde Projection Gradient Analysis<br/>Derive parameter gradient G(y) and memory filtering mechanism"]
    C --> D["Initial Memory Span Requirement Analysis<br/>Vanishing distant MF irrecoverably destroys teacher signal"]
    D --> E["Reservoir Computing Training Strategy<br/>Freeze recurrent eigenvalues and optimize only readout weights"]
    E --> F["Downstream Long-Sequence Modeling & Fast Convergence"]

Key Designs¶

1. Memory Function Equivalence and S4/S4D Theoretical Unification: Eliminating Constraints on Recurrent Eigenvectors When a continuous-time linear time-invariant SSM is discretized via zero-order hold (ZOH), its state recurrence is given by $\bm{x}_{k+1} = \bm{A}\bm{x}_k + \bm{B}\bm{u}_k$. Prior works introduced diagonalized S4D as an empirical approximation to avoid Cauchy kernel computations. This paper demonstrates that any hidden state in a linear recurrent system is strictly a linear combination of past inputs, expressed as $\bm{x}_t = \sum_{\tau=0}^\infty \bm{\alpha}_\tau u_{t-\tau}$. Under IPC theory, the degree to which a delayed input $u_{t-\tau}$ can be reconstructed linearly from the state defines the memory function $M[u_{t-\tau}]$. The exact analytical form satisfies $M[u_{t-\tau}] = \bm{V}_\tau^\top (\bm{V}\bm{V}^\top)^{-1}\bm{V}_\tau$, where $\bm{V}$ is a Vandermonde matrix composed of the eigenvalues of $\bm{A}$, with $\bm{V}_\tau = (\lambda_1^\tau, \dots, \lambda_N^\tau)^\top$. Because total memory capacity and state information depend exclusively on the eigenvalue spectrum $\{\lambda_i\}_{i=1}^N$, the eigenvector transformation matrix $\bm{P}$ and input vector $\bm{B}$ can be chosen arbitrarily without altering the intrinsic information capacity. Consequently, S4 and diagonal S4D sharing identical HiPPO eigenvalues are theoretically equivalent.

2. Vandermonde Projection-Modulated Gradient Dynamics: Memory Attenuation as a Supervision Filter To understand how parameter updates are guided, the authors derive the gradient of the linear readout weight $\bm{\theta}$ under the MSE objective. Given a target output sequence represented as a linear combination of delayed inputs $\bm{y} = \sum_{k=0}^{T-1} \alpha_k \bm{U}_k = \hat{\bm{U}}\bm{\alpha}$, and assuming independent identically distributed random input, the state covariance simplifies to $\bm{X}^\top \bm{X} = T \langle u^2 \rangle \bm{Q}\bm{V}\bm{V}^\top \bm{Q}^\top$. The resulting gradient expression is: $$G(\bm{y}) = 2\langle u^2 \rangle B \bm{V} \left( B \bm{V}^\top \bm{\theta} - \bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V} \bm{\alpha} \right)$$ In this formula, the matrix $\bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V}$ is precisely the normalized Vandermonde projection operator, whose diagonal elements correspond identically to the memory function values $M[u_{t-\tau}]$ at each delay step. This operator acts as a direct weighting filter on the teacher signal coefficients $\bm{\alpha}$: the influence of history at delay $\tau$ on the parameter gradient is fundamentally proportional to the memory retained at that delay. For nonlinear readout layers with activation $f(\cdot)$ (such as GELU or GLU), the gradient maintains the same core structure with modulated teacher coefficients $\bar{\bm{\alpha}}$ corresponding to $\bar{\bm{y}} = \bm{y} \odot f'(\bm{X}\bm{\theta})$.

3. Initial Memory Span Requirement: Prioritizing Long-Range Reach Over Short-Range Precision According to Jaeger's memory capacity bound, the sum of memory capacities is bounded by the state dimension, $\sum_{\tau=0}^\infty M[u_{t-\tau}] \le N$. If an architecture optimizes for high reconstruction accuracy at small delays, memory at large delays inevitably drops to zero. The gradient formulation reveals a critical failure mode: if the target sequence depends on distant historical steps where the model's initial memory is zero ($\bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V} \bm{\alpha} \approx \bm{0}$), the gradient update becomes entirely independent of the target output, causing learning to stall. Conversely, as long as distant memory values remain strictly non-zero—even if their numerical magnitude is small due to capacity constraints—the teacher signal successfully penetrates the projection operator to guide parameter updates. This provides mathematical justification for why HiPPO initializations (S4Dinv and S4Dlin) must distribute eigenvalues along the inner boundary of the unit circle to sustain long, non-vanishing memory tails.

4. Reservoir Computing (RC) Training Paradigm: Freezing Recurrent Weights to Eliminate Overfitting In standard SSM training, recurrent eigenvalues are optimized via gradient descent. However, backpropagation lacks an explicit mechanism to spontaneously create memory in temporal regions where initial memory is zero. Moreover, adapting eigenvalues risks destroying the mathematically optimal orthogonal polynomial spectrum of HiPPO, causing unintended memory collapse elsewhere. In addition, because long-memory eigenvalues must stay extremely close to modulus 1, unconstrained optimization easily causes numerical instability and spectral radius divergence ($\rho > 1$). The authors propose the Reservoir Computing (RC) setting: after initializing recurrent weights with structured HiPPO spectra (S4Dinv/S4Dlin), the diagonal matrix $\bm{\Lambda}$ and discretization step sizes $\Delta$ are frozen throughout training, optimizing only linear projections, feedforward blocks, and readout weights. This eliminates approximately 10% of trainable parameters while preserving the theoretical memory backbone.

A Worked Example¶

Consider an SSM with state dimension $N=4$ and sequence length $T=1024$. Under uniform random eigenvalue initialization ("random"), eigenvalues possess moduli noticeably below 1, causing the memory function to decay to zero for delays $\tau > 40$. Consequently, the lower block of the normalized Vandermonde projection matrix collapses to zero. For a task requiring dependency on an input from $\tau=500$ steps ago ($\alpha_{500} \neq 0$), the driving term $\bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V} \bm{\alpha}$ yields zero at index 500. The gradient vector $G(\bm{y})$ carries only self-correlation terms devoid of supervisory signal, leaving readout parameters unable to learn the relation.

Under S4Dinv initialization, eigenvalues follow $\lambda_m = -\frac{1}{2} + i \frac{N}{\pi} (\frac{N}{2m+1} - 1)$, which map tightly against the unit circle under ZOH discretization. The memory function at delay $\tau=500$ maintains a non-zero value ($M[u_{t-500}] > 0$). While individual reconstruction precision is lower than at $\tau=1$, the non-zero projection component preserves the distant teacher signal, enabling gradient updates to smoothly guide $\bm{\theta}$ toward the analytical least-squares solution $\bm{\theta}^* = \frac{1}{B} (\bm{V}\bm{V}^\top)^{-1}\bm{V}\bm{\alpha}$.

Loss & Training¶

In the analytical single-layer formulation, optimization is conducted under the mean squared error loss: $$L_{\text{MSE}} = \frac{1}{T}\sum_{i=1}^T \left(f(\bm{\theta}^\top \bm{x}_i) - y_i\right)^2$$ For end-to-end multi-layer benchmarks (6 to 8 stacked S4D layers), models are trained with cross-entropy loss using Adam. In the learnable baseline, dedicated learning rates are assigned to layers and step sizes ($\text{lr\_dt} \in [0.0005, 0.01]$). In the proposed RC setting, all recurrent eigenvalue and step-size learning rates are set to 0. Bidirectional sequence processing is applied, doubling the sequence input to readout linear transformations, followed by GLU activation and layer normalizations.

Key Experimental Results¶

Main Results¶

The method was evaluated on the Long Range Arena (LRA) benchmark across all six tasks using six independent random seeds. Minimum and maximum test accuracies across runs are reported alongside the official full S4 baseline.

Task	Sequence Length	Official S4	Learnable S4Dinv	Learnable S4Dlin	RC S4Dinv (Ours)	RC S4Dlin (Ours)
Listops	2,000	0.595	0.594 / 0.589	0.598 / 0.587	0.585 / 0.570	0.593 / 0.577
Text	4,000	0.865	0.853 / 0.839	0.848 / 0.830	0.857 / 0.846	0.850 / 0.843
Retrieval	4,000	0.910	0.863 / 0.760	0.843 / 0.750	0.906 / 0.901	0.903 / 0.898
Image	1,024	0.885	0.864 / 0.859	0.868 / 0.860	0.873 / 0.865	0.877 / 0.871
Pathfinder	1,024	0.940	0.910 / 0.877	0.896 / 0.865	0.899 / 0.889	0.892 / 0.872
PathX	16,384	0.960	0.922 / 0.903	0.926 / 0.912	0.932 / 0.923	0.933 / 0.922

(Note: Bold text denotes cases where the frozen RC setting outperforms the corresponding learnable eigenvalue setting.)

Ablation Study¶

To verify the interaction between eigenvalue structures and weight freezing, the authors compared structured initializations against non-structured configurations (uniform random, real linear, and complex step) under both learnable and RC settings.

Eigenvalues	Training Mode	Retrieval (max / min)	Image (max / min)	Pathfinder (max / min)	PathX (max / min)	Memory Function Profile
random	Learnable	0.903 / 0.899	0.820 / 0.815	0.870 / 0.846	0.877 / 0.841	Medium (minor training extension)
random	RC (Frozen)	0.901 / 0.896	0.733 / 0.698	0.824 / 0.777	0.820 / 0.773	Short (constrained by initial decay)
linear	Learnable	0.899 / 0.894	0.782 / 0.768	0.797 / 0.754	0.820 / 0.503	Short (restricted to real axis)
linear	RC (Frozen)	0.894 / 0.889	0.686 / 0.634	0.769 / 0.506	0.669 / 0.504	Very Short (singular values drop rapidly)
step	Learnable	0.883 / 0.841	0.789 / 0.779	0.858 / 0.824	0.840 / 0.824	Medium (complex conjugate pairs)
step	RC (Frozen)	0.903 / 0.897	0.745 / 0.740	0.832 / 0.793	0.813 / 0.504	Short (lacks polynomial projection basis)
S4Dinv	Learnable	0.863 / 0.760	0.864 / 0.859	0.910 / 0.877	0.922 / 0.903	Longest (optimal polynomial spectrum)
S4Dinv	RC (Frozen)	0.906 / 0.901	0.873 / 0.865	0.899 / 0.889	0.932 / 0.923	Longest & perfectly preserved

Key Findings¶

RC setting prevents overfitting and bolsters seed stability: In the Retrieval task, learnable S4Dinv exhibited substantial run-to-run instability, dropping to 0.760 min accuracy (0.863 max). In contrast, freezing weights via RC S4Dinv stabilized the minimum accuracy at 0.901 while raising the maximum to 0.906. On the challenging PathX task (16k steps), RC consistently beat the learnable counterpart.
Gradient descent fails to spontaneously expand zero-memory regions: Inspecting layer-wise supremum memory functions (MF) before and after training showed that non-structured initializations (random, linear, step) failed to achieve memory lengths comparable to untrained S4D even after converging. This confirms the theoretical claim that backpropagation cannot extend memory into regions where it is initially absent.
Accelerated convergence: On Image, Retrieval, and PathX, test loss under the RC setting dropped substantially faster than under learnable settings, reaching lower asymptotes in fewer epochs by eliminating conflicting gradient updates in the recurrent core.

Highlights & Insights¶

Memory Function as Gradient Filter: Establishes that the memory attenuation curve of a dynamical system is mathematically isomorphic to the projection kernel that filters teacher signals during backpropagation.
Rethinking End-to-End Recurrent Tuning: Shows that when structured initialization provides optimal long-range coverage, gradient adaptation of recurrent matrices is often redundant or counterproductive; fixing recurrent weights acts as an effective regularizer that cuts ~10% of parameters while boosting performance.
Theoretical Capacity Contrast with Transformers: Clarifies that standard attention mechanisms correspond to a linear RNN with a single eigenvalue $\lambda=1$ ($N=1$), whose state Vandermonde matrix collapses to rank 1, offering a theoretical perspective on why multi-eigenvalue SSMs achieve superior sequence memory efficiency.

Limitations & Future Work¶

Reliance on Linear and White-Noise Assumptions: The rigorous mathematical derivation assumes linear readouts and independent white-noise input distributions; extending the formal proof to deep non-linear couplings remains open.
Not Suited for Purely Local Tasks: In tasks requiring solely short-range context, fixed long-range HiPPO spectra needlessly squander memory capacity, where adaptive eigenvalue shortening would be preferable.
Extension to Time-Varying Architectures (Mamba): Mamba introduces input-dependent step sizes, transforming the linear system into a non-linear dynamical model that requires full non-linear IPC analysis.

vs HiPPO & S4 (Gu et al., 2020a, 2022b): HiPPO derived recurrence from function approximation over continuous time; this work explains from an optimization perspective why HiPPO initialization is essential for gradient flow during backpropagation.
vs Vanishing/Exploding Gradients (Zucchet & Orvieto, 2024): Traditional gradient vanishing originates from multiplicative chains across temporal unrollings; the phenomenon identified here occurs within a single sequence layer due to vanishing memory projection terms.
vs Classic Reservoir Computing (Jaeger, 2001b; Nakajima & Fischer, 2021): While classical RC relies on random recurrent weights that struggle on extreme sequence lengths, this work demonstrates that deterministic polynomial spectra in modern SSMs unlock peak performance under frozen RC optimization.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ Establishes a rigorous analytical bridge connecting memory functions with gradient directions, proving S4/S4D equivalence and proposing the RC training paradigm.
Experimental Thoroughness: ⭐⭐⭐⭐☆ Comprehensive evaluation across all 6 LRA benchmarks with multiple random seeds, accompanied by detailed layer-wise MF dynamics.
Writing Quality: ⭐⭐⭐⭐⭐ Rigorous mathematical derivation, crystal-clear conceptual progression, and transparent experimental alignment.
Value: ⭐⭐⭐⭐⭐ Provides fundamental theoretical insights into initialization, structural simplification, and training strategies for state space models.