Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models
Authors: Beidi Chen
Tri Dao
Kaizhao Liang
Jiaming Yang
Zhao Song
Atri Rudra
Christopher Re
Venue: ICLR 2022 Spotlight (2022)
Keywords: Sparse training
butterfly
low-rank
Lottery Tickets
Block sparsity
Hashing
Transformer
ViT
MLP-Mixer

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 $3\times$ 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.3$\times$ faster than the dense MLP-Mixer, Vision Transformer, and GPT-2 small with no drop in accuracy.

Published as a conference paper at ICLR 2022
PIXELATED BUTTERFLY: SIMPLE AND EFFICIENT SPARSE TRAINING FOR NEURAL NETWORK MODELS
Tri Dao"!, Beidi Chen"), Kaizhao Liang Jiaming Yang" 3, Zhao Song " Atri Rudra 5 Christopher Ré
1 Department of Computer Science, Stanford University
2 SambaNova Systems, Inc
8 Department of Probability and Statistics, Peking University
Adobe Research
5 Department of Computer Science and Engineering, University Buffalo, The State University of New York (trid, d, beidic) edwi hewinyjjmad jmpku@gma 1. com, zsongeadobe, com, atri atri@buffalo. edu, chrismre@cs st staffrd.
edu, kai kaizhao. wataao..ianee allaaneammbaoovssytttee.m com,
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 simple and promising approach to achieve this, but there remain challenges as existing methods struggle with accuracy loss, slow training runtime, difficulty in sparsifying all model components. The core problem is that searching for a sparsity mask over a discrete set of sparse matrices difficult and expensive. To address this, our main insight is to optimize over 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 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 accurac--effciinncy tradeoff's 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.
INTRODUCTION
Recent results suggest that overparameterized neural networks generalize well (Belkin et al. 2019), but they are expensive to train (Kaplan et al.) 2020). ideal model should use less compute and memory while retaining the generalization benefits of large models. The simplest and most popular direction is to sparsify these models. This idea has long history in machine learning (LeC'n et al.) 1990) and has driven fundamental progress in other fields such as statistics (Tibshirani 1996). neuroscience (Foldiak) (003), and signal processing (Candes et al. 20066). However, despite significant efforts, speeding up sparse training in wall-clock time without degrading accuracy remains an unresolved problem.
While sparse training an active research area, it has not seen wide adoption. First, it is difficult and expensive to find the sparsity pattern (the possible locations of the nonzeros) that could maintain the same level of accuracy of dense models. Many methods (pruning (Lee et al. 20188). lottery tickets (Frankle and Carbin 2018), hashing (Kitaev et al. (2020) maintain dynamic sparsity masks. However, the large overhead of evolving the sparsity mask can significanlly slow down training and complicate the implementation. Indeed, these methods either require long cycles of pruning and retraining (Frankle and Carbin 2018 or maintain expensive hash tables (Chen al. 2019). Second, most existing methods adopt unstructured sparsity, which may be efficient in theory, but do not take into account the efficiency of training hardware such as GPUs (optimized for dense computation Finally, most methods target single type of operation such as attention (Child et al. 2019 Zaheer et al. 2020). whereas neural network (NN) models often compose different modules (attention, MLP), and in many applications the MLP layers are the main training bottleneck (Wu et al. 2020).
* Equual contribution. Order determined by onin flip.
State-of-the-art sparse training methods require up to 5x more training epochs compared to dense models (vci al.
2020
An unstructured sparse model with 1% nonzero weights can be as slow as a dense model
Hooker 2020

---

Published as a conference paper at ICLR 2022
Model Schema
Pixelated Butterfly
Sparse Masks
Compute Allocation
Attention
Attention Mask
MLP
Flat Block Butterfly
Low-rank
MLP Mask
Figure I: Pixelfly largets GEMM networks (networks whose computation is dominated by matrix multiply), which it views as a series of matrix multiplication. For cach matrix multiply from Model Schema, it (1) allocates compute budget based on dimension and layer type, (2) the budget decides mapping (hyper-parameter) to our proposed flat block butterfly sparsity pattems, (3) outputs a hardware-aware sparse mask. Note since the hardware block device, one memory access to an element a block leads the access to the full block.
A better sparse training method should (i) be simple yet accurate, ideally with static sparsity pattem, (ii) be fast by aligning sparsity pattemn with available hardware, and (iii) have wide coverage of operators that applies to most NN layers. There are three technical challenges. First, we show that given a budget (e.g., total non-zeros in matrix), it NP-hard find the optimal static sparsity pattern for a NN module to minimize the approx- imation error to the dense model. Second, for each sparsity pattern, we need to take into account hardware block-oriented efficiency (accessing each element memory takes the same time as accessing the block of ad- jacent elements (Cook) 2012). illustrated in Fig. 2 Common theoretical measures of efficiency (e.g., number of non-zeros, FL.PP)) do not map well modern hardware designed for block computation. Last, every dif- ferent NN module might require different sparsity patterns, which makes the problem even more complicated.
In our early exploration, we empirically study many sparsity patterns proposed in the literature find those patterns that can closely approximate the dense model (Details in Appendix Kj. We found that one sparsity pattem, namely butterfly low-rank, consistently outperforms the others. This sparsity pattern closely connects to two lines of work in matrix structures: (i) sparse + low-rank matrices, which can capture global and local information (Candes al. 201T (deel and Townsend 2019) Chen et al. 2021], ) and (ii) butterfly matrices (Parker 1995 Dao et al. (2019) whose products can tightly represent any sparse matrix (De Sa al., 2018) Dao al., 2020). Using the fixed sparsity pattern from butterfly matrices, with the addition of low-rank term, would address two of the three challenges above and yield simple way to sparsify most
NN layers (that are based on matrix multiply).
However, butterfly matrices are inefficient on modern hardware: (i) they are difficult to parallelize as they contain sequential products of many factors, and (ii) they are not fedtaar--friedlly because the sparsity patterns are not block-aligned. We propose two simple changes to make Butterfly efficient while retaining their favorable properties. Our proposal, Pixelated Butterfly (Pixelfly), combines flat block butterfly and low-rank matrices to yield a simple and efficient sparse training method.
We design extremely simple sparsity pattern inspired by butterfly + low-rank matrices, which takes into account the hardware's 's block-orienter efficiency. We propose block butterily matrices that are efficient as their sparsity patterss align with hardware blocks. We then introduce flat butterffy, first-order approximation of butterfly with residual connection, that turns the original product of factors into a sum. Flat butterfly matrix multiplications are easy parallelize. Pixelfly, uses the fixed sparsity pattem from flat & block butterfly, along with a low-rank term, to produce a sparse network.
We prove that block butterfly retains the expressiveness of butterfly matrices and can thus tightly capture sparse matrices. We show that flat butterfly matrices can closely approximate large classes of matrices that butterfly matrices capture. Moreover, strictly more expressive than sparse or low-rank matrices alone. Finally, leveraging the recent advance in the neural tangent kemel (NTK), we adapt existing techniques prove the global convergence of gradient descent on training sparse and wide ReL.U networks.
demonstrate that flat block butterfly + low-rank matrices are
Our proposed Pixelfly can be applied to all network modules that rely matrix multiplication (e.g., linear layer, attention, MI.P). To sparsify a full network, we simply need to allocate compute budget for each layer based on matrix and hardware block size.
We empirically validate that Pixelfly can speed up the training of models (Transformers, VIT. MLP-Mixer) without quality drop compared to baselines ImageNet classification, Pixelfly achieve 2.3x training time speedup compared to dense VIT. MLP-Mixer models, and other sparse training baselines, while preserving the same accuracy. On the WikiText- 103 language modeling task, we speed up GPT-2 Medium training by 2.5x and achieve the same perplexity. On the Long Range Arena benchmark, we maintain the same accuracy as Transformer with 5.2x faster training than a dense model, 2x faster than Sparse transformer, and 6x faster than non-block-aligned sparse methods
wide range of domains and tasks. On CIFAR10/100 &

---

Published as a conference paper at ICLR 2022
(Reformer). Our ablation studies highlight the importance of each of our components: our butterfly sparsity improves on existing hand-crafted pattems by up to 2% of accuracy on ImageNet, our hardware-awe aware block-sparsity yields up to 5x speedup, and the balanced compute budget allocation brings 2x speedup compared to baselines that only sparsify attention
2
PROBLEM SETTING
We first define the problem as sparse matrix approximation with simple hardware cost model. Then we briefly introduce butterfly and sparse + low-rank matrices.
Memory Access
Problem Formulation: We focus on the training of GEMM-based models, which can be viewed as a series of matrix multiplies (Given A,B E Raxd compute C== AB").. Speeding up training while maintaining model quality can be mapped to finding an approximation procedure / which reduces the time T of computing C while minimizing enror ElIl AB" [1.-)- Since the hardware a block device, accessing any individual element within a block of
Visualization of memory access for hardware with block size 4: accessing the one (red) location means accessing the full 4 x 4 block (blue).
Figure 2:
memory is the same accessing the full block (Cook) 2012) (Fig. [ A simple cost model of T on hardware with block size b would depend on the number of b-blocks being accessed and compute time (formal definition in Appendix (A]. Our experiment (Appendix L.5) reveals that when the non-zeros are grouped into blocks, picking the smallest block size supported by hardware can speed up operations by 10x compared to sparsity patterss that are not block-aligned.
Low-rank Matrices: Butterfly matrices have been used in numerical linear alge- bra (Parker 1995 Li et al., 2015) and machine learning (Mathicu and LeCun 2014); Jing et al.) 2017 Munkhoeva et al. 2018 Dao et al. [211 Choromanski et al. 2019). They encode the recursive divide-and- conquer structure of the fast Fourier translorm (FFT) algorithm (Cooley and Tukey, 1965) and provably capture any sparse matrix with near-optimal space and time complexity. Sparse and Low-rank structures
Butterfly, Sparse
have been studied in Robust PCA (Candès et al. 201T). graph clustering (Jalali et al. 201T). and co-variance
estimation (Lio. 2011 1b. Recently it has been adopted in attention approximation for Transformers (Chen et al.
2021).
BUTTERFLY MATRICES AND PIXELATED BUTTERFLY
Butterfly matrices (Parker 1995 Dao et al. 2019) are expressive and theoretically efficient. As they contain the set of sparse matrices, we choose to search for the sparsity pattern in this larger class due to their fixed sparsity structure. However, there are three technical challenges. We highlight them here along with our approaches to address them:
Slow speed: butterfly matrices are not friendly to modern hardware as their sparsity patterns are not block-aligned, thus are slow. We introduce a variant of butterfly matrices, block butterfly, which operate at the block level, yielding a block-aligned sparsity pattem..
2. Difficulty of parallelization: the sequential nature of butterfly matrices as products of many factors makes it hard to parallelize the multiplication. We propose another class of matrices, flat butterfly matrices, that are the first-order approximation of butterfly with residual connections. Flat butterfly turns the product of factors into sum, facilitating parallelization.
3. Reduced expressiveness of flat butterfly: even though flat butterfly matrices can approximate butterfly matrices with residual connections, they are necessarily high-rank and cannot represent low-rank matri- ces (Udell and Townsend 2019). We propose to add a low-rank matrix (that is also block-aligned) to flat butterfly to increase their expressiveness.
Combining these three approaches (flat & block butterfly low-rank), our proposal (Pixelated Butterfly) is very simple method to train sparse networks.
3.1 BLOCK BUTTERFLY MATRICES
We propose block version of butterfly matrices, which is more hardware-friente than the regular butterfly. The regular butterfly matrices Dao et al. (2019 2020) will be special case of block butterfly with block size b= 1.
We omit b in the notation if b=
*Pixelfly code is available at httpa///w.itt.............////aa

---

Published as conference paper at ICLR 2022
Butterfly
Block Butterfly
x
X
Bin
15, 21
Flat Butterfly
Flat Block Butterily
Figure
Visualization of Flat, Block, and Flat Block butterfly.
Definition 3.1. block butterfly factor (denoted as BR,b) size kb (where k and block size b is a matrix
Di
D2 where each Dj is &. x block diagonal matrix of block size b of the form 2
of the form B = diag (ong(Da...... Dix112) where DijERAx! We restrict to be a power of
D3
Da
Definition 3.2. A block butterfly factor matrix (denoted (, of size nb with stride and block size is a block diagonal matrix of is (possibly different) butterfly factors of size kb and block size b:
B (m.b)) = diag (B..bl. [x,bb]
Definition 3.3. A block butterfly matrix of size nh with block size (denoted as B( is a matrix that can be expressed as product of butterfly factor matrices: B(n,b) -B
(n.b)
b) B (n.b)
...BB (n.b) Define Bo as the set of
all matrices that can be expressed in the form B(n.b) (for some n))
3.2 FLAT BUTTERFLY MATRICES
In most applications of butterfly matrices neural networks, multiplies the Ollggn) butterfly factors. However, this operation is hard to be efficiently implemented on parallel hardware (e.g.. GPUs) due to the sequential nature of the operation We instead propose to use of butterfly factors that can approximate the products of the factors. This sum of factors results in one sparse matrix with a fixed sparsity pattern, which yields up 3x faster multiplication GPUs (Appendix D)
Residual connections have been proposed to connect the butterfly factors (Vahid et al. 2020). We show that residual products of butterfly matrices have first-order approximation as sparse matnx with fixed sparsity. Let M be a matrix in the set of butterfly matrices In residual form, for some A R:
(n) M= (I+A-BB " (I+AB (/)2 (I+ (I+
(1)
Note that this form can represent the same matrices in the class of butterfly matrices B, since any (m) contains the identity matrix T.
Assuming that À is small, we can expand the residual and collect the terms
( M=I+X( +.. +
(22)).
Definition 3.4. Flat butterfly matrices of maximuan stride le (for power 2) are those of the form 2)
I+ +B(") +B(.
Flat butterfly matrices of maximum stride n are the first-order approximation butterfly matrices residual form (Eq. (D). Notice that flat butterfly of maximum stride is are sparse matrices with (nlogk) nonzeros with a fixed sparsity pattem,, as illustrated in Fig. 3 We call this sparsity pattern the flat butterfly pattern.
Flut block butterfly matrices are block versions of flat butterfly in Section 3.2 empirically validate that flat block butterfly matrices are up faster than block butterfly or regular butterfly (Appendix(1
(shown in Fig. 3 We
Since flat butterfly matrices approximate the residual form butterfly matrices, they have high rank i small (Section 4). This is one of the motivations for the addition of the low-rank term in our method.
'Even with very specialized CUDA kemel, butterfly matrix multiply (O(nlogn) complexity) is only faster than 2019.
dense matrix multiply (O(n²) complexity) for large values of (around 1024) Dao We make the approximation rigorous in Section

---

Published as a conference paper at ICLR 2022
3.3 PIXELATED BUTTERFLY: FLAT BLOCK BUTTERFLY LOW-RANK FOR EFFICIENT SPARSE TRAINING
We present Pixelated Butterfly, an efficient sparse model with simple and fixed sparsity pattern based on butterfly and low-rank matrices. Our method targets GFMMM-based neural networks, which are networks whose computation is dominated by general matrix multiplies (GEMM), such as Transformer and MLP-MMxer. As a result, we can view the network as a series of matrix multiplies.
Given model schema (layer type, number of layers, matrix dimension) and a compute budget, Pixelated Butterfly has three steps: compute budget allocation per layer, sparsity mask selection from the flat butterfly pattern, and model sparsification. We describe these steps in more details:
1. Compute budget allocation: based on our cost model (Appendix [A). given the layer type, number of layers, and matrix dimension, we can find the density (fraction of nonzero weights) of each layer type minimize the projected compute cost. Continuing our goal for simple method, we propose simple rule of thumb: allocate sparsity compute budget proportional to the compute fraction of the layer. For example, if the MLP layer and attention layers are projected to takes 60% and 40% the compute time respectively, then allocate 60% of the sparsity compute budget to MLP and 40% to attention. We verify in model.
Appendix [ that this simple rule of thumb produces similar results to solving for the density from the cost
Sparsity mask selection: given a layer and a sparsity compute budget for that layer, we use one-quarter one-third of the budget for the low-rank part as a simple rule of thumb. We pick the rank as a multiple of the smallest supported block size of the device (c.g.. 32) so that the low-rank matrices are also block-aligned. The remaining compute budget is used to select the sparsity mask from the flat block butterfly sparsity pattern: we choose the butterfly block size as the smallest supported block size of the device (e.g.. 32), and pick the maximum stride of the flat block butterfly (Definition (3.4) to fill up the budget.
3. Model sparsification: The resulting sparse model is simply a model whose weights attention follow the fixed sparsity mask chosen in step 2, with the additional low-rank terms (rank also chosen in step 2). In particular, we parameterize each weight matrix as: W ->=++++ (1-))VVV where B is flat block butterfly matrix (which is sparse), UVT is the low-rank component, and y is leamable parameter. We train the model from scratch as usual.
Our method is very simple, but competitive with more complicated procedures that search for the sparsity pattern (Appendix K). We expect more sophisticated techniques (dynamic sparsity, a better approximation of butterfly) to improve the accuracy of the method.
THEORETICAL ANALYSIS
We characterize the expressiveness of the matrices used in our method. In particular, we prove that block butterfly retains the expressiveness of butterfly, and that flat butterfly can accurately approximate the residual form of butterfly. Moreover, flat block butterfly + low-rank (an instance of sparse low-rank) is more expres- sive than sparse or low-rank matrices alone. Finally, we analyze the training convergence and generalization of networks with sparse weights. All proofs in the Appendix.
4.1 EXPREsSIVENESS OF BLOCK BUTTERFLY
We first prove the expressiveness of block butterfly matrices.
Theorem 4.1. The set B2s of n x block butterfly mutrices with block size 2b contains the set BB, of nxn. block butterfly matrices of block size
By recursive argument, the set of block butterfly matrices whose block size a power of 2 contains the set of regular butterfly matrices.
Dao et al. 2020) show that butterfly matrices tightly represent all structured matrices, such as sparse matrices and many fast transforms. As a result, block butterfly matrices can also represent those structured matrices. In particular,
Corollary 4.2. For any constant block size b that is a power 2, any nbxxbb spare matrix with nonzeros can be written as products of block butterfly matrices with block size b and their transposes, with O (slog)) parameters.
"We describe how to add sparse and low-rank for attention in Appendix

---

Published as a conference paper at ICLR 2022
4.2 EXPRESSIVENESS FLAT BUTTERFLY
We now characterize how the flat butterfly matrices approximate butterfly matrices. In particular, assuming that each butterfly factor has bounded norm, we show that flat-butterfly matrices can accurately approximate the residual form of butterfly with emorr scaling as ()²)).
Theorem 4.3. Let M be a matrix of the form in Definition 3.4 where k= with Biax max, B
ove for some constant 0 1 and some >> O. Then and MIs logu
M A-(1+A(B3) [+ ( +B +B +...
We show that flat butterfly matrices must have high-rank if is small. This is the motivation for the addition of the low-rank term Pixelfly (Section Bp.
Theorem 4,4. Let M be as in Eq. with Buax : max; B and | S
lognEmass for some constant
cVf
0<es 3 and some > O. Let B
max; [Billoc
Assuming B S Bmaxx- Then
[Bmax Bo0 max
logn
rank(I+A(RE) +B("))
clog
Bruax
Boc
max
4.3 EXPRESSIVENESS FLAT BLOCK BUTTERFLY LOW-RANK
Chen et al. (202T) prove that there is natural class of imput sequences (generated by a clustering process) whose attention matrix only be approximated well by sparse + low-rank matrices, and not sparse low-rank matrices alone. We adapt their technique to show a similar result for the class matrices we use in Pixelfly.
We require an extra assumption on the clustering process compared to Chen et al. (2021): the elements in the input sequence form clusters with the same size. Then their attention matrix will have a large block diagonal component well-appe eel-ooppro by flat butterfly, while the rest of the attention matrix of medium size and is well-approximmted by low-rank.
Theorem 4.5 (Informal). There exists a class of input sequences whose attention matrices are well- approximated by flat block butterfly low-rank (a special case of sparse low-rank) but not by sparse or low-rank alone.
The formal theorem statement and proof are in Appendix B3
4.4 CONVERGENCE AND GENERALIZATION OF SPARSE NETWORKS
There are natural questions about the training and generalization of sparse models: do they train similarly dense models, is their generalization close to that of dense models, and can one successfully train them with gradient descent? Our analysis theoretically shows that the answers are yes.
Our analysis relies on the neural tangent kemel (NTK) (Jacot et al. 2018) of the network. The NTK of two data points and y measures the similarity between the gradient of the network when evaluated at compared to the gradient when evaluated at y. This kernel governs the dynamics of the neural network output function f(-0)) throughout the training and its generalization. We build on the great literature of NTK (Li and Liang 2018 Du et al. 2019 Allen-Zhu et al. 20195). The standard result (Song and Yang. 2019) implies the following. if the NTK of the sparse model is close to the NTK of the dense model, then (i) their training convergence speed is similar, (ii) their generalization bounds are similar. For completeness, we state the formal
result in Appendix +
Though this result does not capture the possible regularization effect of sparsity, shows that sparse models with small NTK difference from dense NTK preserve the generalization ability of dense models, subject that has been studied more extensively, both from empirical and from theoretical perspectives. We also show that training wide and sparse networks with gradient descent converges globally, similar to the result for wide dense networks (Du et al. 2019 Allen-Zhi et al. 2019b) in Appendix
EXPERIMENTS
In this section, our goal is to demonstrate that extremely simple fixed sparsity pattern actually speed up sparse model training in wall-clock time without degrading model quality. Specifically, we empirically

---

Published as a conference paper at ICLR 2022
validate three claims that suggest Pixelfly can improve training speed of different model architectures while retaining model quality on a wide range of domains and tasks.
1. Section 5.1 for image classification tasks, we first show the empirical NTK of flat block butterfly low-rank sparsity pattern closer to dense NKT than other baselines. Then we demonstrate our superior end-to-end performance. Specifically, we achieve training speed up on both MLP-Mixer and VIT models by up to 2.3x wall-clock time with no drop in accuracy compared to the dense model and up to 4x compared to Rigl, BigBird and other sparse baselines.
Section [52 for language modeling and text classification tasks, we can speed up GPT-2 small dense model training by 2.1x, achieving perplexity of 22.5 on wikitext- 103. In addition, on Long Range Arena (LRA) benchmark, we maintain the same accuracy but have 5.2x speed-up in training.
3. Section 5.3 we show the necessity of block flat butterfly and low-rank structures, hardware-alignmen and wide coverage of most network layers with ablation studies on these three components of Pixeffly.
5.1 IMAGE CLASSIFICATION
Detoce
0.7
BigDind
We evaluate the quality and efficiency of Pixelfly through three metrics: (1) distance to training dynamic of the dense model: compare the distance between empirical NTK kerne (7 of the models with candidate patterns, including BigBird (Zaheer et al. 2020), Butterfly (Dao et al.] 2020), and that of the dense model, (2) upstream accuracy: compare the accuracy and training time of the Pixelfly, the dense counterpart, and other baselines on same image classification tasks, (3) downstream accuracy: compare the accuracy of our pretrained Pixelfly and dense model fine-tuned on downstream tasks (Appendix [..4]. The empirical NTK of the model with flat block butterffy + low-rank, picked by Pixelfly, is closer to the NTK of the dense model. Pixelfly MIP-mixer and ViT models also retain the same top-1 accuracy of the original dense models while achieving up to 2.3x speed up.
Ficillly
Acmmar
0.
114
NTX Distancel
0.5
NTT Diannoe=
50
500
Foooh
Figure 4: NTK Comparison with Dense Model.
Sctup: We use three popular vision benchmarks, CIFAR-10/100 (Krizhevsky al. 2009) and ImageNet (Deng ct al. 2009). We choose recent popular Vision Transformer (Dosovitskiy et al. (020), T2T-ViT (Yuan et al.
202 and MLP-Mixer (Tlstiikiin et al. (02T) as representative base models. Their major computation
bottlenecks are in different components, e.g. MLP only, attention, or both so we can evaluate the end-to-end
applicability of Pixelfly more clearly.
To character- ize the training dynamic of the sparse networks, we compute the empirical NTK kernels for dense Vision Transformer on CIFAR- 100. Then, we show the rela- tive differences between kernels of models with different sparsity patterns and that of the dense one in Fig. 4 Specifically, we pick a popular sparsity pattern com- bination Bigbird pattern (/a- hherr et al. 2020) for attention layer and random (magnitude-
Empirical NTK:
Figure S: The performance of Pixelfly and ViT MI.-Mixer on CIFARIO, CIFARI00 and ImageNet benchmarks. We measure accuracy and the training time specdup (on ImageNet) compared to the dense model. Model
CIFARIO CIFAR100
ImageNet
Speedup
Mixer-S/16
86.4
58.7
72.4
PixelilyMMixe-SS/
89.8
62.9
72.6
1.7x x
Mixer-B/16
87.6
59.5
75.6
Pixelly-Mixer-B/16
90.6
65.4
76.3
2.3x
VIT-S/16
89.5
65.1
77.7
Pixellly-iT-SS11
91.3
66.8
77.5
1.9x
VIT-B/16
89.9
61.9
78.5
Pixelll-ViT-B/16
92.2
65.1
78.6
2.0x
based sparsity at initialization equals to random) for MLP layer, as a representative baseline. The plot indicates that our designed pattern, flat block butterfly low-rank is the closest one to that of the dense one among all the patterns. Hence, we expect them to enjoy the most benefits of their dense overparameterized counterparts in real tasks. More details on measuring empirical NTK are covered in the Appendix L.3
Training from scratch: We validate that Pixelfly trains up to 2.3x and 2.0x faster than dense MLP-Mixer and ViT models from scratch, with the same accuracy under the same setting (batch size, epochs). Specifically, we sparsify the models with Pixellly and train them on three commonly used vision benchmarking datasets, CIFAR- 10/100 and ImageNet. We measure their Top-1 accuracy wall-clock training time.
Figure Comparison with representative sparse training bascline RigL (Evei et al. 2020)
Model
MaageNer(Aee)
Specdap
Mixer S/32
58.56
Rigl. (Feci et al 20200
56 10
08x
Pixellly (cuns)
59.61
21x
There an emerging consensus that the NTK is an informative measure of how training and convergence behaviors of two models are similar.

---

Published as a conference paper at ICLR 2022
To summarize the general trend, Fig. highlights that our sparse vision models consistently retain the accuracy of their dense counterparts in terms of accuracy and achieve training time speed-up.
Furthermore, we have discussed in Section 1 that current sparse training algorichms aim to dynamic search what could be good sparsity for efficient inference but do not speed up training in wall-clock time. But we still present the comparison results in Fig. 6 for completeness. For a fair comparison, we conduct the experiment on Mixer-S/32 model for 100 epochs because Rigl. aims for sparsity on weights, while we aim for both weights & attention. As expected, Rigl. does not speed up training (the pioneering work has unstructured sparsity and does not achieve speed up on GPU) but surprisingly Pixelfly outperforms both dense and RigL in terms of accuracy while achieving 2.1x x speedup.
Finally, compare Pixelfly with BigBird and Sparse Transformer pattern. For a fair compari- son, we choose ViT as the base model be- cause its major bottleneck is on the T2T attention module (our baselines are efficient attention vari- ants). We can see from Fig. 7 that Pixelfly is the
Figure 7: Comparison with representative sparse attention hase- lines.
Model
ImageNet (Ace)
Specdup
T2T-VIT
81.7
BigBird
81.5
0.9x
Sparse Transforner
81.4
1.3x
only one that can maintain the accuracy and have actual speed up. Further more, Pixelfly speeds up T2T module (large attention) by 1.4x compare to dense.
Pixeltly
81.7
1.4x
5.2 LANGUAGE MODELING AND TEXT CLASSIFICATION
In this section, we aim to evaluate the effectiveness of Pixelfly in the text domain, on a language modeling task and Long Range Arena (L.RA (Tay et al.) (020)) benchmarks. On WikiText-103 (Merity et al. 2016) Pixelfly achieves 22.5 perplexity, which is around the same perplexity GPT-2 small (Radford et al. 2019) but trains 2.1 x faster. On LRA, Pixelfly obtains almost the same accuracy as the full model but gains up to 5.2x speed-up.
Setup: We use WikiText- -03 for language modeling and LRA for classification tasks. We use GPT-2 small and vanilla Transformer as the base dense models. The computational bottleneck of GPT-2 small for moderate sequence length, e.g. 512, would be on both attention and MLP layers, while the bottleneck of transformer on LRA task is on attention since the benchmark is designed to evaluate models under long- context scenarios.
GPT-2-Small, Medium on WikiTexx-1033: We show training GPT-2-Small, Medium and its Pixelfly model from scratch commonly used NLP bench- marking dataset, 103 We measure their perplexity on that dataset, and our training speed up. All setup and finetuning hyperparameters follow the ones in the original paper (Radford al. 2019). We present the results in Fig. [8] It is not hard to that Pixelfly models have great advantages in
Figure 8: The performance of Pixelfly, BigBird and GPT-2- Small, Medium on WikiText- 103. We measure the perplexity and training speed up.
Model
WikiText- 103 (ppl)
Speedup
GPT-2-Small
22.2
BigBird
23.3
0.96x
Pixelfly
22.5
2.1x
GPT-2-Medium
20.9
BigBird
21.5
L.X
accuracy-efficiency tradeoffs since it maintains the
Pixelfly
21.0
2.5x
same perplexity as the dense model but achieve up to 2.5x speed-up in training.
Vanilla Transformer on LRA: We compare vanilla transformer and its Pixelfly models trained from scratch on LRA benchmark. We mea- sure the accuracy, throughput, and training time of both models. Each task has a different sequence length varying between 1024 and 4096. We follow the
Figure 9: The performance of Pixetfly, Reformer and vanilla transformer on Long- Range-Aren benchmarks. We measure the accuracy and training speed.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Avg
Speedup
Transfrmer
36.54
63.12
80.33
41.56
73.49
59.01
Reformer
36.85
58.12
78.36
28.30
67.95
53.90
0.8x
Pixelfly
37.65
66.78
80.55
42.35
72.01
59.86
5.2x
implementation and experimental setting in (Xiong et al. 2021). We compare the performance of Pixelfly against the dense transformer and report the results in Fig, 9 We also include the numbers of other baselines from the same repository in the appendix. We can see Pixelfly cause almost drop in accuracy while achieving 5.2x x speed-up in time.
5.3 ABLATION STUDY
We conduct ablation studies on each component of Pixelfly (Details in Appendix [.S). Specifically, we present (i) how flat block butterfly and low-rank affect the model quality, (ii) how different block size would affect the training speed, (iii) how budget allocation affects the end-to-end speed up.

---

Published as a conference paper at ICLR 2022
Necessity of Flat Block Butterfly and Low-rank: (i) We apply different parameter allocation of flat block butterfly and Low-rank component in Pixelfly Mixer-S model on CIFAR-10 10 under the different density varying in [0.05, 0.1, 0.2]. We found that similar to what was reported in (Chen al. 2021 using around ă budget on Low-rank and Σ on flat block butterfly achieves the best accuracy. (ii) We also compare Pixelfly with baseline sparsity patterns and show it is 2.7x faster than dense, 3x faster than Butterfly, 3.2x faster than BigBird under 10% density.
We study the accuracy- ccuracy-efficiency trade-off for flat block butterfly and random sparsity pattemn with different block sizes from 1-32 ( ()) We found that first, under the same density, the same sparsity patterns covered with different block sizes could have a big difference in efficiency. Under the same block, the pattem with more locality can be more efficient. Last, the density can seem very small, but actually memory access could be up to 100% of the matrix. Therefore, we always want to make full utilization of the smallest block size that the hardware (or compiler) supported.
Block Size:
We sparsify different components of ViT-small separately, including attention and MLP. We show that their compute ratio is approximately 1:2, so if only sparsify one of them, the other one will be the bottleneck preventing end-to-end speed up. Therefore, it is necessary to have an algorithm that can sparsify all layers.
Budget Allocation:
6 RELATED WORK
Lottery Ticket Hypothesis. Models proposed in our work can be roughly seen as a class of manually constructed lottery tickets. Lottery tickets (Frankle and Carbin 2018) are a set of small sub-networks derived from a larger dense network, which outperforms their parent networks. Many insightful studies (Morcos et al.)
2019 Orseau et al. 2020; Frankle et al. 2019 2020 Malach et al. 2020 Pensia et al. (020) are carried out to analyze these tickets. but it remains difficult to generalize to large models due to training cost. In an attempt, follow-up works (Wang et al. 2020 Tanaka et al. 2020) show that one can find tickets without training labels.
We draw inspiration from one of them, Liu and Zenke (2020), which uses the NTK to avoid using labels in sparsifying networks. Other recent works use specialized hardware to accelerate sparse training (Goli and
Aamodt 2020
Raihan and Aamodt 2020).
Neural Pruning. Our work is loosely related to neural network pruning. By iteratively eliminating neurons and connections, pruning has seen great success in compressing complex models. Pioneering work (Han et al., 2015a (b) shows that pruning can produce significantly smaller and faster models for inference. Subsequent methods (Li et al., 2016)] Lin et al. [201 Dong et al. 2017; Sanh et al.. 2020 Lagunas et al.) 2021 Zhu and Gupta 2017 improve on the quality of the pruned models. While both our and the pruning methods aim
produce sparse models, we target training efficiency, whereas pruning mostly focuses on inference efficiency at the cost of sacrificing training speed.
Overparameterized Models and NTK. Our analysis for sparse model convergence relies heavily on recent advance in neural tangent kernel (NTK) (Jacot et al., 2018). NTK is a tool which has been widely
used in analyzing overparameterized models' convergence (Li and Liang] 2018 Du et al.] 2019: Allen-
Zhu et al. 2019b|c (Sog and Yangl 2019) , generalization (Allen-Zhu et al. [2019a). connection data separability (Oymak and Soltanolkotabi 2020). and cost per iteration (Brand et al. (021)). Deep Double Descent (Nakkiran et al. 2019 d' Ascoolli et al. 2020) conjectures that the generalization error improves as the parameter count grows. is not surprising that the community is racing to break the record of the largest
parameter counts (Radford al., 2019 Brown el al.
2020 Dosovitskiy al. 2020 Tolstikhin et al. 2021
Zhang el al. 2021 Naumov et al. 2019 Jumper et al. 2021
We provide extended related work in Appendix M
7 CONCLUSION
In our early exploration of many sparsity patterns with complex training procedures, we found that simple pattern (butterfly + low-rank) consistently (though not always) performed among the best. This motivated us to propose Pixelated Butterfly, a simple and efficient sparse training method. In our quest for simplicity and efficiency, we have chosen to use fixed sparsity that aligns with modern hardware, which was sufficient to yield wall-clock training time speedup without sacrificing accuracy. We are excited about several future directions. Inspired by the remarkable success of model pruning for inference, it is possible that dynamic block sparse mask could be made efficient yet still accurate. Our flat butterfly is simple first order approximation of the rich class of butterfly matrices, and there could be more sophisticated approximations that retain more expressiveness. Our method is first step towards the goal of making sparse models train faster than dense models and make them more accessible to the general machine learning community.

---

Published as a conference paper at ICLR 2022
Ethics Statement. As the amount of data and model size grows, work seeks to understand how to train those large models more efficiently by exploiting sparsity. This potentially connects to energy savings during large-model training. In addition, this allows the general community that has limited access to the computational resources to train and understand those foundation models. Our method is applicable to architectures, which may improve a wide range of
popular models such as MLP-based and Transformer- applications, each with their own potential benefits and harms. For example, making language modeling more efficient might simplify the process of spreading misinformation. Similarly, better image classification models might make automatic surveillance easier. To alleviate the above risks, we need to address application-specte issues like privacy, bias and discrimination, going beyond the accuracy metric we currently considered. Specifically, for image classification task, while our work partially addresses the issue of environmental cost, it does not address other issues such as faimess and bias in model and datasets.
Reproducibilityy Statement. To facilitate the reproducibility of our algorithms and results, (i) we include a link to downloadable source code in supplementary materials, (ii) for our theoretical statements and results, we include clear explanations of any assumptions and a complete proof of the claims from Appendix [A] to Appendix H. for any datasets used in the experiments, complete description of the data processing steps is in Appendix
ACKNOWLEDGMENTS
We thank Laurel Orr, Xun Huang. Sarah Hooper, Sen Wu, Megan Leszczynski, and Karan Goel for their helpful discussions and feedback on early drafts of the paper.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under
No. N000141772266 (Unifying Weak Supervision); ONR N00014-20-1 1-2480 Understanding and Applying Non-Euclidean Geometry in Machine Learnin; N000142012275 (NEPTUNE); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation. American Family Insurance, Google Cloud, Salesforce, Total, the HAI-AWS Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI). and members of the Stanford DAWN project: Facebook, Google, and VMWare. The Mobilize Center a Biomedical Technology Resource Center, funded by the NIH National Institute of Biomedical Imaging and Bioengineering through Grant P41EB027060. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government. Atri Rudra's 's research is supported by NSF grant CCF-1763481.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in neural information processing systems, pages 6155-6166, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pages 242-252. PMLR, 2019b.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. In NeurlPS, 2019c.
Noga Alon. Perturbed identity matrices have high rank: Proof and applications. Combinatorics, Probability and Computing, 18(1-2):3-15 15, 2009.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameteriizatonn. In International Conference on Machine Learning, pages 244-253, PMLR, 2018.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019a.
Sanjeev Aror, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv:1904. /9555 2019b.
10

---

Published as a conference paper at ICLR 2022
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modem machine- practice and the classical bias variance trade-off. Proceedings of the National Academy of Sciences, 116(32): 15849-15854, 2019. 2019.
Rishi Bommasani, Drew Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bemsteenn, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07233 2021.
Jan den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (overparametrized)) neural networks in near-linear time. In ITCS, 2021.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Emmanuel Candes, Justin K Romberg, and Terence Tao. Stable signal recovery from incomplete and inaccurate measurements. Communications on Pure and Applied Mathematics: Journal Issued by the Courant Institute of Mathematical Sciences, 59(8):1007-1223, 2006.
Emmanuel J Candès, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? Journal of the ACM (JACM), 58(3):1-37, 2011. 2011.
Yuan Cao and Quanquan Gu. Generalization error bounds of gradient descent for learning over-parannn over-parameterized deep relu networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 3349-3356, 2020.
Beidi Chen, Tharun Medini, James Farwell, Sameh Gobriel, Charlie Tai, and Anshumali Shrivastava. Slide: In defense of smart algorithms over hardware acceleration for large-scale deep learning systems. arXiv preprint 10X.:::903.03122 2019.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In NeurlPS, 2021.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint
2019.
Krzysztof Choromanski, Mark Rowland, Wenyu Chen, and Adrian Weller. Unifying orthogonal Monte Carlo methods. In International Conference on Machine Learning, pages 1203-1212, 2019.
DC Collins and ES Angel. The diagonal decomposition technique applied to the dynamic programming solution of elliptic partial differential equations. Journal of Mathematical Analysis and Applications, 33(3); 467-481, 1971.
Shane Cook. CUDA Programming: A Developer's 's Guide to Parallel Computing with GPUs. Morgan Kaufmann Publishers Inc., San Francisco, CA, USA, 1st edition, 2012. ISBN 9780124159334,
James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297-30) 1965.
Tri Dao, Albert Gu, Matthew Eichhor, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In International conference on machine learning, pages 1517-1527. PMLR, 2019.
Tri Dao, Nimit S Sohoni, Albert Gu, Matthew Eichhom, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps. International conference representation learning, 2020.
Stéphane d' Ascoli, Levent Sagun, and Giulio Biroli. Triple descent and the two kinds of overfitting: Where & why do they appear? arXiv preprint arXiv:2006.03500 2020.
Christopher De Sa, Albert Gu, Rohan Puttagunta, Christopher Ré, and Atri Rudra. A two-pronged progress in structured dense matrix vector multiplication. In Proceedings of the Twenty-Ninth Annual ACM-SIAM Symposium on Discrete Algorithms, pages 1060-1079. SIAM, 2018
11

---

Published as a conference paper at ICLR 2022
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image databuse. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255 Icec, 2009.
Xin Dong, Shangyu Chen, and Sinno Jialin Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. arXiv preprint 17Xi::7005.05665 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbomn Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition at scale, arXiv preprint arXiv:2010. 1929, 2020.
Simon Du, Xiyu Zhai, Barnabus Poczos, and Aarti Singh. Gradient descent provably optimizes over- parameterized neural networks. In ICLR. https: pecktp:://axiiie http:/axxivv........///// wwt.:////pp.///////////e 1810. 02054 2019.
Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning, pages 2943-2952. PMLR, 2020.
Peter Foldiak. Sparse coding in the primate cortex. The handbook of brain theory and neural networks, 2003.
Dean Foster, Howard Karloff, and Justin Thaler. Variable selection is hard. In Conference on Learning Theory, pages 696-709. PMLR, 2015.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: 03635, 2018.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint 1XXv:1903.01611, 2019.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pages 3259-3269. PMLR, 2020.
Negar Goli and Tor M. Aamodt. Resprop: Reuse sparsified backpropagation. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition (CVPR), June 2020.
Scott Gray, Alec Radford, and Diederik Kingma. Gpu kernels for block-sparse weights. arXiv preprint arXhv:1711.09224, 3. 2017.
Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint vrXiv:1510.00144 2015a.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural networks. arXiv preprint arXiv: 1506.02626, 2015b.
Soufiane Hayou, Amaud Doucet, and Judith Rousseau. Training dynamics of deep networks using stochastic gradient descent via neural tangent kemel. arXiv preprint 1905 13654, 2019.
Sara Hooker. The hardware lottery. arXiv preprint arXiv:2009.06499, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernl:: Convergence and generalization in neural networks. arXiv preprint 1xx::1806.00577 2018
Ali Jalali, Yudong Chen, Sujay Sanghavi, and Huan Xu. Clustering partially observed graphs via convex optimization. In /CML, 2011.
Li Jing, Yichen Shen, Tena Dubcek, John Peurifoy, Scott Skirlo, Yann LeCun, Max Tegmark, and Marin Soljacíć. Tunable efficient unitary neural networks (EUNN) and their application to RNNs. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 1733-1741. JMLR. org, 2017.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figumov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0830 2020.
12

---

Published as a conference paper at ICLR 2022
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
François Lagunas, Ella Charlaix, Victor Sanh, and Alexander M Rush. Block pruning for faster transformers. arXiv preprint arXiv:2109.044880 2021.
Yann LeCun, John S Denker, and Sara A Solla. Optimal brain damage. In Advances in neural information processing systems, pages 598-605, 1990.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32:8572-8583, 2019.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shot network pruning based on connection sensitivity. arXiv preprint 18/0.00/000020 2018.
Hao Li, Asim Kadav, Igor Durdanovie, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. arXiv preprint arXiv: 08700, 2016.
Yingzhou Li, Haizhao Yang, Eileen R. Martin, Kenneth L. Ho, and Lexing Ying. Butterfly factorization. Multiscale Modeling & Simulcaiio,, 13(2):114-332, 2015.
Yuanzhi Li and Yingyu Liang. Learning overparameteriize neural networks via stochastic gradient descent on structured data. In NeurlPS, 2018.
Ji Lin, Yongming Rao, Jiwen Lu, and Jie Zhou. Runtime neural pruming. In Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL httpps://prccceeannnn pd.
neurips. cc/
Tianlin Liu and Friedemann Zenke. Finding trainable sparse networks through neural tangent transfer. In International Conference Machine Learning, pages 6336-6347. PMLR, 2020.
Xi Luo. High dimensional low rank and sparse covariance matrix estimation via convex minimization. arXiv preprint arXiv:/111. 111,113,,, 199, 2011.
Eran Malach, Gilad Yehudai, Shai Shalev-Schwartz, and Ohad Shamir. Proving the lottery ticket hypothesis: Pruning all you need. In International Conference on Machine Learning, pages 6682-6691, PMLR, 2020.
Michael Mathicu and Yann LeCun. Fast approximation of rotations and Hessians matrices. arXiv preprint arXiv:1404.7190 104.1195, 2014.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint wrXivv1609.0788 2016.
Ari S Morcos, Haonan Yu, Michela Paganini, and Yuandong Tian. One ticket to win them all: generalizing lottery ticket initializations across datasets and optimizers. arXiv prepeint arXiv: 1906.0273, 2019.
Marina Munkhoeva, Yermek Kapushev, Evgeny Burnaev, and Ivan Oseledets. Quadrature-based features for kernel approximation. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages 9165-9174. Curran Associates, Inc., 2018
Prectum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. arXiv preprint arXiv: 1912.0222,, 2019.
Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang. Narayanan Sundaraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al. Deep learning recom- mendation model for personalization and recommendation systems. arXiv preprint arXiv: 1906,0009, 2019.
Laurent Orseau, Marcus Hutter, and Omar Rivasplata. Logarithmic pruning is all you need. Advances in Neural Information Processing Systems, 33, 2020.
13

---

Published as a conference paper at ICLR 2022
Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterrizaton:: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, (1):84-105, 2020.
D Stott Parker. Random butterfly transformations with applications in computational linear algebra. 1995.
Ankit Pensia, Shashank Rajput, Alliot Nagle, Harit Vishwakarma, and Dimitris Papailiopoulos. Optimal lottery tickets via subsetsum: Logarithmic pereeeaaatttiiniionnnn sufficient. arXiv preprint arXiv:2006.07990 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 18))9, 2019.
Md Aamir Raihan and Tor M Aamodt. Sparse weight activation training. arXiv preprint arXiv:2001.01669 2020.
Ilya Razenshteyn, Zhao Song, and David P Woodruff. Weighted low rank approximations with provable guarantees. In Proceedings of the sorty-eighth annual ACM symposium on Theory of Compuling, pages 250-263 2016.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115()):11-2522 2015.
Victor Sanh, Thomas Wolf, and Alexander M Rush. Movement pruning: Adaptive sparsity by fine-tuning. arXiv preprint arXiv:2005.0766: 07683, 2020.
Zhao Song and Xin Yang. Quadratic suffices for over-parametriiation via matrix cherooff bound. arXiv preprint arXiv:1906.0333 2019.
Hidenori Tanaka, Daniel Kunin, Daniel LK Yamins, and Surya Ganguli. Pruning neural networks without any data by iteratively conserving synaptic flow. arXiv preprint arXiv:2006.0546 05467, 2020.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers, arXiv preprint PxXv:2001.00000 2020.
Ann Taylor, Mitchell Marcus, and Beatrice Santorini. The penn treebank: an overview. Treebanks, pages 5-22, 2003.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society: Series B (Methodological), 58(1):267-288, 1996.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucie, et Mlp-mixer: An all-mip architecture for vision. arXiv preprint arXiv:2 16600 2021.
Madeleine Udell and Alex Townsend. Why are big data matrices approximately low rank? SIAM Journal on Mathematics of Data Science, 1(1):144-160 2019.
Keivan Alizadeh Vahid, Anish Prabhu, Ali Farhadi, and Mohammad Rastegari. Butterfly transform: An efficient fft based neural architecture design. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12021-12030 IEEE, 2020.
Chaoqi Wang. Guodong Zhang, and Roger Grosse. Picking winning tickets before training by preserving gradient flow. arXiv preprint arXiv:2002.07376, 2020.
Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, and Song Han. Lite transformer with long-short range attention. arXiv preprint 2rXiv:2004.1188 2020.
Yunyang Xiong. Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystromformer: A Nystrom-based algorithm for approximating self-attention, arXiv preprint Stt::2102.03902. 2021.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. arXiv preprint arXiv:2101.11988 2021.
14

---

Published as a conference paper at ICLR 2022
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang. et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 2020.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. arXiv preprint arXiv:2106.04500 2021.
Zhengyan Zhang, Xu Han, Hao Zhou, Pei Ke, Yuxian Gu, Deming Ye, Yujia Qin, Yusheng Su, Haozhe Ji, Jian Guan, al. Cpm: A large-scale generative chinese pre-trained language model. Al Open, 2:93-99, 2021.
Michael Zhu and Suyog Gupta. To prune, or not to prune: exploring the efficacy of pruning for model compression. arXiv preprint arXiv:1710.01878,, 2017.
15

---

Published as a conference paper at ICLR 2022
A PROBLEM FORMULATION
We formulate the problem of sparse model training as sparse matrix approximation with a simple hardware cost model (Section 2
We first describe our simple cost model for sparse matrix multiplication to reflect the fact that parallel hardware such as GPUs are block-oriented (Cook) 2012 Gray et al., [017]: accessing one single element from memory costs the same as accessing one whole block of elements. We then formulate the sparse matrix approximation in the forward pass and the backward pass. The cost model necessitates narrowing the sparsity pattem candidates to those that are block- aalinedd.
Cost model We model the time cost of an operation based on the number of floating point operations and memory access. The main feature is that our cost model takes into account memory coalescing, where accessing a memory location costs the same as accessing the whole block of elements around it (typically b= 16 or 32 depending on the hardware).
Let Costmem be the memory access cost (either read or write) for block of b contiguous elements. Accessing any individual element within that block also costs Cast time, Let Cost be the compute cost of a floating point operation. Let Nolocmmem be the number of block memory access, and Noop be the number of floating point operations. Then the total cost of the operation is Totalcost Coshmem -Nblocmmem Poosttop--Noog
This cost model is a first order approximation of the runtime on modern hardware (GPUs), ignoring the effect of caching.
Block-aligned sparsity pattern, Block cover, and Memory access cost As the memory access cost depends on the number of block of memory being accessed, we describe how the number of nonzero elements in sparse matrix relates to the number of blocks being accessed. We first define a block cover of sparse mask.
Definition A.1. A sparse mask M E (0,1) is (bobbb)-loockk-aligned for any index i.j where Mij- = 1,
we also have May =1 where:
i' b-bii/hb ( jj/ba] for all 0. 1 and = 0..1,...bb 1.
The (b1,b2) b-bockk cover of a sparse mask M E {0.1)mxx is the (b) boobb)--block-alinne mask M' E [0,1)mxx with the least number of nonzeros such that M.j S M for all i.j.
We omit the block size (b1, ) if it is clear from context.
sparse mask M being b2) block-aligned means that if we divide M into blocks of size b, x b2, then each block is either all zeros or all ones. To get the (b,, (by,by) (by)-blockk cover of a sparse mask M, we simply divide M into blocks of size b, x by and set each block to all ones if any location in that block is one.
For a sparse matrix with sparse mask M on a device with block size b the number of block memory access Nutoekmee is the number of nonzero blocks in its (1,b)- cover M' (assuming row-major storage). This corresponds to the fact that to access a memory location on modern hardware (GPUs), the device needs to load a whole block of elements around that location.
Fast sparse matrices mcans alock-aligned sparsity pattern For sparsity patterns that are not block-aligned, such as the random sparse pattern where each location independently zero or nonzero, its (1,b)-block cover might increase the density by a factor of close to b times (we show this more rigorously in the Appendix). As memory access often dominates the computation time, this means that non block-aligned sparsity will often result is b times slower execution than block-aligned ones. In other words, exploiting hardware locality is crucial to obtain speed up.
Therefore, this cost model indicates that instead of searching over sparsity patterns whose total cost is less than some budget C, we can instead search over block-aligned patterns whose number of nonzeros less than some limit k. For our theoretical analysis, we consider sparsity patterns that are blo))-bloc--alignedd In practice, since we need to access both the matrix and its transpose (in the forward and backward pass), we require the sparsity pattern to be both (1,b) 1ao--booo--ligeeeee and (b,1)- (b.1)-loock-alll aloneea This is equivalent to the condition that the sparsity pattern is (b.b) blob))--oo--llllaeee allinned.
Sparse matrix approximation in the forward pass in the forward pass. That is, we have weight matrix A with input B and we would like to sparsify A while minimizing the difference in the output. For easier exposition, focus the case where number of nonzeros in each row is the same.
We now formulate the sparse matrix approximation
16

---

Published as a conference paper at ICLR 2022
Definition A.2 (Forward regression). Given four positive integers m m2n2d2221 matrices A A Rm and B E Rdxx The goal is to find a subly-look--llinnee binary mask matrix M E {0,1) m=d that satisfies
min
|| [A.B-(Ac.M) (AoM)-Bll
(0
s.L. |M
where M, the i-th row of M.
Sparse matrix approximation in the backward pass In the backward pass to compute the gradient wrt the weight matrix A, we would like to sparsify the gradient CB while preserving as much of the gradient magnitude possible.
Definition A.3 (Backward regression). Given four positive integers m mn2n2d d2 k 1, matrices B B g and C E mm The goal to find bllikkee binary mask matrix E (0,1) mxxd such that
min (0.1))) mxd
||C.B" (C-BB)oMM
s.t. |M k.
where M, is the i-th row of M.
Without making any assumptions, such problems are general computationally hard Foster et al. (0155) Razenshteyn et al. (2016).

---

Published as a conference paper at ICLR 2022
B ANALYSIS OF BUTTERFLY VARIANTS
We present formal versions theorems in Section regarding variants of butterfly matrices. We provide full proofs of the results here.
B.I BLOCK BUTTERFLY ANALYSIS
Proof of Theorem 4.7
Let M be an nxx block butterfly matrix with block size b. We want to show that M also has representation as an n x block butterfly matrix with block size 2b.
By Definition
3.3 MM has the form:
8 6
(.0) (1)) .....
M B
B
Notice that we can combine that last two terms to form matrix of the form B
26)
20
(see Fig. 3). Moreover,
other terms in the product of the form B ((b) the form:
can also be written as B
3r
26)
(see Fig.
3 Thus M also has
26
M
=
26
B
28
26
B.
(62))
26
In other words, M is also n block butterfly matrix with block size 2h.
Proof of Corollary 4.2
Dao et al. (2000) Theorem 3) states that any n x sparse matrix with nonzeros can be represented as products of butterlly matrices and their transposes, with ((slogn) parameters.
For constant block size that is power of 2, the set of x block butterfly matrices of block size b contains the set of regular butterfly matrices by Theorem 4.1 Therefore any such sparse matrix also has a representation has products of block butterfly matrices of block size b and their transposes, with (slogg) parameters.
B.2 FLAT BUTTERFLY ANALYSIS
We prove Theorem 43 which relates the first-order approximation the form of flat butterfly matrix with the original butterfly matrix.
Proof of Theorem 4.3
n and let Br,...,,, paxn be the m butterfly factor matrices (we rename them here for simplicity of notation).
Let
E I [I+AB))
I+ ++> XB
Our goal is to show that || Sc. We first recall some properties of Frobenius norm. For any matrices and C, we have ||AC|]=-[|A]]
and || +|CC|]
Expanding the terms of the product E, we have
m E
I
B
18

---

Published as a conference paper at ICLR 2022
Using the above properties of Frobenius norm, bound E:
|
Σ
II
o6 1m,,
Bmax
mm (B
(AmBmax))
(ev7)
<ee (cvE)
i=0
< de
1-cVè
Se,
where in the last step we use the assumption that e 1 2-
We now bound the rank of the first-order approximation.
Proof of Theorem 4.4 Let =++ Em= AB₁. Note that any entry in AB, has absolute value at most
B 1
mAB
Bmax 4'
where we use the assumption that B:00 max < Biax and cs 7.
Thus any diagonal entry in M* has absolute value at least 1- - and the off-diagonal entries are at most eyeBBox. bm
Alon (0099 Theorem 1.1) states that: there exists some e>> such that for any real M ERRx if the diagonal elements have absolute values at least & and the off-diagonal elements have absolute values most where
2/a S then rank(M) Plog1/e
cloga
Applying this theorem to our setting, we have that
Biaax Boo max
rank(M*)
elog
Biax DSS
We just need to show that
B Banax
to satisfy the condition of the theorem.
2even
Indeed, we have that 1 Bauax
< v 2n as each |B:l 2n. Combining the two conditions on
Bmax B maxx.
we have
Bec
shown that 1 Baaux <2c en. This concludes the proof.
mice
B.3 FLAT BLOCK BUTTERFLY + LOW-RANK ANALYSIS
We show that flat butterfly + low-rank (an instance) of sparse + low-rank, is more expressive than either sparse or low-rank alone. We adapt the argument from Chen et al. (202T) show a generative process where the attention matrix can be well approximated by a flat butterfly + low-rank matrix, but not by sparse or low-rank alone.
We describe here generative model of an input sequence to attention, parameterized by the inverse temperature BE R and the intra-cluster distance A E RR
Process 1. Let Q @Rnxx, where 2 2/2lg//3110 (n)), with every now of Q generated randomly as follows:

---

Published as a conference paper at ICLR 2022
1. For C=Q(n), sample C number of cluster centers C1 C1.... E. independently from (NO,Ja/ Vd. For each cluster around C₁ sample n, = number of elements around Co of the form 25j = C + Tij
for j = I,., where Tij ~ (N(
Assume that the total number of elements is π. = cb and
A O(1/log
Let A=QQ" and let Let Q be the matrix whose rows are the vectors Zij where i= i= and j=nj-....... the attention matrix be Mx=cxp(B-A)
Theorem B.1. Let Mg, be the attention matrix in Process ! Fix (0,1). Let R E Rexe be matrix.
Consider low-rank, sparse, and sparse + low-rank approximations M3. Assume (1 (1-A7)logn S 535 Ollggn..
Flat butterfly + low-rank: There exists a allatt butterfly + low-rank R with n +o()) parameters with [M Rl| Sen.
2. Low-rank: R is such that n rn-rank(R)= then \| R|
3. Sparse: If has sparsity o(n"). then M| =
Proof sketch. As the argument is very similar to that of Chen et al. (2021 Theorem 1), describe here the modifications needed to adapt their proof.
The main difference between our generative process and that Chen et al. (021) is that each cluster has the same number of elements, which the same as the block size. The resulting attention matrix will have large block diagonal component, similar to that Chen et al. (202T). However, all the blocks in the block diagonal component has the same block size, which block diagonal component of block size b. Therefore, this flat block butterfly matrix plays the same role the sparse matrix in the proof of Chen al. (202T). The rest of the argument follows that of theirs.
Moreover, a flat block butterfly of block size b contains a
20

---

Published as a conference paper at ICLR 2022
Roadmap The analysis of sparse networks is organized as follows. In Section We list some basic notations that will used. In Section D we consider the problem adding sparsity on W. and we achieve polynomial solving time. In Section E we prove that the gradient descent can be done fast under the sparsity assumption. In Section G we consider the problem of adding spursity on a, and we show that minimizing the dropout loss equivalent with kemel ridge regression problem. In Section H we analyze the dynamics of gradient flow and prove the convergence result.
C NOTATIONS
For a vector æ, we use | to denote its Lp norm, and we mainly consider 1=1 in this paper. For a matrix A, we ||A||o.I|A||a.I|A ||||| For two matrices A, E RRdxm we use Ao B to denote their Hadamard product. We use Taar mma (n.d.m) to denote the time of multiplying n x matrix with another d m matrix. Amin(A) to denote its minimum cigenvalue. We also let vec(A) be the vectorization of matrix A in column first order. We use (+-) to denote standard Euclidean inner product between two vectors.
denote the &o norm, entry-wise & norm and Frobenius norm of A respectively.
a symmetrie matrix A, we use
Moreover, we use N (H,E) to denote the Gaussian distribution with mean and covariance E. We denote the ReLU function by 6(2)= max{2,0) For an event E, use { or g to denote its indicator function.
D SPARSITY ON HIDDEN LAYER WEIGHTS
D.1 APPLYING MASKS BEFORE MULTIPLICATION
Given matrix Raxd B R naively computing AB takes Tmaa (n,d,.... Note that, we can also consider the case where A and B have different size. For simplicity, let us consider the case where matrix A and matrix B have the same size.
Our goal is to find "optimal" binary mask matrix W E (0,1)dxx such that,
min
s.t. B
Remark D.1. In the practical applications care about, the function is the activation function of neural nenwork, e.g., maxU(())mxx(20))
We define a sparse targeted regression problem:
Definition D.2 (Sparse mark regression, & version). Given a matrix e Rdd,, and a vector a E Rd the goal is find k-sparse binary vector E (0,1)d to minimize the following problem:
min| a B B-(a" (aTow) )w)).Blls
Naively, the above problem can be solved in n Olk) via guess all the (+) choices.
Lemma D.3. The targeted sparse mask regression problem can be solved in n O(k)
Proof. We need to guess (+) times, which becomes d0k)) Each time takes nd operations, thus the total time is
(k) -n-)
Definition D.4 (E, version). Given three positive integers m ' 2d matrices TRuxd and Beldxnn We define our problem as finding the binary matrix W E mox that satisfies
min A- (AaW)-Bll)
s.t. |W
where Wis is the i-th row W. W.
Theorem D.5. The problem being defined Definition D.4 can be solved in (lk) time.
Proof. Our problem can be decomposed into m sub-problems as follows:
21

---

Published as a conference paper at ICLR 2022
IA-B-(doW)--ll= (01 ((ADW))--((AAW)BB))))
m
| B (-Aoo) Bl1
m
A B- (A BBI Σ)
where A,e means the i-th row of matrix A. By applying Lemma D33 each sub-problem A.. B B- (A,, W..) Bl11
min
Wi-
can be solved in n-d (00)) time.
Then the problem defined in Definition D.4 can be solved in m-nd0(k) = (k)
time in total. Thus we finish the proof.
In the above Theorem, we show that solving the sparse mask regression problem is NP-hard. However, if we add some mild assumptions and consider minimizing & norm, then we can solve the regression problem in polynomial time, as the following parts show.
Definition D.6 (& version). Given a matrix dxn and vector a a the goal to find a k-sparse binary vector w E (0,1) to solve
min minn|aa . TB- ow )w)).Blk BB11
Lemma D.7. The targeted l version sparse mask regression problem can be solved O(nd+nlogn)
which polynomial time.
Proof. We first consider the situation when a E (0,1)d. In this case, have
||a - (B- ow ) 11|2+11(e ow ) BB]
where T -BI11 is fixed. So we only need to consider the following problem:
max]((a ow ") ).B).
For simplicity we assume a, 1,Vie [d], and we only need to solve
max||ju T. BII
where w has & elements equal to 1 and d elements equal to 0. For E [d]. we compute S ンドー Bij which is the summation of i-th row of B, and sort them as S(1) 2 S(2) ... 2 S(n)- Then we only need let w(a for E [k] and other elements equal to O. Computing all S; takes ((dd) time, sorting S, takes (nlogn) time, thus the total time consumption is Ond
(nd++nloon) nlogn) in this case.
Next, we consider the general case when a E IRd 20-
We let
1,
0; 0=
=a=aBBi.. and a.
Vield]
0,
a==0'
where Bix is the i-th row of B. Then our optimization problem is equivalent to
min||a B (a B-(uu
ow B
w
where B Bedd" and ae (0.1)4, Thus time, thus the total time consumption also +lndn)) this case.
this case into the first case. Constructing B and takes ((nd)
Definition D.8 (e, version). Given three positive integers 2 2 2d2221, matrices A RRaodd and We define our problem as finding the binary matrix E m0,1]mxd that satisfies (A- W
BERdxn 20
min
s.l.
where Wis is the i-th now of W.
Theorem D.9. The problem being defined as Defimition D.8 can solved
malogn))
mmllgg)
time.
22

---

Published as a conference paper at ICLR 2022
Proof. Our problem can be decomposed into sub-problems follows:
7
| B 1 (l-B-(Ao)).Bl,
- ((AoW))BB))
7
>1 A/w (Aco)) B||
f=1
> Al. - Bll-B-AAi.W....
where Ai. means the i-th row of matrix A. By applying Lemma D.7 each sub-problem
min B (w
can be solved in +logn)) time. Then the problem defined in Definition D.8 can be solved in
m- mlognn+ - 0 (mt:
time in total. Thus finish the proof.
D.2 APPLYING MASKS AFTER MULTIPLICATION
Definition D.10. Given matrix B BeRdxy, CERRxn, The goal find mask W [0,1]mxd where each column of W is k-sparse
-CC-B- -CC-B JoWI
min We(0.1 moxd
Remark D.11. The B defined in Definition D.4 is the same as the defined in Definition D.10 is corresponding to the X in the neural network setting.
E GRADIENT COMPUTATION
this section we consider neural network with one hidden layer and m neurons this hidden layer. Suppose E Rd is the input, W = (wi,"" (w) output weight, and M E [0,1]axm is the mask matrix with each column having at most & non-zero entries. The neural network f :Rd > R is defined as
) E Rdxx is the weight matrix of the first layer, a I is the
f(r)=a ((MOW))T..)
For simplicity, only optimize W and fix Consider the mean square loss
1 L(W)= = (1()))) v1) Σ (a (o(MoW) -x:)) -
In the forward computation, for batch of data points on,", EIRd, let X E Rdxx denote the input data points matrix. For convenience, we define
OL(W(())
aw(t)
where 17 is the step size. We define function 9e .R. F.
( v) (AMoW((())) =2)-aa =2)3-
and also denote (x) (=)))))))))) mR
Lemma F.1. We can express AW(() as
Aw(t)=-n(X* ot (X)XXXXTTX)OMM)
and each column of (I has at most k non-zero entries.
23

---

Published as a conference paper at ICLR 2022
Proof. From the definition, know (L)
AW(()= = -1 (W(()
Σ ((a))))) )diigi ((Moow(()) xx
M
mxm
mx1
xd
dxm
= =-
9e (x ToM oM
M
dxn
dxm
Since each column of M has most l non-zero entries, we easily know each column of AW also has most K non-zero entries.
Lemma E.2. Suppose that matrices M E Rd W(t) E Rdxm and AW(t) E R dxx are given and pre-computed, then can compute (x+1(X) in
(mmnk
time. (Here fi+1(X) is the evaluation of at (t++1).)
Proof. The goal is to compute
f++(()= T
MMM (t+1)) -X).
(
dxm
dxm
By using Lemma E.
we have
(M -X (M)) -X
(MOOW() . -X
(X))oM) (MMOW()) (X))))) -X
(M) (TXXX-T(XXX)) (X) T -X
Notice that we have already computed (X to compute (AW(I)) ( -X where (W E IR and X E Rdxn, By using Lemma E. each row of (A) has at most & non-zero entries, thus we can compute (AW(t)) in ((mnk) time.
from previous iteration, so we only need
Lemma E.3. Suppose that matrices M Eg Begdxm,WWee ( and ft(X) are given and pre-computed, then compute OL(W(U) in O(mnk) time.
8W())
Proof. By using Lemma
we have
(t(W(t)) = (X 9. (X))MM
OW(t)
where de (z)= (f(x)=
xx) ae T and ge (9))=(((2))))))))))))))) (z,))) m Rmx We
first compute (t) in ((m) time, then we can construct (x E 1m in n O(mk) time. Given ge (X), since we only need to compute km entries of X X-9. (X), where each entry can computed ((n)
al((W()) in O(mnk) time. (W(
time, thus we can compute
24

---

Published as a conference paper at ICLR 2022
Algorithm The sparse training algorithm
I:
procedure SPARSE TRAINING( (6. (ej)))
2:
Initialization (,) (0))000 ro [m].
3:
fo do
4:
/oorward computation
5: Compute McW(()
Takes (rakk) time.
6: for + do
7:
fx(2,) 6 o (() -xi)
Takes (raak) time.
8:
gr. [z)) (f(x;) 91) diago' ((MoW(t)TT -x;)
Takes (mk) time.
9:
end for /bbacwward computation
11:
(x(X) - (0t (m1), (x (znn).
12:
BL(WO) (X 2e (X)) (F
Takes ((mnk) time.
13:
W (t)
> AW = -7 (W
al:W())
14:
end for end procedure
F NEURAL TANGENT KERNEL, CONVERGENCE, AND GENERALIZATION
Our analysis relies the neural tangent kernel (NTK) (Jacot et al. 2018) of the network.
Definition F11 Let f--0): Rdd + R be the finctionn specified by neural network with parameters 0 R and input dimension d The parameter 9 initialized randomly from a distribution P. Then its neural tangent kernel (NTK) (Jacot et al. 2018) is kernel K: R. xI +R defined by:
(x())) 0f
K(x,y)=
IE 0-P
00
a0
We relate the training and generalization behavior of dense and sparse models through their NTK. The standard result (Song and Yang. 2019) implies the following.
Proposition F.2. Let Scense denote ReLU neural network with L layers with dense weight matrices Otense with NTK Kaensos and let Saparse be the ReLU neural network with the same archûtecture and with weight matrices Osparse whose rows are k-spurse, and with NTK Ksparss Let I1, be the inputs sampled from some distribution Px. Suppose that the empirical NTK matrices Ka Keense (21,315) and K2 (x for ( satisfy [a S..
Training, We knew the the number of iterations of dense network is >uni -2 n log(11c) reach the training loss. For sparse network we need (Amin (Ka)-50 2)-22noog(1/c))
Generalization. We knew the the number of iterations of dense network is Amin (Ka) 1i= to reach the generalization error truining loss. For sparse network we need (Amin (Ka) )) -) 2n 1/e).
These results relate the generalization bound of sparse models to that of dense models.
25

---

Published as a conference paper at ICLR 2022
G DROPOUT NEURAL NETWORK AND KRR
We consider a two layer neural network with ReI.U activation function, and write
f(W = Vm Σ a x)=
V Σ W, x1 =1
(2)
where w,(00)~ (N0)~~N(O,,))) (0,1a) E R., ~ - and all randomnesses are independent. We will fix a, during the training process and use 1 normalization factor, both of which are in the literature of Du et al. v/m
(2019); Song and Yang 2019); Brand et al. (021)
Suppose the training data are (21.31) ....i)............ (Zn. IR x R, we define the classical objective function as follows:
i(W) (f(W,z))
The gradient with respect to loss function L is
aï ow,
a m E
We consider the effect of dropout on network training. For each r E [m], we introduce the mask by defining random variable o, as follows:
0,
with probability 1 1-q:
(1/9.
(1/g.
with probability
is easy to see that ar and o, are independent for any i.j. then B -l,,o)]-[(o, Le Let (0.."..mm)) define our dropout neural net as
We assume
F(W,z,) = Vm Σ 0 ow x)
Vm Σ w. xlw
(3)
Dropout explicitly change the target function, since we need minimize the l distance between F(W,x,or) and y, instead of f(W,z) and y. Formally, we define the dropout loss as
B L(W): = 2
[(P(W..,))
(4)
We first give an explicit formulation which also shows the difference between L and L.
Lemma G.I. The dropout loss defined in Eg. 4 can be expressed as the sum of classical loss L and a regularization tern as
44 m
L(W)= -W)=i(W)++
(w. xr.)2
(5)
2mg
Proof. Since EE|r,] we have
m
a, ol w x)
G w xx = -(WWx)
(6)
ym
m
26

---

Published as a conference paper at ICLR 2022
holds for any iE [n]. Next, show the difference between L and L:
-
a
E
(FW...0))
- (f((.x:)-u)) vi)²
L (F(( -(((W=))uu))
E F[W,2,,0)]
ar, a,,, 2;)(,, (
E on an x:) n..r
72
m
a w 23)
m
w x,,2
(7)
where the first step follows from definition, the second step follows from the linearity of expectation, the third step follows from Eq. (6). the forth step follows from expansion, the fifth step follows from Elo,, Ør.] for T1 F1722 and E B'o,, and the last step follows from a² =. Thus have
a m
L(W)= L(W)=I(W)+
ow x;)²
2mg
and finish the proof.
Before we move on, we introduce some extra notations and definitions. We denote
Wi
yi
W2
V/2
W =vv)<(W)= =
ERnd
and Y=
ER"
:
Wmm
3h
Definition G.2. We define matrix G E Ipaxn which can be viewed as a Gram matrix from a kernel RunUtunntion
associated with ReLU function as follows:
G (X) (0.1)
[x] xj w 2.20, 20, vijjEm]](x)
and assume Ao = Ammin (G
Definition G.3. We define the masked matrix (x,) as
((11,0)
① (22.0)
Pw (X,0)= mm
((Tn
a2o 1
(m
02021 (202
am 1 (um 2
0101
(w) 2002
0202 (2222)20
am (wo 22)) 22000000000
Vm
010 (we a. 20
U202 (wi 2000
amom furm de 2000
and also define the unmasked matrix $ w (X)
aj 1 (m |202
azl (22.
am 1 (wmm 1200
22100
a2 (2100)
Om I (wm 22)) (wm 22002
Tw (X):=
V
0g (wzzen)
On (wm 20
'According to Theorem 3.1 Du et al. 2019) assumption holds when u, not parallel with xy if which is reasonable in reality.
27

---

Published as a conference paper at ICLR 2022
Definition G.4. We define the masked block diagonal matrix Ww (X,
&R as (x(o))):= diag (/1.
where Vrr [m], eh E: Radd is defined as
V= o Σ Diay 1 (2, o,
Iaa; 1 aum,
1=1
We also define the unmasked block diagonal matrix Lw (X) E Rmddddd as
T (x)(x)= == m
diag 2/ (2.
where Vr V. EIRad is defined as
Σ saa 2:,20)200
Lemma G.5. It is easy to verify that
tw (X,r) = (x))).D Do and (V
where
Dg diag(o,,.... 01 Ome pmdxmmd
For convenience, we will simply denote Ωw Lw (x, and ww = (X,,,, Then by using the above notations, we can express our dropout loss as L(W)- = Ee u W - Lemma G.6. we denote A= 1-1 then we have
L(W) - - pow W -V 1 > W W.
Proof. As for the first term, we have
n
m
w W
2
12
a tw a:100 -w, 10,-v11 y1)
m
1
a W. -v1)
As for the second term, since Vw is a block diagonal matrix, we have
1 m
W
W
W
w
11 220) wr))
m
1=1
1 m
(w.) xx, 22)T 1 12u0)))220)
m
pw xz)² 2
1=1==
Thus by using Lemma G.I we have
1 1-1 4 n m
I(W)= P(W)++
(w, xz
2mg 1=1
1
-Y
W V W
and finish the proof.
28

---

Published as a conference paper at ICLR 2022
Remark G.7. A classical kemnel ridge regression problem can be defined
min min
1 2
W
where o:IR. F is a feature map. Note that Lemma G.6 breaks the dropout loss into two parts: the first part an error term, and the second part can seen as a regularization term. Thus the task of minimizing the dropout loss L(W) is equivalent to a kernel ridge regression (KRR) problem.
29

---

Published as a conference paper at ICLR 2022
H DYNAMICS OF KERNEL METHODS (CONTINUOUS GRADIENT FLOW)
The NTK also allows us to analyze the training convergence of sparse networks. We show that gradient descent converges globally when training wide sparse networks. This convergence speed is similar to that of dense models (Du et al. 2019 Allen-Zhu et al. 2019b).
In this section we will discuss the dynamics of kerne method under the mask o, which adds sparsity in the output layer. Our problem will be considered in over-parameterized scheme. First we introduce some additional definitions and notations. We define symmetric Gram matrix G(W) as G(W): w w. E ER For all i.j E {mn] x [ni]. we have
G(W) ky
a
2
m I; Tz
(un .a) (10,,) 2,)
T7
r=1
We define block symmettic matrix H(W) as H(W) = w- -ew E Rmddxmd Then for all i.j E [mz] x [m], the (i-j)-th block of H() is
H(W,, 1
Tk 1 (wn. Raxd a,a,
m
1=1
By using Lemma G.6 we consider the corresponding kemel regression problem:
$ -YW-YR
1 -
W
W
(8)
W
W
2
where E DEER" 1unmnd, W mR Y T and V/ E L.dddmmd The main difference from neural network is that we assume T (related to NTK, e.g., see Definition G.3 and ; (related to regularization term, e.g., see Definition G.4) do not change during the training process.
The gradient of Lę can be expressed as Lw
T ΦW +
(9)
We use W to denote the optimal solution of Eq. B and it satisfies
V WW) w-w-
+ A $ =0
(10)
Since V is positive diagonal matrix, +- exists, thus we have
W ($ -
Next, we consider the question from continuous gradient flow aspect. In time t, we denote W(I)- (t), L wo (t()) w(t) We also denote G(t)=G(W(tt) and ((t)=H(W(()) Following
the literature of Du et al. 2019), we consider the ordinary differential equation defined by
dt,(t) (W)
(11)
dt
(w,
Lemma H.1 (Lemma 3.1 in Du et al. 2019)). Ifm= If
5 log(()) ())), we have with probability at least 1-5, 8,
|[G(O)-G*))) G0
Aa 4 and Amin (G(() Ao
Lemma H.2 (Lemma 3.2 in Du et al. (2019). If W'. are iid generated from (O, then with probability at least 1 - 8, the following holds. For any set of weight vectors wi." ,", Wm.m E RRd that satisfy for any r [mai,]]uu l|,, - w, (0)ll2 < odda for some small positive constant then matrix G E idddd satisfies ([))G(0))22 b and Amin (G) Au
The above lemma shows that for W that is close to W(O), the Gram matrix also stays close to the initial Gram matrix G(D), and its minimal eigenvalue is lower bounded.
Lemma H.3 (Gradient Flow). If we assume Amiia ① 2 Ag > 0, then with probability at least 1 --6 5, for
25 Wi."" Wm E Rd that satisfy Vr [m], - w. (,(())]2 < n²
we have
d क 11
< क $3.12
dt
holds some constant 7 >>00
Proof. By using Eq. 9 and Eq.
dW we can express dr as
dW
(12)
30

---

Published as a conference paper at ICLR 2022
Then we have
BW $W*1
dt
d|ooW $ dW dt -WW (- OW - (
dW
-2( -W
16 DW y
-- $W
TT BW- -VW XWW)
2 $
T
$W
22 का
WW VW*
< 2X $W $ 2 T (W
(13)
where the second step follows from Eq. 12
the fourth step follows from Eq.
10 and the last step follows
from the definition that Àg Amin (G) Amii T)
As for the second term in the Eq. 13
we have
2X(W W
ON W.)
2 T
W
2 2AAu W W
( W-)
2AAo PW ΦW*
(14)
Thus by Eq. (13) and Eq. (14)
we have
@W $W*
+2
dt By letting 2 we finish the proof.
For convenience, we denote u(t) = (t)t))((t)ERR Then it is easy to verify that
1
(t( /m- aa x;)) f(W()))))))) Vie(n],
showing that u(t) is the prediction in time
Lemma H.4 (Convergence rate). we assuaee Amin (G(s)) Ao holds for SSsst, then we have
1. lu(t) - Y (A0+22/m)
-Y
2. Vr [m],] [,, (f) (0)122 < (
Aoymm
Proof. From Eq.
(9) we can express the dynamics by using u(t) as du(t)
dt
-())
(15)
Thus we have
dl|u(()) -Y 2(u(1)- (((())-))(((0000000000 (G)))
dt
= 2((( (((()))))G(()(((())-)))))))))))1 - 22((((t)((((---)2 -Y) DVW
G(()(u(t)-Y),
< - - -a(((t))(tttaaee (T)(I)))YTVWV..... T@WW.
(16)
As for the second term, we have
21
2X(u(t) Y) T
(u(t)- - -Y) T [ W1, Fm wmnl
2
(u(t)- YY) T
T
I:o(u xx). > (Io(wT
m
2
(u(()-)) [ (U,(0)).
(17)
m
31

---

Published as a conference paper at ICLR 2022
where for [n], (t) ER can be expressed
m
Uy(t)=
Or (1m2))200
Exzp(ui mi))
Vm
1
m
n
3. (2)00 ) (10,
m
ar W. 1 (un
(w, 22)
m
We denote U())=
E and have
2\
2A(u(t)-Y) DVW -V) -U(t)
(18)
and our dynamics becomes
d -
21 (u(t)-Y)
- -ollu(tt) r
m
$ - (Do 21
[11u(t--11)
(19)
showing that at, respect to f, and we have
(e(Ao+2A/m)) ((()) -Y -1))=000 Thus (A +22/mJt (((t) - - 1|² is a decreasing function with 5.
l2(t) (u(t)-Y(l -e [[((0) -
As for bounding (t) we use the same method as in Lemma 3.3 of Du et al. (2019). Thus we complete the proof.
Finally, by combining Lemmal
H.l
and H.4 we have the following convergence result. then with probability
Theorem H.S (Convergence gradient flow). Suppose λo > 0, m m
at least 1 - ô over the randomness of initialization, have
|u(t) Y <e (Aa+2A/m)) [u(0) Y
The above theorem shows that in the over-parameterizee setting (when m is large enough), the training loss of the kernel ridge regression problem define in Eq. ( converges
in linear rate. By comparing Theorem HH.5 with Theorem 3.2 in Du et al.) (2019). we can find that the introducing of regularization term makes the convergence speed faster, though the improvement is limited. Further notice that in Section G we prove the equivalence between minimizing the dropout loss and the kernel ridge regression problem. So we conclude our results as:
The introducing of sparsity into neural network makes the convergence speed faster, but the improvement is limited due to the over-parameterized scheme.
32

---

Published as a conference paper at ICLR 2022
I
METHOD DETAILS
We describe some details of our method.
1.1 COMPUTE BUDGET ALLOCATION
We describe here a procedure to compute the budget allocation based on our cost model. This procedure is more complicated than our simple rule of thumb in Section 3.3 and tend to produce the same allocation. For completeness, we include the procedure here for the interested reader.
Given parameter budget B, we find the density of each layer type that minimize the models' total cost of matrix multiplication. For example, in Transformers, let d, and dm be the density the attention and the MLP layers. Let s be the sequence length and d be the feature size. The attention layer with density da will cost de (n²+nd)) and the fully connected layers with density dm, will cost 2dm We then set d, and dm minimize the total cost while maintaining the parameter budget:
minimizes, S., 8 (n² 2 nd) + n2 subject to # of trainable parameters S. B.
(20)
As this is a problem with two variables, we can solve it in closed form.
1.2 LOW-RANK ATTENTION
Section 3.3 we describe how to use the sparsity pattern from flat block butterfly and the low-rank term for
weight matrices. This applies to the linear layer in MLP and the projection steps in the attention.
We also use the sparse low-rank structure in the attention step itself. CCenn al. (2021) describes general method to combine sparse and low-rank attention, where one uses the sparse component to discount the contribution from the low-rank component ensure accurate approximation of the attention matrix.
We follow a simpler procedure, which in practice yields similar performance. We use a restricted version of low-rank of the form "global" sparsity mask (as shown in Fig. 12]. Indeed, sparse matrix whose sparsity pattern follows the "global" pattern is a sum of two sparse matrices, one containing the "horizontal" global components and one containing the "vertical" components. Let w be the width of each of those components, then each of them has rank at most w. Therefore, this sparse matrix has rank at most 2w, and is low-rank (for small w).
We also make the global component block-aligned (i.c., set w to multiple of the smallest supported block size such as 32) for hardware efficiency.
I.3
COMPARISON TO OTHER SPARSITY PATTERNS FOR ATTENTION
In the context of sparse attention, other sparsity patterns such as BigBird and Longformer also contain a "global" component, analogous
low-rank component. Their "local" component is contained in the block diagonal part of the flat block butterfly sparsity pattern.
The main difference that do not the random components (c.g., BigBird), and the diagonal strides from flat block butterily are not found in BigBird or Longformer. Moreover, we apply the same sparsity pattern (+ low-rank) to the linear layers in the MLP and the projection step in attention as well, allowing our method to target most neural network layers, not just the attention layer.
I.4 SPARSITY MASK FOR RECTANGULAR MATRICES
We have described the sparsity masks from flat block butterfly for square matrices. For rectangular weight matrices, we simply "stretch" the sparsity mask. The low-rank component applies to both square and rectangular matrices (as shown in Fig. 10p. We have found this to work consistently well across tasks.
33

---

Published as a conference paper at ICLR 2022
Square Flat Block Butterfly
Rectangular Flat Block Butterfly
Figure 10: Sparsity Mask for Rectangular Matrices.
J BENCHMARKING OF BUTTERFLY MULTIPLY
We validate that flat butterfly matrices (sum of factors) can speed up multiplication GPUs compared to butterfly matrices (products of factors).
Consider the matrix M E = that can be written as products of butterfly factors of strides of up &s (a power of 2), with residual connection:
M ( (3) (T+ h/2 (T
The first-order approximation of M has the form of a flat butterfly matrix with maximum stride is (Section 3.2
Mn - + B (n) (n)) (
4/2
Notice that .M a product of logzk factors, each has 2n nonzeros, so multiplying M by a input vector costs (nloog)) operations (by sequentially multiplying x by the factors of M). The flat version Miac is a sparse matrix with O(nlogk) nonzeros as well, and the cost of multiplying Ma is also O(nlogk). However, in practice, multiplying M is much more efficient on GPUs than multiplying Mr because of the ease of parallelization.
We measure the total time of forward and backward passes of multiplying either Ma and compare to that of multiplying Mr for different maximum strides, shown in Fig. 11 We see that "flattening" the products brings up to x speedup.
Speedup (X)
2
16
32
Maximum stride
Figure 11: Speedup of multiplying Ma compared to multiplying Mr. Flattening the products yields up 3x speedup.
We use matrix size 1024 1024 with block size 32, The input batch size is 2048. We use the block sparse matrix multiply library from /tttps://itthie hwtp://wwwggiiieeeee//oooooooooa The speed measurement done on V100 GPU.
34

---

Published as a conference paper at ICLR 2022
Local
Global (Low-rank)
Butterfly
Random
Figure 12: Sparsity pattern candidate components: Local corresponds to local interaction of neighboring elements; Global (low-rank) involves the interaction between all elements and small subset of elements; Butterfly captures the interaction between elements that are some fixed distance apart; Random is common in the pruning literature.
K EXHAUSTED SEARCHING SPARSITY PATTERNS FOR EFFICIENT SPARSE TRAINING
We describe here our early exploration of searching among different sparsity patterns that has been proposed the literature. We use a metric derived from the NTK, which has emerged as one standard metric predict the training and generalization of the model. We consistently found the butterfly low-rank pattern to perform among the best.
In Appendix K.
we describe the challenges of selecting sparsity pattemss for every model components using the metric derived from the NTK, followed by our approaches. Then in we describe details of empirical NTK computation, which is an important step in our method implementation. Last, in Appendix K.3) we highlight important properties of our method it rediscovers several classical sparsity pattems, and the sparse models can inherit the training hyperparamters of the dense models, reducing the need for hyperparameters tuning.
K.1 CHALLENGES AND APPROACHES
Challenge 1: We seek sparsity patterns for each model components that closely mimic the training dynamics the dense counterpart. As mentioned in Theorem D.9] it is NP-hard to find the optimal sparse matrix approximation. Although NTK provides insights and measurement on the "right" sparse model, bruteforcely computing NTK for one-layer models with all sparsity patterns still infeasible.
Approach 1: Sparsity Pattern Candidates. To address the above challenge, we design our search space to be a limited set of sparsity pattern candidates, each is either a component visualized in Fig. [2 the combination of any two of them. These components encompass the most common types of sparsity pattern used, and can express We provide the intuition behind these sparsity components:
Local: this block- diagonal component in the matrix corresponds to local interaction of neighboring elements. This has appeared in classical PDE discretization (Collins and Angel 197T), and has been rediscovered in Longformer and BigBird attention patterns.
Global: this component involves interaction between all elements and a small subset elements (i.e., "global" elements). This global pattern is low-rank, and this sparse + low-rank structure is common in data science (Udell and Townsend 2019), and rediscovered Longformer and BigBird patterns well. Butterfly: this component corresponds to interaction between elements that are some fixed distance apart. The many divide-and and-conquer algorithms, such as the classical fast Fourier transform (Cooley and Tukey 1965), uses this pattern at each step. Butterfly matrices reflects this sivide-anddoonqueer structure, and hence this sparsity component. The sparse transformer (Child et al. 2019) also found this pattern helpful for attention on image data.
Random: this component is generalization sparsity patterns found in one-shot magnitude, gradient, or momentum based pruning (Lee et al. 2018). Note that at network initialization, they are equivalent to random sparsity.
Challenge 2: Even with fixed pool of sparsity patterns for each layer, if the model has many layers, the number of possible layer- pattern assignments exponentially large.
Approach 2: To further reduce the search space, we constrain each layer type (attention, MLP) to have the same sparsity pattem. For example, if there 10 patterns and 2 layer types, the candidate pool is 10 100 combinations.
35

---

Published as a conference paper at ICLR 2022
Algorithm 2 Model Sparsification
1: Input: model schema S2, compute budget B, dataset subset X, sparsity mask candidate set C.
2: Kdense (fo.X)
Eq.
3: output sparsity mask assignment Sours dmii f- inf
C1 4: for .M.....AA E do
Enumerate all sparsity mask candidate combinations
5: Let be the sparsity mask assignment (f1775)74;;m4) + M,
6:
if (oTooaaoommttes) < then
Eq. 21 Check masks satisfy budget constraint
7:
Let M, be the flattened sparse masks
8:
Kaparac <NNK -NTK(JReMa* (NGM...)
9:
d, 4 D--DISTTNCCE deense, if dymi >d then
Eq.
22
10:
11:
dmmin & Sout 8-88
12:
end if
13:
end if
14:
end for
15:
return Sout
Return sparsity mask assignment
Challenge Computing the empirical NTK on the whole dataset is expensive time and space, as it scales quadratically in the dataset size.
Approach 3:
We compute the empirical NTK randomly chosen subset of the data (i.e., principal submatrix of the empirical NTK matrix). In our experiments, we verify that increasing the subset size beyond 1000 does not change the choices picked by the NTK heuristic. The subsampled empirical NTK can be computed within seconds or minutes.
K.2 ALGORITHM DESCRIPTION
Our method targets GEMM-based neural networks, which are networks whose computation is dominated by general matrix multiplies (GEMM), such as Transformer and MLP-Mixer. As a result, we can view the network a series of matrix multiplies. We first define:
Model schema: a list of layer types t (e.g., attention, linear layers in MLP), number of lay- ers r of that type, and dimension of the matrix multiplies m n. We denote it Ω
A mask M of dimension m x is a binary matrix {0,1)mxx. [ The compute of a mask is the total number of ones in the matrix: compute( (M) -iG
A sparsity pattern Paxn for matrix dimension mxn is set of masks [Mip....MMppb each of dimension
m x n.
A sparsity mask assignment a mapping from model schema Q to masks M belonging some sparsity pattern P: s: (f.r,m.n) + M.
Given a set of sparsity patterns Pt.... the set sparsity mask candidate the union of sparsity masks
in each of P:: C=UP,
A sparsity pattern assignment satisfies the compute budget B if:
TotalCompute(s)-
Σ
compute(((t.rmmmn))5 SB.
(21)
layer type?
Let be the flattened vector containing the model parameters, and let M, be the flattened vector containing the sparsity mask by the sparsity mask assignment Let fa()) be the output of the dense network with parameter / and input x. Then the output the sparse network is fea (z))
The empirical NTK of a network fo on data subset X = (-I..... } is matrix of size xX|x|X]
NTK(fo.X),, < @.f((;)))) 0fo(zy)
09 00
The formal algorithm to assign the sparsity mask to each layer type is described in Algorithm 2 The main idea that, as the set of sparsity mask candidate finite, we can enumerate all possible sparsity mask assignment satisfying the budget and pick the one with the smallest NTK distance to the dense NTK. In practice, we can use strategies to avoid explicitly enumerating all possible sparsity mask, e.g. for each sparsity pattern, we can choose the largest sparse mask that fits under the budget.
K.3
METHOD PROPERTIES: REDISCOVERING CLASSICAL SPARSITY PATTERNS. No ADDITIONAL HYPERPARAMETER TUNING
When applied to the Transformer architecture, among the sparsity components descrihed in Appendix K.1 the -gTdgdd heuristic consistently picks the local and global components for both the attention and
36

---

Published as a conference paper at ICLR 2022
MLP layers. Moreover, the butterfly component is also consistently picked for image data, reflecting the 2D inductive bias in this componem" | surprising that they are also picked for the MLP layers. The most popular type of sparsity pattern in MLP layers top-k (in magnitude or gradient, which at initialization is equivalent to random sparsity). We have proved that lower NTK difference results in better generalization bound for the sparse model. As expected, we observe that this allows the sparse model to use the same hyperparamters (optimizer, learning rate, scheduler) the dense model (Section 5).
While some of these patterss have been proposed for sparse attention, it is
'Convolution (commonly used in image data) can written in terms of fast Fourier transform, which has this same sparse pattern at each step of the algorithm
37

---

Published as a conference paper at ICLR 2022
L EXPERIMENT DETAILS
L.1 DATASETS
Cifar10 (Krizhevsky et al. 2009) consists of 60000 coloured images of resolution 32 x 32. Each of them belong to one of I0 classes, including airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. Among these, 50000 images are allocated to be the training and 10000 images the testing set.
Cifar100 (Krizhevsky et al. 2009) is similar Cifar10. It also consists of images of resolution 32 x 32. In total, there are 60000 images, each of which belongs to one of 100 classes. Each of the 100 classes has 500 images in training set and 100 images testing set.
ImageNet1K (Russakovsky et al. 2015) spans 1000 object classes, containing 1,281,167 training images, 50,000 validation images and 100,000 test images, Although images are collected in different resolutions, in practice they are generally reshaped and cropped into 224 x 224.
WikfText-103 (Merity al. 2016) contains articles from the wikipedia page. extracts verified articles from Wikipedia, which add up to over 100 million tokens. Compared to other datasets, such Penn Treebank (PTB) (Taylor al. 2003), WikiText features larger vocabulary and preserves original upper/lower cases, punctuation and numbers.
I.2 MODEL CONFIGURATIONS AND HYPERPARAMETER
We summarize the details required replicate our experiments below.
Baseline Model: Except for dense model. We choose our baselines for each experiment base on the following. RigL aims sparsify model weights/parameters, BigBird focuses on attention matrices, so we used it as a baseline in Transformer based models (VIT, GPT-2).
use it baseline in MLP-based models (Mixer).
L.2.1
IMAGE CLASSIFICATION
Table 1: Configuration of the Cifar10 experiments.
Model
Optimizer
Weight Decay
Learning Rate
Drop Path
Warnup/Epoch
ViT-Small
AdamW
0.05
0.0005
0.1
5/300
Pixelfly-ViT-Smalll ViTlSmall
AdamW
0.05
0.0005
5/300
ViT-Base
AdamW
0.05
0.0005
0.1
5/300
Pixelfly- ViT-Base
Pixelfly-ViT-Base
AdamW
0.05
0.0005
5/300 0
Mixer-Small
AdamW
0.1
0.0005
0.1
5/300
Pixelty-Mixer-SSall
AdamW
0.1
0.0005
0
5/300
Mixer-Base
Adam
0.1
0.0005
5/300
Pixclfl-Mixer-
AdamW
0.1
0.0005
0
5/300
Model
Optimizer
Weight Decay
Learming Rate
Drop Path
Warmup/Fipoch
ViT-Small
AdamW
0.05
0.0005
0.1
5/300
PixclflyVVTT-Smill
A
0.05
0.0005
5/300
ViT-Base
AdamW
0.05
0.0005
5/300
Pixclfly-VT--uas
AdamW
0.05
0.0005
5/300
Mixer-Small
AdamW
0.1
0.0005
0.1
5/300
Fixelfly-Mixer-Snal
AdamW
0.1
0.0005
0
5/300
Mixer-Base
AdamW
0.1
0.0005
0.1
5/300
Pixelfly- Mixer-Base
AdamW
0.1
0.0005
5/300
Table 2: Configuration of the Cifar100 experiments
We report more details on the models, including number of parameters and FLOPs, in Table
We follow the naming convention in the Vision Transformer paper and MLP-Mixer paper. In particular, VIT-S and VIT-B refers to the small and base ViT models respectively, and 16 refers to the patch size of 16x16. The MLP-Mixer models follows the same convention.
T.2.2 LANGUAGE MODELING
We report more details on the models, including number of parameters and FL.OPs, in Table and Table
38

---

Published as a conference paper at ICLR 2022
Model
Optimizer
Weight Decay
Learning Rate
Drop Path
Warnup/Eipoch
ViT-Small
AdamW
0.05
0.001
0.1
5/300
Vix-SIly-iT-SSmall
A
0.05
0.001
5/300
ViT-Base
Adam
0.05
0.001
5/300
Fixclfl-ViT--Bre
AdamW
0.05
0.001
5/300
Mixer-Small
AdamW
0.1
0.001
5/300
Pixelfy-Miixr--Suee
AdamW
0.1
0.001
5/300
Mixer- Basc
AdamW
0.1
0.001
0.1
5/300
Axceffl-Mixcr-Bie
AdamW
0.1
0.001
5/300
Table 3: Configuration of the ImageNet experiment
Table 4: The performance of Pixelfly and VIT or MLP-Mixer on the ImageNet benchmarks, including the number of parameters and FLOPs. measure the accuracy and the training time speedup (on ImageNet) compared to the dense model.
Model
ImageNet top-1 acc.
Speedup
Params
FLOPs
Mixer-S/16
72.4
18.5M
3.8G
Pixelfl-MMixer-S//
72.6
1.7x
5.9M
1.3G
Mixer-B/16
75.6
59,9M
12.6G
Pixelfy-Mixer-B/16
76.3
2.3x
17.4M
4.3G
ViT-S/16
77.7
48.8M
9.9G
PixelflyViT--S16
77.5
1.9x
16.9M
3.6G
ViT-B/16
78.5
86.6M
17.6G
VixeIfly-ViT-B/16
78.6
2.0x
28.2M
6.1G
Table 5: The performance of Pixelfly, BigBird and GPT-2- --aal WikiText- 103, including the number of parameters and FLOPs. We measure the perplexity and the training speed up.
Model
WikiText-103 (ppl)
Speedup
Params
FLOPS
GPT-2-Small
22.2
117M
48.4G
BigBird
23.3
0.96x
117M
40.2G
Pixelfly
22.5
2.1x
68M
18.5G
GPT-2-Medium
20.9
345 M
168G
BigBird
21.5
1.1x
345 M
134G
Pixelfly
21.0
2.5x
203M
27G
Table 6: Configuration of the WikiText103 experiments
Model
Optimizer
Weight Decay
Learning Rate
Dropout
Warmup/Epoch
GPT----mall
Adam Adam
0.1 0.1
0.0001
0.1
5/100
Pixelfly
0.0001
0.1
5/100
L.3 MEASURING EMPIRICAL NTK
The Empirical NTK is rough estimation of the real NTK, in which the width of the neural net goes to infinity. As the width grows, the kernel gets closer its infinite-width limit. Fortunately, both our models of interest, MLP-Mixer and Vision Transformer, are wide and overly parameterized. Therefore they are only one step away from the real NTK domain. This allows to use the Empirical NTK approximately predict their training behaviors.
As described in equation 2 we first compute the gradient of each data sample, then compute pair-wise product to construct the Empirical NTK. Although we use a relatively small dataset, it's still expensive to build kemel for large models, such as ViTs and M..eMMixer. In practice, we find that it's sufficient compute kernels for a subsampled dataset.
39

---

Published as a conference paper at ICLR 2022
MLP-Mixer and Vision Transformer each represent type of module of interest for sparsification. In MLP-Mixer, we study the sparse behavior of the Linear module, whereas, in Vision Transformer, we mainly focus sparsifying attention. All models are first sparsified to around 10% of the original dense compute. Then we compare their NTK kernels with their original dense kemnel. We run three random seeds to climinate noise, ic., three different initializations for each pair of configurations. We report the mean relative difference between the kernels with respect to the norm of the dense kernel.
L.4 TRANSFER LEARNING EXPERIMENTS
We conduct extended experiments tasks. Specifically, we finetune Pixelfly pretrained model (ImageNet) on CIFAR-I0 and show that it get 99.03% accuracy compared to 98.77% 77% our pretrained dense ViT-B/16 model. In addition, we than 2x speed on downstream task fine-tuning process as well.
the generalization of our pretrained sparse models downstream
L.5 MICROBENCHMARKING
this section, we perform microhenchmarking on 4K 4K sparse matrix multiplication. We aim show that Pixelfly patterns are far more hardware friendly than random patterns. For a 4K matrix, expected density the number of non-zero entries/(4K 4K)): actual density is the number of accessed entries/(4K 44K), e.g. even if there is only one non-zero, 32x32 entries would be accessed because the hardware block size is 32x32.
When random patterns are generated with small block size, (e.g 1 x 1, 2),), the resources, such as memory access and computes(denoted by Actual Density), required to compute random sparse matrix of density 1.25% are equivalent to computing dense matrix multiplication. This is further reflected in the latency: As pattern block size shrinks, deviating from the hardware block size of 32 32, the random patterns' latency worsens, whereas the Pixelfly remains efficient. Vanilla Butterfly is 5x slower than Pixelfly as expected, because (I) does take advantage of the hardware property not structured sparsity(2) a series of products.
Pattern
Block size
Expected Density
Actual Density
Latency(ms)
1.25%
100%
2x2
2.5%
99.84%
9.3
4x4
5%
96.24%
9.04
Random
6x6
10%
93.66%
8.8
8x8
20%
81.89%
7.7
16x 16
40%
34.52%
3.3
32x32
80%
10.15%
Butterfly
1xl
10%
62.50%
5.2
1xl
1.25%
4.62%
0.48
2x2
2.5%
5.38%
0.56
4x4
5%
6.13%
0.63
Pixelfly
6x6
10%
9.64%
0.96
8x8
10%
10.58%
1.05
16x 16
10%
11.30%
1.12
32x32
10%
10.58%
1.04
Table 7: Microbenchmarking of different patterns. Given GPU processes the matrix block by block size 32 x 32, random block pattern's 's latency increases as the block size shrinks, while Pixelfly remains efficient. We measure the latency by averaging 100 runs batch size 4096 for each configuration.
L.6 EFFICIENT IMPLEMENTATION OF PIXELFLY
We run all of our experiments V100 GPUs. We rely on efffcient implementation block sparse matrix multiply and block sparse attention from the libraries Triton (Attps://qithiil com/ openai/ triton) and httpss: //ttttutttttt/// com/ www.huggiinen wwwggggig.ace//torccc.ccomm slotkrrr_bblock_pprss For the low- rank part, we rely on efficient (dense) matrix multiplies from cuBLAS. In particular, to multiply the input r by the low-rank matrix UVT, we multiply U(VTx) z).
We keep the same number of training epochs as that of the dense models (e.g., on ImageNet, the dense model and the Pixelfly model trained for 300 epochs). The training speedup of the Pixelfly models is due faster time per epoch.
40

---

Published as a conference paper at ICLR 2022
We do not use 2:4 sparsity (available on Ampere GPUs such A100). Such fine-grained sparsity is orthogonal to our approach, and we expect that future work incorporating both 2:4 sparsity and block sparsity to yield further speedup.
I..7 ABLATION: SPEED-ACCURACY TRADROFF OF PIXELFLY
We conduct an ablation experiment to examine the speed-accuracy trade of Pixelfly: on the ImageNet dataset and the Mixer-B/16 model, we replace the dense matrices with flat block butterfly low-rank matrices, while varying the compute parameter budget. We plot the speed-accuracy tradeoff in Fig. 13
ImageNetlk top-1 accuracy
76.5
76.0
75.5
Mixer-B/16 (dense)
75.0
Pixelfly Mixer-B/16 (block sparse)
1.0
1.5
2.0
2.5
Training speedup compared to dense model
Figure 13: Speed-accuracy tradeoff' of Pixelfly on ImageNet classification, with Mixer-B/16 as the dense model. Pixelfly maintains exceeds the accuracy of the dense model, up around 2.3 speedup (or around 30% of the number of parameters). Performance degrades when the Pixelfly model has fewer than 30% of the number of parameters.
I..8 COMPARISON AGAINST ORIGINAL. BUTTERFLY
We compare Pixelfly against original Butterfly matrices Dao et al. 2020) on the ImageNet dataset and Mixer-B/ B/16 dense model. We present the results in Table S We notice another benefit of Pixelfly compared to Butterfly: it trains more stably and requires less careful intiialization. Since Butterfly a product of many factors, it requires careful initialization, otherwise the activation and gradient will be very large or very small.
Table 8: The performance of Pixelfly and original Butterfly on MLP-Mixer on the ImageNet benchmarks.
Model
ImageNet top-1 acc.
Speedup
Params
FLOPs
Mixer-B/16
75.6
59 39.9M
12.6G
Butterfly-Mixe Mixer-BB66
76.1
0.8% x
17.4M
4.3G
Exxelfly-Mxee--il
76.3
2.3x
17.4M
4.3G
41

---

Published as a conference paper at ICLR 2022
M EXTENDED RELATED WORK
In this section, we extend the related works referenced in the main paper and discuss them in detail.
M.1 NEURAL PRUNING
Our work is loosely related to neural network pruning. By iteratively eliminating neurons and connections, pruning has seen great success in compressing complex models Han et al. (2015a b) put forth two naive but effective algorithms to compress models up to 49x and maintain comparable accuracy. Li et al. (2016) employ filter pruning to reduce the cost of running convolution models up to 38 %. Lin et al. (0117) prunes the network at runtime, hence retaining the flexibility of the full model. Dong et al. (017)) prunes the network locally in a layer by layer manner. [Sanh et al. (2020) prunes with deterministic first-order information, which more adapttive to pretrained model weights. I Lagunas et al. (202T) prunes transformers models with block sparsity pattern during fine-tuning, which leads to real hardware speed up while maintaining the accuracy. Zhu and Gupta ( finds large pruned sparse network consistently outperform the small dense networks with the same compute and memory footprints. Although both our and all the pruning methods are aiming to produce sparse models, we differ in our emphasis on the overall efficiency, whereas pruning mostly focuses on inference efficiency and disregards the cost in finding the smaller model.
M.2 LOTTERY TICKET HYPOTHESIS
Models proposed in our work can be roughly seen class of manually constructed lottery tickets. Lottery tickets |Frankle and Carbin (2018) are a set of small sub-networks derived from a larger dense network, which outperforms their parent networks in convergence speed and potentially in generalization. A huge number of studies are carried out to analyze these tickets both empirically and theoretically: |Morcos et al.) (2019) proposed
generalized lottery tickets for all vision benchmarks and got comparable results with the specialized lottery tickets; Frankle et al, (2019) improves the stability of the lottery tickets by iterative pruning; Frankle et al. (2020) found that subnetworks reach full accuracy only if they are stable against SGD noise during training: Orseau et al. (2020) provides a logarithmic upper bound for the number of parameters it takes for the optimal sub-networks exist: Pensia al. (2020) suggests way to construct the lottery ticket by solving the subset sum problem and it's a proof by construction for the strong lottery ticket hypothesis. Furthermore, follow-up works (Liu and Zenke 2020) Wang et al. 2020 Tanaka et al. 2020 show that we can find tickets without any training labels.
M.3 NEURAL TANGENT KERNEL
Our work rely heavily on neural tangent kernel in theoretical analysis. Neural Tangent Kemel Dacot Pacot et al.) (2018) is first proposed to analyse the training dynamic infinitely wide and deep networks. The kemnel is deterministic with respect to the initialization as the width and depth go to infinity, which provide an unique mathematical to analyze deep overparameterized networks. Couples of theoretical works are built based upon this: Lee et al. (2019) extend on the previous idea and prove that finite learning rate is enough for the model to follow NTK dynamic. Arora et al. (2019b) points out that there still a gap between NTK and the real finite NNs. Cao and Gu (2020) sheds light on the good generalization behavior of overparameterized deep neural networks. Arora et al. (2019a) is the first one to show generalization bound independent of the network size. Later, some works reveal the training dynamic of models of finite width, pointing out the importance of width in training: HHyou et al. (2019) analyzes stochastic gradient from the stochastic differential equations' point of view; Based on these results, we formulate and derive our theorems on sparse network training.
M.4 OVERPARAMETERNI OVERPARAMETERIZED MODELS
Our work mainly targets overparameterized models. In Nakkiran et al. (2019) the double descendent phenomenon was observed. Not long after that, Ascoli et al.) (2020) discover the triple descendent phenomenon. It's conjectured in both works that the generalization error improves as the parameter count grows. On top of that, Arror et (2018) speculates that overparameterrizaton helps model optimization, and without "enough" width, training can be stuck at local optimum. Given these intuitions, it's not surprising that the practitioning community is racing to break the record of the largest parameter counts: The two large language models, GPT-2 and GPT-3 (Radford et al. 2019 Brown et al. 2220)) are pushing the boundary text generation and understanding: Their amazing zero-shot ability eamn them the title of foundation models (Bommasani al., 2022). On the computer vision side, (Dosovitskiy al. (2000) Tolstikhin al. (2021): Zhai et al. (2021) push the top-1 accuracy on various vision benchmarks to new highs after scaling up to 50 times the parameters: Naumov et al. (2019) shows impressive results recommendation with 21 billion
42

---

Published as a conference paper at ICLR 2022
large embedding: Jumper al. (2011) from DeepMind solve 50 year old grand challenge in protein research with a 46-layer Evoformer. In our work, we show that there a more efficient way to scale up model training through sparsification and double descent only implies the behavior of the dense networks.
43


Reviews:
[{"review_id": "wSvsKqSpSYA", "reviewer": "Reviewer_GfcL", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have a few smaller concerns with respect to the experimental part, where some aspects are not clear to me. Specifically, the selection of models for comparison is not always clear:\n- In the first paragraph of section 5.1 you mention evaluation of Pixelfly against Butterfly by Dao et al., yet I cannot see these results being presented. This part is most interesting as it should show the advantage of your method over the “standard” butterfly approach. \n- It is not always clear in which experiments (NTC or training from scratch, for image classification and language modeling) you are using which model, namely RigL, BigBird, Mixer. The results presented in the tables do not show all comparisons, it appears. \n- Some of the naming schemes/abbreviations used for the models in these sections are not explained in the captions (e.g. what does the S/16 , B/16 stand for). \n- Furthermore there is very little actual results shown on your ablation studies. I know that 9 pages is tight, but this is actually an interesting part, so consider shaving off a bit of space for this elsewhere. \n- For Theorem 4.5, this only relates to “a class of input sequences”: How can you generalize from that? It is a pretty fundamental assumption in all of your work, that flat block butterfly + low-rank performs better than either of them alone.\n- The related work section is kept very short, and the actual advantages/disadvantages of other approaches is not really discussed. I find this however rather important as it should highlight what sets your method apart from those others.\n- You should further check for typos in your appendix ;) (e.g. L.3)", "summary_of_the_paper": "The authors present a new method, Pixelfly, for static sparsity patterns in sparse training of matrix-multiplication-based neural networks. Their methods builds on butterfly matrices as a representation of any kind of sparse matrices, transforms them into a flat block structure in order to a) align them with the hardware structure of GPUs and b) speed up the decomposition by approximating the the butterfly multiplication with a summation (first order approximation), and then combines these flat butterfly patterns with low-rank approximation to increase expressiveness. They prove the validity of the made assumptions in their method and evaluate it on a number of common architectures, showing the actual improvement of their method over other approaches towards sparse training as well as the dense counterparts. ", "summary_of_the_review": "This is an interesting, well explained and soundly tested method. Sparse training for large neural networks is an up-to-date and very interesting research topic, and the paper provides an innovative approach towards this research field. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "_QuDfVplef3", "reviewer": "Reviewer_Ukv2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* A well-motivated work that combines theoretical insights with practical considerations, resulting in a simple yet effective method. \n* The paper is well-written and easy to understand; I especially appreciated the diagrams that illustrate different versions of butterfly matrices.\n* Authors provide rigorous proofs of properties for all Pixelated Butterfly components.\n* In general, the conducted experiments are quite extensive and validate the claims made in the paper.\n\nWeaknesses:\n* Both for image classification and language modeling experiments, it would be great to give explicit parameter/FLOP counts both for Pixelated Butterfly and for the baselines. Also, in my opinion, it is important to consider larger model sizes (1B+ parameters and hidden sizes of 1280+) at least for language modeling: with the current scaling trends in NLP, this is the setting that would benefit the most from this approach, and it is worth verifying that there are no quality losses at this scale.\n* The actual hardware resources used in the main experiments are not described. This affects both the choice of the block size (which might influence the findings) and the availability of fine-grained sparsity patterns, such as 2:4 structured sparsity on NVIDIA Ampere architectures (which should probably be mentioned somewhere in the work).\n* Currently, both from the experiments and the analysis it is not entirely evident whether there are any disadvantages of Pixelated Butterfly compared to regular multiplication of dense matrices: the only apparent one is a slight drop in perplexity on WikiText-103. This might confuse the readers, and thus I would encourage the authors to provide examples (even slightly contrived will do) of matrices in general or parameters of neural networks that are poorly approximated by the proposed procedure. One possible way to do so would be to study how the final task performance changes for a range of compute budgets and a single base architecture with a given size.\n\nQuestions and typos:\n* Theorem 4.1 refers to Equation 1 with $k=n$, but $k$ is not used neither in the theorem nor in the equation. Could you please clarify what does this variable stand for?\n* In Definition 3.3, I believe you refer to *block* butterfly factor matrices instead of butterfly factor matrices.", "summary_of_the_paper": "This work proposes a new sparse parametrization method for layers of neural networks that rely on matrix multiplication. This method aims to improve the computational performance of previous approaches by taking the constraints of hardware into account while maintaining high capacity. Authors develop Pixelated Butterfly — an approach that builds upon the Butterfly family of matrices and leverages the block structure, summation instead of a product and a low-rank term to achieve significant performance improvements compared to the baselines. Also, authors prove several properties regarding the expressive abilities of Pixelated Butterfly and its convergence in the NTK setting.", "summary_of_the_review": "A solid work in general, yet several aspects of the empirical evaluation could be improved and some details are missing.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Cdn1ke3ZvWk", "reviewer": "Reviewer_JNUH", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strong and Weak Points\n\n- (strong) Brief but solid and accessible coverage of the challenges to accelerate sparse matmuls, Butterflies in particular. \n\n- (strong) Results are very compelling.\n\n- (minor weak) Figure 3 is good but It would be better if at the top of each matrix (at least the first row) includes what they are: e.g. $B^{(16,1)}, B^{(8,1)}, B^{(4,1)}, B^{(2,1)} ... B^{(16,2)}, B^{(8,2)}, B^{(4,2)}$. For someone that is not familiar with the recent Butterfly literature it might be not obvious what is going on.\n\n- (minor weak) Could the Author expand L.5? It is not clear to me how the “Expected Density” and “Actual Density” are computed. I think this is an interesting study that supports part of the claims of the paper. It would be great if a new set of columns using vanilla Butterflies could be added.\n\n-  Could the Authors comment upon considerations needed to make the proposed method work for rectangular matrices? (I see Appendix I.4 but adding a bit more detail about what _stretching_ means could be valuable, specially since most matmuls in a normal network would be rectangular) Does this impact the quality of how Flat Butterflies can approximate standard Butterflies?\n\n- I was a bit surprised as I was reading the paper and see the statement “our method targets GEMM-based networks, [..] such as Transformer and MLP-Mixer”, why not mention CNNs? Many popular convolution algorithms follow a GEMM-based implementation (e.g. im2col, im2row, winograd). I’m curious to hear the Authors’ view on why not mention CNNs or having them as a baseline (e.g. a vanilla ResNet34/50)", "summary_of_the_paper": "\nThis paper proposes a method to accelerate training by replacing GEMM-based compute intensive operations with sparse multiplications. Concretely, the Authors focus on a family of sparse matrices known as Butterfly matrices which have well-studies properties and, crucially, have a fixed sparsity pattern with a very low non-zero ratio. However, the speedups that can be attained with Butterfly matrices is often hindered by (1) their sequential (multiplicatively) decomposition which limits parallelisation and, (2) their poor data locality and alignment (due to Butterfly matrices containing zeros around non-zero values most of the times), effectively translating in poor communication-to-compuation ratio. This works addresses both issues by: presenting blocked Butterfly matrices where non-zero values come in the form of $d \\times d$ blocks with d>=2 (instead of 1x1 block); and, by relaxing the construction process of Butterfly matrices and construct them instead as a sum of Butterfly factor matrices (instead of a product). The Authors report large speedups on emerging architectures for Image Classification.", "summary_of_the_review": "### Recommendation\n\nThis is a good paper. Although I wouldn't label as \"significant\" the contribution of going from a Butterfly to its blocked counterpart, the fact that the Authors also demonstrate that the product of Butterfly factors can be replace with a sum of factors, makes it a substantial contribution. Experimental evaluation includes a variate of datasets, that is a big plus. There are other recent works proposing methods to accelerate training with sparsity. However, they require specialise hardware. Authors could consider adding them to their related work or introduction: [1, 2] (note, I’m not an author of these works). I wasn’t familiar with NTK, glad I know about this now. \n\n- [1] https://arxiv.org/abs/2001.01969\n- [2]https://openaccess.thecvf.com/content_CVPR_2020/html/Goli_ReSprop_Reuse_Sparsified_Backpropagation_CVPR_2020_paper.html\n\n\n### Supporting my Recommendation\n\nPlease see points above.\n\n\n### Minor points\n\n- I believe Figure 4 is not ready for colourblind reader", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2lEpuT8uOo6", "reviewer": "Reviewer_Q8VQ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found the paper very interesting and timely, given the huge interest in transformer based networks. The claimed 2$\\times$ training time speedup without loss of accuracy is a very attractive proposition. The first-order approximation of butterfly matrices is a neat trick and there is empirical evidence as well as NTK analysis to show efficacy. However, I do think there are a number of weaknesses in this paper which I will list below:\n\n- Speedup is reported but not discussed in any detail. Do you have the same number of epochs for pixelfly and the baselines? On which device do you measure those speedups? How do you measure wall-clock time (methodology)?\n- Can you control the accuracy-speedup of pixelfly? It seems that a single result was reported for pixelfy for each experiment without any ablations or discussions showing a trade-off between accuracy and sparsity.\n- LRA is renowned for being underspecified, with many of their hyperparameters inconsistently reported between the paper and the github repo. How did you ensure a fair comparison to the results reported in LRA?\n- I did not see any details about the kernel implementation of pixelfly. In python, how did you write the kernel? Is it simply $\\gamma B+(1-\\gamma)UV^T$? and that alone gives you the reported speedups?\n- I am also missing a discussion of speedup in terms of number of FLOPs. It would be nice to compare that to wall-clock speedup.\n", "summary_of_the_paper": "This paper presents a hardware-friendly sparse matrix multiplication that performs competitively on many transformer neural networks. 'Pixelated Butterflies' combine two well-known techniques: butterfly matrices and low-rank matrices. Furthermore, the authors propose to use a first-order approximation of _block_ butterfly matrices to be able to achieve actual speedup on current hardware (I am assuming GPUs, although the device isn't explicitly stated anywhere). Much analysis is presented to show the expressiveness of the proposed sparse matrix representation and experiments are presented on the long-range arena (a transformer benchmark), language modeling tasks and image classification.", "summary_of_the_review": "Good paper with some question marks around the utility, configurability and practicality of the presented approach, and the methodology of speedup measurements.\n\n-----------------\n\nAll of my comments/questions were addressed. I would've liked to see more ablations regarding the accuracy-efficiency tradeoffs but the presented response is more than sufficient. I will therefore increase my score from 6 --> 8.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is an intriguing work that introduces a novel sparse training technique. The core insight is a novel reparametrization or sparsity pattern based on the so-called butterfly matrices that enables fast training and good generalization. The theory is solid and useful. Most importantly, the method is novel and is likely to become impactful. Understanding better what contributes to the excellent performance is an interesting question for future work. In agreement with all the reviewers, it is my pleasure to accept the work.
Decision: Accept (Spotlight)