Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

2025-04-15

AI Summary

Efficiently compressing Attention+SSM hybrid large language models remains challenging due to the intricate interplay between attention and state-space modeling (SSM) components. Method: This paper proposes a group-aware SSM structured pruning technique—the first systematic investigation into the pivotal role of SSM modules in hybrid-architecture compression—combined with multi-granularity pruning across feed-forward networks (FFNs), embedding layers, and Transformer blocks, and MINITRON-style knowledge distillation–based retraining. Contribution/Results: The method preserves sequence modeling capability and architectural integrity while compressing Nemotron-H 8B to 4B parameters, reducing training tokens by 40× and doubling inference speed. Crucially, it achieves higher accuracy than same-scale baselines, significantly advancing the accuracy–latency Pareto frontier.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Data Mining & Knowledge Management: Data Compression

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Compressing Hybrid Attention-SSM LLMs efficiently
• Preserving SSM block integrity during pruning
• Achieving faster inference with higher accuracy

Innovation

• Group-aware pruning for SSM structural integrity
• Combined SSM, FFN, embedding, and layer pruning
• Knowledge distillation-based retraining for efficiency