⚡ electrical engineering

AIRE-Prune: Asymptotic Impulse-Response Energy for State Pruning in State Space Models

AIRE-Prune 是一种针对状态空间模型（SSM）的结构化后训练剪枝方法，它通过基于闭式渐近脉冲响应能量评分进行状态分配与选择来降低状态维度，在多种基准测试中实现了显著的计算量减少且仅造成极小的精度损失。

原作者： Apurba Prasad Padhy, Fernando Camacho, Saibal Mukhopadhyay

发布于 2026-02-03

📖 1 分钟阅读☕ 轻松阅读

原作者： Apurba Prasad Padhy, Fernando Camacho, Saibal Mukhopadhyay

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你拥有一座庞大且超级聪明的图书馆（一个状态空间模型，State Space Model），它能够阅读长篇故事或聆听长段音频。为了理解故事，这座图书馆的大脑内部拥有数以千计的微型“记录员”（称为状态，states）。每个记录员都负责记住过去的一段特定信息。

问题在于，这座图书馆人员过剩了。它雇佣了比实际需要多得多的记录员，这使得图书馆运行缓慢、成本高昂且难以管理，尽管其中大多数记录员只是在无所事事地坐着。

AIRE-Prune 是一种聪明的新方法，它可以在不损害图书馆讲述故事能力的前提下，解雇那些不必要的记录员。以下是它的工作原理，通过简单的类比来解释：

1. 旧方法：“最响亮的尖叫”（最坏情况）

以前的方法试图通过询问：“谁的声音最大？如果把某人推向极限，谁能发出最响亮的尖叫？”来决定保留哪些记录员。

这就像是一场消防演习，你只留下那些能尖叫得最响亮的人。问题在于，在现实生活中，从来没有人会叫得那么大声。因此，你会留下那些擅长尖叫但对日常对话毫无用处的人，同时解雇了那些虽然声音小但确实在做日常工作的实干家。这被称为“最坏情况”（worst-case）方法，通常过于保守。

2. 新方法：“总能量得分”（AIRE-Prune）

论文作者说：“让我们不再担心最响亮的尖叫。相反，让我们衡量每个记录员在漫长的、无限的时间内，实际完成了多少总工作量。”

他们称之为渐近脉冲响应能量（Asymptotic Impulse-Response Energy）。

隐喻： 想象你敲击了一下记录员一次（一个“脉冲”）。这一次敲击会在最终的故事中产生多少能量？
有些记录员就像一面重鼓：你敲击一次，它就会持续震动很久，为故事贡献大量的“能量”。
其他记录员则像一根羽毛：你敲击一下，它几乎纹丝不动。它们几乎对故事没有任何贡献。

AIRE-Prune 根据这种总能量为每一个记录员计算得分。如果一个记录员的得分很低，这意味着他们基本上是“累赘”。

3. “全局记分板”（归一化）

这里有一个棘手的部分：图书馆有不同的房间（层/layers）。在“入口室”，记录员可能天生就很吵；而在“出口室”，记录员可能天生就很安静。如果你只比较原始得分，你可能会不小心解雇掉整个安静房间里的所有人。

AIRE-Prune 通过创建一个归一化记分板来解决这个问题。它观察“入口室”，并问道：“好的，这里的顶尖 10% 的员工是谁？”然后在“出口室”也进行同样的操作。接着，它将所有这些“顶尖员工”放在一张巨大的名单上进行决策。这确保了每个房间都能得到公平对待，无论该房间平时是吵闹还是安静。

4. 结果：精简冗余

论文在著名的挑战集 Long Range Arena（模型需要记忆极长序列数据的测试）上测试了该方法。

神奇数字： 他们能够平均解雇 60.8% 的记录员（状态）。
代价： 图书馆的准确率仅下降了 0.29%。这几乎是不可察觉的。
收益： 因为解雇了这么多记录员，图书馆变得更快，且占用的内存更少。

为什么这比竞争对手更好？

论文将 AIRE-Prune 与之前最好的方法（称为 LAST，它使用的是“最响亮的尖叫”即最坏情况方法）进行了对比。

LAST 像是一个谨慎的经理，为了以防万一而保留了太多人。
AIRE-Prune 则像是一个聪明的经理，他观察实际的每日产出。它发现图书馆携带了许多旧方法未能识别出的“累赘”。

总结

可以将 AIRE-Prune 视为 AI 模型的一个高效人力资源部门。它不是根据谁可能叫得最响来猜测谁很重要，而是衡量谁在一段时间内实际贡献了能量。通过解雇那些低能量的员工，它让 AI 运行得更快、更便宜，同时保持故事（准确性）几乎完全不变。

论文声称，这种方法适用于不同类型的 AI 模型（如 S5、S4D 和 Mamba），并且无需重新训练。你只需计算得分，切掉底部的 60%，大功告成。

技术摘要：AIRE-Prune

问题陈述

深度状态空间模型（SSMs）通过将输入历史压缩到内部状态中，在建模序列数据方面展现出了卓越的效能。然而，状态维度 $n$ 仍然是计算和内存成本的主要驱动因素。虽然现代架构通常利用对角或可对角化的参数化（例如 S5 等多 SISO 或 MIMO 堆叠）来确保稳定性和效率，但这些模型经常面临过度参数化的问题。现有的训练后优化 $n$ 的机制非常有限。以往的方法（如非结构化权重剪枝）并不能直接减少状态维度。其他结构化方法，如最近的 LAST (Gwak et al., 2025)，依赖于最差情况下的频域增益（ $H_\infty$ ）来对状态进行评分。虽然有效，但最差情况度量对于典型工作负载而言过于保守，因为它们强调的是任务可能很少激发的峰值放大，这可能导致次优的压缩比。

方法论：AIRE-Prune

作者引入了 AIRE-Prune（用于状态剪枝的渐近脉冲响应能量，Asymptotic Impulse-Response Energy for State PRUN(E)ing），这是一种针对对角（或可对角化）SSM 的结构化、训练后、层自适应剪枝方法。与以往关注最差情况增益的方法不同，AIRE-Prune 将剪枝与渐近脉冲响应能量对齐，这是一个代表了状态在无限时间范围内对总输出能量贡献的“典型情况”准则。

核心机制

基于能量的排序：
对于由 $(\Lambda, B, C)$ 定义的稳定离散时间对角 SSM 层，其中 $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_n)$ 且 $|\lambda_i| < 1$ ，该方法计算了每个状态（模态） $i$ 在单位脉冲输入下产生的总平方输出能量。
第 $i$ 个模态的闭式解分数为：
$\text{EnergyScore}_{\text{local}}(x_i) = E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1 - |\lambda_i|^2}$
该分数集成了三个因素：
- 可控性： 由 $\|B_{i,:}\|_2^2$ 表示（易于被激发）。
- 可观测性： 由 $\|C_{:,i}\|_2^2$ 表示（易于被测量）。
- 阻尼/记忆： 由极点 $\lambda_i$ 表示；极点靠近单位圆（衰减较慢）的状态会对长期能量做出更多贡献。
跨层归一化：
由于能量尺度在不同层之间可能存在显著差异（例如，由于编码器/解码器增益），作者提出了一种前缀归一化评分策略以实现全局比较：
- 在每个层 $\ell$ 内，按局部能量 $E^{(\ell)}_i$ 降序排列模态。
- 计算排序后能量的前缀和 $S^{(\ell)}_{(i)}$ 。
- 全局得分定义为“风险率”比率：
  $\text{AIREPrune}(x^{(\ell)}_{(i)}) = \frac{E^{(\ell)}_{(i)}}{S^{(\ell)}_{(i)} + \epsilon}$
  这种归一化确保了得分在层内是单调不增的，从而允许在整个模型中应用统一的全局阈值 $\tau$ 。
剪枝执行：
给定目标剪枝率或阈值，算法保留每个层中归一化得分超过 $\tau$ 的最长前缀状态。这实现了层自适应剪枝，即某些层可能会被大幅剪枝甚至完全移除（如果其状态低于全局阈值），而其他层则保持基本不变。

核心贡献

闭式渐近能量评分： 本文提供了一个基于无限时界输出能量的、针对每个学习到的模态的数学推导闭式评分。这避免了昂贵的数值模拟或重新训练。
模态截断的扩展： 该工作将经典的模态截断（通常应用于单个线性系统）扩展到了深层非线性 SSM 堆叠，同时保持了对效率至关重要的对角结构。
典型情况 vs. 最差情况： 通过将重要性指标从最差情况 $H_\infty$ 增益转向总脉冲响应能量，该方法针对的是系统的聚合行为而非罕见的峰值放大。
架构无关性： 该方法既适用于多 SISO 构建和 MIMO 堆叠（如 S5 型），也适用于其他 SSM 系列，如 S4D 和 Mamba。

实验结果

作者在 Long Range Arena (LRA) 基准测试和 Speech Commands 数据集上评估了 S5、S4D 和 Mamba 后端。

压缩效率： 在各种基准测试中，AIRE-Prune 在无需任何重新训练的情况下，实现了平均 60.8% 的状态剪枝，且平均精度下降仅为 0.29%。
与基准对比：
- 在 LRA 系列中，AIRE-Prune 的表现显著优于包括 Random、Uniform Magnitude、Global Magnitude、LAMP、Uniform $H_\infty$ 、Global $H_\infty$ 和 LAST 在内的基准方法。
- 在限制精度损失 $\le 1$ 个百分点 (pp) 的严格约束下，AIRE-Prune 允许比最强基准（LAST）更高的剪枝率。例如，在 Path-X 任务中，相比于 LAST 剪掉 30% 的状态，AIRE-Prune 可以安全地剪掉 70% 的状态。
- 在 Speech Commands 任务中，剪掉了 45% 的状态，且精度损失微乎其微（完整模型的 96.43% 对比剪枝后的 96.40%）。
推理效率： 剪枝直接转化为系统级的收益。在 NVIDIA H100 GPU 上，该方法带来了 1.23 $\times$ 至 2.86 $\times$ 的推理加速，以及 19% 到 65% 的参数减少，具体取决于任务和剪枝比例。
泛化能力： 该方法在 S4D 和 Mamba 架构上展示了稳健的表现，在 S4D 的 Text 任务上实现了高达 90% 的剪枝，在 Mamba 的 Image 任务上实现了 80% 的剪枝，且精度下降极小。

重要性与主张

论文声称 AIRE-Prune 揭示了已训练 SSM 状态空间中的巨大冗余，这些冗余可以在无需重新训练的情况下被移除。通过将剪枝与渐近响应能量而非最差情况增益对齐，该方法实现了激进的压缩，且对任务性能的影响微乎其微。

作者强调，观察到的 AIRE-Prune “阶梯式”精度下降剖面（即精度在达到高剪枝阈值前保持稳定，然后急剧下降）表明重要状态与不重要状态之间存在强烈的分离。这与显示出平滑退化的基准方法形成对比，表明基准方法的评分方法混合了重要和不重要的状态。

在实践中，该方法实现了：

全局预算控制： 设置统一的目标预算，无需进行逐层微调。
块级缩减： 能够当层状态低于阈值时剪掉整个层，将细粒度的稀疏性转化为对硬件更友好的块结构缩减。
免训练优化： 一种无需额外微调的训练后机制，使其能够立即应用于已部署的模型。

研究结论认为，AIRE-Prune 为深度 SSM 的模型阶数缩减提供了一种原则性的、基于能量的方法，在压缩率和精度保持方面均超越了当前的先进剪枝策略。