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

TL;DR

LEAP introduces a learnable, end-to-end unstructured pruning method for large language models that improves accuracy at high sparsity levels by using a novel Bernoulli-via-Gumbel-sigmoid relaxation.

Contribution

LEAP replaces intractable parameterizations with a scalable relaxation, enabling effective end-to-end unstructured pruning of large language models.

Findings

LEAP improves six-task average zero-shot accuracy by +2.59 points over ADMM.

LEAP is effective across five LLM families from 0.5B to 8B parameters.

LEAP achieves 50% and 60% sparsity with improved accuracy.

Abstract

Unstructured sparsity is now natively accelerated by recent GPU kernels and dataflow hardware, shifting the bottleneck from inference execution to the pruning algorithm. State-of-the-art methods for unstructured LLM pruning are layer-wise surrogates derived from the Optimal Brain Surgeon principle, and they sacrifice end-to-end accuracy, especially under aggressive sparsity. End-to-end alternatives such as MaskLLM and PATCH show that learnable masks can close this gap, but their categorical-over-patterns parameterization scales with the number of valid masks per row and does not port to the unstructured setting. We introduce LEAP, which replaces this intractable parameterization with a per-weight Bernoulli-via-Gumbel- sigmoid relaxation that makes end-to-end unstructured mask learning tractable. Across five LLM families from 0.5B to 8B parameters at 50% and 60% sparsity, LEAP improves six-task average zero-shot accuracy by +2.59 points on average over ADMM, the best layer-wise baseline in our sweep.