Ding Zhu, Zhiqun Zuo, Mahdi Khalili

Paper (arXiv)  ·  Code  ·  BibTeX

A weight matrix is written as (S ⊙ A_i) ⊗ B_i. When the small matrix S is sparse, the product is block-wise sparse, and the block size equals the size of B_i. White entries are zero.

Abstract

Large-scale machine learning (ML) models are increasingly being used in critical domains like education, lending, recruitment, healthcare, criminal justice, etc. However, the training, deployment, and utilization of these models demand substantial computational resources. To decrease computation and memory costs, machine learning models with sparse weight matrices are widely used in the literature. Among sparse models, those with special sparse structures (e.g., models with block-wise sparse weight matrices) fit better with the hardware accelerators and can decrease the memory and computation costs during the inference. Unfortunately, while there are several efficient training methods, none of them are designed to train a block-wise sparse model efficiently. As a result, the current methods for training block-wise sparse models start with full and dense models leading to inefficient training. In this work, we focus on training models with block-wise sparse matrices and propose an efficient training algorithm to decrease both computation and memory costs during training and inference. In addition, we will show that our proposed method enables us to efficiently find the right block size for the sparsity pattern during the training process. Our extensive empirical and theoretical analyses show that our algorithms can decrease the computation and memory costs significantly without a performance drop compared to baselines.

Key findings

• Block-wise sparsity through a Kronecker factorization. Each weight matrix is parameterized as W = Σ_i=1^r (S ⊙ A_i) ⊗ B_i, with an ℓ1 penalty on S. A sparse S makes W block-wise sparse, with blocks the size of B_i. Any block-wise sparse matrix with equal-sized blocks can be written in this form (Proposition 1).
• Lower training cost. The factors are trained directly, so training never touches a dense matrix. Group LASSO and pruning, by contrast, start from the full model. For an 8 × 256 matrix with rank 1, the factorization has 128 trainable parameters instead of 2,048. For ViT-tiny on CIFAR-100, it trains 0.16M parameters instead of 5.5M, with 65.37M training FLOPs instead of 2.16G.
• Block size chosen in one training run. Each candidate block size gets its own set of factors, and a group penalty on each set pushes all but one pattern to zero during training. Group LASSO and pruning need one run per candidate. On ViT-tiny, the run kept the 2 × 2 pattern. Trained separately, the three candidate patterns reach 64.56 (2 × 2), 62.99 (4 × 4), and 62.88 (8 × 8) accuracy.
• Accuracy compared with baselines. On CIFAR-100 with 4 × 4 blocks, the method is more accurate than group LASSO, elastic group LASSO, and block-wise RigL on ViT-tiny and Swin-tiny, and than both group LASSO variants on ViT-base (no block-wise RigL result is reported for ViT-base). It is 1.3–3.9 points below the dense models. On a linear MNIST model it is the most accurate method at 2 × 2 blocks, but block-wise RigL is more accurate at larger block sizes.

Results

CIFAR-100, 4 × 4 blocks (Table 3 of the paper):

Bold marks the best result among the sparse methods. On ViT-base, the method trains 11.58M parameters instead of 87.34M and reaches 69.82 accuracy. That is below the dense model (71.34) and above group LASSO (68.41) and elastic group LASSO (66.95).

BibTeX

@article{zhu2025blockwise,
  title   = {An Efficient Training Algorithm for Models with Block-wise Sparsity},
  author  = {Zhu, Ding and Zuo, Zhiqun and Khalili, Mohammad Mahdi},
  journal = {Transactions on Machine Learning Research},
  year    = {2025}
}