LaPrune: Controllable Differentiable Sparsity at Million Scale

2026-08-04

AI Summary

This work addresses the challenge in large-scale sparse models where conventional top-k methods struggle to simultaneously enforce exact sparsity budgets and support differentiable training, while existing continuous relaxation approaches often entangle mask hardness with selection quality, impairing optimization. To resolve this, the authors propose LaPrune, a differentiable sparse layer that adheres strictly to precise sparsity constraints. LaPrune preserves selection quality via the LapSum barrier and introduces a normalized second-moment constraint to enable scale-invariant control over mask hardness, thereby smoothly interpolating between dense equal-magnitude allocation and hard top-k selection. Theoretical analysis establishes guarantees on saturation ratios, near-binary limiting behavior, and worst-case bounds for near-zero components. Experiments demonstrate that LaPrune enables controllable, differentiable, and budget-exact sparse training at million-scale parameter regimes, significantly outperforming fixed-temperature LapSum baselines.

Abstract

Top-$k$ selection determines which components of a sparse model remain active. Hard selection blocks gradients, while continuous relaxations often couple mask hardness to the selected mass. We introduce LaPrune, a mathematically exact-budget differentiable layer that controls the normalized second moment while preserving the selected mass. A LapSum barrier preserves the selection mass, and a normalized second-moment constraint moves the mask from a dense equal-mass allocation toward hard top-$k$ at each budget. We derive a population prediction of the saturated fraction, a near-binary limiting law, and a tight worst-case guarantee on the near-zero fraction. The normalized hardness parameter is invariant to score scale, while a fixed LapSum temperature is not.

Problem

• sparsity
• top-k selection
• differentiable pruning
• mask hardness
• selected mass

Innovation

• differentiable sparsity
• top-k selection
• LapSum barrier
• normalized second moment
• exact-budget control