Efficiënte hybride compressie van taalmodel door groepbewuste SSM-pruning

Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

April 15, 2025

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

Samenvatting

Hybride LLM-architecturen die Attention en State Space Models (SSMs) combineren, bereiken state-of-the-art nauwkeurigheid en runtime-prestaties. Recent onderzoek heeft aangetoond dat het toepassen van compressie en distillatie op Attention-only modellen kleinere, nauwkeurigere modellen oplevert tegen een fractie van de trainingskosten. In dit werk onderzoeken we de effectiviteit van het comprimeren van hybride architecturen. We introduceren een nieuwe groep-gevoelige pruningstrategie die de structurele integriteit van SSM-blokken en hun sequentiemodelleringscapaciteiten behoudt. Bovendien demonstreren we de noodzaak van dergelijke SSM-pruning om verbeterde nauwkeurigheid en inferentiesnelheid te bereiken in vergelijking met traditionele benaderingen. Ons compressierecept combineert SSM-, FFN-, embeddingdimensie- en layer-pruning, gevolgd door hertraining op basis van kennisdistillatie, vergelijkbaar met de MINITRON-techniek. Met deze aanpak comprimeren we het Nemotron-H 8B Hybride model tot 4B parameters met tot 40x minder trainings-tokens. Het resulterende model overtreft de nauwkeurigheid van vergelijkbaar grote modellen terwijl het 2x snellere inferentie bereikt, wat de Pareto-grens aanzienlijk vooruithelpt.

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.