P IXELATED B UTTERFLY : S IMPLE AND E FFICIENT S PARSE TRAINING FOR N EURAL N ETWORK M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Overparameterized neural networks generalize well but are expensive to train.
Ideally, one would like to reduce their computational cost while retaining their generalization beneﬁts.
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 difﬁculty in sparsifying all model components.
The core problem is that searching for a sparsity mask over a discrete set of sparse matrices is difﬁcult and expensive.
To address this, our main insight is to optimize over a continuous superset of sparse matrices with a ﬁxed structure known as products of butterﬂy matrices.
As butterﬂy matrices are not hardware efﬁcient, we propose simple variants of butterﬂy (block and ﬂat) to take advantage of modern hardware.
Our method (Pixelated Butterﬂy) uses a simple ﬁxed sparsity pattern based on ﬂat block butterﬂy and low-rank matrices to sparsify most network layers (e.g., attention,
MLP).
We empirically validate that Pixelated Butterﬂy is 3 × faster than butterﬂy and speeds up training to achieve favorable accuracy–efﬁciency tradeoffs.
On the ImageNet classiﬁcation and WikiText-103 language modeling tasks, our sparse models train up to 2.5 × faster than the dense MLP-Mixer, Vision Transformer, and GPTmedium with no drop in accuracy.
I NTRODUCTION
Recent results suggest that overparameterized neural networks generalize well (Belkin et al., 2019), but they are expensive to train (Kaplan et al., 2020).
An ideal model should use less compute and memory while retaining the generalization beneﬁts of large models.
The simplest and most popular direction is to sparsify these models.
This idea has a long history in machine learning (LeCun et al., 1990) and has driven fundamental progress in other ﬁelds such as statistics (Tibshirani, 1996), neuroscience (Foldiak, 2003), and signal processing (Candes et al., 2006).
However, despite signiﬁcant efforts, speeding up sparse training in wall-clock time without degrading accuracy remains an unresolved problem.
While sparse training is an active research area, it has not seen wide adoption.
First, it is difﬁcult and expensive to ﬁnd 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., 2018), 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 signiﬁcantly slow down training and complicate the implementation.
Indeed, these methods either require long cycles of pruning and retraining (Frankle and Carbin, 2018) 1 or maintain expensive hash tables (Chen et al., 2019).
Second, most existing methods adopt unstructured sparsity, which may be efﬁcient in theory, but do not take into account the efﬁciency of training hardware such as GPUs (optimized for dense computation) 2 .
Finally, most methods target a 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).
A better sparse training method should (i) be simple yet accurate, ideally with a static sparsity pattern, (ii) be fast by aligning sparsity pattern 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 a matrix), it is NP-hard to ﬁnd the optimal static sparsity pattern for a NN module to minimize the approximation error to the dense model.
Second, for each sparsity pattern,
Pixelated
[Figure]
+
Model Schema
[Figure]
we need to take into account hardware block-oriented efﬁciency (accessing each element in memory takes the same time as accessing the block of adjacent elements (Cook, 2012), illustrated in Fig. 2).
Common theoretical measures of efﬁciency (e.g., number of non-zeros, FLOPs) do not map well to modern hardware designed for block computation.
Last, every different 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 to ﬁnd those patterns that can closely approximate the dense model (Details in Appendix K).
We found that one sparsity pattern, namely butterﬂy + 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 (Candès et al., 2011; Udell and Townsend, 2019;
Chen et al., 2021), and (ii) butterﬂy matrices (Parker, 1995; Dao et al., 2019) whose products can tightly represent any sparse matrix (De Sa et al., 2018; Dao et al., 2020).
Using the ﬁxed sparsity pattern from butterﬂy matrices, with the addition of a low-rank term, would address two of the three challenges above and yield
a simple way to sparsify most NN layers (that are based on matrix multiply).
However, butterﬂy matrices are inefﬁcient on modern hardware: (i) they are difﬁcult to parallelize as they contain sequential products of many factors, and (ii) they are not hardware-friendly because the sparsity patterns are not block-aligned.
We propose two simple changes to make Butterﬂy efﬁcient while retaining their favorable properties.
Our proposal, Pixelated Butterﬂy (Pixelﬂy), combines ﬂat block butterﬂy and low-rank matrices to yield a simple and efﬁcient sparse training method.
• We design an extremely simple sparsity pattern inspired by butterﬂy + low-rank matrices, which takes into account the hardware’s block-oriented efﬁciency.
We propose block butterﬂy matrices that are efﬁcient as their sparsity patterns align with hardware blocks.
We then introduce ﬂat butterﬂy, a ﬁrst-order approximation of butterﬂy with residual connection, that turns the original product of factors into a sum.
Flat butterﬂy matrix multiplications are easy to parallelize.
Pixelﬂy, uses the ﬁxed sparsity pattern from ﬂat & block butterﬂy, along with a low-rank term, to produce a sparse network.
• We prove that block butterﬂy retains the expressiveness of butterﬂy matrices and can thus tightly capture sparse matrices.
We show that ﬂat butterﬂy matrices can closely approximate large classes of matrices that butterﬂy matrices capture.
Moreover, we demonstrate that ﬂat block butterﬂy + lowrank matrices are strictly more expressive than sparse or low-rank matrices alone.
Finally, leveraging the recent advance in the neural tangent kernel (NTK), we adapt existing techniques to prove the global convergence of gradient descent on training sparse and wide ReLU networks.
• Our proposed Pixelﬂy can be applied to all network modules that rely on matrix multiplication (e.g., linear layer, attention, MLP).
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 Pixelﬂy can speed up the training of models
(Transformers, ViT, MLPMixer) without quality drop compared to baselines on a wide range of domains and tasks.
On CIFAR10/100 & ImageNet classiﬁcation, Pixelﬂy achieve 2.3 × 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.5 × and achieve the same perplexity.
On the Long Range Arena benchmark, we maintain the same accuracy as Transformer with 5 . 2 × faster training than a dense model, 2 × faster than Sparse transformer, and 6 × faster than non-block-aligned sparse methods (Reformer).
Our ablation studies highlight the importance of each of our components: our butterﬂy sparsity improves on existing hand-crafted patterns by up to 2% of accuracy on ImageNet, our hardware-aware block-sparsity yields up to 5 × speedup, and the balanced compute budget allocation brings 2 × speedup compared to baselines that only sparsify attention.
P ROBLEM S ETTING
We ﬁrst deﬁne the problem as sparse matrix approximation with a simple hardware cost model.
Then we brieﬂy introduce butterﬂy and sparse + low-rank matrices.
[Figure]
Problem Formulation: We focus on the training of GEMM-based models, which can be viewed as a series of matrix multiplies (Given A, B ∈
R n × d , compute C = AB T ).
Speeding up training while maintaining model quality can be mapped to ﬁnding an approximation procedure f which reduces the time T of computing C while minimizing error
E [ (cid:107) f
( A, B ) − AB T (cid:107) 2 F ] .
Since the hardware is a block device, accessing any individual element within a block of memory is the same as accessing the full block (Cook, 2012)
(Fig. 2).
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 deﬁnition 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 10 × compared to sparsity patterns that are not block-aligned.
Butterﬂy, Sparse + Low-rank Matrices: Butterﬂy matrices have been used in numerical linear algebra (Parker, 1995; Li et al., 2015) and machine learning (Mathieu and LeCun, 2014;
Jing et al., 2017; Munkhoeva et al., 2018; Dao et al., 2019; Choromanski et al., 2019).
They encode the recursive divide-and-conquer structure of the fast Fourier transform (FFT) algorithm (Cooley and Tukey, 1965) and provably capture any sparse matrix with near-optimal space and time complexity.
Sparse and Lowrank structures have been studied in Robust PCA (Candès et al., 2011), graph clustering (Jalali et al., 2011), and co-variance estimation (Luo, 2011).
Recently it has been adopted in attention approximation for Transformers (Chen et al., 2021).
B UTTERFLY MATRICES AND P IXELATED B UTTERFLY
Butterﬂy matrices (Parker, 1995; Dao et al., 2019) are expressive and theoretically efﬁcient.
As they contain the set of sparse matrices, we choose to search for the sparsity pattern in this larger class due to their ﬁxed sparsity structure.
However, there are three technical challenges.
We highlight them here along with our approaches to address them:
Slow speed: butterﬂy matrices are not friendly to modern hardware as their sparsity patterns are not block-aligned, thus are slow.
We introduce a variant of butterﬂy matrices, block butterﬂy , which operate at the block level, yielding a block-aligned sparsity pattern.
Difﬁculty of parallelization: the sequential nature of butterﬂy matrices as products of many factors makes it hard to parallelize the multiplication.
We propose another class of matrices, ﬂat butterﬂy matrices, that are the ﬁrst-order approximation of butterﬂy with residual connections.
Flat butterﬂy turns the product of factors into a sum, facilitating parallelization.
Reduced expressiveness of ﬂat butterﬂy: even though ﬂat butterﬂy matrices can approximate butterﬂy matrices with residual connections, they are necessarily high-rank and cannot represent lowrank matrices (Udell and Townsend, 2019).
We propose to add a low-rank matrix (that is also block-aligned) to ﬂat butterﬂy to increase their expressiveness.
Combining these three approaches ﬂat block butterﬂy + low-rank), our proposal (Pixelated Butterﬂy) is a very simple method to train sparse networks.
3. B LOCK B UTTERFLY M ATRICES
We propose a block version of butterﬂy matrices, which is more hardware-friendly than the regular butterﬂy.
The regular butterﬂy matrices Dao et al.
(2019; 2020) will be a special case of block butterﬂy with block size b = 1 .
We omit b in the notation if b = 1 .
Deﬁnition 3.1.
A block butterﬂy factor (denoted as B k,b ) of size kb (where k ≥ 2 ) and block size b is a matrix of the form B k,b =
[Equation]
D [Equation]
(cid:21) where each D i is a k 2 × k 2 block diagonal matrix of block size b of the form diag (cid:0) D i, 1 , . . .
, D i,k/ 2 (cid:1) where D i,j ∈ R b × b .
We restrict k to be a power of 2.
Deﬁnition 3.2.
A block butterﬂy factor matrix (denoted as B ( n,b ) k ) of size nb with stride k and block size b is a block diagonal matrix of nk (possibly different) butterﬂy factors of size kb and block size b :
[Equation]
[Figure]
Deﬁnition 3.3.
A block butterﬂy matrix of size nb with block size b (denoted as B ( n,b ) ) is a matrix that can be expressed as a product of butterﬂy factor matrices: B ( n,b )
=
B ( n,b ) n B ( n,b ) n. . .
B ( n,b ).
Deﬁne B b as the set of all matrices that can be expressed in the form B ( n,b ) (for some n ).
3. F LAT BUTTERFLY MATRICES
In most applications of butterﬂy matrices to neural networks, one multiplies the O (log n ) butterﬂy factors.
However, this operation is hard to be efﬁciently implemented on parallel hardware (e.g., GPUs) due to the sequential nature of the operation 3 .
We instead propose to use a sum of butterﬂy factors that can approximate the products of the factors.
This sum of factors results in one sparse matrix with a ﬁxed sparsity pattern, which yields up to 3 × faster multiplication on GPUs (Appendix J).
Residual connections have been proposed to connect the butterﬂy factors (Vahid et al., 2020).
We show that residual products of butterﬂy matrices have a ﬁrst-order approximation as a sparse matrix with a ﬁxed sparsity.
Let M be a matrix in the set of butterﬂy matrices B .
In residual form, for some λ ∈ R :
[Equation]
Note that this form can represent the same matrices in the class of butterﬂy matrices B , since any B ( n ) k contains the identity matrix I .
Assuming that λ is small, we can expand the residual and collect the terms 4 :
[Equation]
) n ) + (cid:101) O
( λ 2 ) .
Deﬁnition 3.4.
Flat butterﬂy matrices of maximum stride k (for k a power of 2) are those of the form
I + λ ( B ( n ) 2 + B ( n ) 4 + · · · + B ( n ) k ) .
Flat butterﬂy matrices of maximum stride n are the ﬁrst-order approximation of butterﬂy matrices in residual form (Eq.
Notice that ﬂat butterﬂy of maximum stride k are sparse matrices with
O ( n log k ) nonzeros with a ﬁxed sparsity pattern, as illustrated in Fig.
We call this sparsity pattern the ﬂat butterﬂy pattern.
Flat block butterﬂy matrices are block versions of ﬂat butterﬂy in Section 3.2 (shown in Fig. 3).
We empirically validate that ﬂat block butterﬂy matrices are up to 3 × faster than block butterﬂy or regular butterﬂy, as shown in Appendix J.
Since ﬂat butterﬂy matrices approximate the residual form of butterﬂy matrices, they have high rank if λ is small (Section 4).
This is one of the motivations for the addition of the low-rank term in our method.
3. P IXELATED B UTTERFLY : F LAT B LOCK B UTTERFLY + L OW - RANK FOR E FFICIENT S PARSE T RAINING
We present Pixelated Butterﬂy, an efﬁcient sparse model with a simple and ﬁxed sparsity pattern based on butterﬂy and low-rank matrices.
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 as a series of matrix multiplies.
Given a model schema (layer type, number of layers, matrix dimension) and a compute budget, Pixelated Butterﬂy has three steps: compute budget allocation per layer, sparsity mask selection from the ﬂat butterﬂy pattern and model sparsiﬁcation.
We describe these steps in more details:
Compute budget allocation : based on our cost model (Appendix A), given the layer type, number of layers, and matrix dimension, we can ﬁnd the density (fraction of nonzero weights) of each layer type to minimize the projected compute cost.
Continuing our goal for a simple method, we propose to use a 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 Appendix I that this simple rule of thumb produces similar results to solving for the density from the cost model.
Sparsity mask selection : given a layer and a sparsity compute budget for that layer
, we use onequarter to 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 (e.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 ﬂat block butterﬂy sparsity pattern: we choose the butterﬂy block size as the smallest supported block size of the device (e.g., 32), and pick the maximum stride of the ﬂat block butterﬂy (Deﬁnition 3.4) to ﬁll up the budget.
Model sparsiﬁcation : The resulting sparse model is simply a model whose weights or attention follow the ﬁxed 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 5 as: W = γB + (1 − γ ) UV (cid:62) , where B is
a ﬂat block butterﬂy matrix (which is sparse), UV (cid:62) is the low-rank component, and γ is a learnable 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 butterﬂy) to improve the accuracy of the method.
T HEORETICAL ANALYSIS
We characterize the expressiveness of the matrices used in our method.
In particular, we prove that block butterﬂy retains the expressiveness of butterﬂy, and that ﬂat butterﬂy can accurately approximate the residual form of butterﬂy.
Moreover, ﬂat block butterﬂy + low-rank (an instance of sparse + lowrank) is more expressive than sparse or low-rank matrices alone.
Finally, we analyze the training convergence and generalization of networks with sparse weights.
All proofs are in the Appendix.
4. E XPRESSIVENESS OF B LOCK B UTTERFLY
We ﬁrst prove the expressiveness of block butterﬂy matrices.
Theorem 4.1.
The set B 2 b of n × n block butterﬂy matrices with block size 2 b contains the set B b of n × n block butterﬂy matrices of block size b .
By a recursive argument, the set of block butterﬂy matrices whose block size is a power of 2 contains the set of regular butterﬂy matrices.
Dao et al.
(2020) show that butterﬂy matrices can tightly represent all structured matrices, such as sparse matrices and many fast transforms.
As a result, block butterﬂy matrices can also represent those structured matrices.
In particular,
Corollary 4.2.
For any constant block size b that is a power of 2, any nb × nb spare matrix with s nonzeros can be written as products of block butterﬂy matrices with block size b and their transposes, with O ( s log n ) parameters.
4. E XPRESSIVENESS OF F LAT B UTTERFLY
We now characterize how the ﬂat butterﬂy matrices approximate butterﬂy matrices.
In particular, assuming that each butterﬂy factor has bounded norm, we show that ﬂat-butterﬂy matrices can accurately approximate the residual form of butterﬂy with error scaling as (cid:101) O ( λ 2 ) .
Theorem 4.3.
Let M be a matrix of the form in Deﬁnition 3.4 where k = n , with B max := max i (cid:13)(cid:13)(cid:13) B ( n ) i (cid:13)(cid:13)(cid:13) F and | λ | ≤ c √ [Equation] log nB max for some constant 0 < c ≤ 12 and some (cid:15) > 0 .
Then
[Equation]
We show that ﬂat butterﬂy matrices must have high-rank if λ is small.
This is the motivation for the addition of the low-rank term in Pixelﬂy (Section 3).
Theorem 4.4.
Let M be as in Eq.
(1) , with B max := max i (cid:13)(cid:13)(cid:13) B ( n ) i (cid:13)(cid:13)(cid:13) F and | λ | ≤ c √ (cid:15) log nB max for some constant 0 < c ≤ 14 and some (cid:15) > 0 .
Let B ∞ max = max i (cid:107) B i (cid:107) ∞ .
Assuming B ∞ max ≤ B max .
Then
[Equation]
4. E XPRESSIVENESS OF F LAT B LOCK B UTTERFLY + L OW - RANK
Chen et al.
(2021) prove that there is a natural class of input sequences (generated by a clustering process) whose attention matrix can only be approximated well by sparse + low-rank matrices, and not sparse or low-rank matrices alone.
We adapt their technique to show a similar result for the class of matrices we use in Pixelﬂy.
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-approximated by ﬂat butterﬂy, while the rest of the attention matrix is of medium size and is well-approximated by low-rank.
Theorem 4.5 (Informal) .
There exists a class of input sequences whose attention matrices are wellapproximated by ﬂat block butterﬂy + 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 B.3.
4. C ONVERGENCE AND G ENERALIZATION OF S PARSE N ETWORKS
There are natural questions about the training and generalization of sparse models: do they train similarly to 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 kernel (NTK) (Jacot et al., 2018) of the network.
The NTK of two data points x and y measures the similarity between the gradient of the network when evaluated at x compared to the gradient when evaluated at y .
This kernel governs the dynamics of the neural network output function f ( · , θ ) 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., 2019b).
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 F.
Though this result does not capture the possible regularization effect of sparsity, it shows that sparse models with small NTK difference from dense NTK preserve the generalization ability of dense models, a 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-Zhu et al., 2019b) in Appendix H.
E XPERIMENTS
In this section, our goal is to demonstrate that an extremely simple ﬁxed sparsity pattern can actually speed up sparse model training in wall-clock time without degrading model quality.
Speciﬁcally, we empirically validate three claims that suggest Pixelﬂy can improve training speed of different model architectures while retaining model quality on a wide range of domains and tasks.
Section 5.1: for image classiﬁcation tasks, we ﬁrst show the empirical NTK of ﬂat
block butterﬂy +
low
-rank sparsity pattern is closer to dense NKT than other baselines
Then we demonstrate our superior end-to-end performance.
Speciﬁcally, we achieve training speed up on both MLP-Mixer and ViT models by up to 2.3 × wall-clock time with no drop in accuracy compared to the dense model and up to 4 × compared to RigL, BigBird and other sparse baselines.
Section 5.2: for language modeling and text classiﬁcation tasks, we can speed up GPT-2 small dense model training by 2.1 × , achieving a perplexity of 22.5 on wikitext-103.
In addition, on Long Range
Arena (LRA) benchmark, we maintain the same accuracy but have 5 .
2 × speed-up in training.
Section 5.3: we show the necessity of block ﬂat butterﬂy and low-rank structures, hardwarealignment and wide coverage of most network layers by conducting ablation studies on these three components of Pixelﬂy.
5. I MAGE C LASSIFICATION
[Figure]
We evaluate the quality and efﬁciency of Pixelﬂy through three metrics: (1) distance to training dynamic of the dense model: compare the distance between empirical NTK kernel 6 of the models with candidate patterns, including BigBird (Zaheer et al., 2020), Butterﬂy (Dao et al., 2020), and that of the dense model, (2) upstream accuracy: compare the accuracy and training time of the Pixelﬂy, the dense counterpart, and other baselines on same image classiﬁcation tasks, (3) downstream accuracy: compare the accuracy of our pretrained Pixelﬂy and dense model ﬁne-tuned on downstream tasks (Appendix L.4).
The empirical
NTK of the model with ﬂat block butterﬂy + low-rank, picked by Pixelﬂy, is closer to the NTK of the dense model.
Pixelﬂy MLP-mixer and ViT models also retain the same top-1 accuracy of the original dense models while achieving up to 2.3 × speed up.
Setup: We use three popular vision benchmarks, CIFAR-10/100 (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009).
We choose recent popular Vision Transformer (Dosovitskiy et al., 2020),
T2T-ViT (Yuan et al., 2021) and MLP-Mixer (Tolstikhin et al., 2021) 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 Pixelﬂy more clearly.
[Table]
Empirical NTK: To characterize the training dynamic of the sparse networks, we compute the empirical NTK kernels for dense Vision Transformer on CIFAR-100.
Then, we show the relative differences between kernels of models with different sparsity patterns and that of the dense one in Fig. 4.
Speciﬁcally, we pick a popular sparsity pattern combination – Bigbird pattern (Zaheer et al., 2020) for attention layer and random (magnitude-based sparsity at initialization equals to random) for MLP layer, as a representative baseline.
The plot indicates that our designed pattern, ﬂat block butterﬂy + low-rank is the closest one to that of the dense one among all the patterns.
Hence, we expect them to enjoy the most beneﬁts of their dense overparameterized counterparts in real tasks.
More details on measuring empirical NTK are covered in the Appendix L.3.
[Figure]
Training from scratch: We validate that Pixelﬂy trains up to 2.3 × and 2.0 × faster than dense MLP-Mixer and
ViT models from scratch, with the same accuracy under the same setting (batch size, epochs).
Speciﬁcally, we sparsify the models with Pixelﬂy and train them on three commonly used vision benchmarking datasets, CIFAR10/100 and ImageNet.
We measure their Top-1 accuracy wall-clock training time.
To summarize the general trend, Fig.
5 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 algorithms aim to dynamic search what could be good sparsity for efﬁcient 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
Pixelﬂy outperforms both dense and RigL in terms of accuracy while achieving 2 .
1 × speedup.
[Figure]
Finally, we compare Pixelﬂy with BigBird and Sparse Transformer pattern.
For a fair comparison, we choose T2T-ViT as the base model because its major bottleneck is on the T2T attention module (our baselines are efﬁcient attention variants).
We can see from Fig.
7 that Pixelﬂy is the only one that can maintain the accuracy and have actual speed up.
Further more, Pixelﬂy speeds up T2T module (large attention) by 1.4 × compare to dense.
5. L ANGUAGE M ODELING AND T EXT C LASSIFICATION
In this section, we aim to evaluate the effectiveness of Pixelﬂy in the text domain, on a language modeling task and Long Range Arena (LRA (Tay et al., 2020)) benchmarks.
On WikiText-103 (Merity et al., 2016), Pixelﬂy achieves 22.5 perplexity, which is around the same perplexity as GPT-2 small (Radford et al., 2019) but trains 2.1 × faster.
On LRA, Pixelﬂy obtains almost the same accuracy as the full model but gains up to 5 .
2 × speed-up.
Setup: We use WikiText-103 for language modeling and LRA for classiﬁcation tasks.
We use GPT2 small and vanilla Transformer as the base dense models.
The computational bottleneck of GPTsmall 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. Model WikiText-103 ( Speedup
GPT-2-Small 22. BigBird 23. 0.96 × Pixelﬂy 22.
2.1 × GPT-2-Medium 20. BigBird 21. 1.1 ×
Pixelﬂy 21. 2.5 × GPT-2-Small, Medium on WikiText-103: We show training GPT-2-Small, Medium and its Pixelﬂy model from scratch on a commonly used NLP benchmarking dataset, wikiText-103.
We measure their perplexity on that dataset, and our training speed up.
All setup and ﬁnetuning hyperparameters follow the ones in the original paper (Radford et al., 2019).
We present the results in Fig.
It is not hard to see that Pixelﬂy models have great advantages in accuracy-efﬁciency tradeoffs since it maintains the same perplexity as the dense model but achieve up to 2.5 × speed-up in training. Model ListOps [Table]
Pathﬁnder
[Table]
Transformer 36.
[Table]
Vanilla Transformer on
LRA: We compare vanilla transformer and its Pixelﬂy models trained from scratch on LRA benchmark.
We measure the accuracy, throughput, and training time of both models.
Each task has a different sequence length varying between 1024 and 4096.
We follow the implementation and experimental setting in (Xiong et al., 2021).
We compare the performance of Pixelﬂy against the dense transformer and report the results in Fig.
We also include the numbers of other baselines from the same repository in the appendix.
We can see Pixelﬂy cause almost no drop in accuracy while achieving 5.2 × speed-up in time.
5. A BLATION S TUDY
We conduct ablation studies on each component of Pixelﬂy (Details in Appendix L.5).
Speciﬁcally, we present (i) how ﬂat block butterﬂy 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.
Necessity of Flat Block Butterﬂy and Low-rank: (i)
We apply different parameter allocation of ﬂat block butterﬂy and Low-rank component in Pixelﬂy Mixer-S model on CIFAR-10 under the different density varying in [0.05, 0.1, 0.2].
We found that similar to what was reported in (Chen et al., 2021), using around 14 budget on Low-rank and 34 on ﬂat block butterﬂy achieves the best accuracy.
(ii) We also compare Pixelﬂy with baseline sparsity patterns and show it is 2 .
7 × faster than dense, 3 × faster than Butterﬂy, 3 . 2 × faster than BigBird under 10% density.
Block Size: We study the accuracy-efﬁciency trade-off for ﬂat block butterﬂy and random sparsity pattern with different block sizes from 1-32 ( Table 7).
We found that ﬁrst, under the same density, the same sparsity patterns covered with different block sizes could have a big difference in efﬁciency.
Under the same block, the pattern with more locality can be more efﬁcient.
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.
Budget Allocation: 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.
R ELATED W ORK
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., 2020) are carried out to analyze these tickets, but it remains difﬁcult 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 ﬁnd 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 signiﬁcantly smaller and faster models for inference.
Subsequent methods (Li et al., 2016; Lin et al., 2017; 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 to produce sparse models, we target training efﬁciency, whereas pruning mostly focuses on inference efﬁciency at the cost of sacriﬁcing 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; AllenZhu et al., 2019b;c; Song and Yang, 2019), generalization (Allen-Zhu et al., 2019a), connection to data separability (Oymak and Soltanolkotabi, 2020), and cost per iteration (Brand et al., 2021)).
Deep
Double Descent (Nakkiran et al., 2019; d’Ascoli et al., 2020) conjectures that the generalization error improves as the parameter count grows.
It is not surprising that the community is racing to break the record of the largest parameter counts (Radford et al., 2019; Brown et al., 2020; Dosovitskiy et al., 2020;
Tolstikhin et al., 2021; Zhang et al., 2021; Naumov et al., 2019; Jumper et al., 2021).
We provide extended related work in Appendix M.
C ONCLUSION
In our early exploration of many sparsity patterns with complex training procedures, we found that a simple pattern (butterﬂy + low-rank) consistently (though not always) performed among the best.
This motivated us to propose Pixelated Butterﬂy, a simple and efﬁcient sparse training method.
In our quest for simplicity and efﬁciency, we have chosen to use ﬁxed sparsity that aligns with modern hardware, which was sufﬁcient to yield wall-clock training time speedup without sacriﬁcing 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 efﬁcient yet still accurate.
Our ﬂat butterﬂy is a simple ﬁrst order approximation of the rich class of butterﬂy matrices, and there could be more sophisticated approximations that retain more expressiveness.
Our method is a ﬁrst step towards the goal of making sparse models train faster than dense models and make them more accessible to the general machine learning community.
Ethics Statement.
As the amount of data and model size grows, our work seeks to understand how to train those large models more efﬁciently 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 popular models such as MLP-based and Transformer-based architectures, which may improve a wide range of applications, each with their own potential beneﬁts and harms.
For example, making language modeling more efﬁcient might simplify the process of spreading misinformation.
Similarly, better image classiﬁcation models might make automatic surveillance easier.
To alleviate the above risks, we need to address application-speciﬁc issues like privacy, bias and discrimination, going beyond the accuracy metric we currently considered.
Speciﬁcally, for image classiﬁcation task, while our work partially addresses the issue of environmental cost, it does not address other issues such as fairness and bias in model and datasets.
Reproducibility Statement.
To facilitate the reproducibility of our algorithms and results, (i) we include a link to anonymous 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, a complete description of the data processing steps is in Appendix L.
A P ROBLEM F ORMULATION
We formulate the problem of sparse model training as sparse matrix approximation with a simple hardware cost model (Section 2).
We ﬁrst describe our simple cost model for sparse matrix multiplication to reﬂect the fact that parallel hardware such as GPUs are block-oriented (Cook, 2012; Gray et al., 2017): 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 pattern candidates to those that are block-aligned.
Cost model We model the time cost of an operation based on the number of ﬂoating 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 b elements around it (typically b = 16 or 32 depending on the hardware).
Let Cost mem be the memory access cost (either read or write) for a block of b contiguous elements.
Accessing any individual element within that block also costs Cost mem time.
Let Cost ﬂop be the compute cost of a ﬂoating point operation.
Let N blockmem be the number of block memory access, and
N ﬂop be the number of ﬂoating point operations.
Then the total cost of the operation is
Totalcost =
Cost mem · N blockmem + Cost ﬂop · N ﬂop .
This cost model is a ﬁrst 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 a sparse matrix relates to the number of blocks being accessed.
We ﬁrst deﬁne a block cover of a sparse mask.
Deﬁnition A.1.
A sparse mask M ∈ { 0 , 1 } m × n is ( b 1 , b 2 ) - block-aligned if for any index i, j where
M ij = 1 , we also have M i (cid:48) j (cid:48)
= 1 where:
i (cid:48)
=
b
(cid:98)
i
/b 1 (cid:99)
+
r
, j (cid:48)
=
b 2 (cid:98) j/b 2 (cid:99)
+ r 2 for all r 1 = 0 , 1 , . .
, b 1 −
and
r
=
, b
−
The ( b 1 , b 2 ) - block cover of a sparse mask M ∈ { 0 , 1 } m × n is the ( b 1 , b 2 ) -block-aligned mask M (cid:48) ∈ { 0 , 1 } m × n with the least number of nonzeros such that M ij ≤ M (cid:48) ij for all i, j .
We omit the block size ( b 1 , b 2 ) if it is clear from context.
A sparse mask M being ( b 1 , b 2 ) block-aligned means that if we divide M into blocks of size b 1 × b 2 , then each block is either all zeros or all ones.
To get the ( b 1 , b 2 ) -block cover of a sparse mask M , we simply divide M into blocks of size b 1 × b 2 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 N blockmem is the number of nonzero blocks in its (1 , b ) -block cover M (cid:48) (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 b elements around that location.
Fast sparse matrices means block-aligned sparsity pattern For sparsity patterns that are not blockaligned, such as the random sparse pattern where each location is 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 is less than some limit k .
For our theoretical analysis, we consider sparsity patterns that are (1 , b ) - block-aligned.
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 ) -block-aligned and ( b, 1) -blockaligned.
This is equivalent to the condition that the sparsity pattern is ( b, b ) -block-aligned.
Sparse matrix approximation in the forward pass We now formulate the sparse matrix approximation 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, we focus on the case where number of nonzeros in each row is the same.
Deﬁnition A.2 (Forward regression) .
Given four positive integers m ≥ n ≥ d ≥ k ≥ 1 , matrices A ∈ R m × d and B ∈ R d × n .
The goal is to ﬁnd a (1 , b ) -block-aligned binary mask matrix M ∈ { 0 , 1 } m × d that satisﬁes
[Equation]
where M i is the i -th row of M .
Sparse matrix approximation in the backward pass In the backward pass to compute the gradient wrt to the weight matrix
A , we would like to sparsify the gradient CB (cid:62) while preserving as much of the gradient magnitude as possible.
Deﬁnition A.3 (Backward regression) .
Given four positive integers m ≥ n ≥ d ≥ k ≥ 1 , matrices
B ∈ R d × n and C ∈ R m × n .
The goal is to ﬁnd a (1 , b ) -block-aligned binary mask matrix M ∈ { 0 , 1 } m × d such that
[Equation]
where M i is the i -th row of M .
Without making any assumptions, such problems are in general computationally hard Foster et al.
; Razenshteyn et al.
B A NALYSIS OF B UTTERFLY V ARIANTS
We present formal versions of theorems in Section 4 regarding variants of butterﬂy matrices.
We provide full proofs of the results here.
B. B LOCK B UTTERFLY A NALYSIS
Proof of Theorem 4.1.
Let M be an n × n block butterﬂy matrix with block size b .
We want to show that M also has a representation as
an n × n block butterﬂy matrix with block size 2 b .
By Deﬁnition 3.3, M has the form:
[Equation]
Notice that we can combine that last two terms to form a matrix of the form B ( n 2 b , 2 b )(see Fig. 3).
Moreover, other terms in the product of the form B ( nb ,b ) n 2 ib can also be written as B ( n 2 b , 2 b ) n 2 i − 12 b (see Fig. 3).
Thus M also has the form:
[Equation]
In other words, M is also an n × n block butterﬂy matrix with block size 2 b .
Proof of Corollary 4.2.
Dao et al.
(2020, Theorem 3) states that any n × n sparse matrix with s nonzeros can be represented as products of butterﬂy matrices and their transposes, with O ( s log n ) parameters.
For a constant block size b that is a power of 2, the set of n × n block butterﬂy matrices of block size b contains the set of regular butterﬂy matrices by Theorem 4.1.
Therefore any such n × n sparse matrix also has a representation has products of block butterﬂy matrices of block size b and their transposes, with O ( s log n ) parameters.
B. F LAT B UTTERFLY A NALYSIS
We prove Theorem 4.3, which relates the ﬁrst-order approximation in the form of a ﬂat butterﬂy matrix with the original butterﬂy matrix.
Proof of Theorem 4.3.
Let n = 2 m and let B 1 , . .
B m ∈ R n × n be the m butterﬂy factor matrices (we rename them here for simplicity of notation).
Let
[Equation]
Our goal is to show that (cid:107) E (cid:107) F ≤ (cid:15) .
We ﬁrst recall some properties of Frobenius norm.
For any matrices A and C , we have (cid:107) AC (cid:107) F ≤ (cid:107) A (cid:107) F (cid:107) C (cid:107) F and (cid:107)
A + C (cid:107) F ≤ (cid:107) A (cid:107) F + (cid:107) C (cid:107) F .
Expanding the terms of the product in E , we have
E =
[Equation]
(cid:88)
[Equation]
|
[Equation]
(cid:89)
[Equation]
Using the above properties of Frobenius norm, we can bound E :
[Equation]
where in the last step we use the assumption that c ≤ 12 .
We now bound the rank of the ﬁrst-order approximation.
Proof of Theorem 4.4.
Let M ∗
=
I + (cid:80) mi =1 λB i .
Note that any entry in (cid:80) mi =1 λB i has absolute value at most
[Equation]
where we use the assumption that B ∞ max ≤ B max and c ≤ 14 .
Thus any diagonal entry in M ∗ has absolute value at least 1 − 14 = 34 and the off-diagonal entries are at most
[Equation]
Alon (2009, Theorem 1.1) states that: there exists some c > 0 such that for any real M ∈ R n × n , if the diagonal elements have absolute values at least 12 and the off-diagonal elements have absolute values at most (cid:15) where 1 2 √ n ≤ (cid:15) ≤ 14 , then rank( M ) ≥ c log n (cid:15) 2 log1 /(cid:15) .
Applying this theorem to our setting, we have that
[Equation]
We just need to show that B ∞ max B max ≥ 1 2 c √ (cid:15)n to satisfy the condition of the theorem.
Indeed, we have that 1 ≤ B max B ∞ max ≤ √ 2 n as each (cid:107) B
i (cid:107) 0 ≤ 2 n .
Combining the two conditions on B max B ∞ max , we have shown that 1 ≤ B max B ∞ max ≤ 2 c √ (cid:15)n .
This concludes the proof.
B. F LAT B LOCK B UTTERFLY + L OW - RANK A NALYSIS
We show that ﬂat butterﬂy + 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.
(2021) to show a generative process where the attention matrix can be well approximated by a ﬂat butterﬂy + low-rank matrix, but not by a sparse or low-rank alone.
We describe here a generative model of an input sequence to attention, parameterized by the inverse temperature β ∈ R and the intra-cluster distance ∆ ∈ R .
Process 1.
Let Q ∈ R n × d , where d ≥ Ω(log 3 / 2 ( n )) , with every row of Q generated randomly as follows:
For C = Ω( n ) , sample C number of cluster centers c 1 , . . .
, c C ∈ R d independently from N (0 , I d / √ d ) .
For each cluster around c i , sample n i =
b number of elements around c i , of the form z ij =
c i + r ij for j = 1 , . . . , n i where r ij ∼ N (0 , I d ∆ / √ d ) .
Assume that the total number of elements is n = cb and ∆ ≤ O (1 / log 1 / 4 n ) .
Let Q be the matrix whose rows are the vectors z ij where i = 1 , . . .
, C and j = 1 , . .
. , n i .
Let
A = QQ (cid:62) and let the attention matrix be M β = exp( β · A ) .
Theorem B.1.
Let M β , be the attention matrix in Process 1.
Fix (cid:15) ∈ (0 , 1) .
Let R ∈ R n × n be a matrix.
Consider low-rank, sparse, and sparse + low-rank approximations to M β .
Assume (1 − ∆ 2 ) log n ≤ β ≤ O (log n ) .
Flat butterﬂy + low-rank : There exists a ﬂat butterﬂy + low-rank R with n 1+ o (1) parameters with (cid:107) M β − R (cid:107) F ≤ (cid:15)n . 2.
Low-rank : If R is such that n − rank( R ) =
Ω( n ) , then (cid:107) M β − R (cid:107) F ≥ Ω( n ) .
Sparse : If R has sparsity o ( n 2 ) , then (cid:107) M β − R (cid:107) F ≥ Ω( n ) .
Proof sketch.
As the argument is very similar to that of Chen et al.
(2021, Theorem 1), we describe here the modiﬁcations needed to adapt their proof.
The main difference between our generative process and that of Chen et al.
(2021) is that each cluster has the same number of elements, which is the same as the block size.
The resulting attention matrix will have a large block diagonal component, similar to that Chen et al.
However, all the blocks in the block diagonal component has the same block size, which is b .
Moreover, a ﬂat block butterﬂy of block size b contains a block diagonal component of block size b .
Therefore, this ﬂat block butterﬂy matrix plays the same role as the sparse matrix in the proof of Chen et al.
The rest of the argument follows that of theirs.
Roadmap
The analysis of sparse networks is organized as follows.
In Section C we list some basic notations that will be used.
In Section D we consider the problem of 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 sparsity on a , and we show that minimizing the dropout loss is equivalent with a kernel ridge regression problem.
In
Section H we analyze the dynamics of gradient ﬂow and prove the convergence result.
C N OTATIONS
For a vector x , we use (cid:107) x (cid:107) p to denote its (cid:96) p norm, and we mainly consider p = 1 , 2 in this paper.
For a matrix A , we use (cid:107) A (cid:107) 0 , (cid:107) A (cid:107) 1 , (cid:107) A (cid:107) F to denote the (cid:96) 0 norm, entry-wise (cid:96) 1 norm and Frobenius norm of A respectively.
For two matrices A, B ∈ R d × m , we use A ◦ B to denote their Hadamard product.
We use T mat ( n, d, m ) to denote the time of multiplying n × d matrix with another d × m matrix.
For a symmetric matrix A , we use λ min ( A ) to denote its minimum eigenvalue.
We also let vec( A ) be the vectorization of a matrix A in column
ﬁrst order.
We use (cid:104)· , ·(cid:105) to denote standard Euclidean inner product between two vectors.
Moreover, we use N ( µ, Σ) to denote the Gaussian distribution with mean µ and covariance Σ .
We denote the ReLU function by φ ( z )
= max { z, 0 } .
For an event E , we use 1 { E } or 1 E to denote its indicator function.
D S PARSITY ON HIDDEN LAYER WEIGHTS D. A PPLYING MASKS BEFORE MULTIPLICATION
Given matrix A ∈ R n × d , B ∈ R d × n , naively computing AB takes T mat ( n, d, n ) .
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 (cid:62) have the same size.
Our goal is to ﬁnd “optimal" binary mask matrix W ∈ { 0 , 1 } d × n such that,
[Equation]
Remark D.1.
In the practical applications we care about, the function f is the activation function of neural network, e.g., ReLU( z ) =
max { z, 0 } .
We deﬁne a sparse targeted regression problem:
Deﬁnition D.
(Sparse mark regression, (cid:96) 1 version) .
Given a matrix B ∈ R d × n , and a vector a ∈ R d , the goal is to ﬁnd a k -sparse binary vector w ∈ { 0 , 1 } d to minimize the following problem:
[Equation]
Naively, the above problem can be solved in n · d O ( k ) via guess all the (cid:0) dk (cid:1) choices.
Lemma D.3.
The targeted sparse mask regression problem can be solved in n · d O ( k ) .
Proof.
We need to guess (cid:0) dk (cid:1) times, which becomes d O ( k ) .
Each time it takes nd operations, thus the total time is
[Equation]
Deﬁnition D.4 ( (cid:96) 1 version) .
Given three positive integers m ≥ n ≥ d ≥ k ≥ 1 , matrices A ∈ R m × d and B ∈ R d × n .
We deﬁne our problem as ﬁnding the binary matrix W ∈ { 0 , 1 } m × d that satisﬁes
[Equation]
where W i ∗ is the i -th row of W .
Theorem D.5.
The problem being deﬁned as Deﬁnition D.4 can be solved in mnd O ( k ) time.
Proof.
Our problem can be decomposed into m sub-problems as follows:
[Equation]
where A i ∗ means the i -th row of matrix A .
By applying Lemma D.3, each sub-problem
[Equation]
can be solved in n · d O ( k ) time.
Then the problem deﬁned in Deﬁnition D.4 can be solved in
[Equation]
time in total.
Thus we ﬁnish 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 (cid:96) 1 norm, then we can solve the regression problem in polynomial time, as the following parts show.
Deﬁnition D.6 ( (cid:96) 1 version) .
Given a matrix B ∈ R d × n ≥ 0 , and a vector a ∈ R d ≥ 0 , the goal is to ﬁnd a k -sparse binary vector w ∈ { 0 , 1 } d to solve
[Equation]
Lemma D.7.
The targeted (cid:96) 1 version sparse mask regression problem can be solved in
O ( nd + n log n ) which is polynomial time.
Proof.
We ﬁrst consider the situation when a ∈ { 0 , 1 } d .
In this case, we have
[Equation]
where (cid:107) a (cid:62) · B (cid:107) 1 is ﬁxed.
So we only need to consider the following problem:
[Equation]
For simplicity we assume a i = 1 , ∀ i ∈
[ d ] , and we only need to solve
[Equation]
where w has k elements equal to 1 and d − k elements equal to 0 .
For i ∈ [ d ] , we compute S i
=
(cid:80) nj =1 B ij which is the summation of i -th row of B , and sort them as S (1) ≥ S (2) ≥ · · · ≥ S ( n ) .
Then we only need to let w ( i )
= 1 for i ∈
[ k ] and other elements equal to 0 .
Computing all S i takes O ( nd ) time, sorting S
i takes O ( n log n ) time, thus the total time consumption is O ( nd + n log n ) in this case.
Next, we consider the general case when a ∈ R d ≥ 0 .
We let
B i ∗ =
a i B i ∗ and a i = (cid:26) 1 , a i >0 , a i = 0 , ∀ i ∈ [ d ] where B i ∗ is the i -th row of B .
Then our optimization problem is equivalent to
[Equation]
where B ∈ R d × n ≥ 0 and a ∈ { 0 , 1 } d .
Thus we turn this case into the ﬁrst case.
Constructing B and a takes O ( nd ) time, thus the total time consumption is also O ( nd + n log n ) in this case.
Deﬁnition D.8 ( (cid:96) 1 version) .
Given three positive integers m ≥ n ≥ d ≥ k ≥ 1 , matrices A ∈ R m × d ≥and B ∈ R d × n ≥ 0 .
We deﬁne our problem as ﬁnding the binary matrix W ∈ { 0 , 1 } m × d that satisﬁes
[Equation]
where W i ∗ is the i -th row of W .
Theorem D.9.
The problem being deﬁned as Deﬁnition D.8 can be solved in
O ( mnd + mn log n ) time.
Proof.
Our problem can be decomposed into m sub-problems as follows:
[Equation]
where A i ∗ means the i -th row of matrix A .
By applying Lemma D.7, each sub-problem
[Equation]
can be solved in O ( nd + n log n ) time.
Then the problem deﬁned in Deﬁnition D.8 can be solved in m · O ( nd + n log n ) =
O ( mnd + mn log n ) time in total.
Thus we ﬁnish the proof.
D. A PPLYING M ASKS A FTER M ULTIPLICATION
Deﬁnition D.10.
Given matrix B ∈ R d × n , C ∈ R m × n .
The goal is to ﬁnd a mask W ∈ { 0 , 1 } m × d where each column of W is k -sparse
[Equation]
Remark D.11.
The B deﬁned in Deﬁnition D.4 is the same as the B deﬁned in Deﬁnition D.10.
B is corresponding to the X in the neural network setting.
E G RADIENT COMPUTATION
In this section we consider a neural network with one hidden layer and m neurons in this hidden layer.
Suppose x ∈ R
d
is the input,
W = ( w 1 , · · · , w m ) ∈
R
d × m is the weight matrix of the ﬁrst layer, a ∈ R m is the output weight, and M ∈ { 0 , 1 }
d × m is the mask matrix with each column having at most k non-zero entries.
The neural network f : R d → R is deﬁned as
[Equation]
For simplicity, we only optimize W and ﬁx a .
Consider the mean square loss
[Equation]
In the forward computation, for a batch of data points x 1 , · · · , x n ∈ R d , let X ∈ R
d
× n denote the input data points matrix.
For convenience, we deﬁne
[Equation]
where η is the step size.
We deﬁne function g t : R d → R m as
[Equation]
Lemma E.1.
We can express ∆ W ( t ) as
[Equation]
and each column of ∆ W ( t ) has at most k non-zero entries.
Proof.
From the deﬁnition, we know
[Equation]
Since each column of M has at most k non-zero entries, we easily know each column of ∆ W ( t ) also has at most k non-zero entries.
Lemma E.2.
Suppose that matrices M ∈ R d × m , W ( t )
∈ R d × m and ∆ W ( t ) ∈ R
d × m are given and pre-computed, then we can compute f t +1 ( X ) in
O ( mnk ) time.
(Here f t +1 ( X ) is the evaluation of f at W ( t + 1) .)
Proof.
The goal is to compute
[Equation]
By using Lemma E.1, we have
[Equation]
Notice that we have already computed ( M ◦ W ( t )) (cid:62) · X ∈ R m × d from previous iteration, so we only need to compute (∆ W ( t ))
(cid:62) · X where ∆ W ( t ) ∈ R
d × m and X ∈ R d × n .
By using Lemma E.1, each row of (∆ W ( t )) (cid:62) has at most k non-zero entries, thus we can compute (∆ W ( t )) (cid:62) · X in O ( mnk ) time.
Lemma E.3.
Suppose that matrices M ∈ R d × m , W ( t )
∈ R d × m and f t ( X ) are given and precomputed, then we can compute ∂L ( W ( t ))
∂W ( t ) in O ( mnk ) time.
Proof.
By using Lemma E.1, we have
[Equation]
where g t ( x ) =
( f ( x ) − y ) · diag { φ (cid:48)
(( M ◦ W ( t )) (cid:62) · x ) }· a ∈ R m and g t ( X ) =
( g t ( x 1 ) , · · · , g t ( x n ))
∈
R m × n .
We ﬁrst compute M ◦ W ( t ) in O ( mk ) time, then we can construct g t ( X ) ∈ R m × n in n · O ( mk ) time.
Given g t ( X ) , since we only need to compute km entries of X · g (cid:62) t ( X ) , where each entry can be computed in O ( n ) time, thus we can compute ∂L ( W ( t ))
∂W ( t ) in O ( mnk ) time.
1: procedure S PARSE T RAINING ( { x i , y i } i ∈ [ n ] ) 2:
Initialization a r , w r (0) ∼ N (0 , I d ) for r ∈ [ m ] .
3: for t = 1 → T do 4: /*forward computation*/ 5:
Compute M ◦ W ( t ) (cid:46) Takes O ( mk ) time.
6: for i = 1 → n do 7: f t ( x i ) ← a (cid:62) φ (( M ◦ W ( t )) (cid:62) · x i ) (cid:46) Takes O ( mk ) time.
8: g t ( x i ) ← ( f ( x i ) − y i ) · diag φ (cid:48) (( M ◦ W ( t )) (cid:62) · x i ) · a (cid:46) Takes O ( mk ) time.
9: end for 10: /*backward computation*/ 11: g t ( X ) ← ( g t ( x 1 ) , · · · , g t ( x n )) .
12: ∂L ( W ( t )) ∂W ( t ) =
( X · g (cid:62) t ( X )) ◦ M (cid:46) Takes O ( mnk ) time.
W ( t + 1) =
W ( t ) + ∆ W ( t ) (cid:46) ∆ W ( t ) = − η ∂L ( W ( t ))
∂W ( t ) .
14: end for 15: end procedure
F N EURAL T ANGENT K ERNEL , C ONVERGENCE , AND G ENERALIZATION
Our analysis relies on the neural tangent kernel (NTK) (Jacot et al., 2018) of the network.
Deﬁnition F.1.
Let f ( · , θ ): R d
→ R be the function speciﬁed by a neural network with parameters θ ∈ R p and input dimension d .
The parameter θ is initialized randomly from a distribution P .
Then its neural tangent kernel (NTK) (Jacot et al., 2018) is a kernel K : R d × R d
→ R deﬁned by:
[Equation]
We can 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 f dense denote a ReLU neural network with L layers with dense weight matrices θ dense with NTK K dense , and let f sparse be the ReLU neural network with the same architecture and with weight matrices θ sparse whose rows are k -sparse, and with NTK K sparse .
Let x 1 , . . .
, x N be the inputs sampled from some distribution P X .
Suppose that the empirical NTK matrices
K d = K dense ( x i , x j ) and K s =
K sparse ( x i , x j ) for ( i, j ) ∈
[ N ] × [ N ] satisfy (cid:107) K
d − K s (cid:107) ≤ δ .
Training.
We knew the the number of iterations of dense network is λ min ( K d ) − 2 n 2 log(1 /(cid:15) ) to reach the (cid:15) training loss.
For sparse network we need ( λ min ( K d ) − δ ) − 2 n 2 log(1 /(cid:15) ) .
Generalization.
We knew the the number of iterations of dense network is λ min ( K d ) − 2 n 2 log(1 /(cid:15) ) to reach the generalization error (cid:15) training loss.
For sparse network we need ( λ min ( K d ) − δ ) − 2 n 2 log(1 /(cid:15) ) .
These results relate the generalization bound of sparse models to that of dense models.
G D ROPOUT N EURAL N ETWORK AND KRR
We consider a two layer neural network with ReLU activation function, and write
[Equation]
where w r (0) ∼ N (0 , I d ) ∈ R d , a r ∼ unif( {− 1 , +1 } ) and all randomnesses are independent.
We will ﬁx a r during the training process and use 1 √ m normalization factor, both of which are in the literature of Du et al.
; Song and Yang (2019); Brand et al.
Suppose the training data are ( x 1 , y 1 ) , . .
, ( x n , y n ) ∈ R d × R , we deﬁne the classical objective function (cid:98) L as follows:
[Equation]
The gradient with respect to loss function (cid:98) L is
[Equation]
We consider the effect of dropout on network training.
For each r ∈ [ m ] , we introduce the mask by deﬁning random variable σ r as follows:
[Equation]
It is easy to see that E [ σ r ] = 0 · (1 − q ) + (1 /q ) · q = 1 and E [ σ 2 r ] = 0 2 · (1 − q )
+ (1 /q ) 2 · q = 1 /q .
We assume σ i and σ j are independent for any i (cid:54) =
j , then E [ σ i σ j ] =
E [ σ i ] E [ σ j ] = 1 .
Let σ = ( σ 1 , · · · , σ m ) , we deﬁne our dropout neural net as
[Equation]
Dropout explicitly change the target function, since we need to minimize the (cid:96) 2 distance between
F ( W, x, σ ) and y , instead of f ( W, x ) and y .
Formally, we deﬁne the dropout loss as
[Equation]
We ﬁrst give an explicit formulation of L which also shows the difference between L and (cid:98) L .
Lemma G.1.
The dropout loss deﬁned in Eq.
(4) can be expressed as the sum of classical loss (cid:98) L and a regularization term as
[Equation]
Proof.
Since E [ σ r ] = 1 , we have
[Equation]
holds for any i ∈ [ n ] .
Next, we show the difference between L and (cid:98) L :
[Equation]
where the ﬁrst step follows from deﬁnition, the second step follows from the linearity of expectation, the third step follows from Eq.
(6), the forth step follows from expansion, the ﬁfth step follows from E [ σ r 1 σ r 2 ] = 1 for r 1 (cid:54) =
r 2 and E [ σ 2 r 1 ] = 1 q , and the last step follows from a 2 r = 1 .
Thus we have
[Equation]
and ﬁnish the proof.
Before we move on, we introduce some extra notations and deﬁnitions.
We denote
[Equation]
Deﬁnition G.2.
We deﬁne matrix G ∞ ∈ R n × n which can be viewed as a Gram matrix from a kernel associated with ReLU function as follows:
[Equation]
and assume λ 0 = λ min ( G ∞ ) > 0 7 .
Deﬁnition G.3.
We deﬁne the masked matrix Φ W
( X, σ ) ∈ R n × md as
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Equation]
  and also deﬁne the unmasked matrix (cid:98) Φ W ( X ) ∈ R n × md as
[Equation]
  .
Deﬁnition G.4.
We deﬁne the masked block diagonal matrix Ψ W ( X, σ ) ∈ R md × md as Ψ W
[Equation]
:= 1 m diag (cid:16)
[Equation]
(cid:17) .
where ∀ r ∈ [ m ] , ψ r ∈ R d × d is deﬁned as
[Equation]
We also deﬁne the unmasked block diagonal matrix (cid:98)
Ψ W ( X ) ∈ R
md × md as
[Equation]
where ∀ r ∈ [ m ] , (cid:98) ψ r ∈ R d × d is deﬁned as
[Equation]
r
[Equation]
= n (cid:88)
[Equation]
Lemma G.5.
It is easy to verify that
[Equation]
where
[Equation]
For convenience, we will simply denote Φ W =
Φ W ( X, σ ) and Ψ W =
Ψ W
( X, σ ) .
Then by using the above notations, we can express our dropout loss as L ( W ) =
12 E σ [ (cid:107) Φ W W − Y (cid:107) 22 ] .
Lemma G.6.
If we denote λ = 1 − qq ≥ 0 , then we have
[Equation]
Proof.
As for the ﬁrst term, we have
[Equation]
i = ( 1 √ m m (cid:88) r = a r 1 (cid:104) w r ,x i (cid:105)≥ 0 x (cid:62) i · w r − y i )=
n (cid:88)
i =
( 1 √ m m (cid:88) r = a r φ ( w (cid:62) r x i ) − y i )
=
n (cid:88) i =
( f ( W, x i ) − y i )
= 2 (cid:98) L ( W ) .
As for the second term, since (cid:98) Ψ W is a block diagonal matrix, we have
W
[Equation]
Ψ
[Equation]
m m (cid:88)
r = (cid:16) w (cid:62) r · (cid:0) a 2 r n (cid:88) i
=
x
i x (cid:62) i · 1 2 (cid:104) w r ,x i (cid:105)≥ 0 (cid:1) · w r (cid:17) =m m (cid:88) r
= n (cid:88)
i = (cid:0) ( w (cid:62) r x i ) · ( w (cid:62) r x i ) (cid:62) · 1 2 (cid:104) w r ,x i (cid:105)≥ 0 (cid:1)
=
m n (cid:88) i = m (cid:88) r = φ ( w (cid:62)
r x i ) 2 .
Thus by using Lemma G.1, we have
[Equation]
and ﬁnish the proof.
Remark G.7.
A classical kernel ridge regression problem can be deﬁned as
[Equation]
where φ : R d → F is a feature map.
Note that Lemma G.6 breaks the dropout loss into two parts: the ﬁrst part is an error term, and the second part can be seen as a regularization term.
Thus the task of minimizing the dropout loss L ( W ) is equivalent to a kernel ridge regression (KRR) problem.
H D YNAMICS OF K ERNEL M ETHODS (C ONTINUOUS G RADIENT F LOW )
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 kernel method under the mask σ , which adds sparsity in the output layer.
Our problem will be considered in over-parameterized scheme.
First we introduce some additional deﬁnitions and notations.
We deﬁne symmetric Gram matrix G ( W ) as G ( W ) :
= (cid:98) Φ W · (cid:98) Φ (cid:62) W ∈ R n × n .
For all i, j ∈ [ n ] × [ n ] , we have
[Equation]
We deﬁne block symmetric matrix H ( W ) as H ( W ) =
(cid:98) Φ (cid:62) W · (cid:98)
Φ W ∈ R
md × md .
Then for all i, j ∈ [ m ] × [ m ] , the ( i, j ) -th block of H ( W ) is
[Equation]
By using Lemma G.6, we consider the corresponding kernel regression problem:
[Equation]
where (cid:98) Φ ∈ R n × md , W ∈ R md × 1 , Y ∈ R n × 1 and (cid:98) Ψ ∈ R
md × md .
The main difference from neural network is that we assume (cid:98) Φ (related to NTK, e.g., see Deﬁnition G.3) and
(cid:98) Ψ (related to regularization term, e.g., see Deﬁnition G.4) do not change during the training process.
The gradient of L k can be expressed as
[Equation]
We use W (cid:63) to denote the optimal solution of Eq.
(8), and it satisﬁes
[Equation]
Since (cid:98) Ψ is a positive diagonal matrix, (cid:98) Φ − 12 exists, thus we have
[Equation]
Next, we consider the question from a continuous gradient ﬂow aspect.
In time t , we denote W ( t ) =
vec( W ( t )) , (cid:98) Φ( t )
=
(cid:98) Φ W ( t ) , (cid:98) Ψ( t )
= (cid:98) Ψ W ( t ) .
We also denote G ( t ) =
G ( W ( t )) and H ( t )
=
H ( W ( t )) .
Following the literature of Du et al.
(2019), we consider the ordinary differential equation deﬁned by
[Equation]
Lemma H.1 (Lemma 3.1 in Du et al.
If m = Ω( n 2 λ 20 log( nδ )) , we have with probability at least 1 − δ , (cid:107) G (0) − G ∞ (cid:107) 2 ≤ λ 0 4 and λ min ( G (0)) ≥ 34 λ 0 .
Lemma H.2 (Lemma 3.2 in Du et al.
If w 1 , · · · , w m are i.i.d generated from N (0 , I d ) , then with probability at least 1 − δ , the following holds.
For any set of weight vectors w 1 , · · · , w m ∈ R d that satisfy for any r ∈ [ m ] , (cid:107) w r − w r
0) (cid:107) 2 ≤ cδλ 0 n 2 for some small positive constant c , then matrix
G ∈ R d × d satisﬁes (cid:107) G − G (0) (cid:107) 2 < λ 0
and λ min ( G ) > λ 0 2 .
The above lemma shows that for W that is close to W (0) , the Gram matrix G also stays close to the initial Gram matrix G (0) , and its minimal eigenvalue is lower bounded.
Lemma H.3 (Gradient Flow) .
If we assume λ min ( (cid:98) Ψ) ≥ Λ 0 > 0 , then with probability at least 1 − δ , for w 1 , · · · , w m ∈ R d that satisfy ∀ r ∈ [ m ] , (cid:107) w r − w r
0) (cid:107) 2 ≤ cδλ 0 n 2 , we have
[Equation]
[Table]
Proof.
By using Eq.
(9) and Eq.
[Equation]
we can express d W d t as
[Equation]
Then we have
[Equation]
where the second step follows from Eq.
(12), the fourth step follows from Eq.
(10), and the last step follows from the deﬁnition that λ
= λ min ( G ) =
λ min ( (cid:98)
[Equation]
As for the second term in the Eq.
(13), we have
[Equation] [Equation]
Thus by Eq.
(13) and Eq.
(14) we have
[Equation]
By letting γ = 2 λ
+ 2 λ Λ
we ﬁnish the proof.
For convenience, we denote u ( t ) =
(cid:98) Φ( t ) · W ( t )
∈ R n .
Then it is easy to verify that
[Equation]
showing that u ( t ) is the prediction in time t .
Lemma H.4 (Convergence rate) .
If we assume λ min ( G ( s )) ≥ λ 0 2 holds for 0 ≤ s ≤ t , then we have
[Equation] [Equation]
Proof.
From Eq.
(9), we can express the dynamics by using u ( t ) as
[Equation]
Thus we have
[Equation]
As for the second term, we have
[Equation]
where for j ∈ [ n ] , U j ( t ) ∈ R can be expressed as
[Equation]
We denote U ( t ) =
[ U 1 ( t ) , · · · , U n ( t )] (cid:62) ∈ R n and have
[Equation]
18) and our dynamics becomes
[Equation]
19) showing that dd t (cid:0) e ( λ 0 +2 λ/m ) t (cid:107) u
( t ) − Y (cid:107) 22 (cid:1) ≤ 0 .
Thus e ( λ 0 +2 λ/m ) t (cid:107) u ( t ) − Y (cid:107) 22 is a decreasing function with respect to t , and we have
[Equation]
As for bounding (cid:107) w r ( t ) − w r (0) (cid:107) 2 , we use the same method as in Lemma 3.3 of Du et al.
Thus we complete the proof.
Finally, by combining Lemma H.1, H.2, H.3 and H.4, we have the following convergence result.
Theorem H.5 (Convergence of gradient ﬂow) .
Suppose λ 0 > 0 , m = poly( n, 1 /λ 0 , 1 /δ ) , then with probability at least 1 − δ over the randomness of initialization, we have
[Equation]
[Table]
The above theorem shows that in the over-parameterized setting (when m is large enough), the training loss of the kernel ridge regression problem deﬁne in Eq.
(8) converges to 0 in a linear rate.
By comparing our Theorem H.5 with Theorem 3.2 in Du et al.
, we can ﬁnd 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.
I M ETHOD D ETAILS
We describe some details of our method.
I. C OMPUTE 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 a parameter budget B , we ﬁnd the density of each layer type that minimize the models’ total cost of matrix multiplication.
For example, in Transformers, let d a and d m be the density of the attention and the MLP layers.
Let s be the sequence length and d be the feature size.
The attention layer with density d a will cost d a ( n 2 + nd ) , and the fully connected layers with density d m will cost 2 d m nd .
We then set d a and d m to minimize the total cost while maintaining the parameter budget:
[Equation]
# of trainable
[Equation]
As this is a problem with two variables, we can solve it in closed form.
I. L OW - RANK IN A TTENTION
In Section 3.3, we describe how to use the sparsity pattern from ﬂat block butterﬂy 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.
Chen et al.
(2021) describes a general method to combine sparse and low-rank attention, where one uses the sparse component to discount the contribution from the low-rank component to 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 a “global” sparsity mask (as shown in Fig. 12).
Indeed, a 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 2 w , and is low-rank (for small w ).
We also make the global component block-aligned (i.e., set w to be a multiple of the smallest supported block size such as 32) for hardware efﬁciency.
I. C OMPARISON TO O THER S PARSITY P ATTERNS FOR A TTENTION
In the context of sparse attention, other sparsity patterns such as BigBird and Longformer also contain a “global” component, analogous to our low-rank component.
Their “local” component is contained in the block diagonal part of the ﬂat block butterﬂy sparsity pattern.
The main difference that we do not use the random components (e.g., BigBird), and the diagonal strides from ﬂat block butterﬂy 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. S PARSITY M ASK FOR R ECTANGULAR M ATRICES
We have described the sparsity masks from ﬂat block butterﬂy 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. 10).
We have found this to work consistently well across tasks. [Figure] [Equation] [Figure]
J B ENCHMARKING OF B UTTERFLY M ULTIPLY
We validate that ﬂat butterﬂy matrices (sum of factors) can speed up multiplication on GPUs compared to butterﬂy matrices (products of factors).
Consider the matrix M ∈ R n × n that can be written as products of butterﬂy factors of strides of up k (a power of 2), with residual connection:
[Equation]
The ﬁrst-order approximation of M has the form of a ﬂat butterﬂy matrix with maximum stride k (Section 3.2):
[Equation]
Notice that M is a product of log 2 k factors, each has 2 n nonzeros, so multiplying M by a input vector x costs O ( n log k ) operations (by sequentially multiplying x by the factors of M ).
The ﬂat version
M ﬂat is a sparse matrix with O ( n log k ) nonzeros as well, and the cost of multiplying M ﬂat x is also
O ( n log k ) .
However, in practice, multiplying M ﬂat x is much more efﬁcient on GPUs than multiplying
Mx because of the ease of parallelization.
We measure the total time of forward and backward passes of multiplying either M ﬂat x and compare to that of multiplying Mx for different maximum strides, as shown in Fig.
We see that “ﬂattening” the products brings up to 3 × speedup.
[Figure]
We use matrix size 1024 × 1024 with
[Figure]
K E XHAUSTED S EARCHING S PARSITY P ATTERNS FOR E FFICIENT S PARSE T RAINING
We describe here our early exploration of searching among different sparsity patterns that has been proposed in the literature.
We use a metric derived from the NTK, which has emerged as one of the standard metric to predict the training and generalization of the model.
We consistently found the butterﬂy + low-rank pattern to perform among the best.
In Appendix K.1, we describe the challenges of selecting sparsity patterns for every model components using the a 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 patterns, and the sparse models can inherit the training hyperparamters of the dense models, reducing the need for hyperparameters tuning.
K. C HALLENGES AND A PPROACHES
Challenge 1: We seek sparsity patterns for each model components that can closely mimic the training dynamics of the dense counterpart.
As mentioned in Theorem D.9, it is NP-hard to ﬁnd 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 is 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.or 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, 1971), and has been rediscovered in Longformer and BigBird attention patterns.
Global: this component involves interaction between all elements and a small subset of 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 in Longformer and BigBird patterns as well.
• Butterﬂy: this component corresponds to interaction between elements that are some ﬁxed distance apart.
The many divide-and-conquer algorithms, such as the classical fast Fourier transform (Cooley and Tukey, 1965), uses this pattern at each step.
Butterﬂy matrices reﬂects this divide-and-conquer 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 a generalization of 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 a ﬁxed pool of sparsity patterns for each layer, if the model has many layers, the number of possible layer-pattern assignments is exponentially large.
(cid:46)
M
M do combinations 5:
Let s be the sparsity mask assignment ( t i r i , m i , n i
→ M i .
6: if TotalCompute ( s )
B then
, Check if masks satisfy budget constraint 7:
Let M s be the ﬂattened sparse masks 8:
K sparse ← NTK f θ ◦ M s
X ) 9: d s ← D dense
) (cid:46)
10: if d min > d s then 11: d min ← d s s s 12: if 13: end
14: end for 15: return s out
Approach 2:
To further reduce the search space, we constrain each layer type (attention, MLP) to have the same sparsity pattern.
For example, if there are 10 patterns and 2 layer types, the candidate pool is 10 2 = 100 combinations.
Challenge 3: Computing the empirical NTK on the whole dataset is expensive in time and space, as it scales quadratically in the dataset size.
Approach 3: We compute the empirical NTK on a randomly chosen subset of the data (i.e., a 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. A LGORITHM D ESCRIPTION
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 as a series of matrix multiplies.
We ﬁrst deﬁne:
• Model schema: a list of layer types t (e.g., attention, linear layers in MLP), number of layers r of that type, and dimension of the matrix multiplies m × n .
We
denote it as Ω = { ( t 1 , r 1 , m 1 , n 1 ) , . . .
, ( t | Ω | , r | Ω | , m | Ω | , n | Ω | ) } .
• A mask M of dimension m × n is a binary matrix { 0 , 1 } m × n .
The compute of a mask is the total number of ones in the matrix: compute( M ) =
(cid:80) i,j M ij .
• A sparsity pattern P m × n for matrix dimension m × n is a set of masks { M 1 , ..., M | P | } , each of dimension m × n .
• A sparsity mask assignment is a mapping from a model schema Ω to masks M belonging to some sparsity pattern P : s : ( t, r, m, n ) →
M . • Given a set of sparsity patterns P 1 , . .
, P k , the set of sparsity mask candidate C is the union of sparsity masks in each of P
i : C = ∪ P
i • A sparsity pattern assignment s satisﬁes
the compute budget B if:
[Equation]
• Let θ be the ﬂattened vector containing the model parameters, and let M s be the ﬂattened vector containing the sparsity mask by the sparsity mask assignment s .
Let f θ ( x ) be the output of the dense network with parameter θ and input x .
Then the output of the sparse network is f θ ◦ M s ( x ) .
• The empirical NTK of a network f θ on a data subset X = { x 1 , . . .
, x
| X | } is a matrix of size | X | ×
[Equation]
The formal algorithm to assign the sparsity mask to each layer type is described in Algorithm 2.
The main idea is that, as the set of sparsity mask candidate is ﬁnite, 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 ﬁts under the budget.
K. M ETHOD P ROPERTIES : R EDISCOVERING C LASSICAL S PARSITY P ATTERNS , N O A DDITIONAL H YPERPARAMETER T UNING
When applied to the Transformer architecture, among the sparsity components described in Appendix K.1, the NTK-guided heuristic consistently picks the local and global components for both the attention and MLP layers.
Moreover, the butterﬂy component is also consistently picked for image data, reﬂecting the 2D inductive bias in this component 8 .
While some of these patterns have been proposed for sparse attention, it is surprising that they are also picked for the MLP layers.
The most popular type of sparsity pattern in MLP layers is 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) as the dense model (Section 5).
L E XPERIMENT D ETAILS L. D ATASETS
• Cifar
(Krizhevsky et al., 2009) consists of 60000 coloured images of resolution 32 × 32.
Each of them belong to one of 10 classes, including airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks.
Among these, 50000 images are allocated to be the training set and 10000 images the testing set.
• Cifar
(Krizhevsky et al., 2009) is similar to Cifar10.
It also consists of images of resolution 32 × 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 in testing set.
• ImageNet1K (Russakovsky et al., 2015) spans 1000 object classes, containing 1,281, 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.
• WikiText-103 (Merity et al., 2016) contains articles from the wikipedia page.
It extracts veriﬁed articles from Wikipedia, which add up to over 100 million tokens.
Compared to other datasets, such as Penn Treebank (PTB) (Taylor et al., 2003), WikiText features a larger vocabulary and preserves original upper/lower cases, punctuation and numbers.
L. M ODEL C ONFIGURATIONS AND H YPERPARAMETER
We summarize the details required to replicate our experiments below
Baseline Model: Except for dense model.
We choose our baselines for each experiment base on the following. RigL aims to sparsify model weights/parameters, so we use it as a baseline in MLP-based models (Mixer). BigBird focuses on attention matrices, so we used it as a baseline in Transformerbased models (ViT, GPT-2). L.2. I MAGE C LASSIFICATION
[Table]
We report more details on the models, including number of parameters and FLOPs, in Table 4. L.2. L ANGUAGE M ODELING We report more details on the models, including number of parameters and FLOPs, in Table 5 and
Table 6. [Table] [Table] [Table]
[Table]
L. M EASURING E MPIRICAL NTK
The Empirical NTK is a rough estimation of the real NTK, in which the width of the neural net goes to inﬁnity.
As the width grows, the kernel gets closer to its inﬁnite-width limit.
Fortunately, both of our model of interests, MLP-Mixer and Vision Transformer [Table] are wide and overly parameterized.
Therefore they are only one step away from the real NTK domain .
This allows us to use the Empirical NTK to approximately predict their training behaviors
[Table]
As described in equation
[Table]
we ﬁrst compute the gradient of each data sample, then we compute pairwise product to construct the Empirical NTK.
Despite we are using a relatively small dataset, it’s still expensive to construct kernel for large models, such as ViTs and MLP-Mixers.
In practice, we ﬁnd that it’s sufﬁcient to compute kernels for a subsampled dataset.
MLP-Mixer and Vision Transformer each represent one type of module of interest for our sparsiﬁcation.
In MLP-Mixer, we study the sparse behavior of the Linear module, whereas in Vision Transformer we mainly focus on sparsifying attentions.
All models are ﬁrst sparsiﬁed to around 10% of the original dense compute.
Then we compare their NTK kernels with their original dense kernel.
To eliminate noise, we run three random seeds, ie, three different initializations for each pair of conﬁgurations and we report the mean relative difference between the kernels with respect to the norm of the dense kernel.
L. T RANSFER L EARNING E XPERIMENTS
We conduct extended experiments to test the generalization of our pretrained sparse models on downstream tasks.
Speciﬁcally, we ﬁnetune Pixelﬂy pretrained model (ImageNet) on CIFAR-10 and show that it get 99.03% accuracy compared to 98.77% on our pretrained dense ViT-B/16 model.
In addition, we see more than 2 × speed up on downstream task ﬁne-tuning process as well.
L. M ICROBENCHMARKING
In this section, we perform microbenchmarking on a 4K × 4K sparse matrix multiplication.
We aim to show that Pixelﬂy patterns are far more hardware friendly than random patterns.
For a 4K × 4K matrix, expected density is the number of non-zero entries/(4K × 4K) ; actual density is the number of accessed entries/(4K × 4K), e.g. even if there is only one non-zero, 32 × 32 entries would be accessed because the hardware block size is 32 × 32.
When random patterns are generated with small block size, (e.g 1 × 1 , 2 × 2 ), the resources, such as memory access and computes(denoted by Actual Density), required to compute a random sparse matrix of density 1 .
25% are equivalent to computing a dense matrix multiplication.
This is further reﬂected in the latency: As pattern block size shrinks, deviating from the hardware block size of 32 × 32 , the random patterns’ latency worsens, whereas the Pixelﬂy remains efﬁcient.
Vanilla Butterﬂy is 5 × slower than Pixelﬂy as expected, because (1) it does not take advantage of the hardware property – not structured sparsity(2) it is a series of products.
[Table]
L. E FFICIENT I MPLEMENTATION OF P IXELFLY
We run all of our experiments on V100 GPUs.
We rely on efﬁcient implementation of block sparse matrix multiply and block sparse attention from the libraries Triton ( https://github.com/ openai/triton ) and https://github.com/huggingface/pytorch_block_sparse .
For the low-rank part, we rely on efﬁcient (dense) matrix multiplies from cuBLAS.
In particular, to multiply the input x by the low-rank matrix UV (cid:62) , we multiply U ( V (cid:62) x ) .
We keep the same number of training epochs as that of the dense models (e.g., on ImageNet, the dense model and the Pixelﬂy model are trained for 300 epochs).
The training speedup of the Pixelﬂy models is due to faster time per epoch.
We do not use 2:4 sparsity (available on Ampere GPUs such as A100).
Such ﬁne-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.
L. A BLATION : S PEED -A CCURACY T RADEOFF OF P IXELFLY
We conduct an ablation experiment to examine the speed-accuracy trade of Pixelﬂy: on the ImageNet dataset and the Mixer-B/16 model, we replace the dense matrices with ﬂat block butterﬂy + low-rank matrices, while varying the compute / parameter budget.
We plot the speed-accuracy tradeoff in Fig. 13.
[Figure] Training speedup compared to dense model [Figure] L. C OMPARISON A GAINST O RIGINAL B UTTERFLY
We compare Pixelﬂy against original Butterﬂy matrices (Dao et al., 2020) on the ImageNet dataset and Mixer-B/16 dense model. We present the results in Table 8. We notice another beneﬁt of Pixelﬂy compared to Butterﬂy: it trains more stably and requires less careful initialization. Since Butterﬂy is a product of many factors, it requires careful initialization, otherwise the activation and gradient will be very large or very small.
[Table]
M E XTENDED R ELATED W ORK
In this section, we extend the related works referenced in the main paper and discuss them in detail.
M. N EURAL P RUNING
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 ﬁlter pruning to reduce the cost of running convolution models up to 38 % , Lin et al.
(2017) prunes the network at runtime, hence retaining the ﬂexibility of the full model.
Dong et al.
(2017) prunes the network locally in a layer by layer manner.
Sanh et al.
(2020) prunes with deterministic ﬁrst-order information, which is more adaptive to pretrained model weights.
Lagunas et al.
(2021) prunes transformers models with block sparsity pattern during ﬁne-tuning, which leads to real hardware speed up while maintaining the accuracy.
Zhu and Gupta (2017) ﬁnds 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 efﬁciency, whereas pruning mostly focuses on inference efﬁciency and disregards the cost in ﬁnding the smaller model.
M. L OTTERY T ICKET H YPOTHESIS
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 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 to use one 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 to exist; Pensia et al.
(2020) suggests a 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 ﬁnd tickets without any training labels.
M. N EURAL T ANGENT K ERNEL
Our work rely heavily on neural tangent kernel in theoretical analysis.
Neural Tangent Kernel Jacot et al.
(2018) is ﬁrst proposed to analyse the training dynamic of inﬁnitely wide and deep networks.
The kernel is deterministic with respect to the initialization as the width and depth go to inﬁnity, 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 ﬁnite learning rate is enough for the model to follow NTK dynamic.
Arora et al.
(2019b) points out that there is still a gap between NTK and the real ﬁnite NNs.
Cao and Gu (2020) sheds light on the good generalization behavior of overparameterized deep neural networks.
Arora et al.
(2019a) is the ﬁrst one to show generalization bound independent of the network size.
Later, some works reveal the training dynamic of models of ﬁnite width, pointing out the importance of width in training:
Hayou 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. O VERPARAMETERIZED M ODELS
Our work mainly targets overparameterized models.
In Nakkiran et al.
(2019), the double descendent phenomenon was observed.
Not long after that, d’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, Arora et al.
speculates that overparameterization 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., 2020), are pushing the boundary on text generation and understanding; Their amazing zero-shot ability earn them the title of foundation models (Bommasani et al., 2021).
On the computer vision side, Dosovitskiy et al.
; Tolstikhin et al.
; 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 on recommendation with a 21 billion large embedding; Jumper et al.
(2021) from
DeepMind solve a 50 year old grand challenge in protein research with a 46-layer Evoformer.
In our work, we show that there is a more efﬁcient way to scale up model training through sparsiﬁcation and double descent only implies the behavior of the dense networks.