Compresión Eficiente de Modelos de Lenguaje Híbridos mediante Poda SSM con Conciencia de Grupos

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

Resumen

Las arquitecturas híbridas de LLM que combinan modelos de Atención y modelos de Espacio de Estados (SSMs) logran una precisión y rendimiento en tiempo de ejecución de vanguardia. Trabajos recientes han demostrado que aplicar compresión y destilación a modelos basados únicamente en Atención produce modelos más pequeños y precisos con una fracción del costo de entrenamiento. En este trabajo, exploramos la efectividad de comprimir arquitecturas híbridas. Introducimos una novedosa estrategia de poda consciente de grupos que preserva la integridad estructural de los bloques SSM y sus capacidades de modelado de secuencias. Además, demostramos la necesidad de dicha poda en SSM para lograr una mayor precisión y velocidad de inferencia en comparación con enfoques tradicionales. Nuestra receta de compresión combina la poda de SSM, FFN, dimensiones de incrustación y capas, seguida de un reentrenamiento basado en destilación de conocimiento, similar a la técnica MINITRON. Utilizando este enfoque, comprimimos el modelo híbrido Nemotron-H 8B a 4B parámetros con hasta 40 veces menos tokens de entrenamiento. El modelo resultante supera la precisión de modelos de tamaño similar mientras logra una inferencia 2 veces más rápida, avanzando significativamente la frontera 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.