BLADE: Block-Level Adaptive Mask Discovery with Error Compensation
Abstract
N:M sparsity is the pattern hardware accelerates, but selecting its masks well is either cheap or good, not both: one-shot pruners leave a large quality gap to end-to-end learned masks, while the learned reference, MaskLLM, needs thousands of GPU-hours and hundreds of gigabytes of GPU memory at 2:4, and runs out of memory beyond it. We present BLADE (Block-Level Adaptive mask Discovery with Error compensation), which reaches learned-mask quality at a small fraction of end-to-end cost by recasting whole-model N:M compression as a sequence of block-local solves: each block takes its inputs from the already-compressed stream and trains against the dense model's outputs, so upstream error is absorbed rather than compounded, and each solve runs a differentiable Gumbel-softmax search over the block's per-group keep-patterns followed by re-optimization of the surviving weights under the frozen mask. Block-local mask learning runs on a gradient signal that vanishes as the relaxation hardens; we show that with a fair per-optimizer search, every optimizer still converges where it is tuned, but only the sign-based Lion transfers, carrying one configuration, tuned once on a single block of the smallest model, unchanged across blocks and models. Under a unified evaluation harness spanning one-shot, lightweight-learned, and end-to-end-learned baselines on three models, BLADE matches MaskLLM's 2:4 downstream accuracy with roughly 20x less memory and 56x less compute on a single workstation GPU, and at 4:8, where no end-to-end learned masks exist, it attains the best accuracy, perplexity, and KL divergence in every comparison, 2.4% to 4.0% points above the strongest baseline, with its 4:8 models surpassing even MaskLLM's 2:4 accuracy.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.