Aligning Inductive Bias for Data-Efficient Generalization in State Space Models¶
Conference: NeurIPS 2026
arXiv: 2509.20789
Code: https://github.com/Criticality-Cognitive-Computation-Lab/Aligning-Inductive-Bias-for-State-Space-Models
Area: LLM Pretraining
Keywords: state space models, inductive bias, kernel spectral theory, spectral alignment, task-dependent initialization
TL;DR¶
This paper proves from a kernel-theoretic perspective that the inductive bias spectrum of linear time-invariant state space models (LTI SSMs) is governed by their frequency response, and introduces Task-Dependent Initialization (TDI) to align the initial SSM response with task spectra, significantly boosting data efficiency under limited sample regimes without adding parameters or inference overhead.
Background & Motivation¶
The success of modern AI and large sequence models is intimately tied to empirical scaling laws, where performance improves predictably with larger models, dataset sizes, and compute budgets. However, in scientific discovery, healthcare, industrial control, and long-tail domain applications, high-quality labeled sequence data is scarce and costly to collect, making data efficiency—the capability to extract maximum generalizable performance from limited observations—a pivotal research frontier. In finite-sample regimes, generalization capability hinges heavily on a model's inductive bias: when built-in priors coincide with the underlying task structure, the sample complexity decreases substantially; when mismatched, models must consume valuable data samples simply to unlearn or compensate for their unfavorable initial bias.
State Space Models (SSMs), such as S4, S4D, and the convolutional backbones of hybrid selective architectures, have demonstrated exceptional capability in long sequence modeling with sub-quadratic compute. Nevertheless, existing continuous-time discretization and initialization frameworks (such as HiPPO-based or diagonal S4D schemes) are fundamentally task-agnostic, imposing an immutable spectral prior before downstream training. Theoretical parameter-level analyses reveal that standard S4D initializations concentrate the vast majority of transfer function variation in a narrow, low-frequency band; when downstream prediction requires high-frequency temporal features or specific resonant modes, this default low-frequency inductive bias becomes an acute sample-efficiency bottleneck.
To bridge this gap, this paper analyzes SSM filtering through reproducing kernel Hilbert space (RKHS) generalization theory, explicitly formalizing the eigenspectrum of the kernel induced by the LTI SSM convolution operator. Core idea: formalize LTI SSMs as convolutional feature maps whose induced kernel spectrum is asymptotically governed by the squared frequency response magnitude \(|H(\omega)|^2\), and propose Task-Dependent Initialization (TDI) to align initial SSM poles with data-driven task spectral statistics prior to downstream training.
Method¶
Overall Architecture¶
The framework establishes an end-to-end pipeline spanning "induced kernel spectral characterization \(\to\) label-guided task spectrum estimation \(\to\) parameter-level resonant pole alignment". Specifically, an input sequence is processed through the discrete LTI SSM convolution map \(\Phi_{\mathrm{SSM}}(\boldsymbol{u}) = \frac{1}{\sqrt{L}}\boldsymbol{T}_h \boldsymbol{u}\). The singular value spectrum of the causal lower-triangular Toeplitz matrix \(\boldsymbol{T}_h\) governs the kernel eigenvalues, which asymptotically coincide with the squared magnitude of the discrete-time Fourier transform of the SSM impulse response. To align this model prior with the target task, the method computes a task-relevant spectrum (either an input-output cross spectrum or class-stratified Fisher log-power spectrum) from available training data. Finally, resonant poles are strategically allocated to place peak frequency response on task-dominant frequencies, optionally refined via a normalized spectral matching loss.
Downstream supervised training proceeds using standard architectures, loss objectives, and optimizers without any modifications or test-time inference overhead.
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["Training Sequence Data & Labels<br/>D = {(u, Y)}"] --> B["Task Spectrum Estimation<br/>Fisher Log-Power / Cross-Spectrum"]
B --> C["Resonant Pole Peak Initialization<br/>Map dominant task frequencies to stable poles"]
C --> D["Spectral Matching Optimization<br/>Refine parameters via L_spec minimization"]
D --> E["Standard Supervised Training & Inference<br/>Zero extra parameters / Zero inference overhead"]Key Designs¶
1. SSM-Induced Kernel & Frequency Response Equivalence: Connecting State Space Operators to RKHS Spectra
To understand which sequence patterns an LTI SSM favors in finite-sample regimes, the method formalizes the discrete convolution operator \(\boldsymbol{T}_h \in \mathbb{R}^{L \times L}\) as a deterministic linear feature map, defining the SSM-induced kernel as \(K_{\mathrm{SSM}}(\boldsymbol{u}, \boldsymbol{u}') = \frac{1}{L}(\boldsymbol{T}_h \boldsymbol{u})^\top (\boldsymbol{T}_h \boldsymbol{u}')\). Under whitened input assumptions, applying singular value decomposition \(\boldsymbol{T}_h = \sqrt{L}\boldsymbol{U}\boldsymbol{S}\boldsymbol{V}^\top\) yields kernel eigenvalues \(\eta_\rho = s_\rho^2\) with orthonormal eigenfunctions \(\phi_\rho(\boldsymbol{u}) = \boldsymbol{v}_\rho^\top \boldsymbol{u}\). By applying Gray's asymptotic equivalence theorem between Toeplitz operators and circulant matrices as sequence length \(L \to \infty\), the singular values \(\tilde{s}_j\) associated with discrete Fourier frequencies \(\omega_j = 2\pi j / L\) converge asymptotically to the frequency response:
where \(H(\omega) = \sum_{n=0}^\infty h_n e^{-i\omega n}\) is the discrete-time Fourier transform of impulse response \(h_n = \overline{\boldsymbol{C}}\overline{\boldsymbol{A}}^n \overline{\boldsymbol{B}}\). Under Mercer kernel regression theory, target modes corresponding to large kernel eigenvalues \(\eta_\rho\) are learned rapidly in the small-sample regime. Consequently, when task-relevant signals reside in low-gain frequency regions of the SSM filter, the model suffers from severe sample inefficiency.
2. Task Spectrum Estimation: Extracting Label-Conditioned Discriminative Spectral Statistics
To construct an informative prior from limited labeled samples, the initialization pipeline estimates the underlying task spectrum. In theoretical analysis, the input-output cross-covariance \(\mathbf{C}_{\boldsymbol{u}\boldsymbol{Y}} = \frac{1}{P} \sum_{\mu=1}^P (\boldsymbol{u}^\mu - \bar{\boldsymbol{u}})(\boldsymbol{Y}^\mu - \bar{\boldsymbol{Y}})^\top \in \mathbb{R}^{L \times d}\) yields the cross-spectrum \(S_{\mathrm{cross}}(\omega_j) = \frac{1}{d}\sum_{r=1}^d \left|\mathcal{F}[\mathbf{C}_{\boldsymbol{u}\boldsymbol{Y}}^{(:, r)}](\omega_j)\right|^2\). For classification tasks where temporal phase is arbitrary but energy across frequency bands differentiates classes, the method implements the Fisher log-power spectrum \(S_{\mathrm{FLP}}(\omega_k)\).
Extracting log-power features \(z^\mu(\omega_k) = \log(|\mathrm{rFFT}(\boldsymbol{u}^\mu)(\omega_k)|^2 + \epsilon)\), the criterion computes class-prior-weighted between-class variance over within-class variance:
where \(\pi_j\) is the class prior, and \(\mu_j(\omega_k), \sigma_j^2(\omega_k)\) denote class-conditional means and variances. This formulation selectively suppresses background frequencies carrying large marginal power but zero label correlation, isolating the precise bands essential for category discrimination.
3. Resonant Pole Peak Initialization & Spectral Matching Loss: Parameter-Level Alignment
Given normalized task spectrum \(S_{\mathrm{task}}\), TDI modifies the SSM dynamic matrices in two complementary steps. First, Peak Initialization directly selects dominant frequencies \(f_j\) from \(S_{\mathrm{task}}\) and configures continuous-time stable poles. For complex diagonal parameterizations (such as S4D), poles are initialized as \(A_j = -\gamma_j - i 2\pi f_j / \Delta\), where damping \(\gamma_j > 0\) determines resonance sharpness and \(\Delta > 0\) is the discretization step; real-valued SSMs employ corresponding \(2 \times 2\) anti-symmetric resonant blocks. Second, in kernel regression and single-layer settings, parameters are refined by minimizing the normalized spectral matching loss:
Minimizing \(\mathcal{L}_{\mathrm{spec}}\) maximizes the alignment objective \(J = \frac{1}{L}\|\boldsymbol{T}_h \mathbf{C}_{\boldsymbol{u}\boldsymbol{Y}}\|_F^2 = \sum_{\rho=1}^L \eta_\rho \bar{w}_\rho^2\). By the rearrangement inequality, this shifts task projection weights \(\bar{w}_\rho^2\) into leading eigenvalue modes, causing the cumulative power curve \(C(\rho)\) to rise faster and accelerating early generalization.
Loss & Training¶
Following initialization, \(\mathcal{L}_{\mathrm{spec}}\) is discarded and downstream supervised training uses standard cross-entropy or mean squared error loss. In deep multi-layer SSMs (e.g., stacked S4D blocks), the authors demonstrate that applying TDI exclusively to the first layer via constructive peak initialization (without gradient refinement) provides substantial data-efficiency gains while keeping compute minimal.
Key Experimental Results¶
Main Results¶
The method was evaluated on trainable one-layer SSMs (7 sequence datasets, sweeping training ratio \(r \in [0.01, 1.0]\)) and deep S4D models (8 sequence benchmarks, sweeping ratio \(r \in [10^{-3}, 1.0]\)). The table below reports mean final-epoch test accuracy across low- and high-data regimes (compiled from Table 1 and Table 3 of the original paper).
| Architecture | Dataset | Low-Data Base Acc. | Low-Data TDI Acc. | High-Data Base Acc. | High-Data TDI Acc. |
|---|---|---|---|---|---|
| One-Layer SSM | ECG5000 | 62.9 ± 11.8% | 72.8 ± 14.0% | 90.2 ± 1.6% | 90.5 ± 1.1% |
| One-Layer SSM | sMNIST | 61.1 ± 13.3% | 78.4 ± 11.5% | 79.3 ± 26.5% | 88.2 ± 6.2% |
| One-Layer SSM | Wafer | 79.8 ± 9.5% | 87.4 ± 6.2% | 94.9 ± 2.9% | 95.5 ± 3.3% |
| One-Layer SSM | FordA | 52.2 ± 3.5% | 56.8 ± 7.7% | 68.9 ± 5.6% | 70.4 ± 5.1% |
| Deep S4D | sMNIST | 56.3 ± 42.2% | 67.4 ± 38.1% | 99.1 ± 0.2% | 99.0 ± 0.2% |
| Deep S4D | Freq. Cls. | 88.1 ± 18.9% | 91.1 ± 11.5% | 99.2 ± 0.4% | 98.9 ± 0.5% |
| Deep S4D | Speech Commands | 51.8 ± 31.5% | 55.7 ± 29.7% | 87.3 ± 1.6% | 87.9 ± 2.2% |
| Deep S4D | SC10 | 49.2 ± 31.8% | 52.8 ± 31.9% | 92.4 ± 1.1% | 92.7 ± 0.7% |
Paired Improvement & Convergence Analysis¶
To isolate the mechanism across varying training budgets, paired percentage-point improvements over matched dataset-ratio-seed triples for the one-layer SSM are summarized below (from Table 2(a)), along with early-epoch tracking (Early-5 and Early-10):
| Training Ratio \(r\) | Final \(\Delta\)Acc. (%) | Early-5 \(\Delta\)Acc. (%) | Early-10 \(\Delta\)Acc. (%) | Note |
|---|---|---|---|---|
| 0.01 | +5.4 [+1.8, +10.3] | +3.5 [-0.7, +9.7] | +3.2 [+0.8, +6.1] | Strong gain at extreme small samples |
| 0.02 | +7.5 [+0.9, +15.1] | -0.2 [-2.7, +3.1] | +2.1 [-0.4, +5.1] | Peak final accuracy gain |
| 0.04 | +6.5 [+2.3, +11.2] | +3.2 [-1.5, +9.0] | +3.6 [+1.2, +6.8] | Consistent statistical advantage |
| 0.08 | +7.3 [+4.0, +10.8] | +2.7 [+0.4, +5.3] | +4.7 [+3.2, +6.1] | Stable early & final gains |
| 0.16 | +5.6 [+1.5, +10.0] | +1.7 [-1.2, +4.6] | +2.7 [-0.1, +5.7] | Significant improvement over baseline |
| 0.32 | +2.3 [+0.7, +4.0] | +1.5 [+0.0, +3.1] | +1.3 [-0.1, +2.9] | Margin narrows as data increases |
| 0.64 | -1.2 [-6.7, +2.3] | -1.2 [-4.3, +1.6] | -0.6 [-2.2, +1.0] | Confidence interval spans zero |
| 1.00 | -2.3 [-14.1, +7.3] | -0.6 [-4.4, +3.0] | -1.6 [-7.8, +3.3] | Baseline catches up after extended training |
Key Findings¶
- Gains concentrate in limited-data regimes: In small-sample regimes (\(r \le 0.16\)), TDI provides marked generalization gains. When full training sets are provided (\(r \ge 0.64\)), gradient optimization eventually adjusts filter parameters regardless of initialization, eliminating the performance gap.
- Substantially accelerated early learning: On sMNIST with 2.15% data, TDI reaches 80% test accuracy in 20.0 epochs compared to 34.2 epochs for the baseline. Across deep SSM benchmarks, TDI achieves higher Early-10 accuracy in 45/70 evaluated conditions (and 39/50 across tasks suitable for 1D spectral metrics).
- Clear applicability boundaries: On ListOps (hierarchical discrete syntax) and PathFinder (2D spatial path tracing), TDI yields negligible differences (e.g., -0.4% \(\Delta\)Acc on low-data ListOps), reflecting the limitation of 1D Fourier statistics on non-Euclidean or compositional tasks.
Highlights & Insights¶
- Bridging LTI frequency response to RKHS theory: Elegantly connects causal Toeplitz operators and circulant matrix asymptotics to provide an exact mathematical equivalence between \(|H(\omega)|^2\) and kernel eigenvalue distributions.
- Zero-overhead task-adaptive inductive bias: Adapts SSM priors solely during the initialization phase using fast FFT statistics; architecture, parameter counts, and inference compute remain completely identical to the vanilla model.
- Principled, non-universal framing: Authors explicitly characterize TDI not as a universal accuracy booster for large-scale training, but as an inductive bias correction tool specifically engineered for data-constrained, spectrally mismatched regimes.
Limitations & Future Work¶
- Restriction to LTI models: Theoretical derivations rely on linear time-invariance; extending this framework to input-dependent selective models such as Mamba requires local perturbation approximations around an LTI reference.
- Limitation of 1D spectral statistics: One-dimensional Fourier statistics cannot easily encode multi-dimensional spatial topology or discrete token semantics, leaving ListOps and PathFinder outside the target applicability zone.
- Sensitivity at extreme sample sparsity: Below stability transition thresholds (e.g., \(r < 0.001\)), Fisher log-power spectrum estimates suffer from empirical noise, which can degrade alignment quality.
Related Work & Insights¶
- vs. S4 / S4D Initialization (Gu et al., 2022): S4D relies on fixed HiPPO-derived low-frequency priors regardless of the task; TDI diagnoses when this prior causes sample-efficiency bottlenecks and aligns initial poles with task-relevant bands.
- vs. Empirical Frequency Bias Tuning (Yu et al., 2025): Yu et al. tune frequency response heuristically by scaling \(\Delta\); TDI derives a principled spectral matching loss \(\mathcal{L}_{\mathrm{spec}}\) rooted in kernel alignment theory.
- vs. Autocorrelation-Based Initialization (Liu & Li, 2025): Autocorrelation methods rely only on input signals; TDI incorporates label information via cross-spectra and Fisher log-power, preventing high-energy background distractors from hijacking the initial bias.
Rating¶
- Novelty: ⭐⭐⭐⭐⭐ [Establishes the first rigorous theoretical link between LTI SSM frequency response and kernel inductive bias]
- Experimental Thoroughness: ⭐⭐⭐⭐☆ [Span from kernel regression to 1-layer SSMs and deep S4D across diverse datasets, with cluster bootstrap validation]
- Writing Quality: ⭐⭐⭐⭐⭐ [Precise mathematical formalization, self-consistent narratives, and objective discussion of limitations]
- Value: ⭐⭐⭐⭐☆ [Offers an effective, zero-inference-overhead initialization framework for low-data scientific and sequential tasks]