Geometry-aware training of factorized layers in tensor Tucker format

First of all, we would like to thank the reviewer for their feedback.

W1: We agree that our paper can be improved by showcasing stronger empirical evaluations for larger models. We have therefore conducted several new experiments on the popular parameter-efficient fine-tuning using low-rank adapters (LoRA), as suggested by reviewer vGH1. Here, we have tested the proposed training model on larger architectures such as DebertaV3 on Glue benchmark and Stable Diffusion (v1.4) dreambooth. In DebertaV3, we applied our method to matrix layers by interpreting matrices as order-2 tensors. For Stable Diffusion we applied our method to tensor layers in the Unet, while standard LoRA convolution works essentially by reshaping the tensor into a matrix. We believe that these results strengthen the empirical evaluation and we hope these could address your concerns.

Q1. The standard Tucker decomposition is the choice of the parameterization $W(i_1,\dots,i_d) = \sum_{j_1,\dots,j_d} C(j_1,\dots,j_d) U^{(d)}(i_d,j_d)$ to represent the d-mode tensor $W$. This choice is one possible extension of the classical singular value decomposition for tensors. For what concerns the vanilla Tucker decomposition and its training, we compare with the prototype method used often in training low-rank decompositions (such as LoRA), for which the update is simply given by a step of stochastic gradient descent on each factor of the decomposition, i.e. $C(t+1) = C(t)- \lambda_t \nabla_C \mathcal L(t),\, U^{(i)}(t+1) = U^{(i)}(t)-\lambda_t \nabla_{U^{(i)}}\mathcal L(t)$. Apart from the instability with respect to small singular values of vanilla Tucker, we would like to underline that this method introduces additional invariances in the parameter space (since orthonormality of the basis $U^{(i)}$ is not preserved over time). In terms of space complexity, Tucker vanilla and our decomposition are the same. The time complexity of one optimization iteration for the vanilla Tucker is of order $O(b(\prod_{i=1}^d r_i+ \sum_{i=1}^d n_i r_i))$, where $n_i$ are the dimensions of the tensor weight, $r_i$ are the Tucker ranks and $b$ is the batch size. The training of vanilla methods is essentially SGD, so the convergence properties of it are the same of stochastic gradient descent.

Q2 This is a good question and we will add a small discussion of this point in the manuscript. The presented theory and the algorithm is developed in general for Tucker tensors with $d$ modes, that for $d = 2$ covers also the matrix case. In particular then it is possible to apply the method both to matrices and tensors with more than two modes. In the matrix case, the Tucker decomposition would degenerate into a singular value like decomposition of the weight matrix.

Q3 The choice of the name "geometry-aware" refers to the fact that the proposed method is "aware" of the geometric structure underlying the problem. In section 2.1 we formulate the training of tensor-weighted neural networks as a gradient flow $\dot W = -\nabla_W \mathcal L(W)$. When we want to superimpose a low-rank structure, vanilla methods follow a gradient flow $\dot C = -\nabla_C \mathcal L, \, \dot U^{(i)} = \nabla_{U^{(i)}} \mathcal L$. This last gradient flow lies in the (Riemannian) manifold $\mathcal M_r$ of Tucker tensors of rank $\mathbf{r}$, but it is not aware in some sense of the original problem related to the geometry of the constraint. Our proposed method follows the global dynamics $\dot W = -P(W) \nabla_W \mathcal L(W)$, where $P(W)$ is the orthogonal projection on the tangent space $T_W \mathcal M_r$. Notice that the right hand side of this last projected differential equation is a solution of the minimization problem $\underset{\delta W \in T_W\mathcal M_r}{\arg \min}||\delta W+ \nabla \mathcal{L}(W) ||_{F}^2$, basically by definition of orthogonal projection. This variational principle is essentially telling us that locally we're taking the closest direction to following the original unfactorized dynamics. This property is not satisfied by the vanilla gradient flow system, that does not follow locally the original unfactorized dynamics. Moreover, our proposal is aware of certain properties of $\mathcal M_r$, such as high curvature around tensors with small singular values along some mode. As shown in theorems 3.1 and 3.3, the provided theoretical bounds regarding approximation and descent do not depend on the singular values, making the method stable in their presence. This property is also highlighted numerically in figure 2, where we show that vanilla methods seem to suffer from this ill-conditioning.

We will add some details about the questions in the manuscript, and we hope the explaination has clarified any possible doubt. In any case we remain available if any further clarification is needed.

First of all, we would like the thank the reviewer for the insightful feedback. Below, we provide our responses to the main points raised in the review.

Q1+W1 Regarding "How does the one-step integration work? Does it not add significant computational overhead?". With "one-step integration'' we mean a single step of a gradient-based optimization method or, equivalently, any classical time-integration method. In this work, we use stochastic gradient descent (SGD) which in the time integration view is the explicit Euler method. However, the use of other methods is certainly possible. This step does not add any overhead compared to the full-rank training and, in fact, is one of the parts in our method that allows us to have a significant reduction of computational costs and memory footprint. Note that for the full-rank training, the one-step integrate function is called repeatedly using the full gradient, whereas we only need to call this method with the gradient for a low-rank factor. We will add a sentence to our manuscript to clarify what we mean by the one-step integration method.

W2 Thank you for the suggestion which we believe makes a lot of sense. We will add a short description to our manuscript.

Q2 Regarding "Are the conditions and assumptions in Theorem 3.2 and 3.3 reflect the actual DNN training?...". Regarding Theorem 3.3, boundedness and Lipschitz continuity of the gradient are very common assumptions one must make to show most analytic results. Both assumptions are reasonable since one would assume the gradient remains bounded during the optimization and the gradient does not completely change when slightly changing the input. The latter can be achieved using most common activation functions, since it requires a condition on the boundedness of the derivatives. The assumption that the gradient flow remains close to the low-rank manifold is observed empirically, but it is extremely difficult to understand analytically. Numerous experiments have shown that neural networks are usually well-represented by low-rank weights, and some analytic investigations exist for heavily simplified architectures. Regarding Theorem 3.2, the Robbins-Monro conditions are very standard assumptions and can be ensured by the user when picking the learning rate. The stabilization of the spectral distribution over time is something that we observe empirically, though it is hard to show anything on an analytic level. Here, we see that after sufficient iterations of TDLRT, the basis does not change, and most of the change is seen in the core tensor. The drift assumption is meant to play the role of the more standard (but more restrictive) assumption of finite variance, which is common when studying convergence of stochastic methods.

Q3 This is a good question we should have discussed in our manuscript. First, regarding the tolerance parameter for truncating the core tensor, this is, of course, a hyperparameter that needs to be chosen. We commonly use similar values here in different test cases and we never needed an in-depth hyperparameter search to observe good performance. However, it is hard to say if this holds for all architectures and datasets, and most likely, one must adapt here. However, we wish to point out that a single parameter determines individual ranks and compressions for different modes in each tensor in each layer. Here, other approaches, like LoRA, which you mention later, require the user to choose good ranks for each layer, which is certainly a harder task and requires an intensive parameter search if one wants minimal memory overhead. Regarding parameters like learning rate, weight decay etc., we found that our method is relatively robust with respect to their choice, similar to the full baseline. This is aligned with the theoretical findings of e.g. Thm 2.1 and Thm 3.3.

Q4 + W3 Thank you for this excellent question. Certainly, our approach fits LoRA excellently, and we have now provided some results on how TDLRT can be used in such a setting (we refer to the main rebuttal pdf file). The main question is always if these architectures allow for low-rank weights (or in terms of the assumption of Theorem 3.3, if $\varepsilon$ small). If the answer is yes, then our approach should yield good results. Moreover, all the theory presented in the manuscript applies to this case.

L1 Regarding "The utilization of QR decomposition might hinder using the proposed method to train a large model....'' Please note that the QR decomposition only needs to be computed on a small matrix of dimension $n_i\times r_i$, leading to computational costs of $O(n_i\times r_i^2)$. Thus, if the rank is small, the QR decomposition is usually not a limiting factor (and these small ranks are in fact often the case for LoRA style fine-tuning).

L2 Regarding "The number of parameters is determined after training and indirectly adjusted by a hyperparameter. Since a good strategy of choosing is not yet proposed, one might need to train...". We agree, development of a method that takes a parameter budget and determines $\tau$ to obtain the best model for the given budget is a relevant research question. We have been pointed to [1] by reviewer TTwz for an approximative Tucker decomposition (of the core tensor) under a parameter budget. Here the research question is: Given a parameter budget, how to determine the core shapes of the core tensors of all low-rank convolutions in the network to best approximate the full-rank model, if the low-rank dynamics, implicitly given by the data, are representable within this budget. This question is as relevant as it is non-trivial, and we would like to refer to future research to properly address it.

[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.

We wish to thank the reviewer for their feedback. Overall, the reviewer highlights key positive aspects of our work, such as the importance of the topic, the novelty of the approach, and the thorough theoretical analysis. The main weaknesses noted are centered on the need for additional runs and a more detailed discussion of results. We are, however, surprised by the perceived impact of these weaknesses, especially in light of the significant strengths identified. The reviewer mentions that our work is "unclear in many parts" and that "Section 2.1 is hard to follow". However, without further details, it is challenging for us to address these points, particularly since clarity has not been a concern for the other reviewers. In the following, we address all the identified weaknesses (Ws), as we believe several points merit further clarification. We also answer the raised questions (Qs).

W1 Thank you for pointing this out. In that sentence of the abstract, our intention was to refer to the optimality of the Dirac-Frenkel variational principle to approximate the dynamics along with the supporting theoretical results of approximation and convergence. However, this statement is not crucial and can be misinterpreted, therefore we will remove it from the abstract.

W2 We recognize that certain sections of the paper, particularly Section 2.1, contain technical details that may require a background in Riemannian optimization and dynamical model-order reduction theory. We have made significant efforts to make this section as clear and accessible as possible, but we understand that some concepts may still be difficult to follow. However, a conference paper is not the ideal format for providing an extensive introduction to these topics.

We would appreciate further clarification on the specific parts of Section 2.1 or other sections that you found unclear. This will help us address your concerns more effectively and improve the paper's readability. We aim to clarify the main ideas of Section 2.1:

• Our goal is to derive a low-rank gradient flow that closely approximates the standard full gradient flow. To achieve this, we project the gradient flow onto the tangent bundle as described in Equation (3), which represents the standard gradient flow of Riemannian optimization.
• However, this equation forms a stiff system that is challenging to solve directly with a discretized scheme, necessitating a small learning rate.
• By introducing a reparametrization of the weights in Theorem 2.1, we obtain new evolution equations that are well-posed and can be solved with larger learning rates. Despite this improvement, the resulting method remains inefficient as it requires d+1 gradient evaluations.
• In Corollary 2.2, we demonstrate that the basis computation can be significantly simplified, leading to Algorithm 1, which requires only 2 gradient evaluations instead of d+1.

We point out that we have already attempted to convey this structure and our reasoning at the beginning of Section 2. Please let us know if this explanation clarifies your questions, and if not, which specific parts remain unclear.

W3 and Q2 The reason our method is sometimes only shown for large compression values is that one key feature of our method is that the compression rate is somewhat determined by the method itself, rather than being a user-defined input, as with the fixed-rank baselines. Nonetheless, we have adjusted the truncation tolerance to obtain smaller compression rates and have included additional data points in the attached rebuttal PDF file. We emphasize that methods at compression rates below 60% are generally less relevant, so we have focused on the range of large compression rates to maintain the focus on the relevant range of compressions and highlight the strengths of our method. However, if the reviewer considers it valuable, we can include additional runs for smaller compression rates in the appendix.

W4 and Q1 Thank you for your remark. We emphasize that this phenomenon, where networks with a smaller number of parameters outperform their full baselines, is well-documented in the literature, see e.g. [1,2]. One explanation for this behavior is that neural networks are often overparameterized, and parameter reduction or compression can have a regularizing effect enhancing generalization.

Note that in our comparisons to the baseline, we use the same hyperparameters for compression methods and the full-rank baseline. With this approach on VGG16 we indeed obtain better accuracy than the baseline. However, given your comments, we decided to see if we could find better hyperparameters. This gives us results where the baseline outperforms compression methods. We decided to add these results since the overall accuracy is improved this way for the baseline and for all the compression methods.

[1] The lottery ticket hypothesis: Finding sparse, trainable neural networks, 2018 [2] Snip: Single-shot network pruning based on connection sensitivity, 2018

W5 Thank you for pointing out this oversight. We apologize for the confusion.

Upon reviewing Table 1, we realized that an error occurred when bolding the results for AlexNet's compression rate (c.r.) and ResNet's test accuracy. We were updating the table with new competitor results shortly before submission and inadvertently left the bold formatting incorrect. We will correct this in our revised manuscript to accurately reflect the performance of each method.

Regarding the specific metrics, it is true that while the TT-factorized method achieves a higher compression rate for AlexNet, its accuracy is lower compared to TDLRT, although it still outperforms other methods significantly. For ResNet18, while the Tucker RGD method achieves a slightly higher accuracy (0.04%), it does so at the cost of a reduced compression rate.

We will ensure that the table accurately represents these findings and provide a more detailed discussion in the text.

We thank the reviewer for their insightful feedback, below we provide our response to the raised weaknesses and questions.

Q1 Thank you for your feedback regarding the empirical evaluation. In response to your suggestion, we have conducted several new experiments to strengthen our empirical evaluation.

Specifically, we have focused on parameter-efficient fine-tuning using low-rank adapters (LoRA), as suggested by reviewer vGH1. This approach allows us to present additional experimental evidence that complements our previous results. Fine-tuning, rather than compressing a model from scratch during training, enables us to experiment with larger networks while remaining within our time and resource constraints.

In particular, we have tested the proposed training model on larger architectures such as DebertaV3 on Glue benchmark and Stable Diffusion (v1.4) dreambooth. In DebertaV3, we applied our method to matrix layers by interpreting matrices as order-2 tensors. For Stable Diffusion we applied our method to tensor layers in the Unet, while standard LoRA convolution works essentially by reshaping the tensor into a matrix.

These new experiments provide further validation of our method’s performance and demonstrate its applicability to larger models and diverse architectures. We believe these results strengthen the empirical evaluation and could address your concerns about scalability.

We thank the reviewer for their feedback and appreciation of the work. Below we provide our response to the raised questions.

Q1. We appreciate this insightful question. Selecting the appropriate tensor structure a-priori is indeed challenging, as it depends heavily on the specific tensor structures in the neural network and the data used. As shown in Table 1 of our main paper, when using standard direct training of the factors, there is no definitive "winner" among Tucker-factorized, Matrix-factorized, TT-factorized, and CP-factorized approaches across the three test problems. However, when employing our proposed geometry-aware TDRLT scheme for training, we observe notably improved performance.

Therefore, we believe that Tucker's superior performance is primarily due to the training method rather than the decomposition choice alone. While extending this training strategy to CP is challenging due to the absence of Riemannian geometry, we believe that further research into extending the proposed geometry-aware training to TT, HT, and other tree tensor network structures would be highly relevant.

For what concerns the comment in the "Summary" section, we thank the reviewer for bringing this paper to our attention. We agree that the algorithm presented in [1] offers an interesting alternative approach to consider. In particular, their method is well-suited for scenarios where there is a fixed memory budget, such as training on resource-constrained edge devices. In such cases, the objective shifts from "finding the best accuracy-compression trade-off", which is the focus of our method, to "finding the best-performing Tucker decomposition given memory constraints", where the approach from [1] could be highly beneficial.

To explore this aspect further, it would be useful to extend the method in [1], which currently addresses a single Tucker tensor, to an approach that optimizes the shapes of all tensor-valued layers within the neural network given a global network-wise memory constraint. This is an intriguing research direction that could significantly enhance the applicability of tensor decompositions in constrained environments.

We will certainly include a discussion of this paper in our revised manuscript, highlighting it as a potential alternative approach to our proposed method. This addition will provide readers with a broader perspective on the available techniques and their respective advantages and limitations.