Reproduce “TSENOR: Highly-Efficient Algorithm for Finding Transposable N:M Sparse Masks” on OpenResearch

About this paper

Xiang Meng, Mehdi Makni, Rahul Mazumder

Researchers from MIT's Operations Research Center developed TSENOR, a highly efficient algorithm for generating transposable N:M sparse masks, enabling full forward and backward pass acceleration during deep neural network training. This method achieves 1-10% relative error compared to optimal solutions and up to 300x speedup, making it scalable for large language models while preserving accuracy at higher sparsity.

Key insights

• The problem of finding transposable N:M sparse masks can be effectively modeled as a capacitated optimal transport problem, which is amenable to efficient algorithmic solutions.
• Entropy regularization transforms the optimization into a form solvable by Dykstra's algorithm, enabling highly parallelized computation across millions of blocks.
• A refined rounding procedure, combining greedy selection and a novel GPU-optimized local search, is essential for translating fractional solutions into high-quality binary masks.
• Transposable N:M sparsity with larger M values (e.g., 16:32) significantly diminishes accuracy degradation, making it a viable strategy for achieving full training acceleration.

Summary by alphaXiv.