Aligning Inductive Bias for Data-Efficient Generalization in State Space Models

Qiyu Chen; Guozhang Chen

cs.LGSep 25, 2025

Aligning Inductive Bias for Data-Efficient Generalization in State Space Models

Authors: Qiyu Chen, Guozhang Chen

Organizations: School of Computer Science, National Key Laboratory for Multimedia Information Processing, Peking University, China · School of Physics, Peking University, China

Abstract

The remarkable success of modern AI has been closely tied to scaling laws, yet the finite supply of high-quality data makes data efficiency--learning more from less--an increasingly important frontier. A model's inductive bias is a critical lever for data efficiency, but foundational sequence models such as State Space Models (SSMs) often rely on fixed, task-agnostic biases. When this fixed prior is misaligned with the underlying structure of a task, the model may require additional samples to overcome its own bias before learning the relevant signal. In this work, we introduce a principled framework for understanding and aligning the inductive bias of linear time-invariant SSMs. We first formalize this bias through an SSM-induced kernel and show theoretically and empirically that its spectrum is governed by the model's frequency response. This characterization motivates Task-Dependent Initialization (TDI), a fast power-spectrum matching method that aligns the initial SSM bias with the task's spectral characteristics before downstream training. Across controlled synthetic experiments, trainable one-layer SSMs, and deep SSMs on diverse real-world benchmarks, TDI can improve data-efficient generalization primarily when task-relevant spectral structure is present and the default SSM bias is spectrally mismatched. Our results provide both a theoretical lens and a practical tool for task-adaptive inductive bias, suggesting a path toward more data-efficient sequence modeling.

Figures & tables

Figure 1: Graphic overview. (a) A fixed, task-agnostic SSM initialization can place spectral energy in frequency regions that do not match the task-relevant spectrum, analogous to using an unsuitable tool for a job. (b) TDI initializes the SSM so that its frequency response aligns with task-relevant spectral components and assigns more energy to task-important modes. (c) Across trainable one-layer SSMs and deep SSMs, TDI achieves higher mean final-epoch test accuracy under fixed training budgets. A training ratio $r$ denotes training on an $r$ -fraction of the original training split, while keeping the validation and test splits fixed.

Panel of Figure 2 — Figure 2: Sanity check of the SSM-induced kernel eigendecomposition. (a) The empirical eigenvalue spectrum matches the theoretical prediction $\eta_{\rho}=s_{\rho}^{2}$ obtained from the singular values of the Toeplitz convolution operator. (b) The empirical and predicted eigenfunctions are well aligned across most modes. (c) The predicted eigenbasis accurately captures the cumulative task power $C(\rho)$ . (d) Kernel regression generalization predicted from the eigenspectrum matches empirical KRR performance across sample sizes. Error bars indicate standard deviations over 50 independent runs.

Panel of Figure 3 — Figure 3: Kernel regression mechanism experiment. We fix the initial model to be a low-frequency SSM and compare a matched low-frequency task with a mismatched high-frequency task. (a)-(c) In the matched case, the initial model spectrum already overlaps with the task spectrum, so TDI produces little change in the model spectrum, cumulative task power, or generalization. (d)-(f) In the mismatched case, the task-relevant high-frequency components are weakly represented by the initial model. TDI shifts the model spectrum toward these frequencies, moves task power into earlier high-eigenvalue modes, and shifts the kernel regression generalization curve downward across the finite-sample range. This supports the view that TDI improves data efficiency by correcting spectral mismatch rather than by universally improving all initializations.

	Low-data regime				High-data regime
Dataset	Base Acc.	TDI Acc.	Base Loss	TDI Loss	Base Acc.	TDI Acc.	Base Loss	TDI Loss
Binary Freq.	$86.5{\scriptstyle\pm 9.0}$ %	$88.6{\scriptstyle\pm 6.8}$ %	$0.292{\scriptstyle\pm 0.163}$	$0.256{\scriptstyle\pm 0.135}$	$97.5{\scriptstyle\pm 0.5}$ %	$97.3{\scriptstyle\pm 0.6}$ %	$0.070{\scriptstyle\pm 0.012}$	$0.076{\scriptstyle\pm 0.014}$
ECG5000	$62.9{\scriptstyle\pm 11.8}$ %	$72.8{\scriptstyle\pm 14.0}$ %	$1.158{\scriptstyle\pm 0.196}$	$0.981{\scriptstyle\pm 0.615}$	$90.2{\scriptstyle\pm 1.6}$ %	$90.5{\scriptstyle\pm 1.1}$ %	$0.375{\scriptstyle\pm 0.065}$	$0.370{\scriptstyle\pm 0.053}$
FordA	$52.2{\scriptstyle\pm 3.5}$ %	$56.8{\scriptstyle\pm 7.7}$ %	$0.689{\scriptstyle\pm 0.016}$	$0.837{\scriptstyle\pm 0.413}$	$68.9{\scriptstyle\pm 5.6}$ %	$70.4{\scriptstyle\pm 5.1}$ %	$0.569{\scriptstyle\pm 0.065}$	$0.548{\scriptstyle\pm 0.052}$
FordB	$53.5{\scriptstyle\pm 3.5}$ %	$55.0{\scriptstyle\pm 3.9}$ %	$0.729{\scriptstyle\pm 0.054}$	$0.763{\scriptstyle\pm 0.138}$	$60.6{\scriptstyle\pm 2.7}$ %	$61.7{\scriptstyle\pm 2.8}$ %	$0.669{\scriptstyle\pm 0.039}$	$0.640{\scriptstyle\pm 0.014}$
sMNIST	$61.1{\scriptstyle\pm 13.3}$ %	$78.4{\scriptstyle\pm 11.5}$ %	$1.144{\scriptstyle\pm 0.378}$	$0.696{\scriptstyle\pm 0.353}$	$79.3{\scriptstyle\pm 26.5}$ %	$88.2{\scriptstyle\pm 6.2}$ %	$0.375{\scriptstyle\pm 0.144}$	$0.364{\scriptstyle\pm 0.175}$

Table 1: Dataset-level low- and high-data summary for the one-layer SSM baseline and TDI. Low-data averages ratios

\leq 0.16

; high-data averages ratios

\geq 0.32

. Each cell reports the mean across ratios within the corresponding regime, where each ratio is first averaged over 3 random seeds; the standard deviation is computed across these ratio-level means. Accuracy and loss are the final-epoch test metrics. The better value between baseline and TDI is bolded.

	Low-data regime				High-data regime
Dataset	Base Acc.	TDI Acc.	Base Loss	TDI Loss	Base Acc.	TDI Acc.	Base Loss	TDI Loss
CIFAR-10	33.9 $\pm$ 15.7 %	35.4 $\pm$ 14.5 %	2.026 $\pm$ 0.190	2.012 $\pm$ 0.184	73.7 $\pm$ 7.2 %	69.4 $\pm$ 8.5 %	0.948 $\pm$ 0.344	1.094 $\pm$ 0.406
Freq. Classification	88.1 $\pm$ 18.9 %	91.1 $\pm$ 11.5 %	0.385 $\pm$ 0.590	0.331 $\pm$ 0.439	99.2 $\pm$ 0.4 %	98.9 $\pm$ 0.5 %	0.024 $\pm$ 0.014	0.026 $\pm$ 0.012
ListOps	16.6 $\pm$ 7.1 %	16.2 $\pm$ 6.6 %	5.419 $\pm$ 0.754	5.597 $\pm$ 0.941	44.7 $\pm$ 6.9 %	43.2 $\pm$ 6.2 %	1.928 $\pm$ 0.786	1.816 $\pm$ 0.584
sMNIST	56.3 $\pm$ 42.2 %	67.4 $\pm$ 38.1 %	1.145 $\pm$ 1.084	0.897 $\pm$ 0.991	99.1 $\pm$ 0.2 %	99.0 $\pm$ 0.2 %	0.034 $\pm$ 0.011	0.038 $\pm$ 0.011
Pathfinder	51.4 $\pm$ 2.4 %	51.7 $\pm$ 2.6 %	3.482 $\pm$ 0.853	3.343 $\pm$ 0.864	66.6 $\pm$ 10.4 %	68.4 $\pm$ 8.1 %	0.935 $\pm$ 0.567	0.891 $\pm$ 0.407

Table 3: Dataset-level low- and high-data summary for the deep SSM baseline and TDI. Low-data averages ratios

\leq 0.100

; high-data averages ratios

\geq 0.215

. Each cell reports the mean across ratios within the corresponding regime, where each ratio is first averaged over 5 random seeds; the standard deviation is computed across these ratio-level means. Accuracy and loss are final-epoch test metrics. The better value between baseline and TDI is bolded.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure C.1: Stability of the Fisher log-power task spectrum under subsampling. We compute the Fisher log-power spectrum on sequential MNIST using six training-subset ratios while keeping the preprocessing and frequency grid fixed.The estimated spectra become relatively stable once the number of examples reaches the order of $10^{2}$ : the dominant spectral peaks remain largely consistent across larger subsets, while additional examples mainly reduce small fluctuations rather than changing the main task-relevant frequency bands.

Dataset	First stable ratio	Observation
Freq. Cls.	$2.15\%$	Identifiable transition
sMNIST	$2.15\%$	Identifiable transition
Speech Commands	$4.64\%$	Identifiable transition
SC10	$4.64\%$	Identifiable transition
CIFAR-10	$21.5\%$	Negative TDI gains on both sides of the transition
PathFinder	Not reached	No transition identified on the evaluated grid

Appendix

Table C.1: Bootstrap stability of the Fisher log-power spectrum. The first stable ratio satisfies the bootstrap criterion for all five subsets and all subsequent evaluated sizes. Ratios refer to fractions of each dataset’s original training split. Stability does not imply a positive TDI accuracy gain.

Experiment	SSM parameterization	TDI variant
Kernel regression	Complex diagonal SSM	Peak initialization + gradient refinement (Algorithm 3 )
One-layer RNN	Real-valued matrix SSM	Peak initialization + gradient refinement (Algorithm 4 )
Deep SSM	Complex diagonal SSM	Direct construction (Algorithm 3 )

Appendix

Table D.2: TDI variants used in different experiments.

S_{\mathrm{task}}\leftarrow\mathcal{N}\!\left(\operatorname{TaskSpectrum}(\mathcal{D}_{\mathrm{train}})\right).

Appendix

Algorithm 2 Unified Task-Dependent Initialization (TDI)

Panel of Figure 10 — Figure E.2: Sanity check of the SSM-induced kernel eigendecomposition. (a) The empirical eigenvalue spectrum matches the theoretical prediction $\eta_{\rho}=s_{\rho}^{2}$ obtained from the singular values of the Toeplitz convolution operator. (b) The empirical and predicted eigenfunctions are well aligned across most modes. (c) The predicted eigenbasis accurately captures the cumulative task power $C(\rho)$ . (d) Kernel regression generalization predicted from the eigenspectrum matches empirical KRR performance across sample sizes. Error bars indicate standard deviations over 50 independent runs.

Figure E.3: Learning curves of matched case under more parameter settings.

Figure E.4: Learning curves of mismatched case under more parameter settings.

Dataset	Train	Test	Classes	Seq. len.	Input dim.
Binary Freq.	40000	10000	2	256	1
ECG5000	500	4500	5	140	1
FordA	3601	1320	2	500	1
FordB	3636	810	2	500	1
MNIST	60000	10000	10	784	1
TwoPatterns	1000	4000	4	128	1

Appendix

Table E.3: Dataset statistics used in the one-layer trainable SSM experiments. The reported train/test counts correspond to the full dataset split before applying any training ratio.

Dataset	Epochs
Binary Freq.	15
ECG5000	15
FordA	10
FordB	15
MNIST	15
TwoPatterns	15

Appendix

Table E.4: Hyperparameters for the trainable one-layer SSM experiments. The only dataset-dependent training hyperparameter is the number of epochs.

Dataset	Train	Val	Test	Classes	Seq. len.	Input dim.
Freq. Classification	10000	1000	1000	10	1024	1
sMNIST	54000	6000	10000	10	784	1
pMNIST	54000	6000	10000	10	784	1
PathFinder	480001	59999	59999	2	1024	1
CIFAR-10	45000	5000	10000	10	1024	1
ListOps	96000	2000	2000	10	2048	–

Appendix

Table E.5: Dataset statistics used in the deep SSM experiments. The reported train/validation/test counts correspond to the full split before applying any training subset ratio.

Dataset	Depth	Features $H$	State Size $N$	Norm	Pre-norm	Dropout	LR	Batch Size	Epochs	WD	$(\Delta_{\min},\Delta_{\max})$
Freq. Classification	2	64	64	BN	True	0	0.005	64	10	0.01	$(0.001,0.1)$
sMNIST	4	256	64	LN	False	0.1	0.01	128	30	0.05	$(0.001,0.1)$
pMNIST	4	256	64	LN	False	0.1	0.01	128	30	0.05	$(0.001,0.1)$
CIFAR-10	4	256	64	LN	False	0.1	0.01	128	30	0.05	$(0.001,0.1)$
Speech Commands	4	256	64	BN	False	0	0.01	16	30	0.05	$(0.001,0.1)$
SC10	4	256	64	BN	False	0	0.01	16	30	0.05	$(0.001,0.1)$

Appendix

Table E.6: Hyperparameters used in the deep SSM experiments. LR is learning rate and WD is weight decay. BN and LN refer to Batch Normalization and Layer Normalization, respectively; Pre-norm indicates whether normalization is applied before the S4 block.

		Final test accuracy			First-5 test accuracy
Dataset	Ratio	Std. RNN	Baseline	TDI	Std. RNN	Baseline	TDI
Binary Freq.	0.01	$50.4\pm 0.9$ %	$73.0\pm 4.8$ %	$79.1\pm 0.9$ %	$50.2\pm 0.7$ %	$55.1\pm 3.5$ %	$58.6\pm 7.3$ %
Binary Freq.	0.02	$53.6\pm 2.5$ %	$83.6\pm 5.4$ %	$84.2\pm 1.3$ %	$50.8\pm 0.8$ %	$55.9\pm 3.9$ %	$54.1\pm 4.1$ %
Binary Freq.	0.04	$52.9\pm 3.8$ %	$86.6\pm 2.2$ %	$89.9\pm 5.5$ %	$51.0\pm 1.0$ %	$65.2\pm 12.1$ %	$70.9\pm 4.4$ %
Binary Freq.	0.08	$50.4\pm 1.7$ %	$92.5\pm 3.6$ %	$94.9\pm 0.7$ %	$51.2\pm 0.9$ %	$70.7\pm 3.5$ %	$79.1\pm 2.1$ %
Binary Freq.	0.16	$58.9\pm 11.8$ %	$96.7\pm 0.8$ %	$95.0\pm 2.2$ %	$51.8\pm 2.1$ %	$77.6\pm 4.6$ %	$78.5\pm 5.9$ %

Appendix

Table F.7: Complete final performance for Standard RNN, the SSM baseline, and TDI. The left block reports final test accuracy; the right block reports mean test accuracy over the first 5 epochs. Each cell reports mean

\pm

standard deviation over 3 seeds. The best value within each block for each dataset-ratio row is bolded, and the second-best value is underlined.

		Final test accuracy			First-5 test accuracy
Dataset	Ratio	Baseline	TDI	TDI-Frozen	Baseline	TDI	TDI-Frozen
CIFAR-10	0.005	$15.0\pm 1.3$ %	$15.0\pm 1.0$ %	$15.0\pm 0.9$ %	$10.7\pm 0.3$ %	$11.3\pm 0.4$ %	$11.3\pm 0.4$ %
CIFAR-10	0.010	$19.9\pm 1.3$ %	$23.8\pm 2.4$ %	$23.6\pm 2.3$ %	$10.8\pm 1.0$ %	$11.0\pm 0.7$ %	$11.0\pm 0.7$ %
CIFAR-10	0.022	$32.8\pm 3.5$ %	$38.3\pm 1.4$ %	$37.7\pm 1.7$ %	$12.4\pm 0.2$ %	$12.8\pm 1.3$ %	$12.8\pm 1.4$ %
CIFAR-10	0.046	$46.2\pm 1.2$ %	$46.4\pm 1.1$ %	$46.1\pm 0.9$ %	$14.6\pm 0.9$ %	$14.9\pm 1.1$ %	$14.8\pm 1.1$ %
CIFAR-10	0.100	$55.4\pm 1.1$ %	$53.3\pm 0.6$ %	$53.2\pm 0.7$ %	$18.9\pm 1.9$ %	$22.0\pm 3.6$ %	$21.9\pm 3.4$ %

Appendix

Table F.8: Complete final performance for the SSM baseline, TDI, and TDI-Frozen. The left block reports final test accuracy; the right block reports mean test accuracy over the first 5 epochs. Each cell reports mean

\pm

standard deviation over 5 seeds. The best value within each block for each dataset-ratio row is bolded, and the second-best value is underlined.

			Time-to-80% (epochs)		Early-10 (%)		Final $\Delta$ Acc.
Task	Train ratio	Epochs	Baseline	TDI	Baseline	TDI	(pp)
sMNIST	2.15%	100	34.2 ± 4.5	20.0 ± 4.6	10.19	14.15	+0.12 ± 0.47
sMNIST	21.5%	100	4.6 ± 0.9	3.0 ± 0.0	69.71	83.69	−0.08 ± 0.07
Freq. Cls.	2.15%	50	11.6 ± 1.7	9.6 ± 4.0	32.72	63.25	−0.69 ± 0.73
Freq. Cls.	10.0%	50	2.0 ± 0.0	1.2 ± 0.4	92.04	96.03	+0.29 ± 0.29
Freq. Cls.	21.5%	50	1.0 ± 0.0	1.0 ± 0.0	98.04	97.93	+0.13 ± 0.15

Appendix

Table F.9: Early learning and accuracy after extended training. Time-to-80% reports the number of epochs required to reach 80% test accuracy (mean

\pm

SD over five seeds). Early-10 reports mean test accuracy over the first ten epochs, averaged over seeds. Final

\Delta

Acc. reports the paired TDI-minus-baseline final-epoch accuracy difference (mean

\pm

SD), in percentage points (pp). Bold indicates the lower mean Time-to-80% and the higher Early-10 accuracy within each row; ties are both bolded.

Dataset	Train ratio	First-layer TDI	All-layer TDI	$\Delta$ (pp)
sMNIST	0.464%	10.47 ± 0.85	30.39 ± 14.76	+19.92 ± 14.81
sMNIST	1.000%	34.97 ± 23.93	80.41 ± 5.61	+45.45 ± 23.50
sMNIST	2.154%	91.64 ± 3.74	95.22 ± 1.28	+3.58 ± 4.15
CIFAR-10	0.464%	15.85 ± 1.36	17.79 ± 1.65	+1.94 ± 2.35
CIFAR-10	1.000%	22.63 ± 2.37	27.58 ± 3.32	+4.95 ± 3.43
CIFAR-10	2.154%	37.03 ± 1.92	33.93 ± 0.97	−3.09 ± 2.20

Appendix

Table F.10: First-layer versus all-layer TDI. Test accuracy (%) is reported as mean

\pm

standard deviation over five matched seeds.

\Delta

denotes the paired per-seed accuracy difference (all-layer minus first-layer), reported as mean

\pm

standard deviation in percentage points. Bold indicates the higher mean accuracy in each row.

Dataset	Train ratio	TDI-True	TDI-Permuted	Baseline
sMNIST	2.15%	20.0 ± 4.6	28.6 ± 8.7	34.2 ± 4.5
sMNIST	4.64%	11.4 ± 2.6	13.4 ± 2.9	15.4 ± 6.3
sMNIST	10%	5.2 ± 1.3	7.0 ± 2.1	10.6 ± 1.8
Freq. Cls.	2.15%	9.6 ± 4.0	21.6 ± 13.8	11.6 ± 1.7
Freq. Cls.	4.64%	2.8 ± 1.6	7.0 ± 3.7	4.2 ± 0.8
Freq. Cls.	10%	1.2 ± 0.4	3.8 ± 1.9	2.0 ± 0.0

Appendix

Table F.11: Effect of label permutation during task-spectrum estimation. Entries report epochs to reach 80% test accuracy (mean

\pm

SD over five matched seeds; lower is better). Bold indicates the lowest mean in each row.

Asset	Use in this paper	License / Terms
MNIST [ 17 ]	One-layer SSMs and Deep SSMs	MIT License
CIFAR-10 [ 29 ]	Deep SSMs	Official source
UCR Time Series Archive [ 16 ]	One-layer SSMs	Official source
PathFinder [ 43 ]	Deep SSMs	Apache-2.0 License
ListOps [ 32 , 43 ]	Deep SSMs	Apache-2.0 License
Speech Commands / SC10 [ 44 ]	Deep SSMs	CC-BY 4.0 License

Appendix

Table G.12: Existing assets used in the experiments.

Explore similar work

Temporal Task Diversity: Inductive Biases Under Non-Stationarity in Synthetic Sequence Modelling
May 18, 2026Afiq Abdillah Effiezal Aswadi, Oliver Britton, Ross Baker +1Inductive Bias Sequence Modeling
1Independent · University of Oxford
TIDES: Implicit Time-Awareness in Selective State Space Models
May 10, 2026Taylan Soydan, Miguel A. Bessa, Dirk Mohr +1State Space Models Sequence Modeling
AIMM, ETH Zürich Zurich, Switzerland
DSSMs: State Space Models with Explicit Memory via Delay Differential Equations
Jul 11, 2026Yixiao Qian, Song Chen, Jiaxu Liu +2State Space Models Sequence Modeling
College of Control Science and Engineering, Zhejiang University · Department of Mathematics, National University of Singapore · School of Mathematical Sciences, Zhejiang University

May 18, 2026, cs.LGJ/K move · Enter open · S save

Temporal Task Diversity: Inductive Biases Under Non-Stationarity in Synthetic Sequence Modelling

Afiq Abdillah Effiezal Aswadi, Oliver Britton, Ross Baker, Matthew Farrugia-Roberts

1Independent · University of Oxford

Inductive Bias Sequence Modeling Non-Stationarity Temporal Dynamics Deep Learning

Open paper AskPDF arXiv

Modern deep learning science often assumes that neural networks learn from a fixed data distribution. However, many practically important learning problems involve data distributions that change throughout training. How does such non-stationarity impact the inductive biases of deep learning towards models with different structural, generalisation, and safety properties? A fruitful testbed for studying inductive bias is in-context linear regression sequence modelling, where small transformers display strikingly different generalisation patterns depending on the diversity of the (fixed) training task distribution. In this paper, we explore the effect of diversifying the task distribution across training time, finding that such temporal diversity leads to an increased bias towards generalisation over memorisation.

Abstract

Explore similar work

Temporal Task Diversity: Inductive Biases Under Non-Stationarity in Synthetic Sequence Modelling

TIDES: Implicit Time-Awareness in Selective State Space Models

DSSMs: State Space Models with Explicit Memory via Delay Differential Equations