Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Elastic Spectral State Space Models for Budgeted Inference

Dachuan Song, Xuan Wang

TL;DR

Elastic Spectral State Space Models (ES-SSM) address the mismatch between fixed training budgets and heterogeneous real-world deployment by enabling direct runtime truncation without retraining. The method combines a Hankel spectral basis with an input-adaptive spectral gate and budget-aware training (budget dropout) to concentrate predictive information in low-index spectral channels, yielding smooth, predictable performance as the budget $K$ is reduced from the full capacity $\overline{K}=32$. Across long-sequence benchmarks such as Long Range Arena and PG19, ES-SSM matches or exceeds strong baselines at full capacity and exhibits graceful degradation with budget reduction, revealing clear sweet spots where small budgets retain near-maximum accuracy. This elastic approach enables deployment across diverse hardware without multiple training runs or distillation, offering practical impact for scalable, resource-aware sequence modeling.

Abstract

Foundation models are typically trained at a fixed computational capacity, while real-world applications require deployment across platforms with different resource constraints. Current approaches usually rely on training families of model variants or model distillation, which requires additional training and supports only a pre-selected set of sizes rather than fine-grained adaptation at runtime. In this paper, we propose Elastic Spectral State Space Models (ES-SSM), which require only one-time training at full capacity, but can be directly truncated into arbitrary scales for budgeted, runtime inference without retraining. Our ES-SSM builds on Hankel spectral filtering over a state space model (SSM), coupled with a lightweight input-adaptive gate trained under randomized spectral budgets. Using a shared masked normalization rule over the ordered spectral channels, we encourage predictive capability to concentrate in low-index components, while higher-index components act primarily as refinement. We test our algorithm across long-sequence benchmarks spanning text, logic, retrieval, vision, and audio. We demonstrate that a single ES-SSM model trained once can be truncated to provide competitive performance compared with modern Transformer and SSM baselines at similar parameter scales. Furthermore, by testing under various runtime budgets, we observe smooth and stable budget-performance curves over a wide range of truncation levels.

Elastic Spectral State Space Models for Budgeted Inference

TL;DR

Elastic Spectral State Space Models (ES-SSM) address the mismatch between fixed training budgets and heterogeneous real-world deployment by enabling direct runtime truncation without retraining. The method combines a Hankel spectral basis with an input-adaptive spectral gate and budget-aware training (budget dropout) to concentrate predictive information in low-index spectral channels, yielding smooth, predictable performance as the budget is reduced from the full capacity . Across long-sequence benchmarks such as Long Range Arena and PG19, ES-SSM matches or exceeds strong baselines at full capacity and exhibits graceful degradation with budget reduction, revealing clear sweet spots where small budgets retain near-maximum accuracy. This elastic approach enables deployment across diverse hardware without multiple training runs or distillation, offering practical impact for scalable, resource-aware sequence modeling.

Abstract

Foundation models are typically trained at a fixed computational capacity, while real-world applications require deployment across platforms with different resource constraints. Current approaches usually rely on training families of model variants or model distillation, which requires additional training and supports only a pre-selected set of sizes rather than fine-grained adaptation at runtime. In this paper, we propose Elastic Spectral State Space Models (ES-SSM), which require only one-time training at full capacity, but can be directly truncated into arbitrary scales for budgeted, runtime inference without retraining. Our ES-SSM builds on Hankel spectral filtering over a state space model (SSM), coupled with a lightweight input-adaptive gate trained under randomized spectral budgets. Using a shared masked normalization rule over the ordered spectral channels, we encourage predictive capability to concentrate in low-index components, while higher-index components act primarily as refinement. We test our algorithm across long-sequence benchmarks spanning text, logic, retrieval, vision, and audio. We demonstrate that a single ES-SSM model trained once can be truncated to provide competitive performance compared with modern Transformer and SSM baselines at similar parameter scales. Furthermore, by testing under various runtime budgets, we observe smooth and stable budget-performance curves over a wide range of truncation levels.
Paper Structure (27 sections, 1 theorem, 18 equations, 3 figures, 6 tables, 2 algorithms)

This paper contains 27 sections, 1 theorem, 18 equations, 3 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Statement of Contributions.
  3. Related Work
  4. Deployment-oriented model scaling and adaptive inference.
  5. State space models for long-sequence learning.
  6. Hankel spectral filtering and spectral SSMs.
  7. Preliminaries
  8. Setup and notation.
  9. Linear dynamical systems and Hankel spectral basis.
  10. Spectral parameterization of SSM.
  11. Limitations of the base spectral form.
  12. Methods
  13. Budgeted inference through Adaptive spectral gating.
  14. Training ES-SSM with budget dropout
  15. Experiments
  16. ...and 12 more sections

Key Result

Proposition 3.1

Consider the ES-SSM layer, for any budget $K\le \overline{K}$, where $\alpha^{(K)}_k(t)\ge 0$ and $\sum_{k=1}^{K}\alpha^{(K)}_k(t)=1$ for each $t$ (as ensured by eq:gating_masked_softmax). Let $\|\cdot\|_2$ denote the Euclidean norm on vectors, and $\|\cdot\|_{\mathrm{op}}$ the induced operator norm on matrices. Define $\|u\|_{\infty}\triangleq \max_{1\le t\le Then for any input sequence $u_{1:L

Figures (3)

  • Figure 1: Left: pre-norm residual block. Right: ES-SSM layer. Hankel filters produce spectral features, a small gating MLP outputs logits, and a masked softmax activates the first $K$ channels while masking the rest. The output $\hat{y}(t)$ is the gated mixture of the active spectral features.
  • Figure 2: LRA results. Top row: full-capacity accuracy comparison across models on each task (x-axis: accuracy %) at a matched parameter scale of $\sim$20M. Bottom row: ES-SSM budget sweep from the same model trained at $K=32$ (x-axis: runtime spectral budget $K$, y-axis: accuracy %). The star marks the sweet-spot budget $K^\ast$, and the black triangle marks the collapse boundary.
  • Figure 3: PG19 versus runtime budget. ES-SSM exhibits stable sweet spots at moderate budgets and degrades gracefully as $K$ decreases.

Theorems & Definitions (2)

  • Proposition 3.1: BIBO stability under budgeted inference
  • proof