Mask Annealing for Efficient LLM Pruning

• Mask annealing is a technique that progressively hardens soft, learnable masks into binary, hardware-compliant 2:4 sparsity patterns for LLM acceleration.
• The method integrates Hessian-guided groupwise pruning, blending soft and hard gates to optimize mask decisions and maintain model accuracy.
• Annealing schedules for temperature, mask updates, and structure blending ensure a smooth transition from exploratory pruning to final deployable sparse models with reduced retraining cost.

Mask annealing is a class of optimization techniques designed to progressively harden a continuous, learnable soft mask into a hardware-executable semi-structured sparsity pattern, typically for accelerating LLMs. In the context of 2:4 sparsity—where in each group of four weights only two survive—mask annealing manages a transition from exploratory, differentiable pruning to deterministic, group-wise pruning that matches hardware constraints. Recent research, notably the SparseForge framework, demonstrates that mask annealing can maintain or even enhance LLM accuracy post-pruning, while reducing retraining cost by an order of magnitude compared to prior approaches (Hanzuo et al., 7 May 2026).

1. Soft-Mask Representation and Model Integration

Mask annealing begins with the introduction of a soft, learnable mask $m \in [0,1]^{D \times C}$ for each weight matrix $W \in \mathbb{R}^{D \times C}$. During each forward pass in a linear layer, the weights are element-wise multiplied by the mask:

$$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$

This “softly pruned” weight matrix $\widetilde{W}$ allows gradient-based exploration of pruning decisions. Each element $m_{ij}$ can smoothly interpolate between "keep" ($m_{ij} \approx 1$), "drop" ($m_{ij} \approx 0$), or ambiguous/intermediate states ($0 < m_{ij} < 1$), subject to semi-structured group constraints (e.g., 2:4).

2. Optimization Objective and Regularization

The optimization jointly updates the dense weights $(W, B)$ and the mask $m$. The total objective is

$W \in \mathbb{R}^{D \times C}$0

Where:

• $W \in \mathbb{R}^{D \times C}$1 is the main loss function (e.g., next-token prediction).
• $W \in \mathbb{R}^{D \times C}$2 is an optional distillation loss.
• $W \in \mathbb{R}^{D \times C}$3, a binary-preference regularizer maximized at $W \in \mathbb{R}^{D \times C}$4 and minimized at $W \in \mathbb{R}^{D \times C}$5, nudging mask elements towards binary values as training progresses.

3. Hessian-Guided Groupwise Pruning

To direct mask optimization under semi-structured constraints, a second-order, Hessian-based importance score is computed:

$W \in \mathbb{R}^{D \times C}$6

where $W \in \mathbb{R}^{D \times C}$7, estimated via a Hutchinson estimator with Rademacher random vectors. The score $W \in \mathbb{R}^{D \times C}$8 is then standardized layerwise:

$W \in \mathbb{R}^{D \times C}$9

For each 4-element group $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$0, mask annealing constructs:

• Soft gate (for exploration): $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$1, where $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$2 is the $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$3th-largest value in $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$4 ($$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$5 for 2:4) and $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$6 is a learnable or scheduled temperature parameter.
• Hard gate (for structure): $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$7, the definitive groupwise binary pruning pattern.

The sharpness of $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$8 increases as $$\widetilde{W} = m \odot W,\quad Z = X\,\widetilde{W}^T$$9 is annealed downward; $\widetilde{W}$0 enables mask optimization to remain continuous early on, while $\widetilde{W}$1 is the ultimate sparse structure.

4. Annealing Schedules and Mask Updates

Mask annealing interleaves frequent weight updates with less frequent mask updates. At each mask update step $\widetilde{W}$2:

• $\widetilde{W}$3 and $\widetilde{W}$4 are blended:

$\widetilde{W}$5

where $\widetilde{W}$6 is increased from 0 to 1 via a cubic “smooth-step” schedule.

• Mid-penalty correction is injected:

$\widetilde{W}$7

• The mask update (using exponential moving average, EMA):

$\widetilde{W}$8

• Temperature is annealed multiplicatively: $\widetilde{W}$9, with $m_{ij}$0.
• $m_{ij}$1 and $m_{ij}$2 are updated according to their respective schedules.

This multistage schedule ensures that, early in optimization, the mask remains exploratory, while by the end masks converge to near-binary, group-compliant patterns.

5. Final Quenching and Binarization

Upon convergence, a short “hardening window” is used to linearly interpolate $m_{ij}$3 to a binary pattern:

$m_{ij}$4

with $m_{ij}$5 decaying from 1 to 0 over $m_{ij}$6 steps. $m_{ij}$7 is then frozen, yielding a fixed 2:4 mask. Further finetuning absorbs any residual projection error. Progressive annealing minimizes projection error compared to direct thresholding, as mask values are encouraged to settle near 0 or 1 before hard cutoff. Empirical evidence indicates that omitting gradual annealing—e.g., instant thresholding with $m_{ij}$8—introduces larger, often irrecoverable, projection errors (Hanzuo et al., 7 May 2026).

6. Empirical Results and Insights

SparseForge’s implementation of mask annealing on LLaMA-2-7B under 2:4 sparsity achieves:

• 55.96% zero-shot accuracy with 1.25B retraining tokens (vs. 56.43% for dense).
• 57.27% with 5B tokens—surpassing the dense baseline and matching state-of-the-art results achieved with 8× more retraining tokens.

Ablations in (Hanzuo et al., 7 May 2026) show that removing Hessian importance, halving retraining budget, or omitting the low-rank branch each degrades performance by 1.3–2.2 points. Without progressive blending (β-schedule) and temperature annealing, mask values stagnate near 0.5, resulting in a larger soft-to-hard accuracy gap. Fine control of temperature decay ($m_{ij}$9), structure blending ($m_{ij} \approx 1$0), and the mid-penalty regularizer ($m_{ij} \approx 1$1) is each essential for quality recovery; eliminating any element increases perplexity and reduces accuracy by over a point.

7. Summary of Method and Schedules

The core algorithm is succinctly summarized as follows:

Stage	Operation	Purpose
Soft-mask initialization	$m_{ij} \approx 1$2; $m_{ij} \approx 1$3; $m_{ij} \approx 1$4; $m_{ij} \approx 1$5	Begin with identity mask and relaxed gating
Mask update (periodically)	Compute Hessian score $m_{ij} \approx 1$6, standardize; build groupwise soft/hard gates $m_{ij} \approx 1$7; update $m_{ij} \approx 1$8 via EMA and structure blending	Enforce groupwise pruning, explore mask space
Annealing schedules	Anneal $m_{ij} \approx 1$9 ($m_{ij} \approx 0$0), ramp $m_{ij} \approx 0$1 (smooth-step), increase $m_{ij} \approx 0$2 late	Gradual hardening, avoid abrupt transitions
Final hardening	Linearly blend $m_{ij} \approx 0$3 to binary $m_{ij} \approx 0$4 over “hardening window”	Achieve deployable hardware mask with minimal error

This approach allows LLMs to attain high-accuracy semi-structured sparse representations with substantial reductions in retraining cost and minimal accuracy degradation, establishing mask annealing as a central technique in efficient structured LLM compression (Hanzuo et al., 7 May 2026).