Research map / Architectures and efficiency / Pruning and efficient tuning

Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

NeurIPS 2025 · Architectures and efficiency › Pruning and efficient tuning · approach: Layer-wise pruning of LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Compresses the 8B Nemotron-H hybrid (Attention + Mamba2) to a 4B model via group-aware Mamba head/channel pruning, width/depth search, and logit-based knowledge distillation, plus alignment and long-context extension.

Problem

Pruning and distillation have been studied for Transformer models, but compressing hybrid Attention+SSM (Mamba) language models remains underexplored, especially how to prune Mamba layers while preserving SSM group structure.

Approach

A group-aware pruning method for Mamba2 layers that ranks heads only within their B/C groups (preserving broadcast semantics) and shares channel rankings across heads, combined with activation-based FFN/embedding/layer pruning and forward-KL distillation into a hybrid pruning recipe selected by a lightweight-KD architecture search.

Key results

• The compressed Nemotron-H 4B base model averages 64.5 across 16 base benchmarks versus 66.7 for the 8B parent, and outperforms similarly sized community base models (e.g., Qwen-2.5-3B at 60.3) while using 0.38T training tokens versus e.g. 18T for Qwen-2.5.
• The 4B model retains over 96% of the 8B model's accuracy and gives about 1.4x higher throughput than the 8B parent, and about 2.2x higher throughput than Phi-4-Mini-4B.
• Pruned Mamba2 780M from the 1.3B model, trained on 10.5B tokens, achieves average 44.7 versus 42.5 for Mamba2 780M trained from scratch on 300B tokens.

Measured effects reported in the paper

Surprising findings

• Pruning Mamba heads is more favorable than pruning head channels, and the hybrid model tolerates reducing Mamba heads even though attention heads are typically less prunable.
• Depth-only pruning of the hybrid model gives the largest speedups but a 36-layer 4B-plus model underperforms a width-only 4B candidate, showing depth is critical for accuracy.

Limitations

• Attention layers are not pruned since they make up only about 8% of layers, so the attention side of hybrid compression is unexplored.
• Depth pruning combined with width pruning underperformed width-only pruning, and KD for context extension was left as future work.
• Generalization is shown only on a single Mamba2 780M experiment, with fewer evaluation details.

Experimental setup

Datasets: Nemotron-H Phase 3 pretraining data mixture general domain chat data (SFT-KD) math and coding SFT data long-context SFT data derived from chat data

Benchmarks: ARC Challenge ARC Easy CommonsenseQA GSM8K HellaSwag HumanEval HumanEval+ MBPP MBPP+ MMLU OpenbookQA PIQA

Metrics: accuracy pass@1 LM validation loss throughput time to first token safety score

Baselines: Llama-3.2-3B Falcon-3-3B Zamba-2-2.7B Qwen-2.5-3B Phi-4-Mini-4B Mamba2 780M Mamba2 1.3B

Builds on: Minitron FLAP Nemotron-H Mamba2

Closest papers by meaning

• Compact Language Models via Pruning and Knowledge Distillation (NeurIPS 2024 · similarity 0.92)
• You Only Prune Once: Designing Calibration-Free Model Compression With Policy Learning (ICLR 2025 · similarity 0.90)
• Layer as Puzzle Pieces: Compressing Large Language Models through Layer Concatenation (NeurIPS 2025 · similarity 0.89)
• Pruning Large Language Models with Semi-Structural Adaptive Sparse Training (AAAI 2025 · similarity 0.89)
• Týr-the-Pruner: Structural Pruning LLMs via Global Sparsity Distribution Optimization (NeurIPS 2025 · similarity 0.89)
• EvoPress: Accurate Dynamic Model Compression via Evolutionary Search (ICML 2025 · similarity 0.89)