PaperlayerView exact revision on arXiv

Source-linked AI summary

Minitron-SSM: Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

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

arXiv:2504.11409v2cs.CL

TL;DR

Hybrid LLMs combine Attention and SSM strengths but remain large, while pruning has been studied less for hybrids than for Transformer models. This paper introduces group-aware SSM pruning in a broader compression recipe, and reports a 4B model with strong accuracy, faster inference, and substantially fewer training tokens.

  • Problem

    Hybrid LLMs remain billion-parameter models, and pruning methods extensively studied for Transformers remain underexplored for hybrid architectures.

    Source: Introduction · p. 1

  • Method

    The paper combines group-aware Mamba pruning with FFN, embedding-dimension, and layer pruning, then uses knowledge distillation to retrain the compressed model.

    Source: Introduction · p. 2

  • Results

    Nemotron-H 4B retains over 96% of the parent 8B model’s accuracy while improving throughput by ∼1.4×, and achieves state-of-the-art accuracy across task families despite up to ∼40× fewer training tokens.

    Source: Evaluation Summary · p. 10, Conclusions · p. 11

  • Takeaways & Limitations

    Nemotron-H 4B advances the accuracy-efficiency Pareto frontier, with ∼2× faster inference and 2.6% higher accuracy across diverse tasks.

    Source: Conclusions · p. 11

    The analysis factors out sequence length and batch size for simplicity, while stating that it remains valid without them.

    Source: Mamba Pruning · p. 3

Abstract

from arXiv · show

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.

Introduction

Hybrid models combine complementary strengths of Transformers and SSMs, but remain large, and pruning them is underexplored. The paper proposes group-aware Mamba pruning within a broader hybrid compression recipe.

  • Transformers capture global dependencies, while SSMs offer O(N) training scaling and O(1) inference cache size; hybrid models nevertheless remain billion-parameter systems.This motivates compression for deployment in resource-constrained environments.Source: Introduction · p. 1
  • Although structured pruning with knowledge distillation has benefited Transformer models, its application to hybrid models remains underexplored.Prior Mamba pruning approaches include removing the state space module or reducing and merging Mamba dimensions, while a holistic recipe across SSM and other components was lacking.Source: Introduction · p. 1
  • Figure 1 compares 4B model accuracy against inference throughput and compares base-model training budgets with similarly sized community models.Throughput uses input and output sequence lengths of 65536 and 1024, respectively.Source: Figure 1 · p. 2
  • The proposed recipe combines group-aware Mamba pruning with pruning of FFN neurons, embedding channels, and layers.The Mamba method compresses heads and head channels while preserving SSM block structure and sequence modeling capabilities.Source: Introduction · p. 2
  • Nemotron-H 8B is compressed to 4B with up to ∼40× fewer training tokens, state-of-the-art accuracy, and ∼2× throughput speedup versus similarly sized models.The recipe combines pruning with knowledge distillation.Source: Introduction · p. 2

Background

The background describes SSM sequence processing, hybrid SSM-Transformer architectures, and structured pruning as context for the paper’s compression approach.

  • SSMs process inputs through hidden states evolving over time, with hidden state, input, and output related through parameter matrices.The described general equations are linear time-invariant when parameters remain constant across timesteps.Source: Background · p. 2
  • Mamba introduces time-varying selective SSM parameters, while Mamba2 uses Structured State Space Duality to connect SSMs and attention through semi-separable matrix representations.This duality combines linear SSM efficiency with hardware-friendly quadratic computations typical of attention models.Source: Background · pp. 2–3
  • Hybrid SSM-Transformer models pair linear-scaling long-sequence processing with Transformer contextual reasoning; examples include Nemotron-H, Jamba, and Zamba.The cited architectures are associated with reduced KV-cache requirements and increased throughput compared with pure Transformers.Source: Background · p. 3
  • Structured pruning removes blocks of weights, and typical pruning proceeds through importance estimation, model trimming, and accuracy recovery.Knowledge distillation is described as an effective alternative to traditional fine-tuning for accuracy recovery.Source: Background · p. 3

Methodology

The method scores and prunes components across hybrid models while preserving Mamba’s group structure, then uses architecture search and teacher distillation to recover accuracy. It also examines layer importance and the effects of depth pruning.

Methodology

  • Importance estimation uses forward-pass activation scores for Mamba heads and channels, FFN neurons, embedding channels, and layers; low-scoring components are pruned before teacher distillation.The approach follows importance estimation, trimming, and distillation as its pruning workflow.Source: Methodology · p. 3

Mamba Pruning

  • Mamba head rankings preserve group structure, while head channels are pruned consistently across all heads using shared channel rankings.The scoring procedure ranks channels first, then scores heads and sorts them within each group.Source: Mamba Pruning · pp. 4–6

FFN and Embedding Pruning

  • FFN neurons and embedding channels are scored from activations, ranked by importance, and pruned below the target compression ratio.Embedding-channel scores aggregate evidence across layers using the shared embedding channels.Source: FFN and Embedding Pruning · p. 6

FLAP Importance for Hybrid Models

  • The SSM extension of FLAP scores output-projection columns using their weight norms and the variance of their input activations across calibration samples.The resulting scores rank heads within groups and guide trimming of corresponding SSM components.Source: FLAP Importance for Hybrid Models · p. 6

Depth Pruning

  • Pruning 4, 8, 12, 16, or 26 layers left core-knowledge benchmarks largely unaffected but significantly degraded math and coding performance.Layer importance was estimated using Kullback-Leibler divergence between logits from the full model and a model with a layer removed.Source: Depth Pruning · p. 7

Architecture Search

  • The 4B-parameter search combines removal of 4–26 layers with pruning of embedding channels, FFN dimensions, Mamba heads, and head channels.The search generated over a hundred candidate architectures meeting the parameter constraints.Source: Architecture Search · p. 7

Accuracy Recovery with Knowledge Distillation (KD)

  • Accuracy recovery uses logit-based distillation from the original model, with forward KL divergence loss during continued training.The authors adopt distillation in place of conventional fine-tuning for recovery after pruning.Source: Accuracy Recovery with Knowledge… · p. 7

Experiments and Results

The compression experiments identify pruning trade-offs and select a width-only 4B model; evaluation reports strong accuracy and inference results, with a further compression test on Mamba2.

Experiments and Results

  • At 50% compression (8B to 4B), depth-only pruning (#24) provides the highest speedups; math and coding benchmarks show significant accuracy drops below 44 layers.A 36-layer depth-pruned model (#25), despite having ∼1.4× more parameters, performs worse than the least accurate width-only 4B candidate (#23, with 64 Mamba heads).Source: Experiments and Results · p. 8
  • W_x activations often yield the best LM loss in Mamba scoring ablations; after KD, FLAP performs on par with L2 for candidate #1 without a clear advantage.Table 2 compares zero-shot LM loss for scores based on W_x, W_z, and W_O activations.Source: Experiments and Results · pp. 8–9, Table 2 · p. 8
  • At a fixed 4B budget, pruning Mamba components is associated with faster inference and slightly worse LM loss, while pruning FFN and embedding dimensions improves accuracy but slows models.Figure 6 analyzes 125 variants with fixed depth and shows negative correlations among model parameters under the fixed parameter budget.Source: Experiments and Results · pp. 8–9, Figure 6 · p. 10
  • Pruning Mamba heads (m_h) rather than head channels (m_d) yields lower LM loss, reduced latency, and higher throughput when each axis is pruned in isolation.Head pruning also yields better speed improvements than channel pruning within a Mamba layer.Source: Experiments and Results · p. 9, Figure 7 · p. 10

Obtaining the Best Compressed Hybrid Model

  • The final model uses width-only pruning to prioritize accuracy and avoid reducing depth; candidate #1 was chosen for extended KD because its throughput was higher despite matching candidate #2’s LM validation loss.Nemotron-H 8B has 52 layers, fewer than the 64 alternating Attention and FFN layers in Phi-4-4B.Source: Obtaining the Best Compressed Hybrid… · p. 9

Data and Training Hyperparameters; Alignment and Long Context Extension

Evaluation Summary

  • Across five task groups, 4B achieves state-of-the-art accuracy after up to ∼40x fewer tokens; versus Phi-4-4B, throughput is ∼2.2x higher and latency ∼1.8x lower.The 4B model retains over 96% of the original 8B model’s accuracy, and its aligned version leads in math, coding, instruction following, and tool-use tasks.Source: Evaluation Summary · p. 10, Table 4 · p. 11, Table 5 · p. 11

Generalizability to Mamba2

  • With 10.5B training tokens, the compressed 780M model outperforms the 780M model trained from scratch on 300B tokens and achieves an average score comparable to the original 1.3B model.The compression applies SSM and embedding pruning to Mamba2 1.3B.Source: Generalizability to Mamba2 · p. 10

Conclusions

Nemotron-H 4B compresses the 8B hybrid model through group-aware Mamba pruning, other structured pruning, and knowledge distillation. It reports higher accuracy and faster inference, while the compression recipe is open-sourced.

Conclusions

  • Group-aware Mamba pruning combined with structured FFN and embedding-dimension pruning and knowledge distillation compresses Nemotron-H 8B to 4B parameters.The approach uses up to 40× fewer training tokens and preserves over 96% of the original model’s accuracy.Source: Conclusions · p. 11
  • Nemotron-H 4B achieves ∼2× faster inference and 2.6% higher accuracy across a diverse set of tasks.The compressed model retains over 96% of the original 8B model’s accuracy while reducing model size by 50%.Source: Conclusions · p. 11
  • Despite training on 10.5B rather than 300B tokens, the compressed Mamba2 780M model achieves a better average score than the 780M baseline.Table 8 compares the compressed model with Mamba2 780M and 1.3B models trained from scratch.Source: Table 8 · p. 12
  • The authors open-source the compression recipe as a practical blueprint for efficient hybrid model development.Source: Conclusions · p. 12

Related papers

Loading 2504.11409v2…