新技术实现AI模型在训练中变得更轻更快

richlovec 1500_400 (1)
 

训练大型人工智能模型不仅花费巨大资金,还消耗大量时间、能源和计算资源。传统上,要获得更小更快的模型,通常需要先训练一个庞大的模型再进行剪枝,或者从头训练一个小模型,但后者性能较弱。

麻省理工学院计算机科学与人工智能实验室(CSAIL)、马克斯·普朗克智能系统研究所、欧洲学习与智能系统实验室、苏黎世联邦理工学院(ETH)和Liquid AI的研究人员,开发出一种全新方法,能够在训练过程中直接压缩模型,避免了上述权衡。

这项名为CompreSSM的新技术,针对一类被称为状态空间模型(state-space models)的AI架构,这类模型广泛应用于语言处理、音频生成和机器人技术。研究团队借助控制理论中的数学工具,能够识别模型中哪些部分对性能贡献较大,哪些是“负担”,并在训练早期就将不必要的部分剔除。

“这实际上是一种让模型在训练过程中变得更小更快的技术,”CSAIL博士生、论文第一作者Makram Chahine表示,“模型在学习的同时,也会去除对其发展无用的部分。”

关键在于,模型内部不同组件的重要性在训练初期就趋于稳定。研究团队利用Hankel奇异值这一数学量度,衡量每个内部状态对整体行为的贡献,发现只需完成约10%的训练过程,就能可靠地排序各维度的重要性。确定后,较不重要的部分即可安全丢弃,剩余90%的训练则以更小模型的速度进行。

MIT教授、CSAIL主任Daniela Rus指出:“这项工作令人兴奋,因为它将压缩从事后考虑转变为学习过程的一部分。CompreSSM让模型在学习时自行发现高效结构,这是一种根本不同的AI构建思路。”

实验结果令人瞩目。在图像分类基准测试中,压缩后的模型在保持几乎相同准确率的同时,训练速度提升了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's theorem)数学证明了模型状态重要性在训练中平滑变化,并实验证明状态相对排序稳定,确保早期识别为无关紧要的维度不会后期变得关键。

此外,CompreSSM设计了实用的安全机制:若压缩导致性能意外下降,可回退至之前的检查点,给予用户对性能与压缩程度的灵活控制。

该方法适用于内部状态维度与整体性能高度相关的模型,尤其在多输入多输出(MIMO)模型中效果显著。对于单输入单输出架构,性能提升较为有限,因为这类模型对状态维度变化不敏感。

理论上,CompreSSM最适用于线性时不变系统,团队也开发了针对输入依赖、时变架构的扩展。由于状态空间模型涵盖了如线性注意力等新兴架构,应用前景广泛。

Chahine及其合作者视此为迈向更广泛应用的第一步。团队已展示了对线性时变系统如Mamba的扩展,未来计划将CompreSSM推广至线性注意力机制中的矩阵值动态系统,进一步接近当前主流大型AI系统所用的Transformer架构。

“这是第一步,因为这里理论清晰,方法原则性强,”Chahine说,“它是向工业界广泛使用的其他架构扩展的基石。”

未参与本研究的ELLIS图宾根研究所负责人Antonio Orvieto评价:“Chahine及其团队的工作为现代状态空间模型压缩提供了引人注目的理论视角。该方法证明了训练中有效降低状态维度的可能性,且控制理论视角能成功指导这一过程。此项工作为未来研究开辟新方向,所提出算法有望成为预训练大型状态空间模型的标准方法。”

该研究作为论文被国际学习表征会议(ICLR 2026)接收,将于本月晚些时候发表。部分资金支持来自马克斯·普朗克ETH学习系统中心、Hector基金会、波音公司及美国海军研究办公室。


分享:


发表评论

登录后才可评论。 去登录