Compressione Efficiente di Modelli Linguistici Ibridi attraverso il Pruning SSM Consapevole dei Gruppi

Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

April 15, 2025

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

Abstract

Le architetture ibride di LLM che combinano modelli di attenzione e modelli a spazio di stati (SSM) raggiungono precisione e prestazioni di runtime all'avanguardia. Recenti lavori hanno dimostrato che applicare compressione e distillazione a modelli basati esclusivamente sull'attenzione produce modelli più piccoli e accurati a una frazione del costo di addestramento. In questo lavoro, esploriamo l'efficacia della compressione delle architetture ibride. Introduciamo una nuova strategia di pruning consapevole dei gruppi che preserva l'integrità strutturale dei blocchi SSM e le loro capacità di modellazione delle sequenze. Inoltre, dimostriamo la necessità di tale pruning degli SSM per ottenere una migliore precisione e velocità di inferenza rispetto agli approcci tradizionali. La nostra ricetta di compressione combina pruning degli SSM, delle FFN, delle dimensioni di embedding e dei livelli, seguita da un riaddestramento basato sulla distillazione della conoscenza, simile alla tecnica MINITRON. Utilizzando questo approccio, comprimiamo il modello ibrido Nemotron-H 8B fino a 4B di parametri con fino a 40x meno token di addestramento. Il modello risultante supera la precisione di modelli di dimensioni simili, raggiungendo un'inferenza 2x più veloce, avanzando significativamente la frontiera di 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.