Differentiable Learning of Generalized Structured Matrices for Efficient Deep Neural Networks
Authors: Changwoo Lee
Hun-Seok Kim
Venue: ICLR 2024 poster (2024)
Keywords: Structured Matrix
Block Low Rank
Low Rank
Efficient Neural Network
Transformer
Fourier
Dirichlet Kernel
FFT
Boxcar
Pruning
Compression
Primary area: general machine learning (i.e., none of the above)
TL;DR: A differentiable structured matrix learning framework that can discover new types of structured matrices for efficient DNNs.

Abstract: This paper investigates efficient deep neural networks (DNNs) to replace dense unstructured weight matrices with structured ones that possess desired properties. The challenge arises because the optimal weight matrix structure in popular neural network models is obscure in most cases and may vary from layer to layer even in the same network. Prior structured matrices proposed for efficient DNNs were mostly hand-crafted without a generalized framework to systematically learn them. To address this issue, we propose a generalized and differentiable framework to learn efficient structures of weight matrices by gradient descent. We first define a new class of structured matrices that covers a wide range of structured matrices in the literature by adjusting the structural parameters. Then, the frequency-domain differentiable parameterization scheme based on the Gaussian-Dirichlet kernel is adopted to learn the structural parameters by proximal gradient descent. On the image and language tasks, our method learns efficient DNNs with structured matrices, achieving lower complexity and/or higher performance than prior approaches that employ low-rank, block-sparse, or block-low-rank matrices.

Published as a conference paper at ICLR 2024
DIFFERENTIABLE LEARNING OF GENERALIZED STRUCTURED MATRICES FOR EFFICIENT DEEP NEU- RAL NETWORKS
Changwoo Lee, Hun-Seok Kim University of Michigan. Ann Arbor, MI. USA fcwoolee, (anseok]@umiine edu
ABSTRACT
This paper investigates efficient deep neural networks (DNNs) to replace dense unstructured weight matrices with structured ones that possess desired properties. The challenge arises because the optimal weight matrix structure in popular neu- ral network models is obscure in most cases and may vary from layer to layer even in the same network. Prior structured matrices proposed for efficient DNNs were mostly hand-crafted without generalized framework systematically learn them. To address this issue, we propose a generalized and differentiable frame- work learn efficient structures of weight matrices by gradient descent. We first define a new class of structured matrices that covers a wide range of structured ma- trices in the literature by adjusting the structural parameters. Then, the frequency- domain differentiable parameterization scheme based on the Gauusian-Dirichlet kernel is adopted to learn the structural parameters by proximal gradient descent. On the image and language tasks, our method learns efficient DNNs with struc- tured matrices, achieving lower complexity and/or higher performance than prior approaches that employ low-rank, block-sparse, or block-low-rank matrices.
INTRODUCTION
Deep Neural Networks (DNNs) for large language models (LLMM)) (Vaswani et al. 2017 Devlin ct al. 2018: Radford et al. 2019: Brown et al. 2020) and vision tasks (Dosovitskiy et al. 2020 Touvron al. 2021) have shown great success in various domains recent years. The size of the DNNs, however, has increased on an extraordinary scale up to 70 Billion parameters in the single model /hao al. (2023) rrequirnng an unprecedented amount of computing resources and energy to deploy the DNN- based services. Fortunately. under certain conditions, the weight matrices of DNNs trained by stochastic gradient descent (SGD) naturally have well-defined preferred structures, such as low-rank matrices (aaras al. 2023 Huh et al., 2022, Gunasekar et al.) 2017). However, iden- tifying such structures to lower the effective complexity of weight matrices in recent models such Transformers Vaswani et al. (2017) remains challenging problem. Also, it mostly relies on the existing numann-designae
hand-crant craffted structured matrices without a unified systematic approach. classes of structured matrices for DNNs (Li et al.] 2015): [ao et al.) [022] Chen et al., 2022). Notably, each structured matrix is defined dis- jointedly from other formats. For example, neither the block-sparse matrix format nor the low-rank matrix format of the same number of parameters is subset of another, and yet there is no unified representation that describes both well. Moreover, the structure description and the implementation complexity of the structured matrix are in non non-differentiable discrete spaces.
Hence prior works have focused investigating
In this paper, we investigate (locally) optimal structures of the weight matrices as well as a differen- tiable training method to learn them, attempting to answer the following two questions:
1. Is there a universal format that represents wide range of structured matrices?
2. Can the structure of such matrices be learned efficiently, if it exists?
Contributions. Tackling the above two questions, we introduce generalized and differentiable structured matrix format. The main contributions of this work can be summarized as follows.

---

Published as a conference paper at ICLR 2024
1) We propose Generalized Block-low-rank (GBLR) matrix format, which includes many impor- tant structures such as Low-Rank (LR), Block Sparse (BSP), and under some practical conditions. The new structured matrix format consists of two types of param- eters: one guiding the structure of the matrix and the other specifying the content or the values of matrix elements. We show that the LR, BSP and BLR formats are special cases of the GBLR matrix. We also show that the GBLR format is closed under the interpolation between existing GBLR ma- trices in the structural parameter space, which we believe is a strong evidence that the GBLR format is able to capture undiscovered structured matrix formats.
(BLR) matrices
2) We introduce differentiable parameterization the structural parameters widths and locations of blocks of the GBLR format. The structural parameters are defined in the frequency domain, and are processed by the proposed Caaisian-Dirichlet (Gaudi) function followed by inverse Fast Fourier Transform (IFFT) to the format named Ktadi-GBILR, We show that the derivatives of the Gaudi function with respect the structural parameters exist almost everywhere, even when the width is zero.
3) We propose a practical learning algorithm based on proximal gradient descent to train compact neural networks with Gaudi-GBLR matrices. The proposed method is extensively evaluated on the Vision Transformers (ViT) (Dosovitskiy et al. (2200) and MLP-Mixe (Tolstikhin et al. 2021), outperforming prior approaches using hand-designed structured matrices.
2 PRELIMINARIES
Notation. We use @ to indicate the elementwise (Hadamard) product of two matrices. The imag- inary unit is denoted by = V-1. The normalized sine function is defined by sinc(x) = aln where s0nc(0) 1. The (element-wise) floor function is denoted by I-J. The index of elements in a matrix or vector starts from zero (instead of one), following the convention of Cooley & Tukey 1965). Also, I, {0, n....,n n denotes the index set from 0
Assumption. For simplicity, we assume the weights are square matrices. Extension to rectangular matrix cases is discussed in Appendix A.3
2.1 BLOCK-RELATER MATRICES.
A block B E [||[[] of matrix W E maxa is submatrix of W with consec- utive cyclic row and column indices. For example, can be {n 2, 1,0, 1] R = 4. Also, we say two blocks B and B₁ overlap they have shared elements of W.
High-rnnk Blocks
High-rnk Region
Overlao
Two well-known block-structured matri- block sparse (BSP) matrix and a block-low-rank (BLR)AAmettyy et al. 2015) matrix. Informally speaking, ma- trix is block-sparse when non-zero ele- ments gathered in blocks and such blocks sparse in the matrix. A block-low-rand matrix is composed of non-overlapping equally-partitioned low-rank blocks. Figure [ illlstrates the BSP and BLR matrices as well as our proposed generalized biock-low-rank matrix, which introduce in Section 3.1 We present the formal definitions of the block-sparse and block-low-rai matrices in
RSP
BLR
GBLR (Ours)
Figure 1 Comparison of block-sparse, block-low-rank, and our proposed Generalized block-low-rank matri- ces.
Appendix A.I
2.2 STRUCTURED MATRIX
We say matrix W E gaxn is structured if, for any E I", the matrix-vector product (MVP) = Wx requires significantly less number of multiplications than n². For instance, LR, BSP, and BLR matrices are structured because the number of multiplications for MVP is determined by their (low) rank or sparsity. Although a general sparse matrix reduces the complexity MVP to be sub-quadratic function of n, is excluded in discussion of structured matrices because
2

---

Published as a conference paper at ICLR 2024
During Training
After Training
Row Smuciual
Column Stuckand Parana
of.
Addeoss
(12 (5.
Whita
By
B;
rank-1
B₁
Ma
SP
munk-1
Zero podtied Block
Rand
D:
Cropped Content Farons
Figure 2: Left: An example a GBLR matrix with 4 blocks. A block is generated from the structural parameters (w 1,), (w (1) and the content parameters (u, v), where (wR,1%) and (w ,19 form binary masks ( [R( and (( respectively. Note that overlapped regions can have rank higher than one. Right: Efficient Matrix-Vectur Product computations using cropped content parameters and structural parameters. The structural parameters locate the input and output indices/addresses to read and write.
processing moderately sparse matrices (10~50% density) in conventional hardware such GPUs does not reduce the actual run-time due to its unstructured positions of non-zero elements (Chen
al. 2022).
3 PROPOSED METHOD
We are interested in training a deep neural network (DNN) f under a computational cost constraint:
min Σ ((((x).)) s.t. cost(S) < B, 2.
where the first term is the cross-entropy loss for classification task a data point æ and label y. and the constraint cost(f) ) is the number of multiplications to compute the neural network output j (xx) Our method learns weight matrices of DNNs in a generalized structured matrix format in a differentiable manner.
3.1 GENERALIZED BLOCK-LOW-RAND (GBLR) MATRIX
We introduce concept of generalized structured matrix format that explains multiple popular structure matrices used in DNN training. The idea is that a block a matrix can be expressed by sum of rank-1 blocks, i.e., rank-r block is a of rank-1 blocks the same position. In this manner, LR, BSP, and BLR matrices can be expressed under a unified framework (Theorem]) []]
To be specific, proposed structured matrix format is generalized version of Loock-LLw-Randl (BLR) matrices. An n-by-n Generalized Block-Low Srock-LowwRRnd Rank (GBLR) matrix W is obtained by over- lapping multiple rank-1 1 blocks of different sizes arbitrary locations, as depicted in Figure The locations of the blocks as well as their element values are learned simultaneously from training data without explicit size or location restrictions.
Suppose there are K blocks: B1 E xwf Bg E RR xuf B& has two parameter sets: 1) the structural parameters that identify the position of the block in the row and column index set In-1 = (0, I,..., 1}, and 2) the content parameters which specify the actual values of matrix elements.
Bx E RR x Each block
Configuration of structural parameters. The position block given by the indices of the and columns it occupies in the be identified by four numbers: width and location in terms of the row or column indices. Hence, location parameter 7 and width parameter the structural parameters of block.
matrix. Hence, the placement of rectangle block can

---

Published as a conference paper at ICLR 2024
Figure 2 (Left) illustrates block of size w x w. in n X matrix at location (L. 10). The row (column) index set of a block is the sequence of numbers from 1 (19) to 1 + w R (1 + ) where the addition is a cyclic/modulo addition. For each block B; for k = .,..., K, have four parameters: ( (f) for the row and (wg. 1 for the column. We use the notation of = (w )
and o% = ( 1x) to represent the tuple of width and location for the row (of) and column (o
Based on the structural parameter 1 and l one can construct an n-dimensional binary mask that has - E In consecutive ones starting from Ig In in the eyclic order:
mag [i] = (" a ( [1
i S + an < 1 + w
j II a (0,
(2)
otherwise
where a E {0, 1} is necessary to make the order cyclic. We call the mask in Eq. 2 the boxcar mask since the non-zero elements are located consecutively. The boxcar mask is used to select w (cyclic) consecutive non-zero elements of an n-dimensional vector. The mask for the rows m obtained in the same way from wi and If.
Configuration of content parameters. To represent the values of a rank-1 1 block BR, we use two n--dimensional vectors ug and Un as content parameters along with the boxcar masks ma mag. All these parameters (Ø ox uk. Us) are learned during the DNN training simultaneously. Since the boxcar masks mgf and ma guide the location of the block in the n x matrix, U2 and Un element-wine multiplied with the boxcar masks:
ZeroPad B&) (mog ux) O 0z)T,
where the resulting n x n matrix zero- pooo-padddd block. Ideally, we expect the mask to expand / shrink and shift to find the right subset of the elements of content parameter, while the content parameter updates the value of the elements selected by the mask.
Now we formally define the Generalized Block-low-rank (GBLR) format, which is the sum of K sero-padded blocks:
K
K
W
Σ m n m k=1
og) @ (u = Σ
ma O uu
(mag 2
(3)
A GBLR matrix is associated with an average width = 2K Ex= w + u
Definition 1. Let GBLR(n, K) be the set of matrices obtained by Eq. for the average width less than equal 2 0, i.e, w 2K Σk wf R W E GBLR(n, K)s)
Ek=1
7 S A matrix W is (n, K, K)BBR if
We use the notation (WW) tural parameters of W, where w(W) = (wf. wi woR w....., ",. w&} and I(W) = {If 1 18 1 1 ) We simply use & == a if W is clearly inferred in the context.
(wwW), to indicate the collection of the struc-
Efficiency. Once the structural parameters fixed, it is unwise to store and use two n-dimensional content parameters for each block By because only w elements of us and wf elements of ok non-zere according to the boxcar masks in Eq. 3 Henee, one can store and use the cropped content parameters Uin and Dio.0 for MVP between an input x E IR" (which can be also cropped from 10t to 10 + w) and a block B, as described below and in Figure 2 (Right):
ZeroPad( B x = (m (1) O ()( O u
ZeroPad (utR.844m (vic D0C 4m))))
which requires only R + 0 multiplications. Hence, the number multiplications (denoted by FLOPs) for multiplying W E GBLR(n, K, with x E R" is bounded by 2Ks:
FLOPs
w + w 2K 2Ks.
Expressiveness. Low-rank (LR), block sparse (BSP), and block-low-ranh (BLR) matrices pop- ular structured matrices in the DNN literature. and they are special cases of the GLBR matrix under mild conditions. Proofs and formal definitions of LR, BSP, and BLR matrices are in Appendix A.
4

---

Comparison between Boxcar mask and Gaudi masks in the time do- main with different smoothing factors Published as a conference paper at ICLR 2024
Theorem 1. Let n. K,8 be positive integers satisfying Ks 2 Then any n-by-n rank- Ks matri- ces and (n, bbblock-sparse matrices are (n, s) Also, any (n, K, s, ly-bkocck-ow-rrnnk matrices are K. s)-BKLs).BRRRR if K (n/s}:
More importantly a new structured matrix obtained by interpolating the structural parameters of two (n, K, 8) 3 matrices still (n, K, B,BBR, based on Theorem 2 structured matrices can be derived from a set of GBLR matrices.
Therefore, a new type of
Theorem (Closed under structural interpolation). Given two n x n matrices W. Z E GBLR(n, K, s), and [0. 1], consider the following combination between the structural pa- rameters:
w = Laaw(W) (1 ())u(()) 1- Lol(W) (1 -z]
A matrix Y generated by Eq. with the structural parameter (wi 11) is (n, K, - matrix, Y GBLR(n, s)
Theorem D and Theorem tell that (n, K, s-GLBR matrices cover a wide range of popular existing structured matrices and also undiscovered ones. In the following section, we introduce a differentiable tool find/leam structured matrices in GBLR(n, K)
3.2 GAUSSIAN-DIRICHI (GAUDI) MASK
The matrix structure in the GBLR format is determined by the width and location parameters. We aim to extract/learn these structural parameters from training data using stochastic gradient descent. In order to do the non-ciiferentiaabiiiy of the boxcar mask parameters needs to be handled.
To tackle this issue, we introduce a Dirichlet-kermel-buaee itly parameterize the width w and location I in the expression. Consider boxcar mask of length n, width w, and location l. Let (w.1) be the discrete Fourier transform (DFT) the mask (w,t)- The frequency-domain representation of the boxcar mask is given by
parameterization of the boxcar to explic-
sinc (w enk( (==)
(w [A] = e 2n f (w.o) [k] = - -m
2mzj = =--
-rx w
(5)
sinc
==0
e 2 & duw [k].
where d [k] sine ( sxck( 1-
is the Dirichlet kernel of order (Bruckner al. 1997) in
sino &
the discrete domain In-1 {0, 1,....,.. 1}.
Furthermore, propose a smooth version of the Berichlet-kemel-baare mask by convolving the time- domain boxcar mask with the Gaussian-shape function. obtained in the frequency domain by element-wise k2
10-
multiplying the Gaussian function go [=[k] exp 20 with the standard deviation a > 0 to the Dirichlet kernel. We call the resulting function the Houssian-Dirichler (Gaudi) function dg:
d [k] == go [k] dw(kk) [k]]
Figure 3:
m (w, == IDFT(d .e 11)
(6)
The Gaudi mask converges to the Boxcar mask as grows.
where k = 00 1...., 1]. And we call the mask gener- ated by Gaudi function m.(...)) the Gaudi mask, where the parameter a controls the smoothness. Note that as xc, the Gaudi mask converges the boxcar mask (Lemma 5) Figure visualizes the time- domain Gaudi mask approaching the boxcar mask.
A useful property of the Gaudi mask is that parameter, even when the width is zero. To show it, we relax the domain of widths and locations to the continuous interval [0, ni] (see Appendix A.4)
obtain exact derivatives with respect the width
5

---

Published as a conference paper at ICLR 2024
Theorem 3. Let be a finite positive integer. For any [0) and w, [0, n). the derivatives of the Gaudi mask (w.1) with respect to w and are bounded almost everywhere:
@ (w.J)
@m (w
< 00;
<0.
al
Especially when = O, the derivative with respect to neither divergent nor zero. Corollary 4. For any 1 [0, and E (0, o0), the norm of the derivative of the Gaudi mask (w,)) with respect to w at = 0 is well-defined and greater than zero, i.e., 0 Om 8w
Cc.
To allow learning the mask structural parameters in a differentiable manner, we plug the Gaudi mask (Eq. [ into Eq. [ as the mask of GBLR matrices to model Daiisian-Dirichlet GBLR (Gaudi- GBLR) with parameters 0- (Φ) U,.UVV VVV.o)
V wlo.U.V.o)
Σ 1=1
ma @ us (m²g 2 us)) vz)"
In practice, one can use small R at the beginning of the training process to update the corre- sponding (unmasked) content parameters which more than necessary, then gradually increase to adjust the content parameters with a tighter mask. Since the purpose of using a Gaudi mask to learn structural parameters of the GBLR matrix by Gradient Descent, Gaudi-GBLR matrices are later replaced by GBLR matrices once the structural parameters are found/learned. To compute MVP using GBLR matrices, one can use the cropped content parameters and inputs, as we discussed in Section efffciencc without constructing masks at all. Hence, during the inference, there is no overhead to compute Gaudi masks.
3.3 LEARNING GAUDI-GBLR FOR EFFICIENT NEURAL NETWORKS
We now introduce algorithm to learn the structural pa- rameters of Gaudi-GBLR matrices. The goal of the learn- ing algorithm is identify Gaudi-GBLR matrix struc-
Algorithm 1 GBLR Learning by PGD
1: repeat
tures (i.e., their parameters) that allow computationally efficient DNNs. Our discussion is centered around two- layered multi-layer perceptron (MLP) for ease of under- standing. However, the technique can be applied to gen-
2:
(I1, 01) V1),..., (IM, Yad) ~ D +- Li=1 6 (x:)). y:)
3:
4: for trainable p E 0 do
5:
p p 17 AdamW(V!) end for
eral DNNs that incorporate linear layers.
6:
7:
for it 1,
K do Wig t clipo.n (S (wx))
Now let us consider two-layered MI.P 1 with Gaudi- GBLR weight matrix W (xx h i x + b). We initially relax the domain of w E In and / E In-1 of the Gaudi mask to the real-valued space w, E [0, n] as we discuss in Appendix [A.4 Due to the property of Gaudi-GBLR matrices Eq, 4 the computational cost constraint the DNN 1 in Problem 1 can be replaced by a constraint on the sum of the width parameters of W Specifically, we find the width
8:
end for 10: until converge
9:
parameters w {w wf ..... W wx of we satisfying ||w||| < B since Σ w + 3 = | To solve the problem with Gradient Descent, we relax this E -norm constrained problem to unconstrained one using the method of Lagrange multiplier:
min Σ (()((x)) v) + A1/w//1 A O. (a:D)~
The resulting computational budget is implicitly constrained by hyperparameter > 0.
Theorem 3 guarantees the derivatives of the widths and locations of the Faudi-GBLa matrix in fo can be obtained with any positive smoothing parameter > 0 so that we can safely learn the parameters in the continuous domain [0, n). Specifically, we update the width parameter w in the L -oom term in Problem g by Proximal Gradient Descent (PGD):
1 clipo,n (Spa (u (w, [) v))),
(9)

---

Published as a conference paper at ICLR 2024
where S²(2) = (1) (|z] - 14)
i > μ is the element wise soft shrinkage function and otherwise
c. clamps the elements of the inputs to [0, n). In practice, the gradient is calculated with an adaptive optimizer such as A (Loshchilov & Hutter 2017) (see Line 5 in Algorithm [ The overall process is summarized in Algorithm[[] Although Problem is non-linear, our experimental results show that PGD can attain good local optima with an adaptive optimizer and the initializa- tion method we propose in Appendix A.2 A practical learning method for the width and location parameters defined in the discrete spaces Zn and -1 is discussed in Appendix A.5
4 EXPERIMENTS
We evaluate our proposed method by replacing the weight matrices of Vision Transformers (ViT) Dosovitskiy et al. 2020), and MLP-Mixer (Tolstikhin et al. 2021) with Gaudi-GBLR matrices. For the experiment, we set the number of blocks K equal to the number columns the matrix n. We also evaluate alternative schemes for comparisons where the weights are replaced by popular hand- designed structured matrix formats such as Low-Rank (LR), Block-Sparse Saarek--ppseeppus---- -ato--lllss-----eeee Low-Rank (BSP- LR), and Block-low-rank (BLR). For LR matrices, we use singular vector decomposition best rank-8 approximation of the given matrix for the predefined rank of 8. Pixelfly (Chen et al.) 2022) and Monarch (Dao et al., 2022) are schemes that use BSP-LR and BLR matrices respectively. We set the structural parameters of these alternative schemes to exhibit similar computational costs for MVP compared to our proposed scheme. Note that the structural parameter sets for LR, BSP- LR, and BLR do not change across different layers in the neural network. We denote the number of multiplications by FLOPs, and use 8 NVIDIA A40 GPUs in our experiments. The detailed experimental settings are described in Appendix A.7 Our source code is available at https: 7/github. cwt.haa..commooooooooooooo hawmm//hawwwooleeeeeeee Gaadii
find the
4.1 FINE-TUNING RESULTS
Vision Task. We use the pre-trained weights of the ViT-Base on ImageNet and initialize the parameters of the Gaudi-GBLR matrices by Algorithm 2 in Section
VIT-B Finetuned
ImageNat Valication Accurary (%6
79
A.2 The ViTs with Gaudi-GBLR matrices were fine-
78
tuned on ImageNet (Russakovsky et al. 2015). Dur- ing the initialization, we set the computational bud- get of all matrices in the network the same. For a fair comparison, the same set of hyperparameters was used throughout fine-tuning experiments.
Dense
Low Rank
Monarch (BLR) (Dao al.)22)
74
Gaudi GBLR (Durs)
The highest accuracy achieved by Gaudi-GBLR in ViT-Base with patch size of 16 16 the ImageNet validation dataset after fine-tuning for epochs. Figure shows that Gaudi-GBLIR preserves the racy well when the complexity is reduced to 30% of the original 'Dense' model which does use struc- tured matrices (its FLOPs count is normalized to 1).
0.26
0.28
0.30
0.34 Relative FLOPS
0.32
0.36
0.38
Figure 4: ImageNet accuracy after fine- tuning ViT-Base weights replaced with structured weight matrices. Dense: the original ViT-Base model.
The other hand-designed approaches exhibit more significant accuracy degradations for the same complexity reduction. Overall, Gaudi-GBLR strikes better Accuracy-FLOPss trade-off's than LR or Monarch approuches. The higher accuracy for similar FLOPs count quantifies the gain from the learned structured matrices. In addition to VIT, we also tested GBLR matrices on ResNet (He et al.) 2016). Readers interested further details can find the results Appendix A.8
Language Task. We test the GBLR matrix to the weights of the pre-trained GPT-2 (Radford et al. 2019). We compare our method to LR and BLR (Monarch) matrices. We evaluated the perplexity of WikiText103 (Merity et al.) 2016) validation set af- ter 10 epochs of fine-tuning on the training set. The perplexity of each model is in Table | Our method
Table Perplexity by weight type of GPT- 2 after fine-tuning on WikiText103.
Weight Type
Perplexity (L)
Relative FLOPs
Dense
19.36
100%
Low Rank
19.48
43.75%
Monarch
20.56
43.75%
Gaudi-GBLR
19.24
43.7%
7

---

Published as a conference paper at ICLR 2024
-.:.::
Donse
LR
Monarch (BLR) et al.
Pixeltty LEichen at
Goudi GBLR (Ours)
VIT-5, CIFAR-10
VIT-S, CIFAR, 100
S.xP-Mixer-S CIFAR-10
MLP-Mixer S, CIFAR 100
70
(80
87.5
85.0
82.5
0.3
Relative
Relative PLOPs
Relative FLOPs
FLOPS
Figure 5: Accuracy-Cost trade-off' of models trained from scratch CIFAR-10/100 dataset.
Table ImageNet accuracies and GFLOPs on a 224 x 224 image of ViT-Base models trained from scratch with dense matrices and Gaudi-GBLAR matrices.
Query
Key
Value
Invero
MLP FC2
Rank-403
Model
Acc. (%)
GFT.Ps
ViT-Base
78.57
17.2
Layer
Gaudi-GBLR
78.51
5.65
Rankm377
Rank-18
Rericm146
Ranic=205
Rank-408
Table 3: FLOPs type of linear lay- ers with MBLdi matrix ViT-Base trained on ImageNet. Units in 10 FLOPP.
Layer 11
Rank=10
Rank=29
Rank>=89
Rank-618
Layer Type
Min
Max
Avg
Query
4.2
67.6
29.0
Figure 6: Mask patterns of weights for Query. Key, Value, FC1 and FC2 Layer O, 5, and 11 (out of 11). The brighter, the more overlapped blocks. The numbers below indicate the rank the matrix (max: 768).
Key
11.6
87.6
42.1
Value
34.2
135.5
95.0
MLP-FCI
359.8
910.2
597.7
MLP-FC2
349.4
1,175.2
756.6
achieved the lowest perplexity, outperforming the dense baseline and utilizing the smallest number of FLOPs.
4.2 TRAINING FROM SCRATCH
We train VITs (Dosovitskiy et al. 2020) and MLP-Mixers Tolstikhin et al., 2021) with structured weight matrices on CIFAR-1O& (Krizhevsky et al., 2009) by Algorithm [ from randomly- initialized content parameters. We set the training process ends (see Appendix A.7 Figure we study the accuracy-FLOPs trade-off using CIFARIO/100 datasets when the models are trained from scratch (i.c., not fine-tuned from pre-trained weights). As in the ImageNet fine-tuning experiment, Gaudi-GBLR achieves superior accuracy-FLaa -LOPs trade-offs outperforming the other hand designed structured matrices.
for the first epoch and gradually increase it to 100 until
4.3 ANALYSIS ON LEARNED GAUDI-GBLR MATRICES
In this section, we study the learned Gaudi-GBLR matrices in terms of computational budget al- location, and mask patterns. The analysis is based on ViT-Base Gaudi-GBLa matrices trained on ImageNet from scratch by following the same o adaptation rule used in CIFAR10/100 experiments. The accuracy and FL.PPs of the ViT-Base with Gaudi-GBL.R is reported in Table
Learned Computational Budget Allocation. The proposed learning framework automatically allocates the computational budget to all linear layers of ViT-Base during the training process given by Algorithm [ The algorithm finds a well-balanced (not necessarily equal) allocation for each matrix to meet the overall cost budget while minimizing the loss function. As result, Gaudi-GBLR weight matrices in the network have unequal widths requiring different FLOPs for each MVP. Table summarizes the min/max/average FLOPs statistics collected for different types of linear layers. Within an Attention layer, the weights for Values have the highest budget whereas Queries and Keys smaller budget. The smallest layer uses only about 4,200 FLOPSs for MVP involving
8

---

Published as a conference paper at ICLR 2024
768 768 matrix and 768 1 vector. The FLOPs assigned to the linear layer of channel MLPs (MLP-FCI and MLP-FC2 in Table 3 vary significantly. Although the size of the weight matrix used in the MLPs are 4x larger than the ones used for Value, the MLP-FC2-typ layers use 7.96x FL.PPs than the Value-type layers in average.
Visualization. Figure 6 visualizes the locations of the blocks in exemplary Gaudi-GBLR weight matrices of ViT-Base trained from scratch. We select the first, middle, and the last layers of different types: linear layers for Query, Key, Values in Attention modules, and two linear layers (FCI and FC2) in MLPs. Bright colors in Figure 6 highlight regions where masks are overlapped. Interest- ingly, the resulting matrix is neither BSP nor BLR. It is observed that blocks are concentrated in a small number of regions. We believe this is related to the Multi-Head (Vaswani et al. (0177) scheme of the ViT. Each weight matrix of an attention layer is collection of weights for multiple heads. Because some heads are more important than others, more blocks are allocated to those regions. No- tiee the rank of matrices obtained from the GBLR framework differs significantly across different layers and matrix types (Values, Queries, and Keys).
RELATED WORKS
Structured Matrices for DNNs. Prior works have focused on manually designing suitably struc- tured matrices for DNNs since explaining the structure of every weight matrix in practical DNNs such as Transformers (Vaswani et al. 2017) is challenging, Hsu et al. (2022) used weighted low-rank decomposition for the weights of language models. Butterfly matrices (Li al. 2015 Dao et al.) 2019) were adopted in the form of Block-Sparse (BSP) format (Pixelfly) by Chen et al. (2022),
and also in the form of Block-low-rank (BLR) format (Monarch) by Dao et al. (2022). ) In contrast, our method learns the structure of weight matrices from the training data. Layer-wise Low-Rank Rank DNN Compression has been studied in terms of Augmented Lagrangian Multiplier (Idelbayev & Carreira-Perping 2020), subspace grouping (Liebenwein et al., (2021). and quadratic programming (Li Shi 2018). However, they work only on the low-rank format, and involve multiple loops and SVD at every iteration, whereas ours can find layer-wise structure and budget in a single loop.
Mask Learning. For neuron pruning/selection, a binary mask with surrogate gradient has been adopted to learn its pattern. Jang al. (2016) and Maddison et al. (2017) propose alternative distri- butions close to the Bernoulli distribution, which allow the continuous gradient. Movement Pruning Estimator (STE) (Hinton, 2012 [Bngii et al. 2013 Zhou et al. 2016) to pass the gradient through the Top-k function. Lin et al. (2017) adopts deep rein- Forcement learning to select the input-dependent mask. On the contrary, our mask design solves the non-existing gradient problem by fetiiffffntioonhhi parameterization in the frequency domain.
(Sanh et al. 2020) utilizes Straight-Throun
One-shot Neural Architecture Search. Neural Architecture Search (NAS) (Zoph Le 2016
Liu et al., 2018) seeks the optimal neural network structures from training data. To improve search efficiency, Pham et al. (2018) [iu et al. (2018) adopt the one-shot NAS technique that selects sub- network candidates from a super-network. Our method also falls into similar category of finding small-sized neural network in the scope of structured matrix format.
Trequency-doman Learning, DiffStrrde (Riad et al. 2022) learns stride of the pooling operation for images by cropping a rectangular region of the frequency response of an image. They utilize an approximated boxcar mask for differentiable stride. Although DiffStride shares similar compo- nents with Gaudi masks, the fundamental difference is in the design of the mask. We parameterize the mask in the frequency domain where widths and locations inherit the exact gradient.
6 CONCLUSION
We propose generalized and differentiable framework for learning structured matrice for efficient neural networks by gradient descent. We introduce a new generalized format of structured matrices and parameterize the structure in the frequency domain by the Gaussian-Dirichhee (Gaudi) function with a well-defined gradient. Effective learning algorithms are provided for our framework showing flexibility and differentiabilityy to find expressive and efficient structures from training data in end- to-end manner. Evaluation results show that the proposed framework provides the most efficient and accurate neural network models compared to other hand-designed popular structured matrices.

---

Published as a conference paper at ICLR 2024
7 REPRODUCIBILITY STATEMENT
The authors make the following efforts for reproducibility: 1) We release our code at https: //github. com/ http///waloeeeaaaaaee Gaudi BaaauuiiiBBBRR 2) we provide the detailed settings and hyper- parameters in Section 4 A.5 and A.7 and 3) the proofs of all theorems, lemmas, and corollaries are presented in Section
ACKNOWLEDGMENT
We thank Sara Shoouri, Pierre Abillama, Andrea Bejarano-Carbo, and Mingyu Yang for the insight- ful feedback on the paper. This work was supported in part by COGNISENSE, one of seven centers JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA.
REFERENCES
Patrick Amestoy, Cleve Ashcraft, Olivier Boiteau, Alfredo Buttari, Jean- Yves L'Excellent, and Clément Weishecker. Improving multifrontal methods by means of block low-rank represen- tations. SIAM Journal on Scientific Computing, 37(3)A1451-A1444 2015.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 338.3438, 2013.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Andrew M Bruckner, Judith Bruckner, and Brian Thomson. Real analysis. Alassi ysis. com, 1997.
Beidi Chen, Tri Dao, Kaizhao Liang. Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. Interna- tional Conference Learning Representations, 2022. URL https: htttps://openrevieet
forum?id 1Xa y
James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297-301, 1965.
Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In International conference machine learning, pp. 1517-1527. PMLR, 2019.
Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning, pp. 4690- 4721. PMLR, 2022.
Jacob Devlin, Ming- Ming-Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04880. 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An im- age is worth 16x16 words: Transformers for image recognition at scale. In International Confer- ence on Learning Representations, 2020.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Sre- bro. Implicit regularization in matrix factorization. Advances in neural information processing systems, 30, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
10

---

Published as a conference paper at ICLR 2024
Geoffrey Hinton. Neural networks for machine learning coursera video lectures. 2012.
Yen-Chang Hsu, Ting Hua, Sungen Chang, Qian Lou, Yilin Shen, and Hongxia Jin. Language model compression with weighted low-rank factorization. In International Conference on Learning Representations, 2022. URL https
net. ottpppawwweee....ee/eoooereea/ap
Minyoung Huh, Hossein Mobahi, Richard Zhang, Brian Cheung, Pulkit Agrawal, and Phillip Isola. The low-rank simplicity bias in deep networks. Transactions Machine Learning Research, 2022.
Yerlan Idelbayev and Miguel A Carreira-Perpinán. Perpinán. Low-rank compression of neural nets: Learning the rank of each layer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
2020.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzttion with gumbel-softmax. In International Conference on Learning Representations, 2016.
Claude-Pierss Jeannerod, Theo Mary, Clément Pernet, and Daniel S Roche, Improving the complex- ity of block low-rank factorizations with fast matrix arithmetic. SIAM Journal on Matrix Analysis and Applications, 40(4):1478-1 1496, 2019.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Chong Li and CJ Shi. Constrained optimization based low-rank approximation of deep neural net- works. In Proceedings of the European Conference on Computer Vision (ECCV), pp- 732-747, 2018.
Yingzhou Li, Haizhao Yang, Eileen R Martin, Kenneth L Ho, and Lexing Ying. Butterfly factoriza- tion. Multiscale Modeling & Simulation, 13(2):714-732, 2015.
Lucas Liebenwein, Alaa Maalouf, Dan Feldman, and Daniela Rus. Compressing neural networks: Towards determining the optimal layer-wise decomposition. Advances in Neural Information Processing Systems, 34:5328-5344 2021.
Ji Lin, Yongming Rao, Jiwen Lu, and Jie Zhou. Runtime neural pruning. Advances neural information processing systems, 30, 2017.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. In International Conference Learning Representations, 2018.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17X1v:1711.05101, 2017.
Chris J. Maddison, Andriy Mnih, and Yee whye Tch. The concrete distribution: A continuous laxation of discrete random variables. In International Conference on Learning Representations, 2017.
j15Lgg]
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2016.
Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In International conference on machine learning, pp. 4104- PMLR, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Rachid Riad, Olivier Teboul, David Grangier, and Neil Zeghidour. Learning strides in convolutional neural networks. In International Conference on Learning Representations, 2022, URL https:
//openrevi ew. nwt. http://wwwweeeeeetoooooooooo///e
11

---

Published as a conference paper at ICLR 2024
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeey Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJV), 11(())211-255 2015. doi: 10. 1007// 1263-015-0816-y.
Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine- tuning. Advances in Neural Information Processing Systems, 11-20388-20382 2020.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mip-mixer: An all-mlp architecture for vision. Advances in neural information processing systems, 34:24261- 24272, 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Can Yaras, Peng Wang. Wei Hu, Zhihui Zhu, Laura Balzano, and Qing Qu The law of parsimony in gradient descent for learning deep linear networks. arXiv preprint arXiv:2306.011 202.0000.111 / 2023.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. A survey of large language models. arXiv preprint arXiv:2303.18222 2023.
Shuchang Zhou, Yuxin Wu, Zekun Ni, Xinyu Zhou, He Wen, and Yuheng Zou. Dorefa-net: Train- ing low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint arXiv:1606.06160, 2016.
Barret Zoph and Quoe Le. Neural architecture search with reinforcement learning. In International Conference Learning Representations, 2016
12

---

Published as a conference paper at ICLR 2024
A
APPENDIX
A.1
PROOFS
In this section, we provide the missing definitions, propositions, and proofs. The original theorems restated for completeness.
A.1.1 DEFINITIONS
We introduce formal definitions of the block-related matrices discussed in our paper.
Definition (Block-sparse matrix), n-by-n matrix is (n, m, b)-bockk-sparse (BSP) if contains non-overlapping non-zero blocks whose dimension is most X
A (n, m, -locccck-pparse matrix may contain blocks of different sizes, but each of them has maxi- mum s elements.
Definition 3 (Block-low-rank matrix (Amestoy et al., 2015 Jeannerod et al. 2019)). An n-by-n
bo blookkk--oowrr (BLR) if. when it is equally partitioned into p X non- matrix is m, s,
overlapping blocks of dimension 8 8, every block has a rank at most r < 2
The (n, (n),,r)-BB.. matrix has blocks of the same size that tile the entire matrix.
A.1.2 PROOF OF THEOREM
Theorem 1. Let n, K positive integers satisfying K 2 n. Then any n-by-n rank- n matri- ces and (n, K. s)-block-sparse matrices are (n, K, s) Also, any (n, K,8, 8, blobocck-low-ran matrices are (n, 8)BLR if K =mn/s}².
Proof. is sufficient to find the structural parameters of an (n, S,SL matrix for each struc- tured matrix format. Let R wD E IK be the width parameters of rows and columns of the (n, K, s)BLR matrix, respectively, and 1 10 E TK 1 be the locations parameters of rows and columns, respectively, where In {0, 1, n...
Low Rank Matrix. The rank- Ks matrix is equivalent to the GBLR matrix with Ks/n full-sized rank-1 blocks, i.e., wtt wc = n, 0, .. 0] and 1 -==0 0.
Ks/n
Block Sparse Matrix. A block of the (n. K s) s)--block-ppard matrix can have full rank. Consider a rank-1 block of a (n, K, BBBLR matrix with some location parameters where the sum of the row width and the column width is 2s. The maximum dimension of the block B is 8. Then clone and overlap the identical block with the same location and width parameters times to form full--ank block of the dimension most times, any (n, K/s, so---))c---pppaeeeeeeo sparse matrix can modeled by the corresponding width parameters and the location parameters.
Since the full-rank block can be generated K/s
Block Low rank Matrix. Since K n let divide the n-by-n matrix into (n/s)² blocks the dimension of s-by-s. Assign the width and the location parameters correspondingly. The resulting matrix forms the (n, K,s, j--block-low-rann matrix.
A.1.3 PROOF OF THEOREM 2
Theorem (Closed under structural interpolation). Given two n. matrices W. Z GBLR(n, K,s), and rameters:
[0. 1], consider the following combination between the structural pa-
w' = (ww(W) + (1 a)w(Z)), 1'= Lal + (1-))((Z))
matrix Y generated by Eq. with the structural parameter (w',t') is a (n, K, 8) -BLLR matrix, Y GBLR(n, K,s).
13

---

Published as a conference paper at ICLR 2024
Proof. Let w and w be the width of the row and the column of the kthh block of Y, respectively. We use the same style of notation for w (W) ) w w(W), and ((Z) R. w 2) to denote the row and column width of the kth block of W and Z. The interpolated matrix Y is (n, K, s)-GBRR if and only if the mean of the elements of w' is less than or equal to 8, or equivalently the sum of the elements is less than equal to 2 Also, both W and Z have 2Ks sum of the width parameters. From this, the sum of the interpolated width is less than or equal to 2Ks:
Σ w
ft k
w'
k=1
Low W) ft + (1 - -1-a) (Z) + (aw + -1-a) (2)
k=1
K
an W) + (1 a-) aw W) (1-a))o02) 2)
k=
(w (W) k w () + (1 a) Σ w [ o wl (W)
k=
< o2K's + a
2Ks.
Hence the interpolated matrix Y E GBLR(n, K,
A.1.4 PROOF OF LEMMA 5
Lemma 5. Let N be a positive integer. Let 1 In- and w E In. Consider the discrete boxcar
mask (w.)) with the structural parameter (w, !) as in Eq. 2 and the Gaudi mask (w with the same parameters and a >> O as in Eq. 6 Then m (w,1) converges to m as o - 20, where the
bound between two at is given follows:
(m m (1m.1)l22 l12
(1---11
S
20
(w.n)l2
Proof. By Parseval's theorem,
(w.)) mmw
- ((
(10)
.2
du [K]]
20
|d [
1 -
20
12
2
[d [k]]
e
n
(1- 20
1
)2
2
m (w)))) l1
1-
(1I)
where Eq. II is also due to Parseval's theorem.
A.1.5 PROOF OF THEOREM 3
Let us first consider the derivative of the sinc function in eq.
5
14

---

Published as a conference paper at ICLR 2024
dsine
sinc ank
Lemma 6. For any n + O,
du
Proof.
dsine wk
sin auak
dw
dw
aukk
Tuuk
2
72
Twk:
Tuuk
-1
mwk
sin
+
cos
T1
71
71
71
madk
sin Xunk
C08
Buck T1
tok
sinc wk
COS
Note that the derivative in Lemma at w is since lime sinc' (1) 0.
The following lemma is useful for proving Theorem 3
Lemma 7. For any a E (0, oo] and w, E [0, the derivative of the Gaudi function (w with respect to w is as follows:
dd [k]
2m t
dw
sinc E
Proof. Let merge the terms of the Gaudi function de [k] that does not contain w into Cs:
sinc wk
d [
20
sinc
wk vrkw
w sinc
Ck
n
where C₁ = =
20
From Lemma
the derivative of the Gaudi function with respect to
sinc x.
w as follows:
d [k]
C+
a
sinc
wk
e
vades
Ow
Ow
n
C
sinc
wk
-
cas
saank
sinc uk
we
wsinc
wrk e
k (-Tr n n
x
n
w
Twk
Twnk
Ck
(cos
sin
n
= C e
- $
Now we prove Theorem 3
Theorem 3. Let be a finite positive integer. For any (0, and w. [0, the derivatives of the Gaudi mask m. with respect to w and are bounded almost everywhere:
(w
@mg (m.))
am
(w)))
<00
< <0.
Ow
2
15

---

Published as a conference paper at ICLR 2024
Proof. Here we use Parseval's 's theorem again:
dm (w.l)
1
i
dy
-
[k-e
2m
di kiee
2
Ow
71
Zw
Orc
2n
e
2TT E1
sinc
k=0
26
sinc k
Y2
sinc
k=0
72
00,
where the first inequality used the fact that 2= < Similarly, the bound on the nomm of the derivative with respect to the location parameter can be also derived as below:
and
for 0, 1,....... n
sinc
sine
dm o
d
[A)
2
1
σi
01
Øl
sine
12 20
201( x
1 (-2TT1
2xz
e
sinc m
1
2m
sinc
< 00.
A.1.6 PROOF OF COROLLARY 4
Before proving Corollary 4 we introduce another Corollary of Lemma 7
Corollary 8. For any E (0, co] and w., 1 E [0, the derivatives of the Gaudi mask mî with respect and I are given follows:
0mg
(m. [R]
62 2x0 T -
1
Fjk
202
Ow
sinc
k=0
Dmg
(w.)) [k)
R-1
k (-27T)
Fjk
2T
2m
n
n2
where Fjk 22 &. is the (j. k)th element of the inverse discrete Fourier transform (IDFT) matrix.
16

---

Published as a conference paper at ICLR 2024
Proof. By linearity of differentiation and Lemma
0
[k]]
Fjk
2m
d.
Ow
dw
==
n
0
Fjk
2
dalk
72
dw
Fj
e
2
2m w
t
20
sinc
Fjx
2 2 (w-l)
c 20
sinc
A==
The derivative with respect to the location parameter is direct consequence of the derivative of the exponential function.
Corollary 4. For any I [0,n] and E (0, the norm of the derivative of the Gaudi mask
(m with respect to at 0 is well-defined and greater than zero, i.e., 0 <
(w.
Ow
Proof. From Corollary 8
Dmo dw
Fjx
2 e 20
.2
T2
sinc
w=
k=0
Since e 2 e
is not zero for any / E [0,n] and E (0, oo] and the IDFT a matrix F" is unitary (up to scaling factor), the derivative aw (w,w)) + O.
2
sinc
A.2 INITIALIZATION ALGORITIM
Suppose an ny-byn matrix Win is given before the initialization step. We first initialize the structural parameters based on the correlation within the columns of Winn-
Consider the Gram matrix C W. W Wiitt. Cj5 WT Was is is the inner product between the ith column and the jth column of W min That is, the higher |C:yl is, the more correlated the columns are. Thus, for the width and the location parameter of the column of the first block, we pick row of the Gram matrix C. where i, is the index of the row which has the largest norm. Then we apply smoothing filter (e.g.. Gaussian filter with the standard deviation 7 > 0) to the absolute values of the row, which is followed by binarizing the filtered output based on some threshold > 0. Hence, for the kth block, can be summarized as follows:
init
b bl {
Cip [i] 2 Cypy G <T'
C₁2y = F.iiter(CC...llll
i af [k]].
(12)
As result, the binary vector b, consists of the chunks of ones and zeros. We choose the longest chunk of ones and initialize the width and the location parameters (w 1 correspondingly. For the
row parameters of the kth block, the same procedure repeated with the columns of ii selected in the column parameter initialization step.
Next, the content parameters are set by the left and the right singular vectors of W mi scaled by the
singular values, i.c., U + S1/2 Init and V + g1//2 VT where Wint = S mit
init
Finally, we update the structural parameters and content parameters by minimizing the following objective function:
min
Win
A|/w||a wwezz + 1ee
(13)
U.V.
We summarize the initialization steps in Algorithm 2
17

---

Published as a conference paper at ICLR 2024
Algorithm 2 GBLR Parameter Initialization
C W
Target matrix W R
2: 11,12,-..
i'n <- argsort (
3: 5-0
4: for k I...., do
5:
0R + get. get.mask.parrmm Ciall 7.)
6:
= (W m))(W m
7:
6 & argmax( { R 2 h
8:
OF get.mask.parr ams
7T))
9:
b b b + nnz(mp)) + nnz n
if > s then
11:
wI < 0, w t 0: break
end
13:
end for
14:
U.V < WV)
15:
Solve Eq. 13 by Alg.
A.3 RECTENGULAR GAUDI-GBLR MATRICES
When the number of rows is not equal to the number of columns the width and location parameters for rows are defined on I, and Im- respectively, whereas the structural parameters on columns are defined on I, and In- Except for the index sets, the GBLR matrix on the rectangular case is defined exactly the same. Theorem [ and holdd as well. The Gaudi mask also defined in the same manner.
A.4 GAUDI MASK WITH REAL-VALUED WIDTHS AND LOCATIONS
Let us consider the Gaudi mask in Eq. with the continuous-valued widths and locations w, E [0, n). Also, let assume for now that one hand, Eq. 6 still compiles the time-domain signal after IDFT. On the other hand, the output of the IDFT not boxear mask anymore if least one of and [ not an integer. Figure 7 we present the time-domain signal of three cases of the Gaudi mask on n 512: 1) integer- valued width and locations with the smoothing parameter a oc, namely, no Gaussian filter is applied in the Gaudi function; 2) non-integer-valued width and location with integer-valued width and location with o = 100. As illustrated in the figure, the non-integer-valuee parameters generate many spiking errors (middle). Notably, however, the errors alleviated when the smoothing parameter a 100 (right), which almost recovers the shape of the signal with the integer-valued parameters (left).
oc, namely, no Gaussian smoothing is applied. On
and 3) non-
18

---

Published as a conference paper at ICLR 2024
n-232.
0,
location1191
location=191
0.2
0.3
100
200
200
300
300
100
200
300
410
500
100 200
300
400
500
Figure 7: Gaudi Masks in time-domain with different width, location and on n 512. Left: integer width w 128 and integer location / 192 with 3c. i.e., no Gaussian smoothing. Middle: non-integer width parameters, the time-domain signal contains spiking errors. Right: non-integer width w 128.5 and location / = 191.8 with 100. The Gaussian smoothing in the Gaudi function alleviates the spiking errors.
128.5 and location / 191.8 with = ac. Due to the non- non-integer
A.5
MORE DETAILS ON PROXIMAL GRADIENT DESCENT ON GAUDI-GBLR PARAMETERS IN DEEP NEURAL NETWORKS
A.5.1 DISCRETE STRUCTURAL PARAMETERS
Here we discuss solving Problem 8 using the discrete width and location parameters w E I, and
1 In-1- To be specific, given real-valued widths w E [0, n0.n]' K and locations 1 [0, n) | of
Gaudi-GBIl matrix, we apply Straight Through Estimator (Hinton) 2012, Bengio et al. 2013
Zhou et al. 2016) to and before generating the masks. Then, we project each value of and to [0, after the PGD step. We find that this setting is useful when the Gaussian smoothing parameter
a of the Gaudi function very high or infinite, since the non-integer parameters generate spiking errors as discussed in Appendix A.4
A.5.2 MULTIPLE GAUDI-GBLR MATRICES
For DNN with T weight matrices, we suggest adjusting the shrinkage parameter À in eq. based the predefined computational cost budget B for the DNN. For example,
X 0 if the sum of average widths of the weight matrices Usum ((W1) + @(W,) +... + w6 Wr) is below B, and use Ao 0 if Tsum above B where Ao is given as hyperparameter by user. In our experiments, this strategy effectively (not necessarily equally) distributes the computational cost to each matrix meet the overall cost budget constraint during the learning process.
A.6 IMPACT OF SVD
SVD. Since our method involves Singular Vector Decomposition (SVD) during initialization, the time complexity of the SVD discussed in this section. The impact caused by SVD negligible in the training process. The SVD operation performed during initialization decomposes at most 4096 x 1024-sized matrix, which takes less than a second on commercial CPUs or GPUs, whereas the overall training fine-tuning time ranges from few hours to tens of hours. Also, the SVD does not occur during the inference stage, so it does not become a bottleneck in any sense once the network is trained.
19

---

Published as a conference paper at ICLR 2024
A.7 EXTENDED EXPERIMENTAL RESULTS AND DETAILS
A.7.1 FINE-TUNING
In fine-tuning experiments, we use the pre-trained ViT-Base (Dosovitskiy al. 2020). For each pre-trained weight on Query, Key, Value and MLPs, we assign the Gaudi-GBLR matrix where the structural and content parameters of the filter = and the threshold to obtain the binary mask = 0.98 - max(Cs,) in Eq. 12 During the fine-tuning, we use AdamW Loshchilov & Hutter 2017) optimizers.
initialized by Algorithm We use the smoothing parameter
in Eq. 6 Then, all parameters updated by SGD with
For the Block Low-Rank matrices, we use Monarch (Dao et al., 2022) implementation. We project the pre-trained matrix by minimizing the Frobenius norm between the pretrained one and the Monarch matrix by gradient descent for 1000 steps with the learning rate of 0.001.
The initialization learning rate for the structural parameters was set to 0.0025 and decayed to 0.000025. We trained the networks for 35 epochs with the learning rate of 0.0001 for all param- eters, including the structural parameters of the Gaudi-GBLR matrices.
For the WikiText103 experiments, used the publicly available GPT-2 mode mode@for for the pre-trained weights.
A.7.2 TRAINING FROM SCRATCH
For the Gaudi-Git matrices, we initialize location parameters to zero and width parameters the low rank + block sparse structure. The content parameters well other weights are initialized randomly.
During the training with Proximal Gradient Descent in Eq. 9 as in Appendix A.5 the shrinkage parameter À is set to zero if the average width of the Saudi-GBLA matrices of the DNN below the predefined budget, and recovers the predefined value Åo otherwise.
Also, we anneal a = 1 to 100 during the training process. a is fixed to for 5 warm-up epochs, then linearly increases to 100 for 295 epochs. The hyperparameters used in the ImageNet experiment is presented in Table
Table 4: Hyperparameters used in ImageNet experiments using ViT-Base for training from scratch.
Model
Learning
Shrinkage
(iniviinal)
Epochs
Weight Decay
Drop Path Rate
Structural Paroma
Conteol Parana
Ag
Tareet Width
Vil-Base
N/A
DTKil
N/A
N/A
310
0.88
0.1
Vis-Base Caudi-GLLR
0.001
00001
0.04
0.12
708
1,00000
310
0.03
000
A.8 EXTENDED EXPERIMENTAL RESULTS
A.8.1 GBLR MATRIX FOR CONVOLUTION WEIGHTS.
We compare the ImageNet accuracy of ResNet18 (He et al. 2016) with GBLR weight matrices to the Automatic Layer-wise Decomposition Selector (ALDS) (Liebenwein et al. 2021). a simi- lar rank-based layer-wise compression method. Following Liebenwein et al. (2021). we consider the (d, c, ki, ka)-sized weight tensor of convolution layer as d, k1, h2 are the number of output channels, input channels, and the widths and heights of the kernels, respectively. In this manner, we converted the weights of the convolution layers of the ResNet18 model to GBLR matrices, and retrained the model. Table (5oss the accuracy and the relative FLOPPs of the ResNet18 model with dense weights, low-rank weights found by one-shot ALDS, and GBLR weights found by our method. Even without any supervision the structural properties of the convolution kernels, GBLR outperforms ALDS low-rank weights in terms both accuracy and complexity.
(2 matrix, where
20

---

Published as a conference paper at ICLR 2024
Table 5; ImageNet Accuracy on ResNet18.
Weight Type
Accuracy
GFLOPs
Dense
69.62
1.82
ALDS (Liebenwein et al. 2021
69.22
1.03
GBLR (ours)
69.31
1.01
A.8.2 GENERALIZATION GAP OF DNNS WITH GBLR WEIGHTS
In this section, estimate the generalization gap of the DNNs with GBLR weights. We used the VīT-Base model discussed in Section 4.3 which is trained on ImageNet from scratch. The accuracy of the dense and GBLR models on the training and validation sets are evaluated and presented in Table 6 Notably, the training-validate that the dense model. The comparable validation accuracy despite the lower training accuracy provides reasonable empirical evidence that the DNNs with GBLR matrices may exhibit smaller generalization errors than dense models.
accuracy difference of the GBLR model is smaller than
Table 6: Generalization Gap of Dense and GBLR ViT-Base Models.
Weight Type
Training Ace. (%)
Validation Acc. (A) (%)
Training Loss
Validation Loss (A)
Dense
99.03
78.57 (20.46)
0.1640
1.0639 (0. 09.9999
GaLdi-GBLR
94.93
78.51 (16.42)
0.3260
0.9917 (0.6657)
21


Reviews:
[{"review_id": "JjUJ1kPJn1", "reviewer": "Reviewer_g58u", "summary": "The authors introduce a Generalized Block-low-rank  (GBLR) matrix format to construct computationally efficient structures of weight matrices. They also introduce Gaussian-Dirichlet (Gaudi) function to make the structural parameters differentiable and provide an algorithm to learn neural networks with Gaudi-GBLR weight matrices.", "questions": "As I also mentioned in the weakness part, how does the GBLR weight matrices affect the generalization property or the complexity of the neural network?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed GBLR format includes existing important matrix structures. Also, the authors provide a method to make the structural parameters of weight matrices learnable. The idea is interesting and relevant to the community.", "weaknesses": "Providing theoretical investigations of the neural networks learned by the proposed method can improve the quality of the paper. Since the weight matrices are forced to be sparse, I think we need a different analysis from existing analysis for the dense weight matrices. For example, do you have any explanation about the representation power and generalization property of the networks with GBLR weight matrices?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xDOq99ugxL", "reviewer": "Reviewer_f3RM", "summary": "This paper proposes a generalized and differentiable method to learn efficient structures of weight matrices. Moreover, the authors present an effective initialization technique for the proposed method. Some experimental results show the performance of the proposed method.", "questions": "Although the paper is theoretically and experimental sound, there are still some questions need to be discussed in this paper:\n1.\tIn Algorithm 1, what’s AdamW(.), as well as clip?\n2.\tIn Eq. (8), what’s the variable, w or \\theta?\n3.\tThe advantage of the proposed method against existing methods is not clear. \n4.\tThe parameter initialization for the proposed method needs to perform SVD. Thus, the authors should analyze the computational complexity.  \n5.\tThe experimental results are not convincing. The authors should compare the performance of the proposed algorithm and more methods on more models and datasets.\n6.\tThe English language in this paper needs to be improved.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper proposes a generalized and differentiable method to learn efficient structures of weight matrices. \n2. Moreover, the authors present an effective initialization technique for the proposed method. \n3. Some experimental results show the performance of the proposed method.", "weaknesses": "Although the paper is theoretically and experimental sound, there are still some questions need to be discussed in this paper:\n1.\tIn Algorithm 1, what’s AdamW(.), as well as clip?\n2.\tIn Eq. (8), what’s the variable, w or \\theta?\n3.\tThe advantage of the proposed method against existing methods is not clear. \n4.\tThe parameter initialization for the proposed method needs to perform SVD. Thus, the authors should analyze the computational complexity.  \n5.\tThe experimental results are not convincing. The authors should compare the performance of the proposed algorithm and more methods on more models and datasets.\n6.\tThe English language in this paper needs to be improved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xX1n68mhtO", "reviewer": "Reviewer_wuLi", "summary": "The paper presents a new compressed matrix format (parametrization) called GBLR and an optimization method (involving proximal gradient descent, homotopy, and STE) to learn the parameters of the format to make NNs faster and smaller. The proposed parametrization can be understood as a juxtaposition of low-rank matrices of various shapes that needs to be appropriately padded. Such a parametrization is very flexible as contains regular low-rank matrices (if the shape of the blocks are same as the shape of the entire matrix), block low rank matrices, and block sparse. To be more specific, every low-rank submatrix  is represented as a sum of rank-1 matrices; thus, for a given subblock of rank $k$ parameters include shape of submatrix ($w$,$h$), location within orgininal matrix (i,j), and actual values that are $k$ rank-1 matrices of $w \\times h$ stored as $w\\times 1$ and $1 \\times h$ updates.\n\nSuch a parametrization has many non-differentiable parameters, and to make it amenable to SGD based solvers the authors propose several modifications: instead of storing compact $w\\times 1$ and $1\\times h$ rank-1 matrices, the parametrization now involves the whole $n\\times 1$ and $1\\times n$ vectors that are appropirately masked with vector $m$. The mask presents a boxcar filter which is non-differentiable itself, but author approximate it via Gaudi mask in frequency domain with gaussian kernel (with variance $\\sigma^2$); Gaudi mask converges to boxcar with $\\sigma \\to \\infty$ which needs to be driven during training (hence a homotopy). Also, during the training the width $w$, height $h$, location $i,h$, are kept real-valued (for SGD) but apply straight-through estimation for actual matrix recreation. And finally, to drive matrix rank selection, authors propose a cost based selection via an $\\ell_1$ penalty (controlled by $\\lambda$) on certain structure parameters (since the exact equation is not given, I'm guessing on every rank-1 matrix?).\n\nAs for training of the compressed models, the authors plug in GBLR parametrized matrix instead of original weight matrices of nns (the initial values are obtained via an algorithm in A.2), choose appropriate $\\lambda$ (to control for rank), and train or finetune end-to-end on a dataset. The results clearly show that such a scheme gets a better compression-accuracy tradeoff.", "questions": "1. Can you please provide the exact from of $\\ell_1$ penalty used in eq.8? How does FLOPs/parameter counts being represented as F$\\ell_1$ penalty?\n2. What is the value of $\\lambda$ used in experiments? How to choose it properly?\n3. What was the scheme used for $sigma$? The paper says that it was \"gradually increased to 100\", but having exact details is preferred.\n4. Certain computations (Gaudi function) happens in frequency domain that involves DFT and IDFT; how expensive is this operations wrt single training step? No slowdown, 0.75x slowdown, etc..\n5. Please included a more detailed literature review.\n6. Please strengthen the baselines.\n7. I am wondering why section 3.3 discusses the application on the basis of two layer MLP \"for ease of understanding\"? Wouldn't an example on a single layer be much simpler? BTW, there is a typo in this section, multi-layer perception => multi-layer perceptron", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposed a new matrix parametrization that includes many other compressed forms (low rank, block low-rank, and block sparse) as a subset. The parametrization allows to get better error-compression tradeoffs.\n\nThe presentation of the method is very thorough and includes many small details (except for some, see questions) that is definitely a plus for reproducibility.", "weaknesses": "I find two minor weaknesses of the paper, literature review and comparison to other methods.\n\nAlthough I understand that the focus of paper was compression of transformer based models, it seems many relevant low-rank and tensor-decomposition based methods that were used for compression of other networks (CNNs) were left out. Many of those left out papers share similar ideas (e.g., how to parametrize wrt rank) that need to be included. Some of the missed out works include:\n- Factorized Higher-Order CNNs with an Application to Spatio-Temporal Emotion Estimation \n- Low-rank Compression of Neural Nets: Learning the Rank of Each Layer\n- Coordinating filters for faster deep neural networks\n- Constrained optimization based low-rank approximation of deep neural networks\n- Compressing Neural Networks: Towards Determining the Optimal Layer-wise Decomposition\n\nbut there are many others (you can find others by looking within those)\n\nComparison to the relevant baselines. It seems the baselines authors choose are very simple (e.g., using a fixed rank for low-rank compression), and can be considerably strengthened if wanted. Thus I'm asking authors to include stronger baselines (low rank with rank selection, tensor decomposition methods) to compare.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)