Effiziente hybride Sprachmodellkompression durch gruppenbewusstes SSM-Pruning

Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

April 15, 2025

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

Zusammenfassung

Hybride LLM-Architekturen, die Attention- und State-Space-Modelle (SSMs) kombinieren, erreichen state-of-the-art Genauigkeit und Laufzeitleistung. Jüngste Arbeiten haben gezeigt, dass die Anwendung von Kompression und Distillation auf reine Attention-Modelle kleinere, genauere Modelle zu einem Bruchteil der Trainingskosten liefert. In dieser Arbeit untersuchen wir die Effektivität der Kompression hybrider Architekturen. Wir führen eine neuartige gruppenbewusste Pruning-Strategie ein, die die strukturelle Integrität der SSM-Blöcke und ihre Fähigkeiten zur Sequenzmodellierung bewahrt. Darüber hinaus demonstrieren wir die Notwendigkeit eines solchen SSM-Prunings, um im Vergleich zu traditionellen Ansätzen eine verbesserte Genauigkeit und Inferenzgeschwindigkeit zu erreichen. Unser Kompressionsrezept kombiniert SSM-, FFN-, Embedding-Dimensionen- und Layer-Pruning, gefolgt von einer auf Wissensdistillation basierenden Nachschulung, ähnlich der MINITRON-Technik. Mit diesem Ansatz komprimieren wir das Nemotron-H 8B Hybridmodell auf 4B Parameter mit bis zu 40x weniger Trainings-Tokens. Das resultierende Modell übertrifft die Genauigkeit ähnlich großer Modelle und erreicht eine 2x schnellere Inferenz, wodurch die Pareto-Front signifikant vorangetrieben wird.

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.