Minitron-SSM:通过基于组的 SSM 剪枝实现的高效混合语言模型压缩
计算与语言
2025-11-03 v2
摘要
混合 LLM 架构将 Attention 与 State Space Models(SSMs)组合实现最佳的准确率和运行时性能。最近的工作表明,对 Attention-only models 应用压缩和蒸馏可获得更小、更准确的模型,训练成本为原来的一小部分。在本工作中,我们探讨了压缩混合架构的有效性。我们引入一种 novel group-aware pruning strategy,以保留 SSM blocks 及其序列建模能力的结构完整性。此外,我们证明了这种 SSM 剪枝是实现 improved accuracy 和 inference speed 相较于传统方法的必要条件。我们的压缩配方组合 SSM、FFN、embedding dimension 和 layer 剪枝,随后进行 knowledge distillation-based 重训练,类似于 MINITRON 技术。使用此方法,我们将 Nemotron-H 8B Hybrid model 压缩至 4B 参数,训练标记数减少最高可达 40 倍。 resulting model 在准确率方面超过 similarly-sized models,推理速度提升 2 倍,显著推动了 Pareto 前沿的 advance。
引用
@article{arxiv.2504.11409,
title = {Minitron-SSM: Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning},
author = {Ali Taghibakhshi and Sharath Turuvekere Sreenivas and Saurav Muralidharan and Marcin Chochowski and Yashaswi Karnati and Raviraj Joshi and Ameya Sunil Mahabaleshwarkar and Zijia Chen and Yoshi Suhara and Oluwatobi Olabiyi and Daniel Korzekwa and Mostofa Patwary and Mohammad Shoeybi and Jan Kautz and Bryan Catanzaro and Ashwath Aithal and Nima Tajbakhsh and Pavlo Molchanov},
journal= {arXiv preprint arXiv:2504.11409},
year = {2025}
}