LAEP: Adaptive Pruning for MoE Models

• The paper introduces LAEP as a non-uniform, data-driven pruning strategy that adaptively determines which experts to remove at each layer to minimize accuracy loss.
• LAEP leverages methods like genetic search, continuous relaxation, and trajectory-based selection to optimize expert retention using metrics such as calibration loss and routing statistics.
• Its adaptive, layer-specific approach significantly enhances parameter efficiency, pre-training throughput, and inference speed across diverse models including language, vision, and generative architectures.

Layer-Adaptive Expert Pruning (LAEP) encompasses a family of methodologies for non-uniform, data-driven sparsification of Mixture-of-Experts (MoE) architectures through selective retention or pruning of experts at each layer. Unlike uniform-pruning or global-sparsity approaches, LAEP dynamically determines, at each layer, which experts are dispensable with minimal downstream impact, enabling superior parameter efficiency, compute throughput, and—in some modes—even enhanced pre-training efficiency. LAEP is instantiated in both post-hoc and in-training settings across diverse regimes, including LLMs, vision models, and diffusion generative architectures.

1. Conceptual Foundations and Motivation

Mixture-of-Experts models comprise $L$ MoE layers, each with $M_\ell$ experts. For input $x$, a learned router activates a subset (typically top-$k$) of experts for each token. While MoEs provide compute-adaptive capacity and have driven LLMs such as Mixtral-8×7B, Qwen1.5-MoE-A2.7B, and DeepSeek-V2-Lite to state-of-the-art results, their overparameterization at the full model scale presents deployment and training bottlenecks (Yang et al., 2024, ai et al., 20 Jan 2026).

Uniform pruning—eliminating a fixed fraction of experts in every layer—proves suboptimal because early and deep MoE layers exhibit starkly differing redundancy and functional criticality (Yang et al., 2024, Bai et al., 19 Sep 2025, Yang et al., 20 Dec 2025). Shallow layers often contain generic, highly redundant experts, whereas deeper layers may host specialists whose removal disproportionately impacts task accuracy. LAEP addresses this heterogeneity by introducing layer-specific pruning ratios and expert-selection mechanisms, sometimes extending to adaptive determination of the number of active experts per MoE layer (Chitty-Venkata et al., 2 Sep 2025).

2. Algorithmic Frameworks for Layer-Adaptive Pruning

Multiple precise LAEP regimes can be delineated:

Genetic/Blockwise Approaches

MoE-I$^2$ implements a staged LAEP pipeline comprising:

• Per-layer and per-expert importance evaluation: For each expert $e_{i,j}$ in layer $i$, an importance metric $I_{i,j}$ quantifies the calibration loss incurred by removal (cross-entropy or output Frobenius norm perturbation).
• Non-uniform pruning ratio allocation: Higher overall layer importance $I_i$ results in lower allowable prune budgets $P_i$ for layer $i$, tuning global sparsity adaptively.
• Expert selection via layer-level genetic search: Boolean pruning masks are evolved in a GA driven by minimizing output reconstruction error over batches; top $K_i$ candidates are retained.
• Blockwise KT-Receptive Field search: Within windowed blocks of $T$ consecutive layers, brute-force selection of mask combinations across layers jointly optimizes local and cross-layer dependencies, avoiding combinatorial explosion.

This layered, blockwise hybrid paradigm achieves robust parameter reduction while minimizing global loss (Yang et al., 2024).

Differentiable/Continuous Relaxation

DiEP recasts the combinatorial pruning problem into continuous optimization by:

• Defining intra-layer logits $\alpha^{(\ell)}_i$ and inter-layer scalars $\beta^{(\ell)}$, normalized via softmax and gating weights, parameterizing continuous masks.
• Minimizing a combined task/reconstruction loss:

$$\min_{\alpha, \beta} \operatorname{CE}_{\text{task}} + \lambda \|\mathcal{F}'(x;\alpha, \beta) - \mathcal{F}(x)\|_F$$

• Alternating block-coordinate gradient steps for $\alpha$ and $\beta$, jointly enforcing a global expert-sparsity constraint at the final discretization.
• Final pruning by globally sorting all $NL$ expert scores $s^{(l)}_i = \alpha^{(l)}_i \beta^{(l)}$ and removing the quantile corresponding to the overall sparsity.

This single-stage approach discovers optimal, non-uniform sparsity profiles, automatically reflecting redundancy patterns across depth (Bai et al., 19 Sep 2025).

Global Trajectory-Based Selection

MoE Pathfinder casts LAEP as a computation-graph path selection problem:

• The MoE is a DAG where nodes are experts and inter-layer edges represent possible information flow.
• Each expert is scored using a blend of reconstruction-based error, routing probability, and activation strength.
• The pruning objective is to retain those experts traversed by the top $m$ highest-probability paths, as discovered via dynamic programming over calibration data.
• This trajectory-driven method natively induces non-uniform retention: layers traversed by many distinct prominent paths retain more experts; redundant layers retain fewer.

Empirically, this yields highly selective sparsity patterns which mirror cross-layer task salience (Yang et al., 20 Dec 2025).

3. LAEP for Pre-Training Efficiency and Token Routing

Traditional pruning focuses on post-training compression, but recent advances in LAEP enable runtime gains during pre-training as well:

• At pre-training time, LAEP tracks per-expert token loads $T_{i,l}$ and uses explicit constraints on individual and cumulative expert utilization:
  • Prune expert $i$ at layer $l$ if $T_{i,l} \leq \frac{\alpha}{N} T_{\text{all},l}$ and the cumulative least-loaded set remains under a threshold $\beta T_{\text{all},l}$.
• Once token-routing statistics are stable, LAEP prunes the least-utilized experts per-layer and then dynamically reassigns the remaining experts across devices, using a greedy load-balancing heuristic to maintain per-GPU token balance (ai et al., 20 Jan 2026).

These mechanisms lead to significant improvements in training throughput and parameter efficiency.

Model	Params Reduction	Throughput Gain	Test Loss Δ	Notes
1010B (pretrain)	33.3%	+48.3%	–	LAEP vs. unpruned (ai et al., 20 Jan 2026)
20B (various $\alpha$)	Up to 26.5%	Up to +8.3%	<0.02	Variable per-layer pruning

4. Extensions: Inference-Time Allocation and Vision/Generative Models

Layer-adaptive expert selection is not limited to pruning/removal. LAEP also subsumes inference-time optimization of active expert counts per layer:

• LExI/LAEP for inference employs a data-free, weight-only Monte Carlo sensitivity analysis: for candidate top-$k$ values in each layer, output perturbation is measured by the Frobenius distance from the baseline (Chitty-Venkata et al., 2 Sep 2025).
• An evolutionary search optimizes the allocation $\{k_j\}$ to maximize throughput under per-layer and global constraints, subject to a fixed accuracy/resource budget.
• No retraining or calibration data is needed, and memory footprint is unchanged.

For diffusion generation architectures, ALTER generalizes LAEP by coupling per-layer pruning masks (produced via a hypernetwork) with timestep-conditioned expert routing, simultaneously optimizing layer sparsity, expert assignment, and model parameters in a joint end-to-end framework (Yang et al., 27 May 2025).

5. Empirical Performance and Trade-offs

Across a spectrum of LAEP regimes, experiments consistently show that non-uniform, layer-adaptive pruning offers superior trade-off frontiers in terms of accuracy, perplexity, and efficiency compared to uniform or local expert trimming:

• Post-hoc-sparsification: 25–50% parameter reduction with accuracy degradation under 3 points for major language MoEs, e.g., Mixtral-8×7B and Qwen1.5-MoE-A2.7B (zero-shot accuracy drops from 67.53 to 64.55 at 51.8% compression (Yang et al., 2024), 92% of full performance at 50% sparsity for DiEP (Bai et al., 19 Sep 2025), 84% of original MMLU accuracy at 50% sparsity for Pathfinder (Yang et al., 20 Dec 2025)).
• Pre-training acceleration: Up to 48.3% training throughput gain and 33.3% parameter reduction with competitive downstream generalization (ai et al., 20 Jan 2026).
• Inference-time optimization: +5–10% throughput with indistinguishable perplexity, strict dominance over fixed expert count or intra-expert pruning, especially when using optimized GPU backends such as vLLM (Chitty-Venkata et al., 2 Sep 2025).
• Diffusion generation: 25.9% of full model MACs for equal visual fidelity (Alter (Yang et al., 27 May 2025)).

The combination of per-layer importance estimation, adaptive thresholding, and global or trajectory-level coordination underpins these empirical successes.

6. Practical Considerations and Future Directions

Implementation of LAEP incurs modest overheads, typically requiring one-time calibration runs (sometimes amortized by synthetic data), light-weight optimization (genetic or evolutionary search), or, in differentiable approaches, blocked gradient updates over mask parameters. Communication and re-sharding for expert rearrangement, while reported as negligible in published scale, may introduce complexities in extremely large GPU clusters (ai et al., 20 Jan 2026).

Identified limitations include the necessity for hyperparameter tuning for pruning thresholds ($\alpha$, $\beta$), lack of dynamic expert reactivation (in one-shot LAEP), and, in some designs, incompatibility of aggressive pruning with auxiliary load balancing strategies.

Avenues for extension include:

• Integrating adaptive LAEP schedules during training.
• Expanding LAEP to multi-modal or multilingual settings where token distributions are non-uniform across domains.
• Merging LAEP with dynamic routing or online expert splitting.
• Extending LAEP methods for use in fine-tuning and task-specific adaptation scenarios.

7. Representative Algorithms and Their Variants

Method	Pruning Signal	Optimization	Application Stage	Reference
MoE-I$^2$ LAEP	Calibration loss, importance per expert/layer	Layerwise GA, Blockwise joint	Post-training compression	(Yang et al., 2024)
DiEP	Gradient-learned intra/inter-layer scores	Continuous relaxation	Post-training/fine-tuning	(Bai et al., 19 Sep 2025)
Pathfinder	Trajectory-level importance paths	Dynamic programming	Task-specific, post-hoc	(Yang et al., 20 Dec 2025)
LExI	Weight-only, MC perturb. sensitivity	Evolutionary allocation	Inference only	(Chitty-Venkata et al., 2 Sep 2025)
ALTER	Joint layer-mask hypernet + routing	End-to-end co-optimization	Diffusion, generative	(Yang et al., 27 May 2025)
LAEP (pretrain)	Token routing statistics	Prune-then-rearrange	During pre-training	(ai et al., 20 Jan 2026)

This typology illustrates the breadth of the LAEP framework, encompassing both discrete/pruning and continuous/activation strategies, with optimization ranging from combinatorial search to gradient-based relaxation and dynamic-programming trajectories.

---

For further mathematical details, pseudocode, and hyperparameter specifics, see the cited papers: MoE-I$^2$ (Yang et al., 2024), DiEP (Bai et al., 19 Sep 2025), MoE Pathfinder (Yang et al., 20 Dec 2025), LExI/LAEP (Chitty-Venkata et al., 2 Sep 2025), ALTER (Yang et al., 27 May 2025), and LAEP for pre-training (ai et al., 20 Jan 2026).