Research map / Architectures and efficiency / Pruning and efficient tuning

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training

ICML 2025 · Architectures and efficiency › Pruning and efficient tuning · approach: Sparse training from scratch

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes BAME, an N:M sparse training method that keeps both forward and backward passes sparse by pruning-and-regrowing weights within N:M blocks guided by loss-aware criteria and freezing oscillating blocks.

Problem

N:M sparse training typically requires dense backward propagation and dense weight updates to find the optimal N:M mask, incurring high training cost.

Approach

Loss-aware mask adaptation (LMA) prunes and revives weights within each N:M block using first-order Taylor loss-change estimates, computing dense gradients only every ΔT iterations; oscillation-aware block freezing (OBF) uses an EMA of per-block mask oscillation (SAD) to freeze the most oscillating blocks from mask updates.

Key results

• BAME reaches 72.0% top-1 on 1:16 sparse ResNet-50 on ImageNet, exceeding SR-STE by 0.5% at 0.29x training FLOPs versus SR-STE's 0.69x.
• BAME reaches 76.1% top-1 on 1:4 sparse ResNet-50 versus 75.3% for SR-STE, at 0.39x training FLOPs versus 0.74x.
• On CIFAR-10 1:4 ResNet-32, BAME reaches 94.71% top-1, 0.28% above Bi-Mask.

Measured effects reported in the paper

Surprising findings

• Prolonging BAME mask adaptation until the final training iteration causes a further performance drop, as newly regrown weights need substantial training after restoration.

Limitations

• Backward-pass sparsity for hardware N:M tensor cores requires an additional bidirectional mask (Bi-Mask) integration, with results reported separately (BAME†).
• Hyperparameters (ΔT, α, β) were tuned on ResNet-32/CIFAR-10 at 1:16 and the paper notes they are not automatically set.

Experimental setup

Datasets: CIFAR-10 ImageNet-1K

Benchmarks: CIFAR-10 ImageNet-1K

Metrics: Top-1 accuracy training FLOPs test FLOPs parameter count

Baselines: ASP SR-STE LBC Bi-Mask MaxQ RigL

Builds on: SR-STE RigL Bi-Mask gradual pruning bidirectional N:M masks

Closest papers by meaning

• Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks (NeurIPS 2021 · similarity 0.89)
• Bi-directional Masks for Efficient N:M Sparse Training (ICML 2023 · similarity 0.87)
• Training Neural Networks with Fixed Sparse Masks (NeurIPS 2021 · similarity 0.87)
• TSENOR: Highly-Efficient Algorithm for Finding Transposable N:M Sparse Masks (NeurIPS 2025 · similarity 0.86)
• Effective Interplay between Sparsity and Quantization: From Theory to Practice (ICLR 2025 · similarity 0.86)
• Sparser, Better, Deeper, Stronger: Improving Static Sparse Training with Exact Orthogonal Initialization (ICML 2024 · similarity 0.85)