Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning (2504.11409v1)

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.

Efficient Hybrid LLM Compression through Group-Aware SSM Pruning

This paper focuses on optimizing hybrid LLMs by introducing a novel compression method that leverages the unique structural properties of hybrid model architectures, specifically those combining Attention mechanisms and State Space Models (SSMs). Traditionally, such models have shown effectiveness in capturing both global dependencies and sequence processing, though they often remain large and computationally demanding. The paper addresses the challenge of compressing these hybrid architectures efficiently while retaining their accuracy and performance on complex tasks.

The primary contribution is a novel group-aware pruning strategy designed for SSM blocks within hybrid architectures. This method leverages the grouped structure inherent in the SSMs, ensuring that the pruning does not disrupt the sequence modeling capabilities crucial for maintaining model performance. The authors argue that understanding and preserving the structural integrity of SSM blocks is critical to achieving improvements in both model accuracy and inference speed—a key trade-off when compressing models.

The proposed compression recipe includes several advanced techniques:

• SSM Pruning: This maintains group-awareness in pruning strategies to ensure stability in sequence modeling.
• FFN, Embedding Dimension, and Layer Pruning: These are applied to manage the model's scale, while preserving the learning framework.
• Knowledge Distillation-Based Retraining: This follows the pruning process to recover any loss in model accuracy post-pruning, effectively implementing a retraining approach similar to the MINITRON technique.

Experimental validation shows significant improvements. Using the Nemotron-H 8B Hybrid model as a reference, the authors compress it down to 4B parameters while requiring only up to 40 times fewer training tokens. Despite the reduction in model size, the compressed model exceeds the accuracy of similar-sized models and is nearly twice as fast in inference speed. Thus, this method effectively advances the Pareto frontier for hybrid LLMs.

Several practical and theoretical implications arise from this research. Practically, deployments of LLMs in resource-constrained environments can benefit from this compression strategy, especially where computational resources are limited or where latency is a concern. Theoretically, this work helps elucidate the relationships between model architecture, structural pruning, and sequence modeling integrity. It suggests that carefully considered hybrid pruning strategies can unlock new capabilities in LLM deployment and reflect better granularity in model scalability without sacrificing performance.

Future research directions might explore further optimization of sequence modeling capabilities in structured architectures, potentially extending this pruning strategy to other hybrid forms that integrate different types of models. Furthermore, enhanced retraining methods combined with advanced distillation techniques could further enhance model performance post-compression, decreasing token needs without hampering learning processes.

In summary, this paper presents a methodologically robust approach to compressing hybrid LLMs while setting a new standard in maintaining model accuracy and computational efficiency. This work effectively pushes the boundaries on how compression can be achieved across complex hybrid architectures, bringing advancements within reach for real-world applications where resource efficiency is paramount.