Differentiable Learning of Generalized Structured Matrices for Efficient
  Deep Neural Networks
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.
Changwoo Lee, Hun-Seok Kim
# Differentiable Learning of Generalized Structured Matrices for Efficient Deep Neural Networks

###### 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. Finally, we introduce an effective initialization method for the proposed scheme. 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.

## 1 Introduction

Deep Neural Networks (DNNs) for large language models (LLMs) (Vaswani et al., 2017; Devlin et al., 2018; Radford et al., 2019; Brown et al., 2020) and vision tasks (Dosovitskiy et al., 2020; Touvron et al., 2021) have shown great success in various domains in recent years. The size of the DNNs, however, has increased on an extraordinary scale - up to 70 Billion parameters in the single model Zhao et al. (2023) - requiring 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 (Yaras et al., 2023; Huh et al., 2021; Gunasekar et al., 2017). However, identifying such structures to lower the effective complexity of weight matrices in recent models such as Transformers Vaswani et al. (2017) remains a challenging problem. Also, it mostly relies on the existing human-designed / hand-crafted structured matrices without a unified systematic approach. Hence prior works have focused on investigating new classes of structured matrices for DNNs (Li et al., 2015; Dao et al., 2022; Chen et al., 2022). Notably, each structured matrix is defined disjointedly from other formats. For example, neither the block-sparse matrix format nor the low-rank matrix format of the same number of parameters is a 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-differentiable discrete spaces.

In this paper, we investigate (locally) optimal structures of the weight matrices as well as a differentiable training method to learn them, attempting to answer the following two questions:

1. _Is there a universal format that represents a 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 a _generalized_ and _differentiable_ structured matrix format. The main contributions of this work can be summarized as follows.

1) We propose a **Generalized Block-low-rank** (GBLR) matrix format, which includes many important structures such as Low-Rank (LR), Block Sparse (BSP), and Block-low-rank (BLR) matrices under some practical conditions. The new structured matrix format consists of two types of parameters: 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 matrices in the structural parameter space, which we believe is a strong evidence that the GBLR format is able to capture undiscovered structured matrix formats.

2) We introduce a _differentiable_ parameterization of 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 **Gaussian-Dirichlet** (Gaudi) function followed by inverse Fast Fourier Transform (IFFT) to the format named _Gaudi-GBLR_. We show that the derivatives of the Gaudi function with respect to the structural parameters exist almost everywhere, even when the width is zero.

3) We propose a practical learning algorithm based on the 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., 2020) and MLP-Mixer (Tolstikhin et al., 2021), outperforming prior approaches using hand-designed structured matrices.

## 2 Preliminaries

**Notation.** We use \(\odot\) to indicate the elementwise (Hadamard) product of two matrices. The imaginary unit is denoted by \(\imath=\sqrt{-1}\). The normalized \(\operatorname{sinc}\) function is defined by \(\operatorname{sinc}(x)=\frac{\sin\pi x}{\pi x}\) where \(\operatorname{sinc}(0):=1\). The (element-wise) floor function is denoted by \([\cdot]\). The index of elements in a matrix or vector starts from zero (instead of one), following the convention of Cooley & Tukey (1965). Also, \(\mathcal{I}_{n}=\{0,1,\ldots,n\}\) denotes the index set from 0 to \(n\).

**Assumption.** For simplicity, we assume the weights are square matrices. Extension to rectangular matrix cases is discussed in Appendix A.3.

### Block-related Matrices.

**A block \(\mathbf{B}\in\mathbb{R}^{|R|\times|C|}\)** of a matrix \(\mathbf{W}\in\mathbb{R}^{n\times n}\) is a submatrix of \(\mathbf{W}\) with _consecutive cyclic_ row and column indices. For example, \(R\) can be \(\{n-2,n-1,0,1\}\) if \(|R|=4\). Also, we say two blocks \(\mathbf{B}_{1}\) and \(\mathbf{B}_{2}\)**overlap** if they have shared elements of \(\mathbf{W}\).

Two well-known block-structured matrices are a _block sparse_ (BSP) matrix and a _block-low-rank_ (BLR)(Amestoy et al., 2015) matrix. Informally speaking, a matrix is block-sparse when non-zero elements are gathered in blocks and such blocks are sparse in the matrix. A block-low-rank matrix is composed of non-overlapping equally-partitioned low-rank blocks. Figure 1 illustrates the BSP and BLR matrices as well as our proposed generalized block-low-rank matrix, which we introduce in Section 3.1. We present the formal definitions of the block-sparse and block-low-rank matrices in Appendix A.1.

### Structured Matrix

We say a matrix \(\mathbf{W}\in\mathbb{R}^{n\times n}\) is **structured** if, for any \(\mathbf{x}\in\mathbb{R}^{n}\), the matrix-vector product (MVP) \(\mathbf{y}=\mathbf{W}\mathbf{x}\) requires significantly less number of multiplications than \(n^{2}\). For instance, LR, BSP,

Figure 1: Comparison of block-sparse, block-low-rank, and our proposed Generalized block-low-rank matrices.
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 of MVP to be a sub-quadratic function of \(n\), it is excluded in our discussion of structured matrices because processing moderately sparse matrices (10\(\sim\)50% density) in conventional hardware such as GPUs does not reduce the actual run-time due to its unstructured positions of non-zero elements (Chen et al., 2022).

## 3 Proposed Method

We are interested in training a deep neural network (DNN) \(f\) under a _computational cost_ constraint:

\[\min_{f}\sum_{\mathbf{x},\mathbf{y}\sim\mathcal{D}}\mathcal{L}(f(\mathbf{x}),\mathbf{y})\quad \text{s.t.}\quad\mathrm{cost}(f)\leq B, \tag{1}\]

where the first term is the cross-entropy loss for a classification task at a data point \(\mathbf{x}\) and a label \(\mathbf{y}\), and the constraint \(\mathrm{cost}(f)\) is the number of multiplications to compute the neural network output \(f(\mathbf{x})\). Our method learns weight matrices of DNNs in a _generalized_ structured matrix format in a _differentiable_, end-to-end manner.

### Generalized Block-Low-Rank (GBLR) Matrix

We introduce a concept of a generalized structured matrix format that explains multiple popular structure matrices used in DNN training. The idea is that a block in a matrix can be expressed by a **sum of rank-1 blocks**, i.e., a rank-\(r\) block is a sum of \(r\) rank-1 blocks at the same position. In this manner, LR, BSP, and BLR matrices can be expressed under a unified framework (Theorem 1).

To be specific, our proposed structured matrix format is a generalized version of Block-Low-Rank (BLR) matrices. An \(n\)-by-\(n\)_Generalized Block-Low-Rank_ (GBLR) matrix \(\mathbf{W}\) is obtained by _overlapping_ multiple rank-1 blocks of _different_ sizes at _arbitrary_ locations, as depicted in Figure 2 (Left). 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: \(\mathbf{B}_{1}\in\mathbb{R}^{w_{1}^{R}\times w_{1}^{C}},\mathbf{B}_{2}\in\mathbb{R}^{w _{2}^{R}\times w_{2}^{C}},\ldots,\mathbf{B}_{K}\in\mathbb{R}^{w_{K}^{R}\times w_{ K}^{C}}\). Each block \(\mathbf{B}_{k}\) has two parameter sets: 1) the _structural_ parameters that identify the _position_ of the block in the row and column index set \(\mathcal{I}_{n-1}=\{0,1,\ldots,n-1\}\), and 2) the _content_ parameters which specify the actual values of matrix elements.

Configuration of structural parameters.The position of a block is given by the indices of the rows and columns it occupies in the \(n\times n\) matrix. Hence, the placement of a rectangle block can

Figure 2: Left: An example of a GBLR matrix with 4 blocks. A block is generated from the _structural_ parameters \((w^{R},l^{R}),(w^{C},l^{C})\) and the _content_ parameters \((\mathbf{u},\mathbf{v})\), where \((w^{R},l^{R})\) and \((w^{C},l^{C})\) form binary masks \(\mathbf{m}_{(w^{R},l^{R})}\) and \(\mathbf{m}_{(w^{C},l^{C})}\), respectively. Note that overlapped regions can have a rank higher than one. Right: Efficient Matrix-Vector Product computations using cropped content parameters and structural parameters. The structural parameters locate the input and output indices/addresses to read and write.
be identified by four numbers: _width_ and _location_ in terms of the row or column indices. Hence, we use a _location_ parameter \(l\) and a _width_ parameter \(w\) as the structural parameters of a block. Figure 2 (Left) illustrates a block of size \(w^{R}\times w^{C}\) in an \(n\times n\) matrix at location \((l^{R},l^{C})\). The row (column) index set of a block is the sequence of numbers from \(l^{R}\) (\(l^{C}\)) to \(l^{R}+w^{R}\) (\(l^{C}+w^{C}\)) where the addition is a cyclic/modulo addition. For each block \(\mathbf{B}_{k}\) for \(k=1,\ldots,K\), we have four parameters: \((w^{R}_{k},l^{R}_{k})\) for the row and \((w^{C}_{k},l^{C}_{k})\) for the column. We use the notation \(\phi^{R}_{k}=(w^{R}_{k},l^{R}_{k})\) and \(\phi^{C}_{k}=(w^{C}_{k},l^{C}_{k})\) to represent the tuple of width and location for the row (\(\phi^{R}_{k}\)) and column (\(\phi^{C}_{k}\)).

Based on the structural parameter \(w^{C}_{k}\) and \(l^{C}_{k}\), one can construct an \(n\)-dimensional binary mask that has \(w^{C}_{k}\in\mathcal{I}_{n}\) consecutive ones starting from \(l^{C}_{k}\in\mathcal{I}_{n-1}\) in the cyclic order:

\[m_{\phi^{C}_{k}}[j]=m_{(w^{C}_{k},l^{C}_{k})}[j]=\begin{cases}1&\text{if }l^{C}_{k}\leq j+ an<l^{C}_{k}+w^{C}_{k}\\ 0&\text{otherwise}\end{cases},\;j\in\mathcal{I}_{n-1},\;a\in\{0,1\}, \tag{2}\]

where \(a\in\{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^{C}_{k}\) (cyclic) consecutive non-zero elements of an \(n\)-dimensional vector. The mask for the rows \(\mathbf{m}_{\phi^{R}_{k}}\) is obtained in the same way from \(w^{R}_{k}\) and \(l^{R}_{k}\).

**Configuration of content parameters.** To represent the values of a rank-1 block \(\mathbf{B}_{k}\), we use two \(n\)-dimensional vectors \(\mathbf{u}_{k}\) and \(\mathbf{v}_{k}\) as _content_ parameters along with the boxcar masks \(\mathbf{m}_{\phi^{R}_{k}}\), \(\mathbf{m}_{\phi^{C}_{k}}\). All these parameters \((\phi^{R}_{k},\phi^{C}_{k},\mathbf{u}_{k},\mathbf{v}_{k})\) are learned during the DNN training _simultaneously_. Since the boxcar masks \(\mathbf{m}_{\phi^{R}_{k}}\) and \(\mathbf{m}_{\phi^{C}_{k}}\) guide the location of the block in the \(n\times n\) matrix, \(\mathbf{u}_{k}\) and \(\mathbf{v}_{k}\) are element-wise multiplied with the boxcar masks:

\[\mathrm{ZeroPad}(\mathbf{B}_{k})=(\mathbf{m}_{\phi^{R}_{k}}\odot\mathbf{u}_{k})(\mathbf{m}_{ \phi^{C}_{k}}\odot\mathbf{v}_{k})^{T},\]

where the resulting \(n\times n\) matrix is a zero-padded block. Ideally, we expect the mask to expand / shrink and shift to find the right subset of the elements of a 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\) zero-padded blocks:

\[\mathbf{W}=\sum_{k=1}^{K}(\mathbf{m}_{\phi^{R}_{k}}\mathbf{m}^{T}_{\phi^{C}_{k}})\odot(\bm {u}_{k}\mathbf{v}^{T}_{k})=\sum_{k=1}^{K}\left(\mathbf{m}_{\phi^{R}_{k}}\odot\mathbf{u}_{k }\right)\left(\mathbf{m}_{\phi^{C}_{k}}\odot\mathbf{v}_{k}\right)^{T}. \tag{3}\]

A GBLR matrix is associated with an average width \(\bar{w}=\frac{1}{2K}\sum_{k=1}^{K}w^{R}_{k}+w^{C}_{k}\).

**Definition 1**.: _Let \(\mathsf{GBLR}(n,K,s)\) be the set of matrices obtained by Eq. 3 for the average width less than or equal to \(s\), i.e., \(\bar{w}=\frac{1}{2K}\sum_{k=1}^{K}w^{R}_{k}+w^{C}_{k}\leq s\). A matrix \(\mathbf{W}\) is an \((n,K,s)\)-GBLR if \(\mathbf{W}\in\mathsf{GBLR}(n,K,s)\)._

We use the notation \(\mathbf{\phi}(\mathbf{W}):=(\mathbf{w}(\mathbf{W}),\mathbf{l}(\mathbf{W}))\) to indicate the collection of the structural parameters of \(\mathbf{W}\), where \(\mathbf{w}(\mathbf{W})=\{w^{R}_{1},w^{C}_{1},w^{R}_{2},w^{C}_{2},\ldots,w^{R}_{K},w^{C }_{K}\}\) and \(\mathbf{l}(\mathbf{W})=\{l^{R}_{1},l^{C}_{1},l^{R}_{2},l^{C}_{2},\ldots,l^{R}_{K},l^{C }_{K}\}\). We simply use \(\mathbf{\phi}:=\mathbf{\phi}(\mathbf{W})\) if \(\mathbf{W}\) is clearly inferred in the context.

**Efficiency.** Once the structural parameters are fixed, it is unwise to store and use two \(n\)-dimensional content parameters for each block \(\mathbf{B}_{k}\) because only \(w^{R}_{k}\) elements of \(\mathbf{u}_{k}\) and \(w^{C}_{k}\) elements of \(\mathbf{v}_{k}\) are non-zero according to the boxcar masks in eq. 3. Hence, one can store and use the _cropped_ content parameters \(\mathbf{u}_{l^{R}:l^{R}+w^{R}}\) and \(\mathbf{v}_{l^{C}:l^{C}+w^{C}}\) for MVP between an input \(\mathbf{x}\in\mathbb{R}^{n}\) (which can be also cropped from \(l^{C}\) to \(l^{C}+w^{C}\)) and a block \(\mathbf{B}\), as described below and in Figure 2 (Right):

\[\mathrm{ZeroPad}(\mathbf{B})\mathbf{x} =(\mathbf{m}_{(w^{R},l^{R})}\odot\mathbf{u})(\mathbf{m}_{(w^{C},l^{C})}\odot \mathbf{v})^{T}\mathbf{x}\] \[=\mathrm{ZeroPad}\left(\mathbf{u}_{l^{R}:l^{R}+w^{R}}(\mathbf{v}^{T}_{l^ {C}:l^{C}+w^{C}}\mathbf{x}_{l^{C}:l^{C}+w^{C}})\right),\]

which requires only \(w^{R}+w^{C}\) multiplications. Hence, the number of multiplications (denoted by FLOPs) for multiplying \(\mathbf{W}\in\mathsf{GBLR}(n,K,s)\) with \(\mathbf{x}\in\mathbb{R}^{n}\) is bounded by \(2Ks\):

\[\text{FLOPs}=\sum_{k=1}^{K}(w^{R}_{k}+w^{C}_{k})=2K\bar{w}\leq 2Ks. \tag{4}\]
**Expressiveness.** Low-rank (LR), block sparse (BSP), and block-low-rank (BLR) matrices are popular 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.1

**Theorem 1**.: _Let \(n,K,s\) be positive integers satisfying \(Ks\geq n\). Then any \(n\)-by-\(n\) rank-\(\frac{Ks}{n}\) matrices and \((n,\frac{K}{s},s)\)-block-sparse matrices are \((n,K,s)\)-GBLR. Also, any \((n,K,s,1)\)-block-low-rank matrices are \((n,K,s)\)-GBLR if \(K=(n/s)^{2}\)._

More importantly, a new structured matrix obtained by interpolating the structural parameters of two \((n,K,s)\)-GBLR matrices is still \((n,K,s)\)-GBLR, based on Theorem 2. Therefore, a new type of structured matrices can be derived from a set of GBLR matrices.

**Theorem 2** (Closed under structural interpolation).: _Given two \(n\times n\) matrices \(\mathbf{W},\mathbf{Z}\in\mathsf{GBLR}(n,K,s)\), and \(\alpha\in[0,1]\), consider the following combination between the structural parameters:_

\[\mathbf{w}^{\prime}=\lfloor\alpha\mathbf{w}(\mathbf{W})+(1-\alpha)\mathbf{w}(\mathbf{Z})\rfloor, \quad\mathbf{l}^{\prime}=\lfloor\alpha\mathbf{l}(\mathbf{W})+(1-\alpha)\mathbf{l}(\mathbf{Z})\rfloor.\]

_A matrix \(\mathbf{Y}\) generated by Eq. 3 with the structural parameter \((\mathbf{w}^{\prime},\mathbf{l}^{\prime})\) is a \((n,K,s)\)-GBLR matrix, \(\mathbf{Y}\in\mathsf{GBLR}(n,K,s)\)._

Theorem 1 and Theorem 2 tell us that \((n,K,s)\)-GBLR matrices cover a wide range of popular existing structured matrices and also undiscovered ones. In the following section, we introduce a differentiable tool to find/learn structured matrices in \(\mathsf{GBLR}(n,K,s)\).

### Gaussian-Dirichlet (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 so, the non-differentiability of the boxcar mask parameters needs to be handled.

To tackle this issue, we introduce a _Dirichlet-kernel-based_ parameterization of the boxcar to _explicitly_ parameterize the width \(w\) and location \(l\) in the expression. Consider a boxcar mask of length \(n\), width \(w\), and location \(l\). Let \(\hat{\mathbf{m}}_{(w,l)}\) be the discrete Fourier transform (DFT) of the mask \(\mathbf{m}_{(w,l)}\). The frequency-domain representation of the boxcar mask is given by

\[\hat{m}_{(w,l)}[k] =e^{-2\pi\frac{k}{n}l}\hat{m}_{(w,0)}[k]=e^{-2\pi i\frac{k}{n}l} \sum_{j=0}^{w-1}e^{-2\pi ij\frac{k}{n}}=e^{-2\pi i\frac{k}{n}l}w\frac{\text{ sinc}\left(w\frac{k}{n}\right)}{\text{sinc}\left(\frac{k}{n}\right)}e^{\pi k \left(\frac{1-w}{n}\right)} \tag{5}\] \[=e^{-2\pi\frac{k}{n}l}d_{w}[k],\]

where \(d_{w}[k]:=w\frac{\text{sinc}\left(w\frac{k}{n}\right)}{\text{sinc}\left( \frac{k}{n}\right)}e^{\pi k\left(\frac{1-w}{n}\right)}\) is the _Dirichlet kernel of order \(n\)_(Bruckner et al., 1997) in the discrete domain \(\mathcal{I}_{n-1}=\{0,1,\ldots,n-1\}\).

Furthermore, we propose a smooth version of the Dirichlet-kernel-based mask by convolving the time-domain boxcar mask with the Gaussian-shape function. It is obtained in the frequency domain by element-wise multiplying the Gaussian function \(g_{\sigma}[k]=\exp\left(-\frac{k^{2}}{2\sigma^{2}}\right)\) with the standard deviation \(\sigma>0\) to the Dirichlet kernel. We call the resulting function the **Gaussian-Dirichlet** (Gaudi) function \(\mathbf{d}_{w}^{\sigma}\):

\[\mathrm{d}_{w}^{\sigma}[k] :=g_{\sigma}(k)\cdot\mathrm{d}_{n}[k], \tag{6}\] \[\tilde{\mathbf{m}}_{(w,l)}^{\sigma} :=\mathbf{IDFT}(\mathbf{d}_{w}^{\sigma}\cdot e^{-2\pi\frac{k}{n}l}),\]

where \(\mathbf{k}=[0,1,\ldots,n-1]\). And we call the mask generated by a Gaudi function \(\tilde{\mathbf{m}}_{(w,l)}^{\sigma}\) the **Gaudi mask**, where the parameter \(\sigma\) controls the smoothness. Note that as \(\sigma\rightarrow\infty\), the Gaudi mask converges to the boxcar mask (Lemma 5). Figure 3 visualizes the time-domain Gaudi mask approaching the boxcar mask.

Figure 3: Comparison between Boxcar mask and Gaudi masks in the time domain with different smoothing factors \(\sigma\). The Gaudi mask converges to the Boxcar mask as \(\sigma\) grows.
A useful property of the Gaudi mask is that one can obtain exact derivatives with respect to the width parameter, even when the width is zero. To show it, we relax the domain of widths and locations to the continuous interval \([0,n]\) (see Appendix A.4).

**Theorem 3**.: _Let \(n<\infty\) be a finite positive integer. For any \(\sigma\in(0,\infty]\) and \(w,l\in[0,n]\), the derivatives of the Gaudi mask \(\tilde{\mathbf{m}}^{\sigma}_{(w,l)}\) with respect to \(w\) and \(l\) are bounded almost everywhere:_

\[\left\|\frac{\partial\tilde{\mathbf{m}}^{\sigma}_{(w,l)}}{\partial w}\right\|_{2} <\infty,\quad\left\|\frac{\partial\tilde{\mathbf{m}}^{\sigma}_{(w,l)}}{\partial l }\right\|_{2}<\infty.\]

Especially when \(w=0\), the derivative with respect to \(w\) is neither divergent nor zero.

**Corollary 4**.: _For any \(l\in[0,n]\) and \(\sigma\in(0,\infty]\), the norm of the derivative of the Gaudi mask \(\tilde{\mathbf{m}}^{\sigma}_{(w,l)}\) with respect to \(w\) at \(w=0\) is well-defined and greater than zero, i.e., \(0<\left\|\frac{\partial\tilde{\mathbf{m}}^{\sigma}_{(w,l)}}{\partial w}\right\|_{ 2}<\infty\)._

To allow learning the mask structural parameters in a differentiable manner, we plug the Gaudi mask (Eq. 6) into Eq. 3 as the mask of GBLR matrices to model **Gaussian-Dirichlet GBLR** (GauDi-GBLR) with parameters \(\mathbf{\theta}=(\mathbf{\phi},\mathbf{U},\mathbf{V},\sigma)\):

\[\mathbf{W^{\theta}}=\mathbf{W^{(\phi,U,V,\sigma)}}=\sum_{k=1}^{K}\left(\tilde{\mathbf{m}} ^{\sigma}_{\phi_{k}^{R}}\odot\mathbf{u}_{k}\right)\left(\tilde{\mathbf{m}}^{\sigma}_{ \phi_{k}^{C}}\odot\mathbf{v}_{k}\right)^{T}. \tag{7}\]

In practice, one can use small \(\sigma\approx 1\) at the beginning of the training process to update the corresponding (unmasked) content parameters which are more than necessary, then gradually increase \(\sigma\) to adjust the content parameters with a tighter mask. Since the purpose of using a Gaudi mask is 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 3.1-efficiency, without constructing masks at all. Hence, during the inference, there is no overhead to compute Gaudi masks.

### Learning Gaudi-GBLR for Efficient Neural Networks

We now introduce an algorithm to learn the structural parameters of Gaudi-GBLR matrices. The goal of the learning algorithm is to identify Gaudi-GBLR matrix structures (i.e., their parameters) that allow computationally efficient DNNs. Our discussion is centered around a two-layered multi-layer perception (MLP) for ease of understanding. However, the technique can be applied to general DNNs that incorporate linear layers.

Now let us consider a two-layered MLP \(\mathbf{f^{\theta}}\) with a Gaudi-GBLR weight matrix \(\mathbf{W^{\theta}}\): \(\mathbf{f^{\theta}}(\mathbf{x})=\mathbf{h}^{T}a(\mathbf{W^{\theta}}\mathbf{x}+\mathbf{b})\). We initially relax the domain of \(w\in\mathcal{I}_{n}\) and \(l\in\mathcal{I}_{n-1}\) of the Gaudi mask to the real-valued space \(w,l\in[0,n]\) as we discuss in Appendix A.4. Due to the property of Gaudi-GBLR matrices Eq. 4, the computational cost constraint on the DNN \(\mathbf{f^{\theta}}\) in Problem (1) can be replaced by a constraint on the sum of the width parameters of \(\mathbf{W^{\theta}}\). Specifically, we find width parameters of \(\mathbf{W^{\theta}}\) satisfying \(\|\mathbf{w}\|_{1}\leq B\) since \(\sum_{k=1}^{K}w_{k}=\|\mathbf{w}\|_{1}\). To solve the problem with Gradient Descent, we relax this \(\ell_{1}\)-norm constrained problem to a _unconstrained_ one using the method of Lagrange multiplier:

\[\min_{\mathbf{\theta}}\sum_{(\mathbf{x},\mathbf{y})\sim\mathcal{D}}\mathcal{L}(\mathbf{f^{ \theta}}(\mathbf{x}),\mathbf{y})+\lambda\|\mathbf{w}\|_{1},\quad\lambda\geq 0. \tag{8}\]

The resulting computational budget is implicitly constrained by a hyperparameter \(\lambda\geq 0\).

Theorem 3 guarantees the derivatives of the widths and locations of the Gaudi-GBLR matrix in \(\mathbf{f^{\theta}}\) can be obtained with any positive smoothing parameter \(\sigma>0\) so that we can safely learn the
parameters in the continuous domain \([0,n]\). Specifically, we update the width parameter \(\mathbf{w}\) in the \(\ell_{1}\)-norm term in Problem (8) by Proximal Gradient Descent (PGD):

\[\mathbf{w}_{t+1}=S_{\eta\lambda}(\mathbf{w}_{t}-\eta\nabla\mathcal{L}(\mathbf{f}^{\mathbf{\theta }}(\mathbf{x}),\mathbf{y})), \tag{9}\]

where \(S_{\mu}(x)=\begin{cases}\mathrm{sign}(x)\cdot(|x|-\mu)&\text{if }x>\mu\\ 0&\text{otherwise}\end{cases}\) is the element-wise soft shrinkage function. In practice, the gradient is calculated with an adaptive optimizer such as Adam (Kingma & Ba, 2014) or AdamW (Loschnikov & Hutter, 2017). The overall process is summarized in Algorithm 1. Although Problem (8) is non-linear, our experimental results show that PGD can attain good local optima with an adaptive optimizer and the initialization method we propose in Appendix A.2. A practical learning method for the width and location parameters defined in the discrete spaces \(\mathcal{I}_{n}\) and \(\mathcal{I}_{n-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 of columns of 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-plus-Low-Rank (BSP-LR), and Block-low-rank (BLR). For LR matrices, we use singular vector decomposition to find the best rank-\(s\) approximation of the given matrix for the predefined rank of \(s\). 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 FLOPs to denote the number of multiplications, and use 8 NVIDIA A40 GPUs in our experiments. The detailed experimental settings are described in Appendix A.6.

### Fine-tuning Results

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 A.2. The ViTs with Gaudi-GBLR matrices were fine-tuned on ImageNet (Russakovsky et al., 2015). During the initialization, we set the computational budget of all matrices in the network the same. For a fair comparison, the same set of hyperparameters was used throughout our fine-tuning experiments.

The highest accuracy is achieved by Gaudi-GBLR in ViT-Base with a patch size of \(16\times 16\) on the ImageNet validation dataset after fine-tuning it for 35 epochs. Figure 4 shows that Gaudi-GBLR preserves the accuracy well when the complexity is reduced to 30% of the original 'Dense' model which does not use structured matrices (its FLOPs count is normalized to 1). The other hand-designed approaches exhibit more significant accuracy degradations for the same complexity reduction. Overall, Gaudi-GBLR strikes better Accuracy-FLOPs trade-offs than LR or Monarch approaches. The higher accuracy for a similar FLOPs count quantifies the gain from the learned structured matrices.

### Training From Scratch

We train ViTs (Dosovitskiy et al., 2020) and MLP-Mixers (Tolstikhin et al., 2021) with structured weight matrices on CIFAR-10&100 (Krizhevsky et al., 2009) by Algorithm 1 from randomly-initialized content parameters. We set \(\sigma=1\) for the first epoch and gradually increase it to \(100\) until

Figure 4: ImageNet accuracy after fine-tuning ViT-Base weights replaced with structured weight matrices. Dense: the original ViT-Base model.
the training process ends. In Figure 5, we study the accuracy-FLOPs trade-off using CIFAR10/100 datasets when the models are trained from scratch (i.e., not fine-tuned from pre-trained weights). As in the ImageNet fine-tuning experiment, Gaudi-GBLR achieves superior accuracy-FLOPs trade-offs outperforming the other hand-designed structured matrices.

### Analysis on Learned Gaudi-GBLR Matrices

In this section, we study the learned Gaudi-GBLR matrices in terms of computational budget allocation, and mask patterns. The analysis is based on ViT-Base Gaudi-GBLR matrices trained on ImageNet from scratch by following the same \(\sigma\) adaptation rule used in CIFAR10/100 experiments. The accuracy and FLOPs of the ViT-Base with Gaudi-GBLR is reported in Table 1.

**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 1. 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 a result, Gaudi-GBLR weight matrices in the network have unequal widths requiring different FLOPs for each MVP. Table 2 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_ use a smaller budget. The smallest layer uses only about 4,200 FLOPs for MVP involving a \(768\times 768\) matrix and a \(768\times 1\) vector. The FLOPs assigned to the linear layer of channel MLPs (_MLP-FC1_ and _MLP-FC2_ in Table 2) vary significantly. Although the size of the weight matrix used in the MLPs are \(4\times\) larger than the ones used for Value, the _MLP-FC2_-type layers use \(7.96\times\) FLOPs than the _Value_-type layers in average.

\begin{table}
\begin{tabular}{|l|c|c|} \hline Model & Acc. (\%) & GFLOPs \\ \hline ViT-Base & 78.57 & 17.2 \\ \hline w/ Gaudi-GBLR & 78.51 & 5.65 \\ \hline \end{tabular}
\end{table}
Table 1: ImageNet accuracies and GFLOPs on a \(224\times 224\) image of ViT-Base models trained from scratch with dense matrices and Gaudi-GBLR matrices.

Figure 5: Accuracy-Cost trade-off of models trained from scratch on CIFAR-10/100 dataset.

Figure 6: Mask patterns of weights for Query, Key, Value, FC1 and FC2 in Layer 0, 5, and 11 (out of \(0\sim 11\)). The brighter, the more overlapped blocks. The numbers below indicate the rank of the matrix (max: 768).
**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 (FC1 and FC2) in MLPs. Bright colors in Figure 6 highlight regions where masks are overlapped. The rank of each weight matrix is marked under each visualized mask pattern. Interestingly, 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., 2017) scheme of the ViT. Each weight matrix of an attention layer is a collection of weights for multiple heads. It is expected that some heads have more significant impacts on the output while others may not contribute as much. Hence multiple blocks are allocated to regions that correspond to those heads. Notice the rank of matrices obtained from the GBLR framework differs significantly across different layers and matrix types (Values, Queries, and Keys).

## 5 Related Works

**Structured Matrices for DNNs.** Prior works have identified that the weights have simple structures under certain conditions (Yaras et al., 2023; Huh et al., 2021; Gunasekar et al., 2017). However, explaining the structure of every weight matrix in practical DNNs such as Transformers (Vaswani et al., 2017) is still a challenging problem. Hence, prior works have focused on manually designing suitably structured matrices for DNNs. Hsu et al. (2022) used weighted low-rank decomposition for the weights of language models. Butterfly matrices (Li et al., 2015; Dao et al., 2019), inspired by Fast Fourier Transform (FFT), 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). Unlike these prior works that rely on manual designs, our method learns the structure of weight matrices from the training data by stochastic gradient descent.

**Mask Learning** Masking is a popular technique to prune neurons or activations of the DNNs. Since a mask for neuron pruning/selection is a non-differentiable binary vector, its gradient does not exist almost everywhere. Jang et al. (2016) and Maddison et al. (2016) propose alternative distributions close to the Bernoulli distribution, which allow the continuous gradient. Movement Pruning in (Sanh et al., 2020) utilizes Straight-Through Estimator (STE) (Hinton, 2012; Bengio et al., 2013; Zhou et al., 2016) to pass the gradient through the Top-\(k\) function. Lin et al. (2017) adopts deep reinforcement learning to select the input-dependent mask. On the contrary, our mask design solves the non-existing gradient problem by Gaudi-function-based parameterization in the frequency domain without using a surrogate or approximated gradient.

**One-shot Neural Architecture Search.** Neural Architecture Search (NAS) (Zoph and Le, 2016; Liu et al., 2018) seeks the optimal neural network structures from training data. To improve search efficiency, Pham et al. (2018); Liu et al. (2018) adopt the one-shot NAS technique that selects sub-network candidates from a super-network. Our method also falls into a similar category of finding a small-sized neural network in the scope of a structured matrix format while gradually converging to a smaller structure from a super-network given by the baseline network.

**Frequency-domain Learning** DiffStride (Riad et al., 2022) learns a stride of the _pooling_ operation for images by cropping a rectangular region of the frequency response of an image. Riad et al. (2022) utilizes an approximated boxcar mask for a differentiable function. Although DiffStride shares similar components 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 a 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-Dirichlet (Gaudi) function with a well-defined gradient. Effective learning algorithms are provided for our framework showing flexibility and differentiability to find expressive and efficient structures from training data in an 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.
## 7 Reproducibility Statement

The authors make the following efforts for reproducibility: 1) We attach the complete source code used in our experiment as supplemental materials, 2) we provide the detailed settings and hyperparameters in Section 4, A.5, and A.6, and 3) the proofs of all theorems, lemmas, and corollaries are presented in Section A.1.

## Acknowledgment

We thank Sara Shoouri, Pierre Abillama, and Andrea Bejarano for the insightful discussions and feedback on the paper.
