Geometry-aware training of factorized layers in tensor Tucker format

Summary

This paper proposes a method for neural network model parameter compression. It parameterizes weight tensors in the format of Tucker decomposition, and trains the factors instead of the origin weight tensors afterwards. The method is able to adaptively modify the tensor rank. Authors also provide detailed theoretical analysis including the computational steps, and convergence, approximation and the gradient descent guarantees of the method. Experiments show good performance of the method in terms of image classification accuracies, model compressive rates, and the high running efficiency.

Strengths

Writing is easy to follow. Theoretical analysis is self-contained and solid. Experiments are sufficient to demonstrate the claimed good properties of the method.

Weaknesses

I did not observe obvious weakness. However, I found that the studied neural networks seemed to be a bit of out-of-date. As a reviewer, I hope to see applying the method in up-to-date models e.g. transformers, in other fields. I am not sure whether this point would make negative influence on the significance of the work.

Questions

What is a standard Tucker decomposition like? Could you provide more details for the comparison of the proposed method and the vanilla standard Tucker decomposition, especially over the space/time complexity comparison, and their convergency analysis? What are the compressed parameters in the neural networks? Are the parameters merely the convolutional kernel? Can the method be applied on the weights of linear layers? What does ``geometry-aware’’ indicates? What kind of geometry is the method aware of? Could you give more explanation?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

This paper extends the dynamic low-rank neural network training (DLRT) method to the rank-adaptive Tucker tensor format (TDLRT). The proposed reparameterization method greatly reduces the computational complexity and numerical instability of the projected gradient descent. Under certain conditions, TDLRT with SGD converges to a stationary point in expectation, where the tensor found by TDLRT provably approximates the full mode. The experimental results show that TDLRT converges faster with less variance and better accuracy-compression rate tradeoff than other factorization-based and pruning-based methods.

Strengths

- The paper is clearly written. Most parts are supported by sufficient technical details. - The proposed method is an extended version of DLRT. Yet, some unique challenges appear in the Tucker tensor format have addressed with interesting approaches, e.g., the reparameterization method and Corollary 2.2.

Weaknesses

I have some minor concerns on writing and the experimental result - More detailed information about one-step integration methods could be provided for the readers who are not familiar with the concept. - Some sections assume a certain degree of background knowledge on DLRT, e.g., gauge conditions in Line 612. - The evaluation of training time for Tucker Decomposition and TDLRT was conducted on a toy-sized model (LeNet5) and dataset (MNIST). Evaluating training time on a dataset and a model of more practical sizes would provide better understandings of the proposed method.

Questions

- How does the one-step integration work? Does it not add significant computational overhead? - Are the conditions and assumptions in Theorem 3.2 and 3.3 reflect the actual DNN training? Is there any assumption that does not hold in practice? - How sensitive is TLDRT to the choice of hyperparameters like $\tau$, learning rate, weight decay or initialization schemes? Do they require extensive hyperparameter search to find the right ones? - Can TDLRT be used for the low-rank adaptation (LoRA) [1] setting? E.g., fine-tuning the convolution kernels of the U-Nets for the diffusion models by TDLRT with very low tensor ranks. [1] Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models." arXiv preprint arXiv:2106.09685 (2021).

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

- The utilization of QR decomposition might hinder using the proposed method to train a large model on a large dataset, e.g., diffusion model training. - The number of parameters is determined after training and indirectly adjusted by a hyperparameter $\tau$. Since a good strategy of choosing $\tau$ is not yet proposed, one might need to train the DNN multiple times to obtain a desired accuracy and model size.

Summary

The authors present a novel algorithm for training neural network layers using Tucker tensor decomposition. The approach addresses common issues with layer factorisation, including the need for an initial warm-up phase and sensitivity to parameter initialisation. Th method dynamically updates the ranks during training. The authors provide theoretical guarantees on loss descent, convergence, and approximation to the full model, supported by experimental results showing high compression rates and performance comparable to or better than baseline and alternative strategies.

Strengths

- A strong motivation as there is a clear need for further efficiency improvements - Introduces a novel rank-adaptive geometry-aware training method that dynamically updates ranks during training - Proposed to overcome the sensitivity to parameter initialisation and the need for a full-model warm-up phase - Thorough theoretical analysis, including guarantees on loss descent, convergence, and approximation

Weaknesses

- Claims do not match the results. The abstract says "our training proposal proves to be optimal in locally approximating the original unfactorized dynamics" and while there are guarantees, they are not proven to be optimal. - Not very clear in many parts. In particular, Section 2.1 is difficult to follow. - Results seem incomplete. For example, figure 1 does not show compression below 60% and for the proposed method is only shown for 93+% in Figure 1C while the other methods are shown for 60-93%. I could not find any reason for this lack of direct comparison and missing data points. - According to the plots, the proposed method outperforms the full representation at 96% compression. This is a very surprising finding that requires an in-depth discussion, which is lacking. - Table 1: The authors say that "TDLRT outperforms the factorization-based and the pruning-based baselines" and bold their method for Alexnet c.r. and Resnet test acc. but according to the same table, baselines actually outperform the proposed method in those metrics.

Questions

- The proposed method appears to outperform the full representation at 96% compression (Figure 1). Can the authors provide an in-depth discussion and analysis of this finding? What factors contribute directly to this performance, and does it align with theoretical expectations? - Can the authors include data points for the proposed method within the full compression range for a direct comparison with other methods?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

3

Limitations

The paper could benefit from a more detailed discussion on the limitations of the proposed method in different training scenarios and potential strategies to mitigate these limitations.

Summary

The authors study the training of layer factorization models to reduce the number of parameters in deep neural networks. They propose a geometric-aware rank-adaptive training strategy to avoid requiring prior knowledge of ranks and the sensitivity to the weight initializations. Their theoretical results show convergence and approximation error guarantees for the method.

Strengths

The proposed method is quite sensible and is accompanied by good theoretical guarantees.

Weaknesses

The empirical evaluation is slightly weak. While the results could convince me it is better than existing methods (as the proposed method is also quite sensible), from the scale of the model used, it is hard to judge whether it could provide good enough performances on larger models, when model compression during training is more needed.

Questions

Is it possible to conduct experiments on larger models or different architectures (like transformers, even a small one can help) to strengthen the empirical evaluation?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes, the authors discuss about the limitations.

Summary

Reducing the size of neural networks is an important problem for reducing the cost, memory usage, and even inference time. Many works focus on reducing the size after the training phase and use techniques such as sparsification and quantization. This paper, on the other hand, focuses on reducing the size by representing the tensor corresponding to the layers of the network as Tucker decompositions. Perhaps the more important aspect is that the paper considers dynamically changing the Tucker rank while training the model. In other words, the size reduction is simultaneous with training and not after the training is performed. It is argued that due to instability in the gradients, it is required to adopt a geometry-aware training strategy. Essentially, the strategy is to do HOOI based on the gradient of factor matrices. Then, compute a new tensor based on the new factor matrices and update the core tensor based on the gradient of this new tensor. Therefore, for each iteration, two passes are required to update the components: one pass to compute the gradients of factor matrices and one pass to compute the gradient of the core tensor. The paper presents theoretical results about the convergence and reduction of loss. In addition, it presents a sizable empirical study that shows favorable results for the proposed approach. The approach outperforms a variety of factorization and pruning methods in terms test accuracy and compression rate. The paper is generally well-written and provides an appropriate method for a very important problem. I think a relevant paper to discuss is [1], which gives an algorithm to compute the approximately optimal Tucker rank when a size constraint on the size of the Tucker decomposition is given. This could replace the approach based on the tolerance parameter $\tau$ when a hard size constraint is given. It would be interesting to investigate the interaction between that algorithm and the approach presented in this paper. [1] Ghadiri, Mehrdad, Matthew Fahrbach, Gang Fu, and Vahab Mirrokni. "Approximately optimal core shapes for tensor decompositions." In International Conference on Machine Learning, pp. 11237-11254. PMLR, 2023.

Strengths

The paper is generally well-written and provides an appropriate method for a very important problem.

Weaknesses

-

Questions

The results show that Tucker decomposition works better than CP and tensor-train. Is there any intuition for why this is the case? Do you expect Tucker to be better than any other tensor network, or could decompositions like hierarchical Tucker do better? Hierarchical Tucker could be preferable because of the smaller number of required parameters.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

I don't see any direct potential negative societal impact.

Reply to authors' rebuttal

Thank you for answering my questions. I read other reviews and the responses and have decided to keep my score at 7.