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

Mohammad Mozaffari, Younes Hourri, Mohammad Rastegari, Mahyar Najibi

Keywords: Sparsity, LLM Efficiency, LLM Compression, Compression Trinity

TL;DR: LEAP makes end-to-end unstructured LLM pruning tractable by replacing MaskLLM's intractable categorical parameterization with a per-weight Bernoulli-via-Gumbel-sigmoid relaxation, beating ADMM by +2.59 avg zero-shot points.

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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 26