An Efficient Training Algorithm for Models with Block-wise Sparsity
Ding Zhu, Zhiqun Zuo, Mohammad Mahdi Khalili
TL;DR
The paper addresses the inefficiency of training block-wise sparse models by introducing a Kronecker product decomposition-based factorization, enabling training with significantly fewer parameters and FLOPs. It formalizes a decomposed weight representation $W^{[l]}_r = \sum_{i=1}^{r_l} (S^{[l]} \odot A^{[l]}_i) \otimes B^{[l]}_i$ to train block-wise sparse matrices and demonstrates that this representation can express any block-wise sparse matrix while reducing training cost. The authors provide theoretical arguments and empirical results showing substantial reductions in training parameters and FLOPs (up to 97% in certain settings) with minimal accuracy loss, and they introduce a pattern-selection framework to efficiently identify optimal sparsity patterns in one training round. The practical impact is improved scalability of block-wise sparse models on resource-constrained devices and accelerators, with automated pattern selection and applicability to vision transformers and CNNs alike.
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 \textit{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.