🤖 machine learning

SHAPE: Coalition-Aware Expert Pruning for Sparse Mixture-of-Experts LLMs

SHAPE is a task-driven expert pruning framework for Sparse Mixture-of-Experts LLMs that models intra-layer expert cooperation as a cooperative game to identify and retain essential collaborative subsets via Shapley-style attribution, thereby reducing memory footprint while maintaining accuracy without additional training.

Original authors: Yuhao Zhang

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yuhao Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a massive, high-powered kitchen (a Large Language Model) designed to cook up complex answers. Instead of having one giant chef do everything, this kitchen uses a Mixture-of-Experts (MoE) system. Think of it as a team of 100 specialized sous-chefs (the "experts").

When a customer orders a dish (a prompt), a head waiter (the "router") looks at the order and picks the top 3 chefs to work on it together. The magic is that the kitchen doesn't need all 100 chefs to be active at once; it only uses the specific 3 needed for that moment. This makes the kitchen very fast and efficient.

The Problem: The "Memory Wall"
Even though only 3 chefs are cooking at any given second, the kitchen owner has to keep all 100 chefs standing by in the room, fully dressed and ready to go. Why? Because the waiter might need any of them for the next order.

The Issue: If you have a 100-chef kitchen, you need a huge, expensive room (GPU memory) to fit everyone. This makes it impossible to run these powerful kitchens in smaller, cheaper spaces (like a single laptop or a small server).
The Old Fix: Previous attempts to shrink the kitchen were like a manager saying, "Let's fire the chefs who haven't cooked much this week." This is flawed because a chef might not cook often, but when they do cook with a specific partner, they create a masterpiece. Firing them breaks that special teamwork.

The Solution: SHAPE (The Teamwork Detective)
The paper introduces SHAPE, a new way to shrink the kitchen without ruining the food. Instead of just counting how often a chef is called, SHAPE looks at how well they work together.

Here is how SHAPE works, using a simple analogy:

The "Game" Observation: SHAPE watches the kitchen for a short while (using a small "calibration set" of orders). It doesn't just watch who gets called; it watches which groups of 3 chefs get called together and how often those groups succeed.
The "Shapley" Score (Fairness): Imagine a group of 3 chefs creates a perfect dish. How much credit does each chef get?
- If Chef A is great alone but Chef B is useless without Chef C, a simple count might miss this.
- SHAPE uses a mathematical concept called Shapley Value (think of it as a "fairness calculator"). It asks: "If we remove Chef A from this specific trio, does the dish fall apart?"
- If the dish falls apart, Chef A is essential to that team, even if they aren't the most famous chef.
- If the dish tastes the same without Chef A, then Chef A is redundant and can be let go.
The "Quality Coverage" Rule: SHAPE doesn't just fire the bottom 20% of chefs based on a single score. It ensures that every "floor" of the kitchen (every layer of the model) keeps enough chefs to cover the most important teamwork combinations. It uses a "bisection" method (a smart guessing game) to find the perfect number of chefs to keep so the kitchen stays small but still cooks 99% of the same dishes perfectly.

The Results
The researchers tested this on three different modern "kitchens" (Qwen, GPT-OSS, and DeepSeek).

The Outcome: They were able to fire 20% to 40% of the chefs (reducing the memory needed significantly) without the food tasting any worse.
Why it worked: By keeping the chefs who are essential for teamwork rather than just the ones who are busy, the kitchen remained stable.
The Bonus: Because they fired so many chefs, the room (GPU memory) needed to run the kitchen became much smaller, making it possible to run these powerful models on less expensive hardware.

In a Nutshell
SHAPE is like a smart manager who realizes that firing the "quiet" chef who is the secret glue of a winning team is a bad idea. Instead, it fires the "loud" chefs who are easily replaced. This shrinks the team size and the office space needed, while keeping the quality of the work exactly the same.

Technical Summary: SHAPE – Coalition-Aware Expert Pruning for Sparse Mixture-of-Experts LLMs

1. Problem Statement

Sparse Mixture-of-Experts (SMoE) architectures have become pivotal in scaling Large Language Models (LLMs) to hundreds of billions or trillions of parameters while maintaining sub-linear training costs and low per-token inference latency. However, the deployment of these massive models faces a critical "Memory Wall." While conditional computation ensures that only a sparse subset of experts is active per token, the entire expert pool must remain resident in GPU memory to support dynamic, token-dependent routing decisions. This requirement creates prohibitive memory footprints (e.g., terabytes of VRAM for models like DeepSeek-V3.2), limiting their usability in resource-constrained environments.

Existing compression strategies, such as expert pruning, aim to reduce the parameter footprint by removing redundant experts. However, current pruning methodologies suffer from a fundamental modeling gap: they predominantly adopt an expert-wise perspective, scoring experts independently based on metrics like activation frequency or gating magnitude. This approach overlooks the inherently coalitional nature of MoE inference, where the model output is generated by a non-linear composition of a specific top- $k$ coalition of experts. Consequently, independent scoring risks dismantling high-utility expert collaborations, leading to disproportionate performance degradation on complex tasks, even if individual experts appear redundant in isolation.

2. Methodology: SHAPE Framework

The authors propose SHAPE (Shapley-guided, Coalition-Aware Expert Pruning), a task-driven framework that explicitly models intra-layer expert cooperation. SHAPE operates without retraining, relying on a small calibration set to guide pruning decisions. The methodology consists of three core components:

A. Task-Driven Calibration and Routing Traces

SHAPE performs an offline pass on a small, task-specific calibration dataset ( $D_{cal}$ ) using the unpruned model. It records the routing decisions (top- $k$ expert selections) for every token across all MoE layers. This yields a set of observed expert combinations (coalitions) and their activation frequencies, which serve as empirical statistics characterizing task-conditioned expert cooperation.

B. Coalition-Based Utility and Shapley-Style Valuation

Instead of treating experts as independent units, SHAPE formulates the inference process as an empirical cooperative game.

Utility Definition: The utility of a coalition (a specific set of $k$ experts) is defined by its observed activation frequency on the calibration set.
Marginal Contribution: To determine an expert's value, SHAPE estimates its marginal contribution to the coalition's utility. Since removing an expert from a top- $k$ coalition results in a $(k-1)$ -subset that is not directly observed under top- $k$ routing, SHAPE approximates the utility of these sub-coalitions using co-occurrence statistics.
Shapley-Style Approximation: The framework calculates a Shapley-style score ( $\phi$ ) for each expert. This score averages the expert's marginal contribution across all coalitions it participates in, weighted by the coalition size.
Monotonicity Correction: To ensure the approximation respects the monotonicity requirement of cooperative games (i.e., a coalition's utility should not be lower than its sub-coalitions), SHAPE introduces a conservative scaling factor ( $\alpha$ ) derived from observed frequencies. This prevents overestimating expert substitutability.

C. Shapley Quality Coverage Selection

To determine which experts to retain, SHAPE employs a quality-coverage selection rule rather than a simple count-based threshold.

Layer-wise Coverage: For each layer, the framework sorts experts by their non-negative Shapley scores. It selects the minimal subset of experts such that their cumulative Shapley mass covers a fraction $\alpha$ of the total layer contribution.
Global Budget Matching: A binary search (bisection) is used to find the optimal coverage threshold $\alpha$ that satisfies a global target keep-rate (e.g., retaining 60% or 80% of experts).
Structural Preservation: This approach ensures that layers with highly concentrated value retain fewer experts, while layers with diffuse contributions retain more, preventing "layer collapse" where pruning disproportionately affects specific layers.

3. Key Contributions

Coalitional Expert Valuation: The paper reframes MoE pruning from independent importance estimation to coalition-aware contribution modeling. It introduces a principled mechanism using Shapley values to attribute expert contributions conditioned on routed coalitions, distinguishing between frequently selected experts and those essential for high-utility collaborations.
Structure-Preserving Pruning Strategy: SHAPE proposes a topology-aware selection mechanism that preserves layer-wise cooperative capacity. By enforcing a quality-coverage constraint based on cumulative Shapley value, it avoids the structural imbalances caused by global ranking or fixed per-layer counts.
Empirical Validation: Extensive experiments demonstrate that preserving expert cooperation is critical for effective compression. The framework achieves significant pruning (up to 40%) with negligible performance degradation on reasoning-intensive benchmarks, outperforming both heuristic baselines and state-of-the-art training-free pruning methods.

4. Experimental Results

The authors evaluated SHAPE on three modern MoE backbones: Qwen3-30B-A3B, GPT-OSS-20B, and DeepSeek-V2-Lite. The evaluation covered diverse benchmarks including math reasoning (GSM8K, MATH-500), code generation (HumanEval), knowledge QA (GPQA-Diamond), and medical QA.

Performance Retention: Under a 20% pruning ratio, SHAPE maintained performance comparable to, and in some cases slightly exceeding, unpruned baselines. At a 40% pruning ratio, SHAPE consistently outperformed baseline methods (including random pruning, frequency-based, gating-based, RAEP, and EASY-EP), particularly on reasoning-intensive tasks where expert cooperation is critical.
Robustness: While reasoning tasks (e.g., GSM8K, GPQA-Diamond) showed higher sensitivity to pruning, SHAPE's coalition-aware approach preserved accuracy significantly better than expert-wise heuristics. Larger models (Qwen3, GPT-OSS) demonstrated greater resilience, but SHAPE remained effective even on the more compact DeepSeek-V2-Lite.
Memory Efficiency: Pruning directly reduced the peak GPU memory footprint. The study highlighted that serving the full expert pool incurs substantial memory overhead regardless of activation sparsity. SHAPE enabled clear reductions in VRAM usage (proportional to the pruning ratio) without additional training or architectural changes.
Ablation Studies: Comparisons with global ranking and layer-wise independent pruning confirmed that the quality-coverage constraint is vital. Global pruning led to structural imbalance, while layer-wise pruning failed to account for task-specific capacity distribution. SHAPE's balanced approach yielded the most stable performance across all pruning ratios.

5. Significance and Claims

The paper claims that SHAPE addresses a critical limitation in current MoE compression: the failure to account for the cooperative nature of expert inference. By shifting the focus from individual expert metrics to coalitional utility, SHAPE provides a robust, training-free solution for deploying massive MoE models in memory-constrained environments.

The authors emphasize that their work does not require retraining or architectural modification, making it a practical, deployment-oriented optimization. The results suggest that a significant portion of the expert pool in modern MoE models contributes limited marginal utility when considered in isolation, but the structure of their collaboration is what sustains model performance. SHAPE successfully identifies and preserves these critical collaborative structures, enabling a favorable trade-off between model capacity and resource footprint. The open-source code is provided to facilitate further research and application in this domain.