Elastic Spectral State Space Models for Train-Once Budgeted Inference

Dachuan Song; Junyu Yin; Zechen Hu; Xuan Wang

cs.LGJan 30, 2026

Elastic Spectral State Space Models for Train-Once Budgeted Inference

Authors: Dachuan Song, Junyu Yin, Zechen Hu, Xuan Wang

Organizations: Department of Electrical and Computer Engineering, George Mason University, Fairfax, VA, USA · Department of Computer Science, George Mason University, Fairfax, VA, USA

Abstract

Modern sequence models are typically trained at a fixed computational capacity, while real-world applications require deployment across platforms with different resource constraints. Existing approaches either train or distill separate compact models for selected budgets, or use elastic architectures that shrink dimensions such as width, depth, Feed-Forward Network (FFN) size, or SSM state dimension. In this paper, we propose Elastic Spectral State Space Models (ES-SSM), a train-once, export-many sequence modeling framework that gains elasticity through spectral approximation of the SSM sequence operator. ES-SSM builds on Hankel spectral filtering for state space models, where long-range token mixing is represented through fixed Hankel spectral channels that define an operator-level approximation resolution. To make this spectral elasticity reliably deployable, ES-SSM combines input-adaptive channel-wise gates with budget dropout, training the same spectral prefixes that are used by compact models at inference time. This encourages low-index channels to become predictive on their own, while higher-index channels provide refinements for larger budgets. We evaluate ES-SSM across byte-level language modeling, Long Range Arena, Speech Commands V2, and offline reinforcement learning benchmarks. Across these settings, a single trained ES-SSM can be truncated to competitive compact models compared with modern Transformer and SSM baselines at similar parameter scales. Furthermore, by testing under various runtime budgets, we observe smooth and stable quality-cost curves over a wide range of truncation levels. Code is available at https://github.com/GMU-MICO/ES-SSM.

Figures & tables

Figure 1: ES-SSM architecture overview. [A] ES-SSM is embedded in a pre-norm residual block with an ES-SSM spectral sublayer followed by an FFN; residual paths carry the block input around each sublayer. [B] The ES-SSM spectral sublayer

\mathcal{S}_{K}

combines fixed Hankel spectral channels with a tokenwise Gate MLP that produces input-adaptive, channel-wise budgeted gates. For a runtime budget

K\leq\overline{K}

, direct prefix truncation keeps channels

1,\ldots,K

in the exported compact model and removes tail channels

k>K

. Each retained channel is projected by

M_{k}

, modulated by the scalar gate

g_{k}^{(K)}(t)

, and added to the retained-channel sum. Thus, one full-capacity training can be truncated into compact models at multiple runtime budgets.

Figure 2: Budget dropout sampler and prefix-truncation behavior. Left: ES-SSM samples a log-scale target

K_{\mathrm{target}}

and snaps it to the nearest element of the runtime-budget set

\mathcal{K}

, placing more training pressure on small-

K

prefixes while preserving coverage of larger budgets. Right: budget dropout trains the exported prefixes directly, yielding stable validation bits-per-byte (BPB) across runtime budgets; without budget dropout, a model trained only at

\overline{K}

degrades sharply under direct prefix truncation.

Figure 3: SlimPajama main-table visualization for the

1.3

B-family,

30

B-token scaling experiment. [A] Test BPB (lower is better) vs. runtime budget

K

(log scale). [B] Test BPB vs. inference throughput (kt/s, log scale; sequence length

L{=}2048

, batch

1

, single B200 GPU). The two panels share the same legend on the right. Solid lines denote methods with elastic capabilities, whereas dashed lines denote non-elastic baselines trained independently at each shown capacity. We note that these dashed non-elastic baselines require significantly more training time and therefore, should not be used for direct comparison with our method. Instead, they are presented as oracle results for reference. Markers are the data points in Appendix G.1 , Table 11 ; for ES-SSM, the ten

K

values instantiate exported parameter budgets spanning roughly

330

M to around

1.3

B. The light-pink band highlights the high-BPB degradation region (Test BPB

>1.5

). Inline values annotate the headline operating points. For readability, both panels use a piecewise-linear

y

-axis that stretches the competitive band

[0.90,1.10]

and compresses the collapse range. Per-family numerical results are reported in Appendix G.1 .

Figure 4: D4RL quality and training-cost comparison. [A] Quality (avg. normalized return, larger means better) vs. runtime budget

K

. ES-SSM (solid blue) is a single

\overline{K}{=}32

training exported at five budgets via prefix truncation; Transformer (dashed red) is five independently trained models, with the dashed connector serving only as a visual guide. The bottom strip shows the per-

K

delta (ES-SSM

-

Transformer);

\checkmark

\approx

\times

denote win, tie, and loss (tie threshold

0.5

pts). [B] Total training wall-clock on the full D4RL suite (9 datasets

\times

3 seeds, RTX 3090). The Transformer bar is segmented into its five separate training runs, while ES-SSM trains a single model. The teal arrow reports the savings. Full numerical tables are reported in Appendix G.4 .

Figure 5: LRA budget sweep at the matched

\sim\!50

M scale. Each curve is obtained by direct prefix truncation from the same

\overline{K}{=}32

ES-SSM checkpoint. Gold stars mark the smallest

K

retaining at least

98\%

of the full-budget score; black triangles mark the smallest tested budget

K

that retains at least

90\%

of the full-budget score. Across all six tasks, ES-SSM degrades smoothly as

K

decreases rather than collapsing abruptly. Full LRA tables are reported in Appendix G.2 .

Appendix figures & tables32 assets

Supplementary material from the paper’s appendix.

Appendix

Config	main	spec.	Val BPB K=2	Val BPB K=8	Val BPB ${}_{K=\overline{K}}$	$\Delta_{K=\overline{K}}$
B. naive bf16	bf16	low	1.7826	1.6064	1.6089	$+0.0063$
C. naive fp16+scaler	fp16	low	1.7841	1.6088	1.6114	$+0.0087$
D. bf16 + fp32 spectral ( ours )	bf16	fp32	1.7786	1.6002	1.6027	—
E. fp16 + fp32 spectral	fp16	fp32	1.7778	1.5997	1.6023	$-0.0004$

Appendix

Table 1: Precision-policy ablation at

160

M parameters on SlimPajama,

1.5

B byte-tokens,

\overline{K}{=}32

. All four configurations share architecture, optimizer, schedule, data, and random seed; only the AMP main-network dtype and whether the spectral branch is forced to fp32 differ. “low” in the spec. column means the spectral multiply runs in the AMP main dtype (i.e. the FFT itself is auto-cast to fp32 by PyTorch’s AMP fallback for FFT, but the complex multiply runs in bf16 / fp16 ).

\Delta_{K=\overline{K}}

is val BPB minus our policy (D) at

K{=}\overline{K}{=}32

Figure 6: $\Delta$ Val BPB relative to our policy (row D in table 1 ) at runtime budget $K\in\{2,4,8,16,\overline{K}\}$ . Positive entries (red) are worse than ours; negative (blue) are better. Both naive low-precision rows (B, C) are uniformly positive ( $10/10$ cells), and the fp16-naive row C is darker than the bf16-naive row B at every $K$ . The forced- fp32 spectral row (E) is statistically indistinguishable from D ( $|\Delta_{K}|<0.001$ BPB at every $K$ ).

Figure 7: Per-

K

validation-BPB difference relative to our precision policy (D) after

1.5

B SlimPajama tokens. Positive values are worse than our policy. The naive low-precision spectral configurations are worse at every runtime budget, while the fp16 main-network run with an fp32 spectral branch stays within the

\pm 1

mBPB band.

Figure 8: Training dynamics for the precision-policy ablation. Top: validation BPB at

K=\overline{K}=32

over training. Bottom: validation-BPB difference relative to our policy (D), in mBPB. The two fp32-spectral configurations track each other within a small band, while the naive low-precision spectral configurations keep a positive offset after early training.

Hyperparameter	Value
Optimizer	AdamW, $(\beta_{1},\beta_{2})\!=\!(0.9,0.95)$ , $\varepsilon\!=\!10^{-8}$
Weight decay	$0.1$ on non-LayerNorm, non-embedding, non-bias parameters
Peak learning rate	$3\!\times\!10^{-4}$
Warmup	Linear, first $2\%$ of total optimization steps
Decay	Cosine to $10\%$ of peak at the final step
Gradient clipping	Global $\ell_{2}$ norm clipped to $1.0$

Appendix

Table 2: Training recipe.

Figure 9: Full-budget anchor cadence sweep. Each cell shows validation BPB above the best cadence for that runtime budget, in milli-BPB and clipped at $120$ for readability. The black row marks the default $n_{\mathrm{full}}\!=\!8$ , which forces $K_{\mathrm{train}}\!=\!\overline{K}$ once every eight post-warmup optimizer steps; the green box marks the empirically stable band $n_{\mathrm{full}}\!\in\!\{4,6,8,12,16\}$ . The endpoint $n_{\mathrm{full}}\!=\!1$ degenerates to full-budget training and collapses small- $K$ prefixes, while very large cadences undertrain the tail channels. Exact means and standard deviations are reported in table 3 .

$n_{\mathrm{full}}$	$K\!=\!2$	$K\!=\!3$	$K\!=\!4$	$K\!=\!6$	$K\!=\!8$	$K\!=\!12$	$K\!=\!16$	$K\!=\!24$	$K\!=\!32$
$1$ ( $\equiv$ no_bd)	$4.245\pm 0.190$	$4.072\pm 0.109$	$4.049\pm 0.069$	$3.485\pm 0.060$	$3.005\pm 0.046$	$2.506\pm 0.065$	$2.316\pm 0.049$	$2.208\pm 0.032$	$2.187\pm 0.027$
$2$	$2.527\pm 0.032$	$2.331\pm 0.020$	$2.272\pm 0.021$	$2.227\pm 0.024$	$2.216\pm 0.024$	$2.209\pm 0.025$	$2.206\pm 0.025$	$2.205\pm 0.025$	$2.206\pm 0.025$
$3$	$2.486\pm 0.037$	$2.317\pm 0.023$	$2.267\pm 0.024$	$2.225\pm 0.023$	$2.215\pm 0.024$	$2.210\pm 0.024$	$2.208\pm 0.023$	$2.209\pm 0.024$	$2.213\pm 0.024$
$4$	$2.466\pm 0.038$	$2.308\pm 0.021$	$2.262\pm 0.022$	$2.222\pm 0.022$	$2.213\pm 0.023$	$2.209\pm 0.023$	$2.208\pm 0.023$	$2.211\pm 0.024$	$2.215\pm 0.024$
$6$	$2.457\pm 0.035$	$2.302\pm 0.021$	$2.257\pm 0.021$	$2.221\pm 0.022$	$2.213\pm 0.023$	$2.210\pm 0.023$	$2.210\pm 0.024$	$2.213\pm 0.024$	$2.218\pm 0.025$
$8$ (reference)	$2.454\pm 0.034$	$2.301\pm 0.021$	$2.256\pm 0.022$	$2.221\pm 0.022$	$2.213\pm 0.021$	$2.210\pm 0.022$	$2.211\pm 0.022$	$2.215\pm 0.022$	$2.222\pm 0.023$

Appendix

Table 3: Effect of the full-budget anchor cadence

n_{\mathrm{full}}

on validation BPB across runtime budgets, mean

\,\pm\,

std over

3

random seeds using a

287.1

M-parameter model trained for

80

M tokens per run (

\approx\!0.28

tokens per parameter). Two competing effects make the curve U-shaped: as

n_{\mathrm{full}}\!\downarrow

the

K\!=\!\overline{K}

channel gets more gradient and

K\!=\!\overline{K}

BPB falls monotonically; as

n_{\mathrm{full}}\!\downarrow

the small-

K

operating points get less probability mass and small-

K

BPB rises. The reference choice

n_{\mathrm{full}}\!=\!8

sits in a wide flat optimum spanning

n_{\mathrm{full}}\!\in\![4,16]

at every

K

, inside which the trade-off is empirically balanced —

K\!=\!2

val BPB lies in

[2.447,2.466]

and

K\!=\!\overline{K}

val BPB in

[2.215,2.226]

, all within two standard errors of one another. The catastrophic collapse at

n_{\mathrm{full}}\!=\!1

(

K\!=\!2

val BPB

4.245\!\pm\!0.190

) recovers the no-budget-dropout failure mode of table 23 as a degenerate limit. The sweep does not identify a single best

n_{\mathrm{full}}

— neither

4

nor

16

is statistically distinguishable from

8

on this seed set — but it rules out both endpoints and licenses

8

as a defensible round-number default rather than a hand-tuned one.

Training budget set	Mean val. BPB	Max. gap	$K=2$	$K=32$
Dense-31	2.2307	–	2.4838	2.2244
Paper-10	2.2319	0.0030	2.4848	2.2253

Appendix

Table 4: Sensitivity to the set of budgets sampled during training. Both checkpoints are evaluated on all 31 integer runtime budgets. Lower BPB is better.

Runtime budget $K$	40M-token warmup	Zero warmup	Absolute difference
$2$	$1.9552$	$1.9465$	$0.0087$
$8$	$1.8221$	$1.8182$	$0.0039$
$32$	$1.8205$	$1.8204$	$0.0001$

Appendix

Table 5: Validation BPB under a 40M-token full-budget warmup and zero warmup.

Method	Full-budget val. BPB $\downarrow$	Steps and byte-tokens	Analytic compute	Measured wall-clock
Always-full control	$1.839901$	$100.0\%$	$100.0\%$	$100.0\%$
ES-SSM	$\mathbf{1.834882}$	$117.3\%$	$\mathbf{69.8\%}$	$\mathbf{60.6\%}$

Appendix

Table 6: Training resources required to reach matched full-budget validation quality. Resource columns are normalized to the always-full control (

100\%

); lower values indicate lower total cost. The ES-SSM row corresponds to the first measured checkpoint that meets the target quality.

Component	Expression	Scales with $K$ ?
Channel projections $\{M_{k}\}_{k\leq K}$	$K\,d^{2}$	yes
Gate MLP output row $\mathbf{W}_{g,2}[:K,:]$	$K\,(d/2)$	yes
Gate MLP output bias $\mathbf{b}_{g,2}[:K]$	$K$	yes
Gate MLP hidden layer $\mathbf{W}_{g,1},\mathbf{b}_{g,1}$	$d\,(d/2)+d/2$	no
FFN (widening $4$ )	$8\,d^{2}+5\,d$	no
$2\times$ LayerNorm (pre-spectral, pre-FFN)	$4\,d$	no

Appendix

Table 7: Per-layer parameter ledger. “Buffer” entries are cached spectral quantities and are not learnable parameters.

Stage	FLOPs / token	Scales with $K$ ?
Gate MLP hidden	$d^{2}/2$	no
Gate MLP output (sliced)	$K\,d/2$	yes
Spectral rFFT (input)	$\mathrm{O}(d\log L)$	no
Spectral complex multiply	$\mathrm{O}(K\,d)$	yes
Spectral irFFT (output)	$\mathrm{O}(K\,d\log L)$	yes
Channel projection einsum	$K\,d^{2}$	yes

Appendix

Table 8: Per-token FLOPs ledger. The spectral FFT is amortized across the sequence; we report per-token cost by dividing the

\mathrm{O}(Ld\log L)

sequence-level FFT cost by

L

Knob	Value
Sequence length $L$	$2048$
Vocabulary	$258$ (bytes + PAD + BOS)
Token budget	$30$ B byte-tokens
Effective batch	$16$ sequences $\times\,2$ grad accum. $\times\,2048=65{,}536$ tokens / step
Optimizer	AdamW, $\beta{=}(0.9,0.95)$ , $\varepsilon{=}10^{-8}$
Weight decay	$0.1$ on tensor parameters; $0$ otherwise

Appendix

Table 9: SlimPajama scaling: training hyperparameters. Identical across all six families except for the capacity schedule: elastic families use a capacity-sampling schedule over their runtime-budget set, while non-elastic families train at a single fixed

K

Knob	Value
$d_{\mathrm{model}}$	$1262$
$n_{\mathrm{layers}}$	$20$
Maximum spectral budget $\overline{K}$	$32$
Runtime-budget set $\mathcal{K}$	$\{2,3,4,5,6,8,12,16,24,32\}$
Gate MLP hidden width	$d_{\mathrm{model}}/2=631$
Gate output bias init.	$-2.0$

Appendix

Table 10: SlimPajama main ES-SSM architecture. This is the full-budget configuration used to train the ES-SSM rows in table 11 ; smaller ES-SSM entries are exported compact models obtained by slicing this checkpoint along the spectral-channel axis.

		Quality			Inference cost
Family	$K$	Params (M)	Val BPB	Test BPB	Lat. (ms)	Tput (kt/s)	Mem (GB)
ES-SSM (ours)	2	335.00	1.40	1.41	6.18	331.40	0.87
ES-SSM (ours)	3	366.90	1.22	1.23	7.50	273.10	0.93
ES-SSM (ours)	4	398.70	1.10	1.10	8.65	236.80	1.00
ES-SSM (ours)	5	430.60	1.05	1.06	9.98	205.20	1.06
ES-SSM (ours)	6	462.50	1.02	1.03	11.34	180.60	1.13

Appendix

Table 11: Byte-level SlimPajama main scaling results (1.3B family, 30B byte-token training). Elastic families (ES-SSM, MatFormer, MatMamba) train once at

\overline{K}=32

and are evaluated at all ten runtime budgets in

\mathcal{K}

; non-elastic families (Transformer++, Pythia-style, GLA) contribute four independently trained reference budgets. These non-elastic rows are high-resource reference baselines rather than equal-training-cost competitors, since supporting multiple budgets requires multiple separate training runs. Quality is bits-per-byte (lower is better); inference cost is measured at sequence

L{=}2048

, batch

1

, on a single B200 GPU. For ES-SSM, the ten

K

values instantiate ten exported parameter budgets, spanning roughly

330

M to around

1.3

B. A small black downward triangle marks a collapsed test result whose Test BPB is more than

90\%

higher than the best Test BPB within the same family. All BPB values are rounded to two decimal places.

Model	$d_{\text{model}}$	$n_{\text{layers}}$	Specific	Params
Transformer (FA-2)	$512$	$16$	$h{=}8$ , RoPE	$50.8$ M
Mamba-2	$512$	$13$	$d_{\text{state}}{=}128$ , $\text{headdim}{=}64$	$50.1$ M
STU	$512$	$5$	$K{=}32$ , fixed Hankel basis	$52.8$ M
ES-SSM (ours)	$512$	$5$	$\overline{K}{=}32$ , adaptive gate, budget dropout	$53.5$ M

Appendix

Table 12: Per-architecture configurations at the matched

\sim\!50

M parameter scale. All models share

d_{\text{model}}{=}512

. Layer counts are tuned so each backbone stays in the same effective parameter range; embedding and classifier head are tiny in comparison.

Task	Seq. len.	Batch	Grad accum	Epochs	LR (peak)
ListOps	$2{,}000$	$32$	$4$	$6$	$3{\times}10^{-4}$
Text	$4{,}096$	$16$	$4$	$5$	$2{\times}10^{-4}$
Retrieval	$4{,}000\,{\times}\,2$	$8$	$4$	$4$	$2{\times}10^{-4}$
sCIFAR	$1{,}024$	$64$	$4$	$10$	$4{\times}10^{-4}$
Pathfinder	$1{,}024$	$64$	$4$	$4$	$3{\times}10^{-4}$
Path-X	$16{,}384$	$4$	$8$	$2$	$1{\times}10^{-4}$

Appendix

Table 13: Per-task training recipe. Effective global batch is batch

\times

grad_accum .

Model	ListOps	Text	Retrieval	sCIFAR	Pathfinder	Path-X
Transformer (FA-2)	$48.70{\pm}0.50$	$76.38{\pm}0.92$	$70.84{\pm}0.42$	$61.21{\pm}0.56$	$84.12{\pm}0.30$	FAIL
Mamba-2	$49.60{\pm}0.41$	$85.05{\pm}0.59$	$81.31{\pm}1.10$	$78.08{\pm}0.66$	$81.03{\pm}0.52$	$74.02{\pm}0.53$
STU	$60.10{\pm}0.03$	$89.25{\pm}1.15$	$89.68{\pm}0.26$	$91.24{\pm}0.14$	$94.28{\pm}0.27$	$88.01{\pm}0.32$
ES-SSM (ours)	$\mathbf{61.70{\pm}0.68}$	$\mathbf{92.77{\pm}0.60}$	$\mathbf{91.49{\pm}0.72}$	$\mathbf{92.88{\pm}0.24}$	$\mathbf{94.44{\pm}0.20}$	$\mathbf{93.08{\pm}0.13}$

Appendix

Table 14: LRA full-capacity results at the matched

\sim\!50

M parameter scale (

\overline{K}{=}32

). Top-

1

test accuracy (%), mean

\pm

std over

3

seeds for every task. Bold marks the per-task best.

Task	$2$	$3$	$4$	$5$	$6$	$8$	$12$	$16$	$24$	$32$
ListOps	$46.77$	$56.09$	$59.73$	$60.28$	$\underline{60.77}$	$61.08$	$61.33$	$61.51$	$61.64$	$61.70$
Text	$86.46$	$87.67$	$89.52$	$90.45$	$\underline{91.10}$	$91.75$	$92.21$	$92.49$	$92.68$	$92.77$
Retrieval	$84.90$	$88.38$	$\underline{90.30}$	$90.48$	$90.85$	$91.12$	$91.31$	$91.40$	$91.44$	$91.49$
sCIFAR	$89.18$	$\underline{91.63}$	$92.17$	$92.41$	$92.37$	$92.55$	$92.69$	$92.73$	$92.81$	$92.88$
Pathfinder	$85.94$	$88.58$	$90.57$	$91.23$	$91.61$	$91.80$	$92.55$	$\underline{93.68}$	$94.16$	$94.44$
Path-X	$81.91$	$85.17$	$87.50$	$88.89$	$89.82$	$90.75$	$\underline{91.68}$	$92.34$	$92.80$	$93.08$

Appendix

Table 15: ES-SSM budget sweep on a single

\overline{K}{=}32

training run,

\sim\!50

M scale. Test top-

1

accuracy (%), mean over

3

seeds for every task. Underlined entries mark the smallest

K

that retains

\geq 98\%

of the

\overline{K}

score (sweet spot

K^{\star}

Model	Params	Top- $1$ (%)	Latency (ms)	ECE
Conformer	$10.31$ M	$\mathbf{97.67\pm 0.26}$	$13.82\pm 0.19$	$0.067\pm 0.014$
Transformer++	$10.64$ M	$91.92\pm 1.71$	$\phantom{0}4.23\pm 0.07$	$0.084\pm 0.014$
Mamba (causal)	$10.20$ M	$97.19\pm 0.52$	$\phantom{0}9.15\pm 0.10$	$0.086\pm 0.025$
Audio-Mamba (bidir.)	$10.20$ M	$97.59\pm 0.26$	$\phantom{0}9.34\pm 0.18$	$\mathbf{0.059\pm 0.005}$
AR-STU	$10.51$ M	$95.37\pm 0.07$	$\phantom{0}\mathbf{2.34\pm 0.04}$	$0.072\pm 0.006$
ES-SSM (ours, $\overline{K}{=}32$ )	$10.66$ M	$97.40\pm 0.29$	$\phantom{0}2.97\pm 0.02$	$\mathbf{0.059\pm 0.003}$

Appendix

Table 16: SC-10 full-capacity accuracy, inference cost, and calibration. Top-

1

accuracy and Expected Calibration Error [ 45 ] use mean

\pm

std over

5

seeds. Per-sample latency is measured at batch

1

on an A100 80GB with bfloat16, using the median over

30

warmed-up runs. All models are matched at

\sim 10

M parameters. Bold marks the best mean in each column; the ES-SSM row is the single full-capacity training run from which table 17 is exported.

$K$	Params	Top- $1$ Acc. (%)	Retention (%)	Budget (%)	Latency (ms)
$\phantom{0}2$	$2.78$ M	$96.25\pm 0.35$	$\phantom{0}98.81$	$\phantom{0}6.25$	$2.95$
$\phantom{0}3$	$3.04$ M	$96.96\pm 0.24$	$\phantom{0}99.55$	$\phantom{0}9.38$	$2.94$
$\phantom{0}4$	$3.30$ M	$97.44\pm 0.28$	$100.05$	$12.50$	$2.97$
$\phantom{0}5$	$3.56$ M	$97.61\pm 0.20$	$100.22$	$15.62$	$2.97$
$\phantom{0}6$	$3.83$ M	$97.66\pm 0.25$	$100.28$	$18.75$	$2.96$
$\phantom{0}8$	$4.35$ M	$\mathbf{97.73\pm 0.15}$	$\mathbf{100.34}$	$25.00$	$2.98$

Appendix

Table 17: SC-10 budget sweep on a single

\overline{K}{=}32

ES-SSM training run. Compact ES-SSM

(K)

models are exported by direct prefix truncation of the channel projection stack

\{M_{k}\}

and Hankel spectral channels to the prefix

K

, so the Params column reflects a deployable standalone model. Top-

1

test accuracy is mean

\pm

std over

5

seeds; retention is relative to

K{=}\overline{K}{=}32

. Bold marks the peak accuracy across all

K

Protocol	Trainings / cell	Time / training	Suite time / cell	Total time	Relative total
ES-SSM exported suite	$1$	$53.3$ min	$53.3$ min	$23.99$ h	$1.00{\times}$
Transformer 5- $K$ suite	$5$	$47.9$ min	$239.6$ min	$107.84$ h	$4.50{\times}$
ES-SSM reduction	$5{\times}$ fewer	–	$186.3$ min saved	$83.85$ h saved	$77.8\%$ lower

Appendix

Table 18: D4RL training-cost accounting. A cell denotes one dataset–seed pair; the full suite contains

9

datasets and

3

seeds (

27

cells). ES-SSM trains one full-capacity model per cell at

\overline{K}{=}32

and exports all five runtime budgets by direct prefix truncation. The matched Transformer suite trains one independent model for each budget

K\in\{2,4,8,16,32\}

. Times are single-RTX-3090 training wall-clock; rollout evaluation time is excluded.

Dataset	DT-paper	DT (ours)	ES-SSM ( $K{=}32$ )	$\Delta$ vs. DT (ours)
halfcheetah-medium	$42.6$	$42.22\!\pm\!0.23$	$\mathbf{42.30\!\pm\!0.09}$	$+0.08$
halfcheetah-medium-replay	$36.6$	$34.64\!\pm\!0.80$	$33.64\!\pm\!2.37$	$-1.00$
halfcheetah-medium-expert	$86.8$	$91.12\!\pm\!0.94$	$\mathbf{91.31\!\pm\!0.77}$	$+0.19$
hopper-medium	$67.6$	$50.91\!\pm\!2.51$	$\mathbf{52.59\!\pm\!2.74}$	$+1.68$
hopper-medium-replay	$82.7$	$47.96\!\pm\!19.93$	$\mathbf{79.50\!\pm\!4.02}$	$\mathbf{+31.54}$
hopper-medium-expert	$107.6$	$82.83\!\pm\!17.59$	$\mathbf{88.06\!\pm\!9.74}$	$\mathbf{+5.23}$

Appendix

Table 19: D4RL MuJoCo: normalized returns. Higher is better. Mean

\pm

std over

3

seeds,

100

rollout episodes per cell. “DT-paper” is the originally reported number [ 36 ] ; “DT (ours)” is the same architecture trained and evaluated under our recipe (

150

k steps,

100

-episode rollout) for paired comparison. ES-SSM at

K{=}32

matches DT in parameter count (

711

k vs.

730

k). Bold ES-SSM cells:

\mathrm{ES{\text{-}}SSM}\geq\mathrm{DT(ours)}

. Bold

\Delta

|\Delta|>5

Family	$K$	Params	Norm. return	Lat. (ms)	Tput (kt/s)	Mem (MB)
ES-SSM exported compact models (one $\overline{K}{=}32$ training)
ES-SSM	$32$	$711$ k	$\mathbf{70.23}$	$2.75$	$21.9$	$16$
ES-SSM	$16$	$461$ k	$\mathbf{65.94}$	$2.74$	$21.9$	$13$
ES-SSM	$8$	$335$ k	$\mathbf{66.55}$	$2.73$	$22.0$	$12$
ES-SSM	$4$	$273$ k	$66.43$	$2.69$	$22.3$	$11$
ES-SSM	$2$	$241$ k	$65.02$	$2.74$	$21.9$	$10$

Appendix

Table 20: D4RL MuJoCo inference table. ES-SSM rows are exported compact models from a single

\overline{K}{=}32

training run; Transformer rows are independently trained baselines (one row

=

one training). Norm. return averaged over

9

datasets,

3

seeds,

100

rollout episodes. Inference cost measured at batch

1

K_{\mathrm{ctx}}{=}20

on RTX 3090. Bold marks the higher normalized return within each matched budget pair.

Family	Params	Effective batch	Final val. BPB ( $\downarrow$ )	Wall
Transformer++ (full)	$304$ M	$65{,}536$ tokens/step	$1.0800$	$\sim\!62$ h
MatFormer (full width)	$304$ M	$65{,}536$ tokens/step	$1.1026$	$\sim\!56$ h
ES-SSM ( $\overline{K}{=}32$ )	$335$ M	$65{,}536$ tokens/step	$\mathbf{1.0771}$	$\sim\!66$ h

Appendix

Table 21: Byte-level generality replication at around

330

M parameters and

10

B byte-tokens on SlimPajama. Validation BPB on the shared

1000

-segment cache; lower is better. All three runs reach the full

10

B-token budget cleanly and write a final best.pt at the last validation gate.

Tokens (B)	$1$	$2$	$3$	$4$	$5$	$6$	$7$	$8$	$9$	$10$
Transformer++	$1.338$	$1.256$	$1.214$	$1.189$	$1.164$	$1.140$	$1.119$	$1.102$	$1.090$	$1.080$
MatFormer (full width)	$1.358$	$1.278$	$1.238$	$1.212$	$1.188$	$1.165$	$1.143$	$1.126$	$1.115$	$1.103$
ES-SSM ( $\overline{K}{=}32$ )	$1.296$	$1.251$	$1.211$	$1.181$	$1.156$	$1.135$	$1.117$	$1.102$	$1.089$	$\mathbf{1.077}$

Appendix

Table 22: Full-capacity validation BPB versus tokens seen. Each row is one architecture; columns step the byte-token budget from

1

B to

10

B in

1

B increments. All three trajectories descend monotonically with no observed instability or recovery spike.

Figure 10: Byte-level generality training trajectories. Full-capacity validation BPB versus tokens seen for the three architectures at the around

330

M-parameter /

10

B-byte-token replication scale. Same data as table 22 ; lower is better.

Training regime	$K\!=\!2$	$K\!=\!3$	$K\!=\!4$	$K\!=\!5$	$K\!=\!6$	$K\!=\!8$	$K\!=\!12$	$K\!=\!16$	$K\!=\!24$	$K\!=\!32$
ES-SSM (reference)	1.965	1.873	1.847	1.833	1.827	1.822	1.821	1.822	1.824	1.826
No budget dropout ( $\overline{K}$ )	3.974	4.583	4.596	4.432	3.555	2.797	2.079	1.895	1.810	1.796

Appendix

Table 23: Without budget dropout, the exported compact model at budget

K

is a post hoc direct prefix truncation of a full-capacity model and is not optimised for any operating point below

\overline{K}

. At full capacity, the reference scores

1.826

BPB versus

1.796

without budget dropout, a

0.030

BPB penalty in this single-seed comparison. In exchange, budget dropout improves BPB by

0.975

K\!=\!8

2.710

K\!=\!3

, and a peak

2.749

K\!=\!4

. These low-budget gains are substantially larger than the full-capacity penalty and the reported seed noise of this protocol. The non-monotonicity of the no-budget-dropout row, where

K\!=\!3

scores worse than

K\!=\!2

(

4.583

vs.

3.974\,

BPB) and

K\!=\!4

even worse still, makes the same point structurally: without budget dropout, channel ordering under direct prefix truncation is not a controlled accuracy–compute knob, and adding one more channel at a small runtime budget can hurt because that channel was never trained to be used in isolation.

Gating variant	$K\!=\!2$	$K\!=\!3$	$K\!=\!4$	$K\!=\!5$	$K\!=\!6$	$K\!=\!8$	$K\!=\!12$	$K\!=\!16$	$K\!=\!24$	$K\!=\!32$
Adaptive sigmoid gate (reference)	1.965	1.873	1.847	1.833	1.827	1.822	1.821	1.822	1.824	1.826
Uniform $g^{(K)}_{k}\!=\!1/\sqrt{K}$	2.246	2.131	2.091	2.068	2.058	2.049	2.043	2.041	2.045	2.052
Per-channel scalar, input-independent	2.104	2.000	1.965	1.947	1.939	1.933	1.929	1.929	1.933	1.940

Appendix

Table 24: Both degeneracies cost BPB at every

K

. At

K\!=\!\overline{K}

the uniform gate costs

0.226\,

BPB and the per-channel scalar

0.114\,

BPB relative to the input-adaptive reference; the corresponding gaps at

K\!=\!2

are

0.281

and

0.139\,

BPB. All four main gaps are an order of magnitude above the seed noise floor. The per-channel scalar variant sits between the uniform gate and the input-adaptive reference at every

K

and accounts for

0.112

of the

0.226\,

BPB total gating gain at

K\!=\!\overline{K}

— almost exactly half — suggesting that learned per-channel reweighting is mildly useful in isolation but that the dominant share of the gain comes from conditioning on

u(t)

. The split should be read with the caveat that the adaptive and per-channel variants differ in capacity (a two-layer MLP versus

K\!=\!32

scalars), so the comparison changes both the gate’s functional form and its parameter count; a parameter-matched test would require a wider architectural search beyond this ablation. Direct evidence for input-conditionality of the reference gate appears in table 26 , and the order-of-magnitude difference in effective gate support between adaptive and non-adaptive variants appears in table 27 .

$n_{\mathrm{full}}$ (steps)	$K\!=\!2$	$K\!=\!3$	$K\!=\!4$	$K\!=\!5$	$K\!=\!6$	$K\!=\!8$	$K\!=\!12$	$K\!=\!16$	$K\!=\!24$	$K\!=\!32$
$\infty$ (never force $\overline{K}$ )	1.955	1.869	1.845	1.832	1.826	1.823	1.823	1.825	1.829	1.834
$16$	1.960	1.871	1.847	1.834	1.828	1.825	1.824	1.825	1.828	1.831
$8$ (reference)	1.965	1.873	1.847	1.833	1.827	1.822	1.821	1.822	1.824	1.826
$4$	1.974	1.877	1.851	1.836	1.828	1.823	1.821	1.821	1.822	1.824
$1$ (no budget dropout)	3.974	4.583	4.596	4.432	3.555	2.797	2.079	1.895	1.810	1.796

Appendix

Table 25: Coarse full-budget anchor cadence ablation. The endpoint

n_{\mathrm{full}}\!=\!1

collapses budget dropout to pure full-capacity training and shows the same direct-truncation failure mode as table 23 , as expected. Among the non-endpoint budget-dropout settings (

\infty,16,8,4

), the curve is flat to within

0.02\,

BPB at every

K

, so this sparse one-seed table should not be used to rank these cadences. The default

n_{\mathrm{full}}\!=\!8

is justified by the finer three-seed validation in Appendix D ; this table only shows that reasonable non-endpoint cadences preserve the budget dropout behavior while

n_{\mathrm{full}}\!=\!1

does not.

Channel index $k$	Mean $\gamma_{k}$	Std $\gamma_{k}$	Min $\gamma_{k}$	Max $\gamma_{k}$
$2$	$0.326$	$0.334$	$0.000$	$1.000$
$3$	$0.150$	$0.231$	$0.000$	$0.999$
$4$	$0.068$	$0.142$	$0.000$	$0.995$
$5$	$0.032$	$0.091$	$0.000$	$0.987$
$6$	$0.024$	$0.061$	$0.000$	$0.969$
$8$	$0.007$	$0.029$	$0.000$	$0.957$

Appendix

Table 26: Selected-channel sigmoid gate statistics, averaged over all positions and held-out sequences and across layers. The trained model assigns substantial gate mass to the low-index channels (

\bar{\gamma}_{k}\!\approx\!0.33

k\!=\!2

and

\bar{\gamma}_{k}\!\approx\!0.15

k\!=\!3

), and the per-channel standard deviation matches or exceeds the per-channel mean at every listed channel (most visibly

0.231

vs.

0.150

k\!=\!3

), so the gate is strongly input-conditional rather than saturating at a constant. Several high-index channels that are nearly closed on average (

k\!\geq\!8

\bar{\gamma}_{k}\!\leq\!0.01

in our seq-

2048

measurements) still reach maximum gate values from about

0.4

to nearly

1.0

on specific token contexts; the gate is opportunistically opening these channels rather than learning a hard mask. This pattern of input-conditional, channel-specific use is the load-bearing property table 24 identifies.

Variant	$K_{\mathrm{eff}}^{\,\mathrm{PR}}$	$K_{\mathrm{eff}}^{\,(0.05)}$	$K_{\mathrm{eff}}^{\,(0.01)}$
Adaptive sigmoid gate (reference)	$\phantom{0}3.84$	$\phantom{0}4$	$\phantom{0}7$
Adaptive, no budget dropout ( no_bd )	$\phantom{0}4.35$	$\phantom{0}4$	$\phantom{0}7$
Per-channel scalar ( gate_static )	$28.25$	$32$	$32$
Uniform $1/\sqrt{K}$ ( gate_uniform )	$32.00$	$32$	$32$

Appendix

Table 27: Effective gate support for the gate variants, computed from the per-channel mean unscaled gate activities measured on the held-out validation cache. The adaptive gate concentrates gate activity into a small number of channels: the reference model has participation ratio

3.84

, and only four channels exceed the

0.05

mean-activity threshold. Removing budget dropout changes this count only modestly, while the non-adaptive variants fail to sparsify: the uniform gate uses all channels by construction, and the per-channel scalar remains broadly spread across the channel set. This supports the interpretation that input-adaptive gating, not merely a static channel weight, concentrates gate activity onto a compact subset of channels. This diagnostic measures gate activity only, rather than the full channel contribution

g_{k}^{(K)}(t)M_{k}(\Phi_{k}*u)(t)

; low mean-activity channels may still provide input-specific high-budget refinements.

Sequence Modeling

Sequence modeling studies how to represent, predict, classify, and generate ordered data such as text, time series, events, and biological sequences.

See all 261 papers

Sep 26, 2026

Elastic Selective Spectral Hybrids for Train-Once, Export-Many Budgeted Inference

Deploying a language model under different computing and latency budgets calls for compact models with different quality-cost trade-offs. Towards this end, elastic spectral state space models provide ordered and temporally decomposed channels that can be truncated, but they are associated with linear time-invariant filters that cannot selectively preserve relevant past information or forget irrelevant information as the context evolves. To address this, we introduce the Elastic Selective Spectral Hybrid (ESSH), which realizes each Hankel spectral channel as an independent recurrent unit using fitted damped rotation modes. It also features an input-dependent decay and write/read gates that make temporal retention and state update input-dependent while preserving channel-wise truncation and a structured recurrence for efficient execution. ESSH combines these selective spectral mixers with sliding-window attention and jointly trains multiple capacities by reducing spectral-channel count and feed-forward width at different rates through a two-rate capacity map with full-model distillation. The resulting models support chunked parallel training and fused recurrent decoding while avoiding computation for discarded channels. At full capacity, ESSH achieves language-modeling quality comparable to similarly sized independently trained models, while smaller exports exhibit a smooth quality-cost trade-off. We validate the effectiveness of the proposed framework using language understanding, retrieval, cross-domain text, and DNA experiments, by assessing quality retention and the trade-off against independently trained and elastic baselines. At the 1.53B model configuration, fused batch-one decoding takes 1.37 ms per token on a B300, providing a 2.14-2.80x speedup over the tested Mamba-2 and Mamba-3 implementations and 3.03x over Transformer++ at matched parameter counts.

Dachuan Song, Chuchu Chen, Xuan Wang

George Mason University · George Washington University

Selective SSMs LLM Inference Acceleration Language Modeling LLM Compression Efficient Language Model Training

Ask PDF

Jun 29, 2026

MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

State space models (SSMs) have emerged as efficient linear-time alternatives to attention for long-sequence modeling. However, existing SSMs often suffer from instability and memory degradation over extended horizons due to poorly conditioned first-order updates and unbalanced update geometry. We introduce MuonSSM, a general framework that stabilizes SSM training by explicitly conditioning the geometry of memory updates rather than the recurrent transition matrix. MuonSSM augments SSMs with a momentum-based pathway and a lightweight Newton Schulz transformation on low-rank input injections, yielding bounded and spectrally conditioned updates while preserving parallel scan complexity. Theory shows that MuonSSM improves gradient propagation, mitigates spectral amplification, and enriches memory representations over long horizons. Extensive experiments across language, vision, and time-series benchmarks show consistent gains in accuracy, robustness, and long-context performance when integrated into diverse SSM backbones. These results establish geometric conditioning of updates as a principled pathway to stable, scalable sequence modeling.

Thai-Khanh Nguyen, Ngoc-Bich-Uyen Vo, Thieu N. Vo +2

Faculty of Information Technology, Dainam University, Hanoi University of Science and Technology, Hanoi 10000, Vietnam · Faculty of Artificial Intelligence, Posts and Telecommunications Institute of Technology, Hanoi 10000, Vietnam · Department of Computer Science, University of Bath, Bath BA2 7AY, UK +1

Long-Context Modeling Sequence Modeling State Space Models

Ask PDF

May 8, 2026

Star Elastic: Many-in-One Reasoning LLMs with Efficient Budget Control

Training a family of large language models (LLMs), either from scratch or via iterative compression, is prohibitively expensive and inefficient, requiring separate training runs for each model in the family. In this paper, we introduce Star Elastic, a novel LLM post-training method that adds N nested submodels to a given parent reasoning model using the compute of one run (N-fold savings) via a single post-training job. Beyond reducing training costs, Star Elastic also addresses a fundamental limitation of efficient reasoning: the rigidity of static architectures, which forces the allocation of constant resources regardless of token difficulty. By unlocking elastic budget control, Star Elastic enables a novel inference scheme that uses different submodels for each reasoning phase (thinking and answering). Star Elastic supports (1) nesting along the SSM, embedding channel, MoE, and FFN axes, (2) learning nested submodels via an end-to-end trainable router, and (3) curriculum-based knowledge distillation. Building on the Nemotron Elastic framework, we apply Star Elastic to the NVIDIA Nemotron Nano models, with a particular focus on hybrid Mixture-of-Experts (MoE) architectures: from Nemotron Nano v3 (30B/3.6A), we generate 23B (2.8A) and 12B (2.0A) variants with 160B training tokens. All nested models match or outperform independently trained baselines of comparable size and achieve a 360x reduction versus pretraining from scratch and a 7x reduction over state-of-the-art compression. Crucially, elastic budget control advances the accuracy-latency Pareto frontier, achieving up to 16% higher accuracy and 1.9x lower latency via dynamic per-phase model selection. We further extend Star Elastic to quantized regimes via Quantization-Aware Distillation (QAD), producing nested NVFP4 and FP8 elastic checkpoints that preserve zero-shot slicing while delivering smaller deployment footprints.

Ali Taghibakhshi, Ruisi Cai, Saurav Muralidharan +17

NVIDIA

LLM Compression LLM Post-Training

Ask PDF

Abstract

Sequence Modeling

Explore similar work

Elastic Selective Spectral Hybrids for Train-Once, Export-Many Budgeted Inference

MuonSSM: Orthogonalizing State Space Models for Sequence Modeling

Star Elastic: Many-in-One Reasoning LLMs with Efficient Budget Control