Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models

TL;DR

Pixelated Butterfly introduces a simple, hardware-aware sparse training approach by combining flat block butterfly and a low-rank component to sparsify GEMM-based networks. The method uses a fixed sparsity pattern and a learnable gamma parameter to balance a sparse butterfly term with a low-rank correction, enabling efficient training across attention and MLP layers. Theoretical results establish expressiveness and convergence properties under neural tangent kernel analyses, while extensive experiments on ImageNet, WikiText-103, and LRA demonstrate substantial wall-clock speedups (up to 2.5x) with negligible accuracy loss. Overall, the work offers a practical, broadly applicable framework for accelerating sparse training in transformer- and MLP-based architectures, with clear hardware-alignment benefits and solid theoretical grounding.

Abstract

Overparameterized neural networks generalize well but are expensive to train. Ideally, one would like to reduce their computational cost while retaining their generalization benefits. Sparse model training is a simple and promising approach to achieve this, but there remain challenges as existing methods struggle with accuracy loss, slow training runtime, or difficulty in sparsifying all model components. The core problem is that searching for a sparsity mask over a discrete set of sparse matrices is difficult and expensive. To address this, our main insight is to optimize over a continuous superset of sparse matrices with a fixed structure known as products of butterfly matrices. As butterfly matrices are not hardware efficient, we propose simple variants of butterfly (block and flat) to take advantage of modern hardware. Our method (Pixelated Butterfly) uses a simple fixed sparsity pattern based on flat block butterfly and low-rank matrices to sparsify most network layers (e.g., attention, MLP). We empirically validate that Pixelated Butterfly is 3x faster than butterfly and speeds up training to achieve favorable accuracy--efficiency tradeoffs. On the ImageNet classification and WikiText-103 language modeling tasks, our sparse models train up to 2.5x faster than the dense MLP-Mixer, Vision Transformer, and GPT-2 medium with no drop in accuracy.

Figures (13)

• Figure 1: Pixelfly targets GEMM-based networks (networks whose computation is dominated by matrix multiply), which it views as a series of matrix multiplication. For each matrix multiply from Model Schema, it (1) allocates compute budget based on dimension and layer type, (2) the budget decides a mapping (hyper-parameter) to our proposed flat block butterfly sparsity patterns, (3) outputs a hardware-aware sparse mask. Note since the hardware is a block device, one memory access to an element in a block leads to the access to the full block.
• Figure 2: Visualization of memory access for a hardware with block size 4: accessing the one (red) location means accessing the full $4\times4$ block (blue).
• Figure 3: Visualization of Flat, Block, and Flat Block butterfly.
• Figure 4: NTK Comparison with Dense Model.
• Figure 5: The performance of Pixelfly and ViT or MLP-Mixer on CIFAR10, CIFAR100 and ImageNet benchmarks. We measure the accuracy and the training time speedup (on ImageNet) compared to the dense model.

Theorems & Definitions (57)

• Definition 3.1
• Definition 3.2
• Definition 3.3
• Definition 3.4
• Theorem 4.1
• Corollary 4.2
• Theorem 4.3
• Theorem 4.4
• Theorem 4.5: Informal
• Definition A.1