Эффективное сжатие гибридных языковых моделей с помощью группового SSM-прунинга

Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

April 15, 2025

Авторы: Ali Taghibakhshi, Sharath Turuvekere Sreenivas, Saurav Muralidharan, Marcin Chochowski, Yashaswi Karnati, Raviraj Joshi, Ameya Sunil Mahabaleshwarkar, Zijia Chen, Yoshi Suhara, Oluwatobi Olabiyi, Daniel Korzekwa, Mostofa Patwary, Mohammad Shoeybi, Jan Kautz, Bryan Catanzaro, Ashwath Aithal, Nima Tajbakhsh, Pavlo Molchanov

cs.AI

Аннотация

Гибридные архитектуры больших языковых моделей (LLM), сочетающие механизмы внимания (Attention) и модели пространства состояний (State Space Models, SSMs), достигают наивысшей точности и производительности. Недавние исследования показали, что применение сжатия и дистилляции к моделям, основанным исключительно на механизме внимания, позволяет создавать более компактные и точные модели при значительно меньших затратах на обучение. В данной работе мы исследуем эффективность сжатия гибридных архитектур. Мы представляем новую стратегию группового прореживания, которая сохраняет структурную целостность блоков SSM и их способность к моделированию последовательностей. Кроме того, мы демонстрируем необходимость такого прореживания SSM для достижения улучшенной точности и скорости вывода по сравнению с традиционными подходами. Наш метод сжатия объединяет прореживание SSM, FFN, размерности эмбеддингов и слоев, за которым следует повторное обучение на основе дистилляции знаний, аналогично технике MINITRON. Используя этот подход, мы сжимаем гибридную модель Nemotron-H 8B до 4 миллиардов параметров, используя до 40 раз меньше токенов для обучения. Полученная модель превосходит по точности модели аналогичного размера, обеспечивая при этом двукратное ускорение вывода, что значительно продвигает границу Парето.

English

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.