Compressão Eficiente de Modelos de Linguagem Híbridos por meio de Poda SSM com Consciência de Grupo

Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

April 15, 2025

Autores: 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

Resumo

Arquiteturas híbridas de LLM que combinam modelos de Atenção e modelos de Espaço de Estados (SSMs) alcançam precisão e desempenho de tempo de execução de última geração. Trabalhos recentes demonstraram que a aplicação de compressão e destilação em modelos baseados apenas em Atenção resulta em modelos menores e mais precisos, com uma fração do custo de treinamento. Neste trabalho, exploramos a eficácia da compressão em arquiteturas híbridas. Introduzimos uma nova estratégia de poda consciente de grupos que preserva a integridade estrutural dos blocos SSM e suas capacidades de modelagem de sequências. Além disso, demonstramos a necessidade dessa poda SSM para alcançar melhorias na precisão e velocidade de inferência em comparação com abordagens tradicionais. Nossa receita de compressão combina poda de SSM, FFN, dimensão de embedding e camadas, seguida de retreinamento baseado em destilação de conhecimento, semelhante à técnica MINITRON. Utilizando essa abordagem, comprimimos o modelo híbrido Nemotron-H 8B para 4B de parâmetros, com até 40x menos tokens de treinamento. O modelo resultante supera a precisão de modelos de tamanho similar enquanto alcança inferência 2x mais rápida, avançando significativamente a fronteira de Pareto.

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.