Minitron-SSM: Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2025
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866908621722353664 |
|---|---|
| author | Taghibakhshi, Ali Sreenivas, Sharath Turuvekere Muralidharan, Saurav Chochowski, Marcin Karnati, Yashaswi Joshi, Raviraj Mahabaleshwarkar, Ameya Sunil Chen, Zijia Suhara, Yoshi Olabiyi, Oluwatobi Korzekwa, Daniel Patwary, Mostofa Shoeybi, Mohammad Kautz, Jan Catanzaro, Bryan Aithal, Ashwath Tajbakhsh, Nima Molchanov, Pavlo |
| author_facet | Taghibakhshi, Ali Sreenivas, Sharath Turuvekere Muralidharan, Saurav Chochowski, Marcin Karnati, Yashaswi Joshi, Raviraj Mahabaleshwarkar, Ameya Sunil Chen, Zijia Suhara, Yoshi Olabiyi, Oluwatobi Korzekwa, Daniel Patwary, Mostofa Shoeybi, Mohammad Kautz, Jan Catanzaro, Bryan Aithal, Ashwath Tajbakhsh, Nima Molchanov, Pavlo |
| contents | Hybrid LLM architectures that combine Attention and State Space Models (SSMs) achieve state-of-the-art accuracy and runtime performance. Recent work has demonstrated that applying compression and distillation to Attention-only models yields smaller, more accurate models at a fraction of the training cost. In this work, we explore the effectiveness of compressing Hybrid architectures. We introduce a novel group-aware pruning strategy that preserves the structural integrity of SSM blocks and their sequence modeling capabilities. Furthermore, we demonstrate the necessity of such SSM pruning to achieve improved accuracy and inference speed compared to traditional approaches. Our compression recipe combines SSM, FFN, embedding dimension, and layer pruning, followed by knowledge distillation-based retraining, similar to the MINITRON technique. Using this approach, we compress the Nemotron-H 8B Hybrid model down to 4B parameters with up to 40x fewer training tokens. The resulting model surpasses the accuracy of similarly-sized models while achieving 2x faster inference, significantly advancing the Pareto frontier. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2504_11409 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Minitron-SSM: Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning Taghibakhshi, Ali Sreenivas, Sharath Turuvekere Muralidharan, Saurav Chochowski, Marcin Karnati, Yashaswi Joshi, Raviraj Mahabaleshwarkar, Ameya Sunil Chen, Zijia Suhara, Yoshi Olabiyi, Oluwatobi Korzekwa, Daniel Patwary, Mostofa Shoeybi, Mohammad Kautz, Jan Catanzaro, Bryan Aithal, Ashwath Tajbakhsh, Nima Molchanov, Pavlo Computation and Language Hybrid LLM architectures that combine Attention and State Space Models (SSMs) achieve state-of-the-art accuracy and runtime performance. Recent work has demonstrated that applying compression and distillation to Attention-only models yields smaller, more accurate models at a fraction of the training cost. In this work, we explore the effectiveness of compressing Hybrid architectures. We introduce a novel group-aware pruning strategy that preserves the structural integrity of SSM blocks and their sequence modeling capabilities. Furthermore, we demonstrate the necessity of such SSM pruning to achieve improved accuracy and inference speed compared to traditional approaches. Our compression recipe combines SSM, FFN, embedding dimension, and layer pruning, followed by knowledge distillation-based retraining, similar to the MINITRON technique. Using this approach, we compress the Nemotron-H 8B Hybrid model down to 4B parameters with up to 40x fewer training tokens. The resulting model surpasses the accuracy of similarly-sized models while achieving 2x faster inference, significantly advancing the Pareto frontier. |
| title | Minitron-SSM: Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning |
| topic | Computation and Language |
| url | https://arxiv.org/abs/2504.11409 |