DarwinLM: Evolutionary Structured Pruning of Large Language Models

Technical Summary: DarwinLM

Problem Statement
Large Language Models (LLMs) have achieved significant success but face massive computational costs that hinder deployment, particularly in real-time applications. While structured pruning offers a solution by compressing models to provide end-to-end speed improvements on mainstream hardware, existing methods often fail to account for two critical factors:

1. Non-uniform Sensitivity: Different model components (layers, attention heads, MLP dimensions) exhibit varying sensitivities to pruning. Uniform compression often leads to performance collapse.
2. Post-Compression Training: A pruning method must not only identify a capable substructure but also ensure that the structure is amenable to recovery via post-compression fine-tuning. Prior methods, such as ZipLM, often optimize for local layer-wise errors or one-shot performance, which does not correlate well with downstream task performance or fine-tuning recovery.

Methodology
The authors propose DarwinLM, a training-aware structured pruning method that utilizes an evolutionary search process to find optimal non-uniform sparsity allocations. The pipeline consists of two main stages:

1. Evolutionary Search with Training-Aware Selection:

   • Initialization: The search begins with a "parent" model generated via one-shot structured pruning using second-order information (Optimal Brain Surgeon formulation).
   • Mutation: In each generation, the algorithm generates "offspring" models by copying the parent and applying a "level switch" mutation. This involves shifting sparsity levels between sub-blocks (e.g., increasing sparsity in one layer while decreasing it in another) to maintain the global target sparsity or speedup constraint.
   • Training-Aware Selection: A core innovation is the integration of lightweight fine-tuning into the selection process. Instead of selecting offspring based solely on one-shot metrics (like KL-divergence on a small calibration set), the method employs a multi-step selection process. Offspring are trained on progressively larger token budgets (e.g., 10K $\to$ 50K $\to$ 200K tokens). Poorly performing models are eliminated at each stage. This ensures the selected "fittest" model is not just accurate in a static state but is also robust and recoverable after continued training.
   • Termination: Once the search converges, the final candidate undergoes a more extensive fine-tuning stage (e.g., 10B tokens) before final evaluation.

2. Pruned Layer Database:

   • To facilitate efficient search, the method pre-computes a database of sparse sub-blocks (MLP and Attention modules) at various discretized sparsity levels.
   • For Attention modules, pruning is performed per-head. For MLPs, dimensions are pruned in multiples of 32.
   • For models with Group-Query Attention (GQA), the K and V matrices are not pruned directly; instead, corresponding heads are removed during the forward pass to maintain computational compatibility.

3. Extension to MoE Architectures:

   • The method is extended to Mixture-of-Experts (MoE) models (e.g., Qwen3-30B-A3B). In this setting, the search optimizes sparsity within the expert MLPs while maintaining uniform sparsity across MoE blocks. Attention modules are generally left unpruned as experts contain the majority of parameters.

Key Contributions

• Training-Aware Evolutionary Search: DarwinLM introduces a novel evolutionary framework that explicitly accounts for the model's ability to recover performance during fine-tuning, rather than optimizing solely for one-shot accuracy.
• Non-Uniform Structured Pruning: The method enables fine-grained, non-uniform compression of LLMs, leveraging the varying sensitivity of different layers to achieve higher compression ratios without accuracy loss.
• MoE Adaptation: This work represents the first exploration of non-uniform structured pruning specifically for MoE architectures, demonstrating that structured sparsity is effective even in complex expert-based models.
• Efficiency: The search process is highly efficient, completing in approximately 8 hours on 4 consumer-grade GPUs, with the final fine-tuning stage taking roughly half a day on a standard cluster.

Experimental Results
The authors validate DarwinLM on Llama-2-7B, Llama-3.1-8B, Qwen-2.5-14B-Instruct, and Qwen3-30B-A3B (MoE).

• Performance: DarwinLM achieves state-of-the-art performance in one-shot structured pruning. For example, when pruning Llama-3.1-8B to half its size, it achieves a 5.9% higher average zero-shot accuracy compared to the best prior method (ZipLM).
• Data Efficiency: The method significantly reduces the training budget required for recovery. DarwinLM surpasses ShearedLlama (which uses 50B tokens for fine-tuning) while requiring only 10B tokens. In the same 10B token setting, DarwinLM outperforms ShearedLlama.
• MoE Results: On the Qwen3-30B-A3B model, DarwinLM produces a 16B-A2B variant that retains $\ge$90% of the original model's accuracy after 10B token fine-tuning.
• Hardware Efficiency: The pruned models demonstrate direct speedups and memory reductions. For instance, the 2.7B DarwinLM model achieves 1.98$\times$ throughput and uses 2.43$\times$ less memory compared to the dense 7B baseline on L40s GPUs.
• Comparison: The method outperforms uniform pruning baselines and other structured pruning methods (including ShortGPT, Shortened-Llama, and EvoPress) across various sparsity levels, particularly at high compression rates where other methods degrade rapidly.

Significance
The paper claims that DarwinLM represents a significant advancement in LLM compression by bridging the gap between structural search and post-training recovery. By demonstrating that non-uniform structured pruning can be effectively applied to both dense and MoE architectures with minimal training data, the work challenges the notion that high-performance compression requires massive retraining budgets. The authors position DarwinLM as a practical, hardware-agnostic solution that enables the deployment of efficient LLMs without sacrificing accuracy.