训练中压缩状态空间模型,高效保持性能。
The Curious Case of In-Training Compression of State Space Models
- 用控制理论识别关键状态维度,在训练时动态压缩模型。
- 压缩后模型优化更快,性能优于直接小尺寸训练的模型。
- 适合追求效率与性能平衡的序列建模研究者。
状态空间模型(SSMs)为高效处理长序列任务而设计,兼具可并行训练和快速推理的特点。其核心是保持隐藏状态的递归动力系统,更新开销随状态维数增长。关键挑战在于在表达能力与计算负担间取得平衡。控制理论中的汉克尔奇异值分析提供了一种衡量各状态能量的方法,并可通过平衡截断将原系统压缩至更小表示,同时保证性能。我们利用汉克尔矩阵的特征值稳定性,将此方法应用于训练过程中的SSMs,仅保留高影响力维度。所提方法CompreSSM适用于线性时不变的SSMs(如线性循环单元),也可扩展至选择性模型。实验表明,训练中压缩显著加速优化,同时保持表达能力;压缩模型能保留直接小维度训练所丢失的任务关键结构。即:从大模型开始、训练中压缩的方案,在实现计算效率的同时维持更高性能。项目代码见github.com/camail-official/compressm。
原文摘要 · Abstract (English)
State Space Models (SSMs), developed to tackle long sequence modeling tasks efficiently, offer both parallelizable training and fast inference. At their core are recurrent dynamical systems that maintain a hidden state, with update costs scaling with the state dimension. A key design challenge is striking the right balance between maximizing expressivity and limiting this computational burden. Control theory, and more specifically Hankel singular value analysis, provides a potent framework for the measure of energy for each state, as well as the balanced truncation of the original system down to a smaller representation with performance guarantees. Leveraging the eigenvalue stability properties of Hankel matrices, we apply this lens to SSMs \emph{during training}, where only dimensions of high influence are identified and preserved. Our approach, \textsc{CompreSSM}, applies to Linear Time-Invariant SSMs such as Linear Recurrent Units, but is also extendable to selective models. Experiments show that in-training reduction significantly accelerates optimization while preserving expressivity, with compressed models retaining task-critical structure lost by models trained directly at smaller dimension. In other words, SSMs that begin large and shrink during training achieve computational efficiency while maintaining higher performance. Project code is available at github.com/camail-official/compressm.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。