Elastic Spectral State Space Models for Train-Once Budgeted Inference¶

会议: NeurIPS 2026
arXiv: 2601.22488
领域: LLM效率
关键词: 状态空间模型, 弹性推理, 汉克尔谱滤波, 算子近似, 预算丢弃

一句话总结¶

针对传统序列模型无法低成本跨算力预算部署的难题，提出基于汉克尔谱滤波算子分解的 ES-SSM，结合自适应通道门控与对数预算丢弃（Budget Dropout）实现单次训练即可通过前缀截断直接导出任意紧凑模型。

研究背景与动机¶

现代序列模型通常在固定的计算容量下完成训练，但在工业落地的真实场景中，模型往往需要同时部署于高吞吐云端集群、算力有限的边缘服务器以及超低功耗终端设备。针对多平台算力异构的问题，主流工业界方案主要依赖针对每个目标预算独立训练一组规模各异的模型，或是通过大模型向小模型蒸馏。这种范式虽然能为选定预算获得良好性能，但每扩展一个算力档位都需要经历高昂的再训练、数据校准或微调流程，不仅严重阻碍动态算力调度，也带来了成倍的算力与存储开销。

以可瘦身网络（Slimmable Networks）、Once-for-All 以及近期面向序列模型的 MatFormer 和 MatMamba 为代表的弹性神经网络，尝试通过单次超网训练来支持多预算推理。然而，现有弹性方法在本质上属于“架构级弹性”——其收缩机制主要依赖于缩小物理宽度、裁减网络深度、截断 FFN 隐层或裁剪 SSM 隐状态维度。由于这类粗粒度的物理维度裁剪与序列算子的功能退化之间缺乏直接的数学映射，模型在小算力截断时极易陷入表示崩塌，往往需要依赖复杂的超网搜索、渐进式收缩策略或繁琐的校准才能勉强维持性能。

本文另辟蹊径，从算子逼近（Operator Approximation）这一全新视角审视序列模型的弹性问题。在长程序列建模（尤其是状态空间模型 SSM）中，核心计算在于因果卷积算子对时序 token 的长距离混合。基于汉克尔矩阵（Hankel Matrix）构造的固定谱通道为序列算子提供了从低到高天然有序的逼近分辨率，但若直接在全容量下训练，梯度优化会将有效信号任意分散在各个通道或形成跨通道抵消，导致直接前缀截断时性能急剧崩溃。本文的核心 idea 是：将固定汉克尔谱基与输入自适应通道门控相结合，并在训练期引入匹配推理截断行为的对数预算丢弃（Budget Dropout），迫使低索引谱通道独立承载核心序列预测能力、高索引通道作为增量细化项，实现单次训练即可免搜索、免微调无缝截断导出紧凑模型。

方法详解¶

整体框架¶

ES-SSM 的核心是以算子逼近精度为控制手柄构建弹性序列模型。网络主体采用标准的 Pre-Norm 残差块架构，由 ES-SSM 谱子层 $\mathcal{S}_K$ 与随后的两层前馈网络（FFN）交替串联而成。输入序列 $u_{1:L}$ 经过嵌入层与首个残差块前归一化后，送入 ES-SSM 谱子层：该子层利用预计算好的固定汉克尔因果卷积核组提取跨时序长程特征，同时轻量门控 MLP 根据当前输入动态产生输入自适应的通道激活门控值。对于任意运行时预算 $K \le \overline{K}$，网络直接保留前 $K$ 个谱通道及其对应的通道投影参数，其余高阶通道直接物理丢弃，从而无需任何重排序与额外微调即可独立导出对应参数量的紧凑模型。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入序列 u"] --> LN1["Pre-LN 归一化"]
    LN1 --> GBF["汉克尔固定谱基因果卷积<br/>共享 2L 长度 real-FFT 频域滤波"]
    LN1 --> MLP["输入自适应通道门控<br/>轻量 MLP 动态预测标量激活"]
    GBF --> BPT["直接前缀截断与尺度缩放<br/>截断保留前 K 个通道并按 1/√K 归一化"]
    MLP --> BPT
    BPT --> CP["通道投影与聚合<br/>门控特征沿通道融合投影至模型维度"]
    CP --> Res1["残差叠加 + DropPath"]
    Res1 --> FFN["两层 FFN 变换 + Pre-LN 残差"]
    FFN --> Out["块输出 r(out)"]

关键设计¶

1. 基于汉克尔谱基的因果算子谱分解：以数学逼近分辨率取代架构物理维度裁剪

传统状态空间模型 $x(t) = A x(t-1) + B u(t), y(t) = C x(t) + D u(t)$ 的输入-输出映射表现为因果卷积 $y(t) = D u(t) + (G * u)(t)$，其中卷积核 $G(\tau) = C A^\tau B$。在对称稳定设置下，矩阵 $A$ 的正交特征分解使得卷积核退化为几何标量序列 $\nu(\lambda) = [1, \lambda, \dots, \lambda^{L-1}]^\top$ 与系数矩阵的乘积之和。本设计直接利用对连续稳定几何轨迹加权积分得到的半正定汉克尔矩阵 $Z_{ij} = \int_0^1 (1-\rho)^2 \rho^{i+j-2} d\rho = \frac{2}{(i+j)^3 - (i+j)}$，提取其前 $\overline{K}$ 个最大特征向量作为通用的固定离散因果谱通道 $\Phi_k = \{\phi_k(\tau)\}_{\tau=0}^{L-1}$。这使得卷积核被严格展开为 $G(\tau) \approx \sum_{k=1}^K M_k \phi_k(\tau)$，其中 $M_k \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ 为唯一可学习的通道投影矩阵。由于汉克尔特征值呈指数级衰减，数学定理保证了截断至前 $K$ 阶谱通道具有理论上的最小均方逼近误差界，赋予了截断操作极强的数学可解释性，从根本上克服了以往弹性网络剪切层宽或剪切层数时与任务表征损失脱节的缺陷。

2. 输入自适应通道门控与保持前缀一致性的直接截断：实现动态表征与免微调导出

若仅使用静态投影 $M_k$，模型面对不同上下文时通道贡献完全刚性，无法动态抑制噪声或突出瞬时重要频段。为此，ES-SSM 在谱子层中引入了输入自适应的通道门控机制。对于输入 $u(t)$，轻量双层 MLP 首先计算通道未归一化 logits： $$s(t) = W_{g,2} \text{GELU}(W_{g,1} u(t) + b_{g,1}) + b_{g,2}$$ 其中隐藏宽度设为 $d_g = d_{\text{in}} / 2$，输出偏置初始化为 $b_{g,2} \equiv -2$（使初始门控激活值较小，保护早期残差流）。在运行时预算 $K$ 下，作用于通道 $k \le K$ 的门控系数显式定义为： $$g_k^{(K)}(t) = \frac{1}{\sqrt{K}} \operatorname{sigmoid}(s_k(t))$$ 该设计具有两个关键性质：第一，针对通道 $k$ 的 logit $s_k(t)$ 仅由 $W_{g,2}$ 的第 $k$ 行决定，与后续高阶通道 $j > K$ 的权重和偏置严格解耦，保证了截断在机械结构上的绝对前缀一致性（Prefix Consistency）；第二，引入 $1/\sqrt{K}$ 确定性缩放因子，平衡了从低预算到高预算时谱累加项的范数方差。在推理时，只需对 $W_{g,2}$、$b_{g,2}$、$\{M_k\}$ 与 $\{\Phi_k\}$ 切取前 $K$ 个切片即可导出完全独立的紧凑模型。

3. 对数预算丢弃训练算法与全容量锚定：强制前缀学习与防止高阶通道欠拟合

即使谱基天然有序，在全容量下联合优化依然会导致高低索引通道之间发生梯度的协同依赖或抵消，使得直接截断在极小预算下失效。ES-SSM 提出了对数预算丢弃（Budget Dropout）策略：在前 $E_{\text{warmup}}$ 个 epoch 采用全预算 $K_{\text{train}} = \overline{K}$ 预热以稳定全局参数；随后在常规训练步中，以概率 $1/n_{\text{full}}$（默认 $n_{\text{full}} = 8$）强制执行全容量锚定更新（Anchor Update），而在其余步中通过对数均匀采样动态抽取当前训练预算： $$K_{\text{train}} = \arg\min_{k \in \mathcal{K}} \Big| k - K_{\min} \left(\overline{K} / K_{\min}\right)^\xi \Big|, \quad \xi \sim \text{Unif}[0, 1]$$ 对数采样将预算视为倍乘容量梯级，极大增加了对小预算 $K$ 前缀的采样优化压力；同时，前缀嵌套性保证了低索引通道 $k$ 参与反向传播的边缘概率满足 $p_i \ge p_j \, (i < j)$，迫使低频基底通道具备强大的独立预测能力。而周期性的全容量锚定步为尾部通道提供了 $1/n_{\text{full}}$ 的梯度更新保底率，保证了高阶通道能稳定习得高阶细化特征而不至于退化。

损失函数 / 训练策略¶

ES-SSM 采用任务端到端目标函数进行单阶段训练，无需教师-学生蒸馏损失，也不需要复杂的子网采样排名惩罚。在每个优化步骤中，根据对数预算丢弃采样的 $K_{\text{train}}$ 仅执行对应前缀的紧凑前向传播与反向传播： $$\mathcal{L}_{\text{train}}(\theta) = \mathcal{L}_{\text{task}}(F_\theta(u; K_{\text{train}}), y)$$ 前缀计算中的长程因果卷积全部通过长度为 $2L$ 的实数 FFT（rFFT）与频域复数相乘、逆变换（irFFT）批处理完成。为消除汉克尔奇异值跨越数个数量级对小模长通道造成的精度截断偏差，ES-SSM 采用了混合精度策略：主干网络（Embedding、FFN、残差变换）在 bfloat16 下运行，而谱分支的 FFT、复数相乘与逆 FFT 计算显式锁定在 float32 单精度下执行。优化器采用 AdamW，配合预热余弦退火学习率调度与梯度范数裁剪。

实验关键数据¶

主实验¶

在 SlimPajama 627B 字节级语言建模任务中，全容量模型规模约为 1.3B 参数，训练序列长度 $L=2048$，总计训练 30B byte-tokens。ES-SSM 仅通过单次训练即导出涵盖 $K \in \{2, 3, 4, 5, 6, 8, 12, 16, 24, 32\}$ 的 10 种预算紧凑模型，并与代表性弹性架构（MatFormer、MatMamba）以及独立训练的非弹性模型进行系统对比。

模型族	预算/状态 $K$	参数量 (M)	验证 BPB ↓	测试 BPB ↓	推理延迟 (ms)	吞吐量 (kt/s)	显存 (GB)
ES-SSM (Ours)	2	335.0	1.40	1.41	6.18	331.4	0.87
ES-SSM (Ours)	4	398.7	1.10	1.10	8.65	236.8	1.00
ES-SSM (Ours)	8	526.2	0.99	0.99	14.13	144.9	1.25
ES-SSM (Ours)	16	781.1	0.96	0.97	28.60	71.6	1.76
ES-SSM (Ours)	32 (全容量)	1291.0	0.95	0.96	59.83	34.2	2.78
MatFormer	2	341.8	5.32	5.35	4.52	453.1	0.84
MatFormer	8	577.4	3.76	3.78	7.19	284.8	1.40
MatFormer	32	1310.7	0.98	0.99	14.32	143.0	3.20
MatMamba	2	262.0	1.63	1.65	3.02	678.1	0.58
MatMamba	8	689.0	1.02	1.04	7.27	281.7	1.52
MatMamba	32	1310.0	1.00	1.00	11.96	171.2	2.89
Transformer++ (独立训练)	2	354.4	1.04	1.06	5.11	400.6	0.75
Transformer++ (独立训练)	32	1323.8	0.93	0.94	11.49	178.3	2.63
GLA (独立训练)	2	335.0	1.07	1.08	3.94	519.8	0.62
GLA (独立训练)	32	1300.0	0.96	0.97	8.27	247.6	2.39

在跨模态基准 Long Range Arena（LRA, ~50M 参数规模）评测中，ES-SSM 展现了对超长因果建模的极强泛化性，在全容量下全面超越基线，并在极端长度任务 Path-X（序列长 16,384）中取得突破性提升：

模型	ListOps	Text	Retrieval	sCIFAR	Pathfinder	Path-X (16k)
Transformer (FA-2)	48.70	76.38	70.84	61.21	84.12	失败 (OOM/崩溃)
Mamba-2	49.60	85.05	81.31	78.08	81.03	74.02
STU (静态汉克尔谱)	60.10	89.25	89.68	91.24	94.28	88.01
ES-SSM (Ours, 全容量)	61.70	92.77	91.49	92.88	94.44	93.08

消融实验¶

在 770M 参数规模、400M tokens 的 SlimPajama 诊断配置下，系统验证了预算丢弃机制（Budget Dropout）与输入自适应通道门控的核心价值：

配置 / 变体	$K=2$ 验证 BPB ↓	$K=4$ 验证 BPB ↓	$K=8$ 验证 BPB ↓	$K=32$ 验证 BPB ↓	说明
ES-SSM 完整机制 (参考)	1.965	1.847	1.822	1.826	自适应门控 + Budget Dropout
去除 Budget Dropout (仅在 $\overline{K}$ 训练)	3.974	4.596	2.797	1.796	极端过拟合全容量，小预算崩溃且非单调
均匀门控 ($g_k^{(K)} = 1/\sqrt{K}$)	2.246	2.091	2.049	2.052	丧失输入动态调制能力，全预算大幅退化
静态通道可学标量门控 (输入无关)	2.104	1.965	1.933	1.940	无法响应序列时序模式，性能明显弱于自适应
锚定更新频次 $n_{\text{full}} = 1$ (等同无丢弃)	3.974	4.596	2.797	1.796	蜕变为纯全容量训练，小预算性能崩塌
锚定更新频次 $n_{\text{full}} = \infty$ (无全容量保底)	1.955	1.845	1.823	1.834	高阶通道更新概率偏低，导致大预算略微欠拟合

关键发现¶

低预算崩塌防线：没有 Budget Dropout 时，仅在全容量下训练的模型其低预算截断性能极度恶化（$K=4$ 时 BPB 从 1.847 暴跌至 4.596），且呈现非单调异常震荡（$K=4$ 反而劣于 $K=2$）。这表明在未经正则约束的联合优化中，高阶与低阶通道会演化出剧烈的相互抵消机制；而 Budget Dropout 强行打破了这一依赖。
自适应门控的有效支撑集稀疏化：对训练后门控激活的诊断表明，自适应门控将绝大多数激活质量自发压缩在低索引通道中（前 4 个通道的平均参与度最高），其参与比（Participation Ratio）仅为 3.84，远低于静态门控的 28.25，证明了网络自主学会了将核心信息归并至前缀。
D4RL 离线决策的算子级复用：在强化学习离线决策测试中，ES-SSM 单次训练导出的 5 组策略模型与独立训练的 Decision Transformer 表现持平甚至在 $K=8, 32$ 反超，同时将整体训练时间从 107.84 小时骤降至 23.99 小时，节省了 77.8% 的训练算力。

亮点与洞察¶

算子逼近视角开辟弹性新范式：不同于以往暴力裁减神经网络物理维度的弹性结构，ES-SSM 首次将序列算子的泛函谱分解与系统级弹性部署统一起来，使“降低计算量”在数学上严格等价于“以更低阶的谱基逼近原序列微分算子”。
精妙的机械前缀一致性设计：门控 MLP 的权重组织方式使得前 $K$ 个通道的未归一化响应与后续未激活通道完全独立，避免了在部署阶段执行任何架构搜索或激活再标定。
对数采样的算力调度泛化：对数尺度预算采样策略与周期性全容量锚定更新的结合，不仅彻底解决了多重子网络同时训练时的梯度失衡难题，其设计原则可广泛借鉴于混合专家模型（MoE）路由及弹性注意力机制的设计。

局限与展望¶

算子实现的工程底层优化空间：由于缺乏工业级高度优化的 Triton/CUDA 专用频域算子核函数，在长序列但小通道数场景下，ES-SSM 相比于拥有成熟 FlashAttention 或 Mamba 硬件级流式扫描内核的基线，在纯延迟表现上仍有工程瓶颈，后续需开发针对谱投影收缩的专用融合内核。
离散预设预算的覆盖限制：当前实验依赖预先离散化设定的运行时预算集合 $\mathcal{K}$，对于完全连续的浮点级算力限制调度及在线动态自适应退火预算仍有待进一步探索。
更大模型规模的验证：论文主干实验最高验证至 1.3B 参数与 30B 字节 token 规模，未来需在数十亿至千亿级大语言模型预训练阶段验证谱弹性的扩展法则（Scaling Law）。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从算子近似而非架构物理裁剪的视角构建弹性序列模型，数学基础坚实且思想独特。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1.3B 规模语言建模、LRA 超长序列、语音分类及离线强化学习决策，跨模态验证非常详尽。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨，数学推导详尽完备，消融实验针对性强。
价值: ⭐⭐⭐⭐⭐ 为边缘端与异构硬件的大规模部署提供了极具工业落地潜力的单次训练多端导出方案。