Differentiable Learning of Generalized Structured Matrices for Efficient Deep Neural Networks

We appreciate the constructive suggestions! We revised and extended the manuscript based on your comments.

Q1: 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 parameterize wrt rank) that need to be included.

A1: We carefully went through the suggested methods used in CNNs. We added a new comparison for the ImageNet accuracy of ResNet18 with GBLR weight matrices vs. the Automatic Layer-wise Decomposition Selector (ALDS) [1], which is a similar rank-based layer-wise compression method. Following [1], we consider the $(d, c, k_1, k_2)$-sized weight tensor of a convolution layer as a $d \times (c k_1 k_2)$-sized matrix, where $d, c, k_1, k_2$ 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 1 shows the accuracy and the relative FLOPs of the ResNet18 model with dense weights, low-rank weights found by one-shot ALDS [1], and GBLR weights found by our method. Even without any supervision on the structural properties of the convolution kernels, GBLR outperforms ALDS low-rank weights in terms of both accuracy and complexity.

Table 1.

Q2: Can you please provide the exact form of $\ell_1$ penalty used in eq.8? How does FLOPs/parameter counts being represented as $\ell_1$ penalty?

A2: The FLOPs / parameter count is directly represented by the sum of the width parameters as in Eq. 4 (restated below):

$$\mathrm{FLOPs}=\sum_{k=1}^K (w_k^R + w_k^C),$$

which is the $\ell_1$ norm of the vector containing width parameters $\boldsymbol{w}=\{w_1^R, w_1^C,\ldots,w_K^R,w_K^C\}$. Hence, minimizing the $\ell_1$ norm of this width parameter vector $\boldsymbol{w}$ of a GBLR matrix reduces the FLOPs for matrix-vector product.

Q3: What is the value of $\lambda$ used in experiments? How to choose it properly?

A3: We used $\lambda=0.02-0.04$ for the ViTs, $\lambda=0.1-0.2$ for the MLP-Mixers, and $\lambda=0.001-0.005$ for ResNet-18. The range of the optimal choice for the $\lambda$ hyperparameter depends on the model, learning rate, and dataset. We used a grid search enumerating the search space to find the $\lambda$ hyperparameter in our experiments. Automatically selecting an optimal $\lambda$ from the budget is left as a topic for future work.

Q4: What was the scheme used for $\sigma$? The paper says that it was "gradually increased to 100", but having exact details is preferred.

A4: We initially set $\sigma$ to 1. At each epoch, we start increasing $\sigma$ linearly to 100 until the end of the training process, except during the learning rate warm-up and cool-down steps. For example, if the training process consists of 110 epochs including 5 warm-up epochs at the beginning and 10 cool-down epochs at the end, $\sigma$ remains at 1 until epoch 5 and then increases linearly to 100 until epoch 100. We provided the $\sigma$ scheduling used in our experiment in Appendix A.7.2.

Q5: 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..

A5: FFT/IFFT does not impose significant overhead in our methodology. Let us consider ViT-Large for example. At each fully-connected layer, 1-D Inverse Fast Fourier Transform (IFFT) is performed over a $4096 \times 1024$-sized weight matrix, which requires $4096 \times 2.5 \times 1024 \log_2 1024 = 104,857,600 \approx 105\times 10^6$ FLOPs.

In contrast, the matrix-matrix multiplication (MMM) between a $(B, N, 4096)$-sized input tensor and a $4096\times 1024$-sized weight matrix requires $B\times N \times 4096 \times 1024$ FLOPs where $B$ and $N$ are the batch size and sequence length, respectively. For $B=128$ and $N=197$, the number of FLOPs for MMM is $106\times 10^9$, which is three orders of magnitude larger than the number of FLOPs required for IFFT.

Reference

[1] Liebenwein, Lucas, et al. "Compressing neural networks: Towards determining the optimal layer-wise decomposition." Advances in Neural Information Processing Systems 34 (2021): 5328-5344.

Thank you for the comprehensive suggestions! The below provides answers to some of your questions.

Q1: In Algorithm 1, what’s AdamW(.), as well as clip?

A1: $\mathrm{AdamW}(\cdot )$ stands for the AdamW optimizer [1]. Line 5 in Algorithm 1 ($p \gets p - \eta \cdot \mathrm{AdamW}(\nabla \ell)$) corresponds to the conventional parameter update to minimize the loss $\ell$ by using the AdamW optimizer.

The function $\mathrm{clip}(\mathbf{w}, 0, n)$ clamps each element of $\mathbf{w}$ to the interval $[0, n]$, which is the domain of the width parameter. We included this operation in Algorithm 1 to ensure that the width parameter stays within $[0,n]$ after the gradient descent.

We have added a detailed explanation of Algorithm 1 in Section 3.3.

Q2: In Eq. (8), what’s the variable, w or $\theta$?

A2: In Eq. (8), the variable $\theta$ denotes the set of the structural and content parameters, namely, $\theta=(\boldsymbol{\phi}, \boldsymbol{U}, \boldsymbol{V}, \sigma)$, as stated above Eq. (7). The variable $\boldsymbol{w}$ denotes the width parameters $\boldsymbol{w}=\{w_1^R, w_1^C,\ldots,w_K^R,w_K^C\}$ of $\boldsymbol{W}^{\boldsymbol{\theta}}$. We updated Section 3.3 to clarify the definition of $\boldsymbol{w}$.

Q3: The advantage of the proposed method against existing methods is not clear.

A3: The key advantage of our method lies in its ability to learn the structural parameters of weight matrices by stochastic gradient descent, overcoming the non-differentiable nature of typical matrix structures --e.g., non-zero locations of the sparse matrix or the rank of the low-rank matrix. Prior approaches used hand-crafted structures because of this issue. Our approach allows treating structural parameters (widths and locations of the masks in Eq. 2) as additional trainable parameters to automatically identify/learn efficient structures in practice. This enables the end-to-end, layer-wise learning of the optimal structured matrix for each weight matrix given the computational budget for each layer.

Most existing rank-based compression methods are confined to either a single type of structured matrix (e.g., low-rank or block low-rank) [2,3] relying on a heuristic rule for budget allocation across layers [3,4], or segregating the structural parameter optimization from the training process of the other parameters [2,4]. This is primarily because, in those prior approaches, the structural formats are defined in a non-differentiable and combinatorial manner.

Using the learned structures and parameters from our approach generally provides higher accuracy for the same complexity (or comparable accuracy with lower complexity) compared to the model in prior works.

Q4: The parameter initialization for the proposed method needs to perform SVD. Thus, the authors should analyze the computational complexity.

A4: Please note that our objective is to minimize the inference complexity. And the SVD overhead during training is negligible. The SVD during initialization takes less than one minute, whereas the overall training or fine-tuning time ranges from a few hours to tens of hours. Also, the SVD is unnecessary during the inference stage once the network is trained. We have included an analysis of the impact of SVD on training time in the Appendix.

References

[1] Loshchilov, Ilya, and Frank Hutter. "Decoupled Weight Decay Regularization." International Conference on Learning Representations. 2018.

[2] Liebenwein, Lucas, et al. "Compressing neural networks: Towards determining the optimal layer-wise decomposition." Advances in Neural Information Processing Systems 34 (2021): 5328-5344.

[3] Chen, Beidi, et al. "Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models." International Conference on Learning Representations. 2021.

[4] Dao, Tri, et al. "Monarch: Expressive structured matrices for efficient and accurate training." International Conference on Machine Learning. PMLR, 2022.

Thank you for the insightful feedback!

Q1: How does the GBLR weight matrices affect the generalization property or the complexity of the neural network?

A1: We hypothesize based on empirical observations that the GBLR weight matrices reduce the generalization gap of neural networks. We observed that when the weights of DNNs are initialized to the GBLR matrices and trained from scratch, the training accuracy is generally lower than that of the dense model, whereas the validation accuracy is similar. For instance, Table 1 shows the ImageNet accuracy of dense and GBLR ViT-Base models.

Table 1.

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 the dense models. We added this discussion to the paper.

However, given the low-rankness and sparsity of GBLR matrices, we agree with the reviewer that we need a different technique to fully understand the generalization and approximation bounds of the DNNs with GBLR matrices. Due to time constraints, unfortunately, we could not conclusively determine the theoretical generalizability of the GBLR matrices. Nonetheless, we believe the generalized structured format has potential to serve as a good tool for a better understanding of the generalization bounds of DNNs using the low-rank matrices [1,2,3] and the approximation bounds of DNNs with sparse matrices [4,5]. Conceptually, they are based on the low-dimensional representations of large and complex functions. Since our method expands the low-dimensional matrix representations from the hand-designed models to a general parametric space, we are eager to study the GBLR matrix theoretically in future work.

References

[1] Arora, Sanjeev, et al. "Stronger generalization bounds for deep nets via a compression approach." International Conference on Machine Learning. PMLR, 2018.

[2] Suzuki, Taiji, Hiroshi Abe, and Tomoaki Nishimura. "Compression based bound for non-compressed network: unified generalization error analysis of large compressible deep neural network." International Conference on Learning Representations. 2020.

[3] Baykal, Cenk, et al. "Data-Dependent Coresets for Compressing Neural Networks with Applications to Generalization Bounds." International Conference on Learning Representations. 2018.

[4] Klusowski, Jason M., and Andrew R. Barron. "Approximation by combinations of ReLU and squared ReLU ridge functions with $\ell^ 1$ and $\ell^ 0$ controls." IEEE Transactions on Information Theory 64.12 (2018): 7649-7656.

[5] Domingo-Enrich, Carles, and Youssef Mroueh. "Tighter Sparse Approximation Bounds for ReLU Neural Networks." International Conference on Learning Representations. 2021.