💻 computer science

LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

The paper introduces LEAP, a learnable end-to-end adaptive pruning method that utilizes a per-weight Bernoulli-via-Gumbel-sigmoid relaxation to enable tractable unstructured sparsity learning in large language models, significantly outperforming existing layer-wise baselines in zero-shot accuracy at high sparsity levels.

Original authors: Mohammad Mozaffari, Younes Hourri, Mohammad Rastegari, Mahyar Najibi

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mohammad Mozaffari, Younes Hourri, Mohammad Rastegari, Mahyar Najibi

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 you have a massive, incredibly smart library (a Large Language Model) that knows almost everything. But this library is so huge that it takes up an entire warehouse, requires a giant power plant to run, and is too slow for anyone to actually use on a regular computer.

To fix this, you want to throw away 50% to 60% of the books (the "weights" or connections in the model) to make it smaller and faster. The goal is to keep the library just as smart, even with half the books.

This is where the paper LEAP comes in. Here is the story of how they solved the problem, explained simply:

The Problem: The "Too Many Choices" Trap

For a long time, scientists tried to cut out books using two main strategies:

The "Layer-by-Layer" Approach (Old Way): Imagine a librarian looking at one shelf at a time and deciding, "This book looks useless, throw it out." They do this for every shelf without talking to the other shelves.
- The Flaw: Sometimes a book looks useless on its own shelf but is actually crucial for a story that spans the whole library. By looking at shelves in isolation, this method accidentally throws away important connections, making the library less smart.
The "Pattern" Approach (The Newer, Broken Way): Recently, some researchers tried to look at the whole library at once. They said, "Let's learn exactly which books to keep." But they tried to do this by grouping books into small sets (like 4 books at a time) and asking, "Which pattern of these 4 should we keep?"
- The Flaw: This works great for small groups. But if you try to do this for a whole row of 4,000 books (which is what modern AI models have), the number of possible "patterns" becomes a number so huge it's impossible to write down. It's like trying to list every possible combination of lottery numbers for a billion tickets. The computer chokes; the math breaks.

The Solution: LEAP (The "Individual Vote" System)

The authors of this paper, LEAP, realized they needed a new way to vote on which books to keep without getting overwhelmed by the math.

Instead of asking, "Which group pattern do we keep?" (which is impossible for huge groups), they asked a much simpler question for every single book: "Should this specific book stay or go?"

The Analogy: Imagine a massive election where every single voter (every weight in the model) gets a tiny ballot. Instead of voting for a complex "team strategy," they just vote "Yes" (keep) or "No" (discard).
The Magic Trick: To make this work, they used a mathematical trick (called "Gumbel-sigmoid") that lets the computer "guess" the votes smoothly at first, and then slowly make those guesses harder and sharper until every vote is a clear "Yes" or "No."

How They Did It (The Recipe)

They didn't start from scratch. They used a smart starting point:

The Warm Start: They first used a quick, simple method (called Wanda) to get a rough idea of which books were likely useless. They used this as a "head start" so the computer didn't have to guess blindly.
The Training: They let the computer "learn" the best set of votes by reading a small amount of text (a calibration stream) and adjusting the "Yes/No" votes to keep the library's intelligence high.
Freezing the Books: Crucially, they didn't change the words inside the books (the model weights). They only changed the decision of which books to keep. This keeps the library's original "personality" and knowledge intact, just in a smaller package.

The Results: Smarter, Faster, Leaner

They tested this on five different famous AI models (ranging from small to very large) and cut them down by 50% and 60%.

The Scorecard: They compared LEAP against the best existing methods.
The Win: LEAP was consistently better. On average, it improved the model's accuracy by 2.59 points over the previous best method. In some cases, the improvement was as high as 5.40 points.
The Speed: Because they cut the model in a way that fits perfectly with new, fast computer chips (GPUs) designed for this exact job, the resulting model runs much faster without losing its smarts.

Why This Matters

Before this, if you wanted to make a huge AI model smaller and faster, you had to choose between:

Option A: Keep it accurate but huge and slow.
Option B: Make it small and fast, but dumb.

LEAP shows you can have your cake and eat it too. It provides a practical way to shrink these massive AI brains by half, keeping them just as smart, so they can run on regular computers instead of needing a supercomputer.

In short: LEAP is a new, smarter way to edit a giant AI model by letting every single connection vote on whether it should stay, resulting in a smaller, faster, and still very intelligent AI.

Technical Summary: LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

Problem Statement
The deployment of Large Language Models (LLMs) is increasingly bottlenecked by memory and compute constraints. While weight sparsity offers a solution, the landscape is divided into structured, semi-structured (e.g., 2:4), and unstructured regimes. Recent advances in GPU kernels (e.g., FlashLLM, SpInfer) and dataflow hardware have shifted the bottleneck from executing unstructured sparsity to inducing it with minimal accuracy loss.

Current state-of-the-art methods for unstructured pruning, such as Wanda, SparseGPT, Thanos, and ADMM, rely on layer-wise surrogates derived from the Optimal Brain Surgeon (OBS) principle. These methods minimize local reconstruction errors rather than the end-to-end model loss. Consequently, they accumulate local errors that compound in deep networks, leading to significant accuracy degradation, particularly under aggressive sparsity (50–60%).

Conversely, end-to-end learnable mask methods like MaskLLM and PATCH successfully optimize masks directly against language modeling loss but are restricted to semi-structured patterns. Their parameterization relies on a categorical distribution over valid patterns within a group. While tractable for small groups (e.g., $\binom{4}{2}=6$ patterns for 2:4 sparsity), this approach becomes combinatorially intractable for unstructured sparsity. For a row of width $d=4096$ at 50% sparsity, the number of valid masks is $\binom{4096}{2048} \approx 10^{1229}$ , making the storage and indexing of pattern logits impossible.

Methodology: LEAP
The authors introduce LEAP (Learnable End-to-End Adaptive Pruning), a framework designed to make end-to-end unstructured mask learning tractable by replacing the intractable categorical-over-patterns parameterization with a per-weight approach.

Per-Weight Bernoulli Relaxation: Instead of assigning a logit to every valid pattern, LEAP assigns a learnable logit to every individual weight. For a weight matrix $W \in \mathbb{R}^{m \times n}$ , a parallel logit matrix $P \in \mathbb{R}^{m \times n}$ is introduced. The stochastic mask $M$ is generated via the Gumbel-sigmoid relaxation:
$M = \sigma\left(\frac{\alpha P + g}{\tau}\right)$
where $g$ is Gumbel noise, $\sigma$ is the sigmoid function, $\alpha$ is a scale factor, and $\tau$ is a temperature parameter. This formulation scales linearly with the number of weights ($O(mn)$) rather than the number of valid patterns.
Optimization Strategy:
- Frozen Weights: The pretrained weights $W$ are kept frozen. Only the mask logits $P$ are trained. This preserves the calibration of the pretrained model and simplifies deployment.
- Initialization: $P$ is initialized using a one-shot Wanda mask (setting selected entries to $+s$ and others to $-s$ ) to provide a reasonable starting point, avoiding cold-start issues.
- Schedules: A scale schedule ( $\alpha$ ) ramps up to push the sigmoid toward binary values, while a temperature schedule ( $\tau$ ) decays to sharpen the mask.
- Loss Function: The objective combines the language modeling loss on a calibration stream with two regularization terms:
  1. Global Sparsity Regularizer: Enforces the target density globally across all layers, allowing individual layers to adjust their density based on end-to-end importance.
  2. Magnitude-Aware Term: Biases the optimization toward retaining higher-magnitude weights to prevent degenerate minima where critical weights are dropped.
Training Efficiency: The method trains for approximately 2,000 iterations on a small general-text calibration stream (SlimPajama).

Key Contributions

Identification of Combinatorial Obstruction: The paper formally identifies that categorical-over-patterns parameterizations (used by MaskLLM/PATCH) cannot scale to unstructured sparsity due to the exponential growth of valid patterns.
Tractable Reformulation: The authors propose a per-weight Bernoulli-via-Gumbel-sigmoid reformulation that enables end-to-end differentiable learning for unstructured masks at LLM scale.
LEAP Framework: A lightweight pruning framework that operates on frozen weights, requiring only a small calibration dataset and roughly 2,000 training steps.
Empirical Validation: Extensive evaluation across five LLM families (Qwen-2.5, Gemma-3, LLaMA-3.2, LLaMA-3.1) ranging from 0.5B to 8B parameters at 50% and 60% sparsity.

Results
LEAP was evaluated against layer-wise baselines (Wanda, SparseGPT, Thanos, ADMM) and the semi-structured learnable baseline MaskLLM.

Accuracy Gains: Across ten (model, sparsity) settings, LEAP improved the six-task average zero-shot accuracy by +2.59 points over ADMM, the best layer-wise baseline.
Specific Improvements: The largest gain was +5.40 points on LLaMA-3.2 1B at 60% sparsity (50.39 vs. 44.99).
Perplexity: LEAP consistently achieved lower perplexity (PPL) than baselines. For example, on LLaMA-3.1 8B at 50% sparsity, LEAP achieved 7.66 PPL compared to ADMM's 9.12.
Sparsity Allocation: Experiments showed that under end-to-end optimization with a global budget, LEAP converges to a near-uniform per-block density allocation, contrasting with the non-uniform variations seen in semi-structured methods.
Compute Cost: While more expensive than one-shot methods (taking hours on H100s vs. minutes), LEAP is significantly cheaper than MaskLLM (e.g., 276 GPU-hours for LLaMA-3.1 8B vs. thousands of A100 hours for MaskLLM on comparable scales).

Significance and Claims
The paper claims that LEAP resolves the fundamental obstruction preventing end-to-end mask learning in unstructured sparsity. By decoupling mask learning from weight updates, it offers a practical path to compressing foundation models without sacrificing the calibration of pretrained weights.

The authors emphasize that LEAP produces masks in the 50%–60% unstructured regime, which are directly compatible with existing accelerated kernels (e.g., FlashLLM, SpInfer, MACKO). The significance lies not in creating new inference kernels, but in enabling the deployment of these kernels on high-quality, pruned models that were previously unattainable via layer-wise methods. The work suggests that unstructured sparsity, when induced via end-to-end learning, can offer a favorable trade-off between compression and accuracy for commodity hardware deployment.

Limitations and Scope

Frozen Weights: The current scope keeps weights frozen. While joint weight-and-mask optimization is noted as a natural extension, it is not implemented in this work.
Uniform Allocation: The observation of uniform per-block density is presented as a regime-specific finding for the studied models and sparsity levels, not a general rule against non-uniform allocation.
Safety: The authors caution that compressed models may exhibit different failure modes on long-tail inputs and recommend re-evaluation for safety-critical deployments.