Efficient Hybrid Language Model Compression through Group-Aware SSM
  Pruning
Authors: name: Ali Taghibakhshi; name: Sharath Turuvekere Sreenivas; name: Saurav Muralidharan; name: Marcin Chochowski; name: Yashaswi Karnati; name: Raviraj Joshi; name: Ameya Sunil Mahabaleshwarkar; name: Zijia Chen; name: Yoshi Suhara; name: Oluwatobi Olabiyi; name: Daniel Korzekwa; name: Mostofa Patwary; name: Mohammad Shoeybi; name: Jan Kautz; name: Bryan Catanzaro; name: Ashwath Aithal; name: Nima Tajbakhsh; name: Pavlo Molchanov
Published: 2025-04-15 17:26:29

Abstract: 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.