Efficient Hybrid Language Model Compression through Group-Aware SSM
  Pruning
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.
67fff37f995d6bb5fe72e2bb
Ali Taghibakhshi
67fff37f995d6bb5fe72e2bc
Sharath Turuvekere Sreenivas
67fff37f995d6bb5fe72e2bd
Saurav Muralidharan
67fff37f995d6bb5fe72e2be
Marcin Chochowski
67fff37f995d6bb5fe72e2bf
Yashaswi Karnati
67fff37f995d6bb5fe72e2c0
Raviraj Joshi
67fff37f995d6bb5fe72e2c1
Ameya Sunil Mahabaleshwarkar
67fff37f995d6bb5fe72e2c2
Zijia Chen
67fff37f995d6bb5fe72e2c3
Yoshi Suhara
67fff37f995d6bb5fe72e2c4
Oluwatobi Olabiyi
67fff37f995d6bb5fe72e2c5
Daniel Korzekwa
67fff37f995d6bb5fe72e2c6
Mostofa Patwary
67fff37f995d6bb5fe72e2c7
Mohammad Shoeybi
67fff37f995d6bb5fe72e2c8
Jan Kautz
67fff37f995d6bb5fe72e2c9
Bryan Catanzaro
67fff37f995d6bb5fe72e2ca
Ashwath Aithal
67fff37f995d6bb5fe72e2cb
Nima Tajbakhsh
67fff37f995d6bb5fe72e2cc
Pavlo Molchanov
claimed_verified
2025-04-20T15:03:39.466Z
646d0c1c534e52f8c30500a6
https://cdn-avatars.huggingface.co/v1/production/uploads/646d0c1c534e52f8c30500a6/75VH8ClbRaP75BU2ONfXE.png
Pavlo Molchanov
user
pmolchanov
Efficient Hybrid Language Model Compression through Group-Aware SSM
  Pruning
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.
A novel group-aware pruning strategy is introduced to compress Hybrid LLM architectures, enhancing accuracy and inference speed while reducing training costs and parameters.
Hybrid LLM architectures Attention State Space Models SSM FFN embedding dimension layer pruning knowledge distillation MINITRON Nemotron-H parameter compression training tokens