CompreSSM:在训练过程中压缩状态空间模型的新技术

训练大型人工智能模型不仅成本高昂,还极其耗时,并消耗大量能源和计算资源。以往若想得到更小、更快的模型,通常有两条路:先训练一个庞大模型再进行剪枝压缩,或者从一开始就训练一个小模型,但往往性能不足。

麻省理工学院计算机科学与人工智能实验室(CSAIL)、马克斯·普朗克智能系统研究所、欧洲学习与智能系统实验室、苏黎世联邦理工学院(ETH)以及 Liquid AI 的研究人员提出了一种新方法,试图打破这一传统权衡:在训练过程中就对模型进行压缩,而不是训练完成后再动手。

在训练中压缩状态空间模型

这项技术被称为 CompreSSM,面向一类称为状态空间模型(State Space Models, SSMs)的人工智能架构。SSM 被广泛应用于语言处理、音频生成和机器人等任务。研究团队借助控制理论中的数学工具,识别模型内部哪些部分真正起作用、哪些是冗余,并在训练早期有针对性地移除不必要的组件。

“本质上,这是让模型在学习的同时变得更小、更快的一种方法。”论文第一作者、MIT 电气工程与计算机科学博士生、CSAIL 成员 Makram Chahine 表示,“在训练过程中,模型会主动剔除对其发展无用的部分。”

该研究论文已被国际学习表征会议(ICLR 2026)接收,将于本月晚些时候发表,目前已在 arXiv 预印本平台公开。

关键洞见:早期即可判断重要性

研究的核心发现是:模型内部不同组件的重要性在训练初期就会趋于稳定。团队使用一种称为 Hankel 奇异值 的数学量度,来衡量每个内部状态对整体模型行为的贡献。他们证明,只需在训练进度约 10% 时,就能可靠地对各个维度的重要性进行排序,区分出关键与非关键部分。

一旦完成排序,重要性较低的组件就可以被安全地丢弃,后续约 90% 的训练过程便可以在一个更小的模型上进行,从而显著提升训练速度。

MIT 教授、CSAIL 主任、论文资深作者 Daniela Rus 指出:“这项工作的特别之处在于,它把压缩从‘事后处理’变成了学习过程的一部分。CompreSSM 让模型在训练时自行发现高效结构,而不是先把大模型训完再想办法缩小。这代表了一种构建人工智能系统的全新思路。”

基准测试中的训练加速与性能表现

在多项关键基准测试中,CompreSSM 展现出显著优势。在图像分类任务上,压缩后的模型在几乎不损失准确率的前提下,训练速度最高提升约 1.5 倍。

例如,在 CIFAR-10 数据集上,一个被压缩到原始状态维度约四分之一的模型,仍然取得了 85.7% 的准确率;而从头直接训练同样小尺寸的模型,准确率只有 81.8%。在广泛使用的状态空间架构 Mamba 上,该方法实现了约 4 倍 的训练加速:将 128 维模型压缩到约 12 维的同时,依然保持具有竞争力的性能。

Chahine 解释说:“你仍然获得了大模型的性能,因为在预热阶段已经捕捉到了大部分复杂动态,然后只保留最有用的状态。模型的表现依旧优于从一开始就训练的小模型。”

与剪枝和知识蒸馏的差异

CompreSSM 的一个显著特点是其坚实的理论基础。传统剪枝方法需要先完整训练大模型,再在事后删除部分参数,训练阶段仍要承担大模型的全部计算成本。知识蒸馏则需要先训练一个大型“教师”模型,再训练一个较小的“学生”模型,整体训练工作量几乎翻倍。

CompreSSM 则在训练中期做出压缩决策,从而避免了上述两种方法的高成本。研究团队将 CompreSSM 与这些主流方法进行了直接对比。

与近期提出的一种基于 Hankel 核范数正则化的光谱方法相比,CompreSSM 在训练速度上快了 40 倍以上,同时准确率更高。该正则化方法由于在每一步梯度更新中都需要进行昂贵的特征值计算,训练速度大约慢了 16 倍,且最终模型性能也不理想。

在 CIFAR-10 上与知识蒸馏方法对比时,CompreSSM 在高压缩率场景中优势尤为明显:当状态维度被大幅缩小时,蒸馏模型的准确率明显下滑,而 CompreSSM 压缩后的模型仍能保持接近完整模型的性能。此外,由于知识蒸馏在训练每一步都要同时前向传播教师和学生模型,即便学生模型较小,其训练速度仍慢于直接训练完整基线模型。

理论保证与实用“安全网”

在理论层面,研究人员利用 Weyl 定理 证明了模型各个状态的重要性在训练过程中是平滑变化的,并通过实验表明这些状态的重要性排序在训练中保持稳定。这为实践者提供了信心:那些在早期被判定为不重要的维度,不会在训练后期突然变得关键。

此外,该方法还设计了一个实用的“安全网”机制。如果压缩步骤意外导致性能明显下降,使用者可以回滚到压缩前保存的检查点。

Chahine 表示:“这种机制让用户可以直接控制愿意牺牲多少性能,而不必去设定一些不直观的能量阈值。”

适用场景与局限

CompreSSM 并非对所有模型都同样有效。它最适用于那些内部状态维度与整体性能高度相关的模型,而这一点会随任务和架构的不同而变化。

该方法在 多输入多输出(MIMO) 模型上效果尤佳,因为这类模型中状态大小与表达能力之间的关系最为紧密。相对地,对于每通道为单输入单输出的架构,收益较为有限,因为这类模型对状态维度变化本身不太敏感。

在理论上,CompreSSM 最适合 线性时不变系统。研究团队也提出了针对输入依赖的时变架构的扩展,以适应日益流行的现代 SSM 变体。值得一提的是,状态空间模型家族还包括 线性注意力 等架构,而线性注意力正作为传统 Transformer 的替代方案受到越来越多关注,应用前景广阔。

未来方向:迈向更大规模架构

Chahine 及其合作者将这项工作视为进一步研究的起点。团队已经展示了 CompreSSM 在诸如 Mamba 等线性时变系统上的扩展,下一步计划是将该方法推广到用于线性注意力机制的矩阵值动力系统,从而进一步接近当前支撑大多数大型 AI 系统的 Transformer 架构。

“我们必须从这里起步,因为在这一设定下理论最为清晰,方法也最具原则性。”Chahine 说,“这为向工业界当前使用的其他架构扩展提供了跳板。”

未参与该研究的 ELLIS 图宾根研究所首席研究员、马克斯·普朗克智能系统研究所独立组组长 Antonio Orvieto 评论道:“Chahine 及其同事的工作,为现代状态空间模型(SSMs)的压缩提供了一个引人注目且有理论支撑的视角。

该方法展示了在训练过程中有效降低模型状态维度的可行性,并证明控制理论视角可以成功指导这一过程。这项工作为未来研究打开了新的方向,所提出的算法有潜力成为预训练大型 SSM 模型的标准方法。”

本文经 MIT 新闻(web.mit.edu/newsoffice/)授权转载,该网站专注报道 MIT 的研究、创新与教学相关内容。


分享:


发表评论

登录后才可评论。 去登录