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

TL;DR

This work provides a theoretical framework for gradient-based learning in linear state space models via memory function (MF), showing that initialization must establish the longest possible memory to propagate teacher information, even at the cost of memory precision. It proves that MF—and thus learning dynamics—are governed primarily by the eigenvalues of the internal weight matrix, with S4 and S4D effectively sharing memory structure, while Vandermonde-based analyses reveal numerical instability for certain eigenvalue sets. The authors propose a reservoir computing (RC) approach with fixed eigenvalues, demonstrating faster convergence and reduced overfitting on long-range tasks, and show that, in many cases, RC rivals or surpasses learnable-eigenvalue setups. Overall, MF offers a principled lens to compare SSMs with Transformer-based models and suggests practical training strategies for long-memory tasks, while outlining avenues for extending the theory to nonlinear SSMs and attention mechanisms.

Abstract

State space models (SSMs) have gained attention by showing potential to outperform Transformers. However, previous studies have not sufficiently addressed the mechanisms underlying their high performance owing to a lack of theoretical explanation of SSMs' learning dynamics. In this study, we provide such an explanation and propose an improved training strategy. The memory capacity of SSMs can be evaluated by examining how input time series are stored in their current state. Such an examination reveals a tradeoff between memory accuracy and length, as well as the theoretical equivalence between the structured state space sequence model (S4) and a simplified S4 with diagonal recurrent weights. This theoretical foundation allows us to elucidate the learning dynamics, proving the importance of initial parameters. Our analytical results suggest that successful learning requires the initial memory structure to be the longest possible even if memory accuracy may deteriorate or the gradient lose the teacher information. Experiments on tasks requiring long memory confirmed that extending memory is difficult, emphasizing the importance of initialization. Furthermore, we found that fixing recurrent weights can be more advantageous than adapting them because it achieves comparable or even higher performance with faster convergence. Our results provide a new theoretical foundation for SSMs and potentially offer a novel optimization strategy.

Figures (15)

• Figure 1: The MF of five eigenvalue realizations. The left (right) panel shows the absolute eigenvalues (MF). The horizontal axis is the index of eigenvalues (delay $\tau$ of input). The system size $N$ is $32$, and the results were obtained by calculating Eq. \ref{['eq:original MF']}, where $T=1024$ was typically used in previous studies' models.
• Figure 2: Test loss and accuracy of three tasks. The top (bottom) row indicates test loss (accuracy), and each column corresponds to a different task. The horizontal axes represent epochs. Lines represent averages across seeds, with shaded regions indicating standard deviations. Solid lines represent models in the RC setting, while dotted lines indicate models with trainable eigenvalues. Colors indicate different eigenvalue realizations.
• Figure 3: The supremum MF of each layer before and after training. The top (bottom) panels show MFs before (after) training. From left to right, each column represents one layer. The horizontal (vertical) axis shows the input delay $\tau$ (MF). The five colors indicate the eigenvalue realizations. The state size $N$ is $32$, and the results are obtained by calculating Eq. \ref{['eq:original MF']} with $T=1024$. As an example, we show the results for the Pathfinder task.
• Figure 4: The linear readout weights $\bm{W}^{(\ell)}$ and $\bm{b}^{(\ell)}$ at each layer and epoch. From top to bottom, the rows correspond to epochs 0, 10, and the final epoch, respectively. From left to right, the columns show the weights of layers 1 through 6. Following Figure \ref{['fig:S4Dinv_params_kernel']}, we present an example from the Pathfinder task.
• Figure 5: The average MF of five eigenvalue realizations. The horizontal axis is the delay $\tau$ of input. The system size $N$ is $32$, and the results are obtained by calculating Eq. \ref{['eq:original MF']}, where $T=1024$.