Geometry-aware training of factorized layers in tensor Tucker format

Authors: emanuele  zangrando
steffen  schotthöfer
gianluca  ceruti
jonas  kusch
francesco  tudisco
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:50:51

Abstract: Reducing parameter redundancies in neural network architectures is crucial for achieving feasible computational and memory requirements during train and inference of large networks. Given its easy implementation and flexibility, one promising approach is layer factorization, which reshapes weight tensors into a matrix format and parameterizes it as the product of two rank-r matrices. However, this family of approaches often requires an initial full-model warm-up phase, prior knowledge of a feasible rank, and it is sensitive to parameter initialization.
In this work, we introduce a novel approach to train the factors of a Tucker decomposition of the weight tensors. Our training proposal proves to be optimal in locally approximating the original unfactorized dynamics and stable for the initialization. Furthermore, the rank of each mode is dynamically updated during training.
We provide a theoretical analysis of the algorithm, showing convergence, approximation and local descent guarantees. The method's performance is further illustrated through a variety of experiments, showing remarkable training compression rates and comparable or even better performance than the full baseline and alternative layer factorization strategies.

Main paper: 
 introduction: 
 The memory footprint and computational cost for inference and training are major limitations of modern deep learning architectures. A variety of techniques have been developed to address this issue, aiming to reduce model size and computational complexity. Popular approaches include weight sparsification [26, 59, 30] and quantization [83, 15] . However, pruning via sparsification struggles to take advantage of GPU hardware designed for dense matrices, and it is difficult to provide error estimates on model performance when using quantization [24] . Moreover, while able to reduce resource requirements for inference, these methods struggle to achieve memory reduction during training without affecting performance. As pointed out in [20, 24] , training accurate sparse or quantized neural networks from the start is particularly challenging. At the same time, the training phase of modern architectures can require several days on several hundreds of GPUs [7] , requiring huge memory storage and energy consumption overheads. Thus, being able to reduce the resource demand of both inference and training while maintaining model performance is of critical importance.Along with sparsification and quantization, layer factorization is another popular and successful model compression approach. Representing layer weights using different types of matrix and tensor factorizations can yield huge memory reduction while retaining model performance and robustness [67, 13] . A wealth of recent work has shown theoretical and experimental evidence suggesting that layer weights of over-parametrized networks tend to be low-rank [3, 5, 22] and that removing small singular values may even lead to increased model performance while dramatically reducing model size [70, 68] . One significant advantage of low-rank factorizations is that a low-parametric factorized model can be used throughout the entire training [79, 36, 68] and fine-tuning phase [33, 80] . The most direct way of doing this is by representing each layer's weight tensor W as the product of small factors and then directly training each factor independently in a block-coordinate descent. In the matrix case, this boils down to imposing the rank-r parametrization W = U V ⊤ for each layer W , with U, V rectangular matrices with only r columns. The network is then trained on the set of rank-r matrices M r = {U V ⊤ : U, V have r columns}, interpreting the loss as a function of U, V alone. This is the approach taken also by popular fine-tuning strategies such as LoRA [32, 86, 50] . When W is a higher-dimensional tensor, as in the case of convolutional kernels for example, the same approach can be implemented using different types of tensor low-rank factorizations, including canonical polyadic (CP) and Tucker formats [43, 49, 63, 71, 42, 21] . While direct training of a layer's factors is widely used in deep learning, this approach has two major disadvantages:1. The rank r of the factorization needs to be chosen a-priori, and the performance of the compressed model can highly depend on it [70, 79] ; 2. The training flow is highly sensitive with respect to the choice of the initialization, which may result in a high-oscillatory slow-converging loss and sub-optimal performance and may require a warm-up phase during which the full model is trained prior to the rank reduction [79, 36, 68] .Point 2, in particular, is directly related to the geometry of the constraints set. In the matrix case, it is well-known that M r is a Riemannian manifold with points of very high curvature near small singular values [19] . These points give rise to stiffness and result in ill-conditioning. This intrinsic poor conditioning can be overcome by projecting the gradient flow on the tangent bundle of M r as presented in [68] .In the higher-dimensional tensor case, we face the same problems. However, trying to adapt the approach for matrices to other tensor factorizations is not trivial, as it may lead to prohibitive computational costs and memory requirements scaling with the order of the tensor. Moreover, not all the tensor factorizations have the required Riemannian structure to design projections and tangent planes. This paper introduces a rank-adaptive geometry-aware training algorithm that trains factorized tensor layers in Tucker format taking advantage of the underlying Riemannian structure, yielding strictly better performance than direct factorizations and overcoming both Points 1 and 2 above. Our main contributions are:• We design an algorithm for training tensor layers in Tucker format that is rank-adaptive, as the ranks of the layers are dynamically updated during training to match a desired compression rate; • We provide theoretical guarantees of loss descent, convergence to stationary points in expectation, and approximation to the ideal full model; • We provide extensive experimental evaluation showing that the proposed method yields remarkable training compression rates (e.g., more than 95% for VGG16 on CIFAR10), while achieving comparable or even better performance than the full baseline and alternative baselines.Related work on network compression methods differs structurally by the mathematical object of consideration, i.e., matrix-or tensor-valued parameter structures, as well as the type of parameter reduction. Weight pruning [28, 61, 73, 59, 81] enables parameter reduction by enforcing sparsity, i.e., zero-valued weights, whereas low-rank compression imposes parameter reduction by factorization of weight matrices [34, 49, 79, 66, 84] and tensors [44, 71, 4, 63, 42, 38, 42, 72, 35] . On top of approaches that transform tensor layers into compressed matrices [68, 34, 49, 79] , different tensor decompositions have been used to compress convolutional layers. Such approaches include CP decomposition [44, 71, 4, 63] , Tucker [42, 38] , tensor trains [42, 72] or a combination of these [21] . Other works focus on performing efficient optimization on Stiefel manifolds to preserve orthonormality, with methods based on regularization or landing [57, 8, 13, 1] , cheap parametrizations of orthogonal groups [46, 48] and Riemannian schemes [58, 67, 2] . Other methods consider only the floating point representation of the weights, e.g. [75, 25, 27, 14, 76] , or a combination of the above [51] . From the algorithmic point of view, related work can be categorized into methods that compress networks entirely in a postprocessing step after full-scale training [60, 56, 44, 38, 21, 4] , iterative methods where networks are pre-trained and subsequently compressed and fine-tuned [28, 34, 79] , and methods that directly compress networks during training [68, 61] . As no full-scale training is needed, the latter approach offers a better potential reduction of the overall computational footprint. Only a few of these methods propose strategies for dynamically choosing the compression format during training or fine-tuning, e.g., by finding the ranks via alternating, constraint optimization in discrete [47] and discrete-continuous fashions [34] . However, both these approaches require knowledge of the full weights during training and overall are more computationally demanding than standard training. In [68] , a rank-adaptive evolution of the gradient flow on a low-rank manifold was proposed to train and compress networks without using the full-weight representation; however, only for matrix-valued layers. While a direct extension of this method to Tucker tensors is possible, the resulting algorithm exhibits a prohibitive memory footprint and computational complexity. The development of rank-adaptive training methods for tensor-valued layers poses non-trivial challenges that may prevent loss descent and performance of the compressed net. For example, numerical instabilities arising from the CP decomposition during training have been observed in [44] and [63] . 
geometry-aware training in tucker format: 
 For a tensor W , we write Mat i (W ) to denote the matrix obtained by unfolding W along its i-th mode.The tuple ρ = (r 1 , r 2 , . . . , r d ) is called Tucker rank of W if r i = rank(Mat i (W )). Every d-order tensor W with Tucker rank ρ = (r 1 , . . . , r d ) can be written in Tucker form (or Tucker decomposition) W = C × d i=1 U i , entry-wise defined asW (i 1 , . . . , i d ) = r1,...,r d α1,...,α d =1 C(α 1 , . . . , α d )U 1 (i 1 , α 1 ) • • • U d (i d , α d )where C ∈ R O(r 1 • • • r d + n 1 r 1 + • • • + n d r d ), as opposed to the O(n 1 • • • n d )cost required by the standard full representation. When n i ≫ r i , e.g., n i > 1.5r i , the latter is much larger than the former.In the following, we develop a rank-adaptive algorithm that trains layers in Tucker form in a robust and efficient manner. Our derivation follows five points:1. We introduce the dynamical low-rank approximation framework in Section 2.1, which provides gradient flows for layers in Tucker format. However, the direct use of these evolution equations to train the network will require prohibitively small learning rates due to the high curvature of the manifold of Tucker tensors. 2. We introduce a reparameterization in Theorem 2.1 that allows us to formulate robust dynamics for the reparametrized factors. 3. By integrating numerically the resulting gradient system (with e.g. SGD as explicit Euler) along with a basis augmentation step, we propose a geometry-aware rank-adaptive training strategy for the network in Tucker format. This approach, however, requires d + 1 forward and backward evaluations of the network, resulting in significantly increased computational costs. 4. We show in Corollary 2.2 that computational costs can be substantially reduced by noting that the computation of the augmented basis can largely be simplified. This leads to our proposed training scheme in Algorithm 1. The algorithm is equivalent to the integration of the gradient system but requires only two instead of d + 1 gradient evaluations. 5. Due to its equivalence to the approach constructed in point 3, we can show that Algorithm 1 indirectly updates weights along straight lines on the manifold, thus leading to three main theoretical properties: loss descent (Theorem 3.1), convergence in expectation (Theorem 3.2), and a robust bound showing approximation of the full model (Theorem 3.3).For ρ = (r 1 , . . . , r d ), the setM ρ = {W : rank(Mat i (W )) = r i , i = 1, . . . , d}is a manifold with the following tangent space at any pointEQUATIONwhere S j is the Stiefel manifold of real n i × r i matrices with orthonormal columns. To design a strategy that computes layer weights within M ρ using only the low-rank Tucker factors C and {U i } i , we formulate the training problem as a continuous-time gradient flow projected onto the tangent space (1) . As shown in Section 3, the continuous formulation will allow us to derive a modified backpropagation pass which uses only the individual small factors C, {U i } i and that does not suffer from a slow convergence rate due to potential ill-conditioned tensor modes (see also Section 4.2).Let f be a neural network and let W be a weight tensor within f . Consider the problem of minimizing the loss function L with respect to just W while keeping the other parameters fixed. This problem can be equivalently formulated as the gradient flowEQUATIONwhere, for simplicity, we write the loss as a function of only W and where "dot" denotes the time derivative. When t → ∞, the solution of (2) approaches the desired minimizer. Now, suppose we parametrize each tensor layer in a time-dependent Tucker formW (t) = C(t) × d i=1 U i (t) ∈ M ρ . Then Ẇ (t) ∈ T W (t) M ρ ,the tangent space of M ρ at W (t). Thus, (2) boils down toEQUATIONwhere P (W ) denotes the orthogonal projection onto T W M ρ . Using standard derivations from dynamical model order reduction literature [40] , the projected gradient flow in (3) leads to the following evolution equations for the individual factors C(t) andEQUATIONwhere † denotes the pseudoinverse and where we omitted the dependence on t for brevity. Even though (4) describes the dynamics of the individual factors, the equations for each factor are not fully decoupled. A direct integration of (4) would still require taping the gradients ∇ W L with respect to the full convolutional kernel W . Moreover, the pseudoinverse of the matrices Mat i (C) † adds a stiffness term to the differential equation, making its numerical integration unstable. The presence of this stiff term is actually due to the intrinsic high-curvature of the manifold M ρ and is well understood in the dynamic model order reduction community [39, 53, 37, 54, 10, 9] . As observed in [68] , an analogous term arises when looking at low-rank matrix parameterizations, and it is responsible for the issue of slow convergence of low-rank matrix training methods, which is observed in [79, 36, 68] .To overcome these issues, we use the following change of variables. Let Mat i (C) ⊤ = Q i S ⊤ i be the QR decomposition of Mat i (C) ⊤ . Note that S i is a small square invertible matrix of size r i × r i . Then, the matrix K i = U i S i has the same size as U i and spans the same vector space. However, the following key result holds forEQUATIONwhere Ten i denotes "tensorization along mode i", i.e. the inverse reshaping operation of Mat i .The supplementary material provides the proof in Appendix D. The theorem above allows us to simplify (4), obtaining a gradient flow that only depends on the small matrices K i and the small core tensor C. Moreover, it eliminates a stiffness term; this added regularity appears reasonable as no inversion is now involved in the differential equations. A rigorous regularity statement can be found in Theorem 3.3. We would like to underline the importance of the careful construction of K i to arrive at this theorem. A naive extension of [68] to Tucker tensors can be constructed by a reshaping ofW into matrices Mat i (W ) = U i S i V ⊤ i with S i = Mat i (C) and V i = j̸ =i U j .Then, K i = U i S i can be used to update U i into all directions i ≤ d which directly inherits the robustness properties presented in [68] . However, this construction of K yields a prohibitive memory footprint of O(n i Π j̸ =i r j ) and computational costs of O(n i Π j̸ =i r 2 j ) rendering the resulting method impractical. Based on the numerical integration of ( 5), we propose a robust, memory-efficient, and rank-adaptive method to update the network parameters by using only the core tensor C and the basis matrices K i . The proposed approach is as follows: Fix an approximation tolerance τ > 0, then, first update the basis matrices performing for all i = 1, . . . , d:1. Form K i = U i S i, where S i is the square r i × r i matrix from QR decomposition of Mat i (C) ⊤ 2. Compute K new i with one step integration of (5) starting from K i 3. Form the new augmented matrix U new i by orthonormalizing the columns of [U i , K new i ] Second, update the core tensor and truncate: 4. Lift the core tensor C = C × d i=1 (U new i ) ⊤ U i using the new augmented basis matrices 5. Compute C new with one step integration of (5) starting from C using fixed basis matrices U new i 6. Perform a rank adjustment step to the prescribed tolerance by computing the best Tucker approximation of C new , i.e. solving the following optimization (rounding) task:EQUATIONwhere ρ = (r 1 , . . . , r d ) and M ≤2ρ denotes the set of tensors with component-wise Tucker rank lower than 2ρ.In practice, the final rank adaptive step is done by performing a high-order SVD (HOSVD) [16] on C new . The parameter τ is responsible for the compression rate of the method, as larger values of τ yield smaller Tucker ranks and thus higher parameter reduction. The computed tensorC ∈ M ρ ′ has the form C = C ′ × d i=1 U ′ i ∈ M ρ ′ and the computed U ′ i ∈ R 2ri×r ′ i with r ′i ≤ 2r i are then pulled back to the initial dimension by the change of basisU i = U new i U ′ i ∈ R ni×r ′ i, and the new core tensor C is then assigned C ′ . This implementation, however, comes at the expense of evaluating the network and gradient tapes d + 1 times for an order d tensor.The next key result will overcome this issue and will allow us to reduce the necessary network and gradient tape evaluations to two. Corollary 2.2.Let W = C × d i=1 U i ∈ M ρ be such that (6) holds. Let Mat i (C) ⊤ = Q i S ⊤ i be the QR decomposition of Mat i (C) ⊤ and let K i = U i S i . Then, span U i , Ki = span U i , ∇ Ui L W .The supplementary material provides the proof in Appendix E. Using Corollary 2.2, one can replace the individual forward evaluation and descend steps for K i by a single backpropagation. All available new information is given by the gradients ∇ Ui L, which can be evaluated from the same tape.Combining the above strategy with Corollary 2.2 we obtain Algorithm 1: an efficient rank-adaptive geometry-aware training method for tensor in Tucker format. Note that without the explicit computation of K i = U i S i we compute all basis gradients ∇ Ui L in a single network evaluation and we use ∇ Ui L to augment the basis. Note that stochastic gradient evaluations can be done in practice and that Algorithm 1: TDLRT: Efficient Tensor Dynamical Low-Rank Training in Tucker format.Input :Initial low-rank factors C ∼ r 1 × • • • × r d ; U i ∼ n i × r i ;adaptive: Boolean flag that decides whether or not to dynamically update the ranks; τ : singular value threshold for the adaptive procedure.G i ← ∇ Ui L(C × d i=1 U i ) /* Single-sweep grad evaluation */ for each mode i do U new i , _ ← QR([U i , G i ]) /* Augmentation and orthonormalization */ C ← C × d i=1 (U new i ) ⊤ U i /* Projection of Tucker core onto new basis */ C ← descent step with direction ∇ C L C × d i=1 U new i (C, U 1 , . . . , U d ) ←Tucker decomposition of C up to relative error τ as in ( 6)U i ← U new i U i , for i = 1, . . . , d /* Rank adjustment */momentum methods are applicable for the descent step on line 5 of Algorithm 1. In case the rank decreases after the retraction step, we only use the corresponding subset of the old basis functions to form the momentum term. In case of rank increase, the momentum term of the new basis vectors is initialized as zero. As a side note, the rank truncation proposed in Algorithm 1 allows to maintain the nice theoretical guarantees of the algorithm, but in practice any tensor-completion/factorization method such as [23] could be used for this step.The computational costs for the full network training come from back and forward passes through each layer. For a layer with weight tensorW ∈ R n1×•••×n d , they require O(b i n i ) operations, where b is the batch size. When using TDLRT, these computational costs reduce to O(b i r i + b i n i r i ) operations, yielding a significant reduction in computational costs to determine the gradient. However, performing low-rank updates also adds computational costs due to several factorizations. Here, the QR and SVD on Mat i (C), which are needed in the updates of U i and the truncation step, require O( i r i j r j ), and the QR on K i requires O( i n i r 2 i ) operations. Hence, in total, we have for every layer a cost ofO(b i r i + b i n i r i + d i=1 (n i r 2 i + r i d j=1 r j )) operations for TDLRT, vs. the O(b i n i ) required by the full baseline. Thus, TDLRT scales linearly with the dimensions n i , and for r i ≪ n i , which is typically the case, see Appendix C.2, it has advantageous computational cost. 
convergence and approximation analysis: 
 In this section, we present our main theoretical results. First, we show Algorithm 1 guarantees descent of the training loss, provided the compression tolerance is not too large. Second, we show that when Algorithm 1 is implemented with SGD with a decaying learning rate, the method converges to a stationary point in expectation. Third, we prove that the compressed network computed via the rank-adaptive TDLRT scheme in Algorithm 1 well-approximates the full model that one would obtain by standard training, provided the gradient flow of the loss is, at each step, approximately low-rank. The latter result shows that if a high-performing subnetwork of low Tucker rank exists, then the proposed TLDRT will probably approximate it. For brevity, some statements here are formulated informally, and all proofs and details are deferred to Appendix F in the SM. Suppose that for each convolution W , the gradient ∇ W L, as a function of W , is locally bounded and Lipschitz, i.e.,∥∇ W L(Y)∥ ≤ L 1 and ∥∇ W L(Y 1 ) -∇ W L(Y 2 )∥ ≤ L 2 ∥Y 1 -Y 2 ∥ around W . Then, Theorem 3.1 (Descent). Let W (λ) = C × dj=1 U j be the Tucker low-rank tensor obtained after one training iteration using Algorithm 1 and let W (0) be the previous point. Assuming the one-step integration from 0 to λ is done exactly, it holds L W (W (λ)) ≤ L W (W (0)) -αλ + βτ , where α, β > 0 are constants independent of λ and τ , and where L W denotes the loss as a function of only W .We now prove that the rank-adaptive training method in Algorithm 1 converges to a stationary point in expectation if implemented with SGD and decaying learning rate. Tiny-Imagenet Figure 1 : Comparison of compression performance for different models against the full baseline for the Cifar10 (a-c) and Tiny-Imagenet (d) benchmark. The mean accuracy of 20 weight initializations is displayed. TDLRT achieves higher compression rates at higher accuracy with lower variance between initializations. Theorem 3.2 (Convergence). Denote with W (t) is the weight tensor after t passes of Algorithm 1 before the rank truncation step, and W (t) the one obtained after the rank truncation. Assume Algorithm 1 is implemented using SGD as a descent method with learning rate sequence {λ t } satisfying the Robbins-Monro conditions:t λ t = +∞ t λ 2 t < +∞ .Suppose also that the spectral distribution stabilizes fast enough over time, i.e., t≥0 E[∥ W (t) -W (t)∥] < +∞ and that the projected stochastic gradient has a controlled drift, namelyE ∥∇L(t -1) × j P Uj (t) ∥ 2 | t -1 ≤ µ + ν∥∇L(t -1) × j P Uj (t-1) ∥ 2for some µ, ν > 0, where P U is the orthogonal projection onto the range of U . Then, the following convergence condition holds lim inf t→∞ E∥∇L(t -1) × j P Uj (t-1) ∥ 2 = 0 Details of the proof are contained in the appendix. Theorem 3.3. For an integer k, let t = kλ, and let W (t) be the full convolutional kernel, solution of (2) at time t. Let C(t), {U i (t)} i be the Tucker core and factors computed after k training steps with Algorithm 2, where the one-step integration from 0 to λ is done exactly. Finally, assume that for any Y in a neighborhood of W (t), the gradient flowEQUATIONwhere the constants c 1 , c 2 , c 3 depend only on L 1 and L 2 .In particular, both bounds in the above theorems do not depend on the higher-order singular values of the exact nor the approximate solution, which shows that the method does not suffer instability and slow convergence rate due to potential ill-conditioning (small higher-order singular values). Note that this result is crucial for efficient training on the low-rank manifold and is not shared by direct gradient descent training approaches, as we will numerically demonstrate in the following section. Moreover, we emphasize that (7) provides a sufficient condition for the computation of a high-performing subnetwork. In fact, for smooth enough network models f , condition (7) implies that f (W (t)) ≈ f (C(t) × d j=1 U j (t)), i.e. the computed subnetwork approximates the full model. 
experiments: 
 In the following, we conduct a series of experiments to evaluate the performance of the proposed method as compared to the full model and to standard layer factorization and model pruning baselines. The full baseline is the network trained via standard implementation. In order to test the method on tensor layers, we consider here convolutional networks and apply the decomposition to the convolutional kernel W . In terms of layer factorization, we compare against different baseline approaches: direct training of the factors in the low-rank matrix factorization format [34, 49, 79, 36] , the low-rank tensor Canonic-Polyadic format [44, 71, 4, 63] , the low-rank tensor Tucker format [42, 38] , and the low-rank Tensor-Train format [62] . The convolution operation can then be written completely in terms of the small factors using the factorization ansatz. While the above papers propose different initialization and rank selection strategies, all the referenced literature trains the factors in the chosen layer factorization format by implementing forward and backward propagations simultaneously and independently on each factor in a block-coordinate fashion. This way of training on the low-rank manifold ignores the geometry of the manifold, whereas TDLRT directly exploits the underlying geometry to avoid points of high curvature. We compare the training strategy alone. Thus, we ignore any model-specific initialization and regularization addition. We also compare with Riemannian gradient descent (RGD) for tensors in Tucker format implemented using the HOSVD retraction [74, 17] and with the matrix dynamical training algorithm [68] , where the standard forward and backward passes are endowed with a rank-adaptive QR projection step, similar to the proposed Algorithm 1. In terms of pruning techniques based on sparsification, we compare with methods from two of the most popular strategies: iterative magnitude pruning (IMP) [20] , and single-shot pruning at initialization, single-shot network pruning (SNIP) [45] and Gradient Signal Preservation (GraSP) [78] . The experiments are performed on an Nvidia RTX3090, Nvidia RTX3070 and one Nvidia A100 80GB. The code is available in the supplementary material.The compression performance of TDLRT is evaluated on CIFAR10 and tiny-imagenet. The typical data augmentation procedure is employed for this dataset: a composition of standardization, random cropping, and a random horizontal flip. All methods are trained using a batch size of 128 for 70 epochs each, as done in [79, 36] . All the baseline methods are trained with the SGD optimizer; the starting learning rate of 0.05 is reduced by a factor of 10 on plateaus, and momentum is chosen as 0.1 for all layers. The rank r of each tensor mode for the fixed-rank baseline methods is determined by a parameter κ, i.e., we set r = κr max . The proposed TDLRT method employs Algorithm 2, where SGD is used for the descent steps at lines 4 and 9, with momentum and learning rate as above. Dynamic compression during training is governed by the singular value threshold τ , see Equation (6). Figure 1 (a-c ) shows the mean accuracy of TDLRT as compared to competing factorization baselines. TDLRT achieves higher compression rates at higher accuracy with lower variance between weight initializations than the competing methods. In the case of the VGG16 benchmark, TDLRT is able to maintain baseline accuracy for compression rates over 90% and exceeds the baseline on average for τ = 0.03, i.e., 95.3% compression. Alexnet has 16.8% of the parameters of VGG16. Thus, compression is naturally more challenging to achieve. Nevertheless, TDLRT outperforms the baselines and remains close to the full network performance. Similar behavior is observed on ResNet18.Table 1 shows a comparison of the best-performing compression between all the factorization-based and pruning-based baseline methods as well as TDLRT in the CIFAR10 benchmark for Alexnet, ResNet18, and VGG16. In the proposed evaluation, TDLRT is on par or outperforms all the alternatives, including pruning based on sparsity (implemented without warmup for the sake of a fair To further highlight the advantages of Algorithm 2 as compared to standard simultaneous gradient descent on the factors of the decomposition, we show in Figure 2 the accuracy history of LeNet5 on MNIST using TDLRT as compared to standard training on Tucker and CP decompositions. In the case of TDLRT, an optimization step denotes the evaluation of Algorithm 2 for all convolutional layers for one batch of training data, while for the other methods, we refer to a standard SGD batch update for all factors of the tensor decompositions of all layers. All linear layers of the network are trained with a traditional gradient descent update and are not compressed. In this experiment, we initialize the network weights to simulate a scenario where the rank is overestimated. To this end, we employ spectral initialization with singular values decaying exponentially with powers of ten.Integrating the low-rank gradient flow with the TDLRT Algorithm 2 leads to faster and more robust convergence rates of the network training process.The computational performance in inference and training of convolutional layers in Tucker decomposition depends on their current tensor ranks, see Section 2. We evaluate the inference time of 120K RGB images and memory footprint of VGG and AlexNet in Tucker factorization as used in Algorithm 2 and compare them to the non-factorized baseline models in Figure 2 . As a result, for realistic compression rates, see also Figure 1 , the computational footprint of TDLRT is significantly lower than the corresponding baseline model.Rank adaptive TDLRT training comes at the additional expense of QR and SVD operations per optimization step compared to standard fixed rank training without orthonormalization. However, the increased robustness of the optimization and faster convergence reduces the computational overhead of TDLRT as demonstrated in Figure 2 . In order to provide a fair comparison between the methods, we report the time to achieve 60% accuracy target at a compression rate of 90% for all methods, on LeNet5 MNIST with the setting as in Section 4.2. We refrain from measuring time to convergence since the standard Tucker decomposition is not able to reach similar accuracy levels at the same compression ratio as TDLRT.In this section, we are presenting another application of our method, namely fine-tuning pre-trained models through the use of low-rank adapters. In particular, our approach Algorithm 1 is completely agnostic between model compression and adaptation, the approach is the same. More precisely, given a pre-trained model f W * with tensor pre-trained weights W * , it is possible to add a low-rank Tucker correction W . Then, given a task loss L(f ), by defining L(W ) = L(f W * +W ) we report ourselves to the original formulation of the problem. Moreover, we would like to stress that this approach can be applied also to matrices, which are a particular case of tensors with d = 2 modes. To showcase these two settings, we present two different settings in which we test our method. In Table 2 (left) we show the fine-tuning of Deberta V3 [29] on the GLUE benchmark [77] . In this test case here, the low-rank adapters have been applied to all matrices in attention layers, and the final performance against LoRA [33] fine-tuning is reported. Apart from our additional hyperparameter τ , all the other hyperparameters had been kept as reported in [85] .In Table 2 (right), we report the results for fine-tuning stable diffusion [64] with Dreambooth [65] .To adapt the model, we applied Tucker tensor corrections to each convolution of the Unet, and a matrix adapter to each attention layer of the text encoder network. We applied our method on all these correctors, while keeping the same hyperparameters setting as in [55] . We would like to observe that the kind of adapters they propose for convolutions consist in a matrix factorization of a reshaping of the convolutional tensor. This results in a potentially bigger number of parameters needed, as a matrix factorization would be O((Cd1 d 2 + F )r) against a O(Cr 1 + F r 2 + d 1 r 3 + d 2 r 4 + r 1 r 2 r 3 r 4 )for a plain Tucker factorization, where F is the number of output features, C is the number of input channels, d 1 and d 2 are the spatial dimensions of the convolutional kernel. This observation is in fact reflected in the numbers in Table 2 , in which we can observe a higher compression potential for tensor decompositions compared to matrix ones.  
discussion and limitations: 
 This work leverages the geometry of the Tucker tensor factorization manifolds to construct a robust and efficient training algorithm for neural networks in compressed Tucker format. The proposed method has a theoretical backbone of approximation error bounds to the full model and guarantees of loss descent and convergence to stationary points in expectation. The method is superior to standard factorization approaches with alternating or simultaneous gradient descent, as demonstrated in the compression-to-accuracy ratio for various benchmarks and models. The method provides a significant reduction of hyperparameters to a single parameter τ . This parameter has a clear interpretation as compression rate per layer depending on the layer's importance on the overall network training dynamics. We, however, note that further strategies to pick an adequate τ are possible and remain to be investigated. Further, we note that an efficient implementation on GPUs requires an efficient tensor rounding algorithm in Algorithm 1. Finally, the proposed method assumes well-performing low-rank Tucker sub-nets exist in the reference network. While we observe this empirically, further investigations are required to provide theoretical evidence supporting this assumption, similar to the case of fully-connected linear layers, see [3, 6] .This paper presents work whose main goal is to reduce the training costs of tensor-based architecture, while maintaining mathematically sound guarantees of performance in terms of convergence and approximation. As in the majority of deep learning research, there are potential societal consequences of our work, none of which we feel must be specifically highlighted here. A positive societal impact is the reduces carbon emissions by more efficient neural networkt training and inference. Regarding ethical aspects, we feel nothing has to be added.B Algorithm for tensor taping based on Theorem 2.1For the sake of completeness, we provide the algorithmic description of the method to generate gradients of the Tucker factors as obtained by Theorem 2.1.The training algorithm for tensor-valued neural network layers in Tucker format is presented in Algorithm 2. Through the lens of computational efficiency, the main difference to Algorithm 1 is the following:Each time we back-propagate through a convolutional layer W = C × d i=1 U i , we form the new variable K i = U i S i as in Theorem 2.1, we integrate the ODE in (5) from K i (0) = K i to K i (λ), λ > 0, and then update the factors U i by forming an orthonormal basis of the range of K i (λ). This strategy directly follows from Theorem 2.1.Similar to Algorithm 1, we implement the orthonormalization step via the QR factorization while we perform the integration of the gradient flow via stochastic gradient descent with momentum and learning rate λ, which coincides with a stable two-step linear multistep integration method [69] . Once all the factors U i are updated, we back-propagate the core term by integrating the equation for C in (5), using the same approach.Similar to Algorithm 1, the Tucker rank of the new kernel can be adaptively learned with a key basis-augmentation trick. The implementation for Algorithm 2 works as follows: Each time we backpropagate K i ∈ R ni×ri , we form an augmented basis K i by appending the previous U i to the new K i (λ),K i = [K i |U i ].We compute an orthonomal basis U new i ∈ R ni×2ri for K i and we form the augmented2r 1 × • • • × 2r d core C = C × d i=1 (U new i ) ⊤ U i .We then backpropagate the core C integrating (5) starting from C(0) = C. Finally, we perform a rank adjustment step by computing the best Tucker approximation of C to a relative tolerance τ > 0. This step corresponds to solving the following optimization (rounding) task:Find C ∈ M ≤2ρ of smallest rank ρ ′ = (r ′ 1 , . . . , r ′ d ) such that ∥ C -C∥ ≤ τ ∥ C∥ where ρ = (r 1 , . . . , r d ) and M ≤2ρ denotes the set of tensors with component-wise Tucker rank lower than 2ρ. In practice, this is done by unfolding the tensor along each mode and computing a truncated SVD of the resulting matrix. The tensor C ∈ M ρ ′ is then further decomposed in its Tucker decomposition yielding a factorizationC = C ′ × d i=1 U ′ i ∈ M ρ ′ .The parameter τ is responsible for the compression rate of the method, as larger values of τ yield smaller Tucker ranks and thus higher parameter reduction. To conclude, the computed U ′ i ∈ R 2ri×r ′ i with r ′ i ≤ 2r i are then pulled back to the initial dimension of the filter by settingU i = U new i U ′ i ∈ R ni×r ′ i, and the new core tensor C is then assigned C ′ . This implementation shares the robust error bound of Algorithm 1. However it comes at an increased computational cost due to d + 1 necessary evaluations of the network and gradient tape, where the first d are due to the basis updates K i and the last for the coefficient update S. C ∼ r 1 × • • • × r d ; U i ∼ n i × r i ;adaptive: Boolean flag that decides whether or not to dynamically update the ranks; τ : singular value threshold for the adaptive procedure. for each mode i do Q i S ⊤ i ← QR decomposition of Mat i (C) ⊤ K i ← U i S i K i ← descent step; direction ∇ Ki L(Ten i (Q ⊤ i ) × j̸ =i U j × i K i ); starting point K i if adaptive then /* Basis augmentation step */ K i ← [K i | U i ] U new i ← orthonormal basis for the range of K i C ← C × d i=1 (U new i ) ⊤ U i C ← descent step; direction ∇ C L C × d i=1 U new i ; starting point C if adaptive then /* Rank adjustment step */ (C, U 1 , . . . , U d ) ← Tucker decomposition of C up to relative error τ U i ← U new i U i , for i = 1, . . . , d else U i ← U new i , for i = 1, . . . , dThe rank evolution over the optimization steps of VGG16 on Cifar10 are displayed in Figure 3 . The color gradients indicate the position of the respective tensor basis in the network, where lighter green denotes bases near the input and darker green denotes bases near the output layer of VGG16. Higher singular value cutoff tolerance τ results in faster rank decay; however, across different choices of tau, a monotonous decrease in the ranks to a steady state is observable.The proposed method Algorithm 1 allows us to choose any Tucker ranks at initialization In Figure 3 , we initialize the network layers with full rank. The decay of the layers' ranks over time is typical for Alg1 and, indeed observed for other architectures as well. This is aligned with theoretical [5] and empirical [18] findings stating that neural networks trained with SGD exhibit a low-rank structure. Tab.4 contains the Tucker ranks of Alexnet compressed with TDLRT as a supplement to the results presented in section4.1. All the test cases were run with the rank adaptive version of the integrator.For Cifar10, a standard random crop and random flip are used for data augmentation at training time.All methods are trained for 70 epochs using a batch size of 128. All methods are trained using the SGD optimizer, with a starting learning rate of 5 × 10 -2 with a scheduler that reduces it by a factor of 10 whenever validation loss reaches a plateau. Polyak momentum was 0.1 for all layers but batch normalizations, which was set to 0.9.where † denotes the pseudoinverse. We assume that the tensor W admits a differentiable Tucker tensor representation, i.e., W (t) = C(t) × d i=1 U i (t) ∈ M ρ for t ∈ [0, λ] satisfying the associated gradient-flow tensor differential equationẆ (t) = -∇ W L(W (t)) .For sake of brevity, the parameter is going to be omitted. Next, we perform a QR-factorizationMat i (C) ⊤ = Q i S ⊤i . Thus, we observe that the matrix Mat i (W ) admits an SVD-like decomposition as followsMat i (W ) = U i S i V ⊤ i with V ⊤ i := Q ⊤ i j̸ =i U ⊤ j ,where the matrix Q i possess othornormal columns, i.e., Q ⊤ i Q i = I. We introduce the quantityK i := U i S i where S i = Mat i (C)Q i .By construction, we observe thatS i = U ⊤ i Mat i (W )V i .The tiny matrix S i satisfies the following differential equationṠi = U ⊤ i Mat i (W )V i + U ⊤ i Mat i ( Ẇ )V i + U ⊤ i Mat i (W ) Vi = U ⊤ i U i =0 S i V ⊤ i V i + U ⊤ i Mat i ( Ẇ )V i + U ⊤ i Mat i (W ) Vi =0 = -U ⊤ i Mat i (∇ W L)V i .The first null identity follows from the gauge condition on U i . The second null identity follows by the initial assumption Vi = 0. Therefore, the matrix K i satisfies the differential equationKi = (U i S i ) = Ui S i + U i Ṡi = -(I -U i U ⊤ i ) Mat i (∇ W L × j̸ =i U ⊤ j ) Mat i (C) † S i -U i U ⊤ i Mat i (∇ W L)V i = (U i U ⊤ i -I) Mat i (∇ W L × j̸ =i U ⊤ j )Q i -U i U ⊤ i Mat i (∇ W L)V i = (U i U ⊤ i -I) Mat i (∇ W L) • ( ⊗ j̸ =i U j )Q i -U i U ⊤ i Mat i (∇ W L)V i = (U i U ⊤ i -I) Mat i (∇ W L)V i -U i U ⊤ i Mat i (∇ W L)V i = -Mat i (∇ W L)V i . (9)To conclude, we setZ = K i V ⊤ i and let W = Ten i (K i V ⊤ i ) = Ten i (Z). We obtain ∇ Ki L(W ) = ∇ Z L(Ten i (Z))∇ Z K i = ∇ Z L(Ten i (Z))V i .where we remind thatEQUATION)The first K i -differential equation is then obtained combining ( 11) and ( 10)EQUATION)The right-hand side can be further reduced using standard tensorization formulas [41] Teni (K i V ⊤ i ) = Ten i (Q ⊤ i ) × j̸ =i U j × i K i .The second differential equations follows by observing that∇ W L = ∇ C L × i U i + j L × j Uj × i̸ =j U i .The tensor C satisfies the differential equationĊ = -∇ W L(W ) × i U ⊤ i = -∇ C L(W ) × i U i × i U ⊤ i = -∇ C L(W ) × i U ⊤ i U i =I = -∇ C L(C × i U i ) .where the extra terms disappear due to the imposed gauge conditionsEQUATION)Proof. From Eq. ( 9) of the proof in Theorem 2.1 it is apparent, thatEQUATION) Moreover, we observe that using the chain rule,∇ Ui L = Mat i (∇ W L)∇ Ui (U i S i V ⊤ i ) = Mat i (∇ W L)V i S ⊤ iThen, with (13) we have∇ Ui L = Mat i (∇ W L)V i S ⊤ i = -Ki S ⊤ i .Using full-rankness of S i concludes the proof.Theorem. Let W (λ) = C × d j=1 U j be the Tucker low-rank tensor obtained after one training iteration using Algorithm 2 and let W (0) be the previous point. Then, for a small enough learning rate λ, it holds L W (W (λ)) ≤ L W (W (0)) -αλ + βτ , where α, β > 0 are constants independent of λ and τ , and where L W denotes the loss as a function of only W .Proof. Let W (t) = C(t) × i U 1 i .Here, W (t) and C(t) denote the augmented solutions for t ∈ [0, λ] arising from the intermediate steps of the TDLRT Algorithm 2. We observe thatEQUATIONIntegrating ( 14) from t = 0 until t = λ, we obtainL( W (λ)) ≤ L( W (0)) -αλ.Because the augmented subspaces U i contain by construction the range and co-range of the initial value, we have that W (0) = W (0). Furthermore, the truncation is such that∥W (λ) -W (λ)∥ ≤ τ . Therefore L(W (λ)) ≤ L( W (λ)) + βτ where β = max 0≤τ ≤1 ∥∇L τ W (λ) + (1 -τ ) W (λ) ∥.Lemma F.1. The following estimate holds∥W (λ) -W (λ) × j U j (λ)U j (λ) ⊤ ∥ ≤ Θ = O(h(h + ϵ)) ,where the hidden constants depend only on L 1 and L 2 .Proof. It has been shown in [68, Appendix] that there exists a constant θ ∝ O(h(h + ϵ)) such that∥U j U ⊤ j Mat j (W (λ)) -Mat j (W (λ))∥ ≤ θ ∀j = 1, . . . , d ,where the value θ has been shown to depend only on the constants L 1 , L 2 and λ. The proof of the Lemma follows a recursive constructive argument∥W (λ) -W (λ) × d j U j (λ)U j (λ) ⊤ ∥ ≤ ∥W (λ) -W (λ) × d-1 j U j (λ)U j (λ) ⊤ ∥ + ∥W (λ) × d-1 j U j (λ)U j (λ) ⊤ -W (λ) × d j U j (λ)U j (λ) ⊤ ∥ ≤ ∥W (λ) -W (λ) × d-1 j U j (λ)U j (λ) ⊤ ∥ + ∥ W (λ) × d U d (λ)U d (λ) ⊤ -W (λ) × d-1 j U j (λ)U j (λ) ⊤ ∥ ≤ ∥W (λ) -W (λ) × d-1 j U j (λ)U j (λ) ⊤ ∥ + ∥W (λ) × d U d (λ)U d (λ) ⊤ -W (λ)∥ ≤ ∥W (λ) -W (λ) × d-1 j U j (λ)U j (λ) ⊤ ∥ + θ .The conclusion is obtained iterating the provided argument.Theorem F.2. For an integer k, let t = kλ, and let W (t) be the full convolutional kernel, solution of (2) at time t. Let C(t), {U i (t)} i be the Tucker core and factors computed after k training steps with Algorithm 2, where the one-step integration from 0 to λ is done exactly. Finally, assume that for any Y in a neighborhood of W (t), the gradient flow-∇L W (Y ) is "ε-close" to M ρ . Then, ∥W (t) -C(t) × d j=1 U j (t)∥ ≤ c 1 ε + c 2 λ + c 3 τ /λwhere the constants c 1 , c 2 and c 3 depend only on L 1 and L 2 .Proof. First, we provide a bound for the local error, i.e., the error obtained after one training epoch. If we apply Lemma F.1, we obtain that∥W (λ) -C(λ) × d j=1 U j (λ)∥ ≤ ∥W (λ) -W (λ) × d j=1 U j (λ)U j (λ) ⊤ ∥ + ∥W (λ) × d j=1 U j (λ)U j (λ) ⊤ -C(λ) × d j=1 U j (λ)∥ ≤ Θ + ∥ W (λ) × d j=1 U j (λ) ⊤ -C(λ) × d j=1 U j (λ)∥ ≤ Θ + ∥W (λ) × d j=1 U j (λ) ⊤ -C(λ)∥ .It suffices to study the latter term. For t ∈ [0, λ], we define the quantityC(t) := W (t) × d j=1 U j (λ) ⊤ It satisfies the differential initial value problem ˙ C = -∇ W L(W ) × d j=1 U j (λ) ⊤ , C(0) = W (0) × d j=1 U j (λ) ⊤ .The term W (t) can be written as a perturbation of CW (t) = C × d j=1 U j (λ) + R(t) , where R(t) = W (t) -W (t) × d j=1 U j (λ)U j (λ) ⊤ . Then, we observe that ∥W (t) -W (λ)∥ ≤ λ 0 ∥ Ẇ (s)∥ds = λ 0 ∥-∇ W L(W (s))∥ds ≤ C 1 λ .The remainder can be estimated as follows∥R(t)∥ ≤∥R(t) -R(λ)∥+∥R(λ)∥ ≤ 2L 1 λ + 2Θ .Furthermore, the full gradient can be re-written as∇ W L(W (t)) = ∇ W L( C(t) × d j=1 U j (λ) + R(t)) = ∇ W L( C(t) × d j=1 U j (λ)) + D(t), where the defect D(t) is defined asD(t) := ∇ W L( C(t) × d j=1 U j (λ) + R(t)) -∇ W L( C(t) × d j=1 U j (λ)) . Because of the Lipschitiz assumption, we have that∥D(t)∥ ≤ L 2 ∥R(t)∥ ≤ 2L 2 (L 1 λ + Θ) .Next, we compare the two differential equations˙ C(t) = -∇ W L( C(t) × d j=1 U j ) × d j=1 U ⊤ j + D(t), Ċ(t) = -∇ W L(C(t) × d j=1 U j ) × d j=1 U ⊤ j ,where C(0) = C(0), by construction. The solution C(λ) of the second tensor-differential equation coincides with the solution of the last training step of the TuckerDLRT algorithm. The first differential equation has been constructed such that its solution isC(λ) = W (λ) × j U ⊤ j .Therefore, the study of the local error follows by a direct application of the Gronwall inequality∥C(λ) -C(λ)∥ ≤ exp(C 2 λ)2L 2 (L 1 λ + Θ)λ .To conclude, the global error in the training epochs follows by using the Lipschitz continuity of the gradient flow: We move from the local error in time to the global error in time by a standard ODEs argument of Lady Windermere's fan [82, §II.3] .In this section, we provide the details of the proof of convergence to stationary points in the stochastic setting, Theorem 3.2. The proof extends the approach of [31] to the tensor case and relaxes some of the assumptions there made on the matrix case, while following the same overall structure.Theorem G.1 (Convergence). Let W (t) be the weight tensor after t ∈ N iterations of Algorithm 1 before the rank truncation step, and W (t) as the one obtained after the rank truncation. Assuming that • Algorithm 1 is implemented using SGD as the descent method.• The loss function is assumed to be positive, locally bounded, and differentiable with a Lipschitz gradient.• The learning rate sequence λ t satisfies the Robbins-Monro conditions, i.e.t λ t = +∞ t λ 2 t < +∞ .• The spectral distribution stabilizes fast enough over time, i.e. EQUATION• The projected stochastic gradient has a controlled drift, namelyEQUATIONwhere P U = U U T is the orthogonal projection onto the range of U , andE [•|t] = E •|W (t), {U i (t)} d i=1denotes the conditional expectation. Then the following convergence condition holds lim inf t→∞ E ∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 = 0The convergence proof of Theorem 3.2 is based on the following technical lemmas. Lemma G.2. Let L be a differentiable loss function, and assume that its gradient ∇L(W ) is onesided Lipschitz continuous with constant L 2 . Then, for any W and W ′ , the following inequality holdsL(W ) ≤ L(W ′ ) + ⟨∇L(W ′ ), W -W ′ ⟩ + L 2 2 ∥W -W ′ ∥ 2Proof. By using Cauchy-Schwartz inequality, we have thatL(W ) = L(W ′ ) + 1 0 d dt L W ′ + t(W -W ′ ) dt = L(W ′ ) + ⟨∇L(W ′ ), W -W ′ ⟩ - 1 0 ⟨∇L(W ′ ) -∇L (W ′ + t(W -W ′ ) , W -W ′ ⟩ dt ≤ L(W ′ ) + ⟨∇L(W ′ ), W -W ′ ⟩ + L 2 ∥W -W ′ ∥ 2 1 0 t dt = L(W ′ ) + ⟨∇L(W ′ ), W -W ′ ⟩ + L 2 2 ∥W -W ′ ∥ 2 .Lemma G.3. Let U j,1 = [U j,1 |U j,0 ] be a given basis set with orthonormal columns, L denote the loss function computed on the whole dataset, and L B denote the loss calculated on a batch B. Then, for any j * ∈ {1, . . . , d} and W , it holds thatEQUATIONProof. We introduce first the functionϕ( U ) := ∇L(W ) × j̸ =j * P Uj,0 × j * P U , ∇L B (W ) × j̸ =j * P Uj,0 × j * P U .Let m be the the infimum on the set U as defined belowm = inf U ∈U ϕ( U ) with U = {U ∈ R n j * ×r j * | rank(U ) ≤ r j * } .Because the term ∇L(W ) can be decomposed into a sum of ∇L B (W ) × j̸ =j * P U j * ,0 × j * P Uj,1 and its orthogonal component, we observe that the infimum satisfies m ≤ ϕ(U j * ,1 ) = ∇L(W ), ∇L B (W ) × j̸ =j * P U j * ,0 × j * P Uj,1 .Hence, by taking the expectation on both sides, we can conclude thatE ∇L(W ), ∇L B (W ) × j̸ =j * P Uj,0 × d j * P Uj,1 ≥ inf U ∈U E ϕ( U ) ≥ inf U ∈U ∥∇L(W ) × j̸ =j * P Uj,0 × j * P U j * ,1 ∥ 2The conclusion follows by observing that the most right term in the inequality is positive.With the technical lemmas established above, we are now in the position to prove the theorem 3.2.Proof. (Thereom 3.2) To simplify the notation, we will denote byE t [•] the conditional expectation E t [ • ] := E[ • | W (t -1)]. We first remind the reader of two properties of the conditional expectation. Specifically, for any deterministic function ψ and any random variable X, we have thatE t ψ(W (t -1)) = ψ(W (t -1)) , E E t [X] = E[X] .We will begin by examining an upper bound for the one-step drift of 2. We denote by W (t) = C(t) × j U j (t) the weight tensor before truncation at step t ∈ N. As per the assumption, the optimization in the C-step of 2 is defined using an SGD update. Therefore, we haveC(t) = C(t-1)× j U j (t-1)-λ t ∇L(W (t-1))) × j U j (t) ⊤ = W (t-1)-λ t ∇L(W (t-1)) × j U j (t) ⊤ .HenceW (t) = C(t) × j U j (t) = W (t -1) -λ t ∇L(W (t -1)) × j P Uj (t).where P U = U U T is the orthogonal projector onto the range of arbitrary matrix U . Applying G.2 we have thatE t L( W (t)) -L(W (t -1)) ≤ -λ t E t ⟨∇L(W (t -1)), ∇L B (W (t -1)) × j P Uj (t) ⟩ + λ 2 t L 2 2 E t ∥∇L B (W (t -1)) × j P Uj (t) ∥ 2We notice that for the augmented basis U j (t) = [U j (t-1)|U j (t)], it holds P Uj (t) = P Uj (t-1) +P Uj (t) .When we expand L B (W (t -1)) × j P Uj (t) using its sum representation and apply Lemma G.3 to the mixed terms, we obtainEQUATIONThe locally bounded loss computed on the truncated approximation W (t) is bounded viaEQUATIONBy combining equations (17) and (18) , we arrive at the following boundEQUATIONFollowing assumption (16) , we haveE t L(W (t)) -L(W (t -1)) ≤ -λ t ∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 + λ 2 t L 2 2 µ + ν∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 + CE t ∥W (t) -W (t)∥ = -λ t (1 - 1 2 λ t L 2 ν)∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 + λ 2 t L 2 µ 2 + CE t ∥W (t) -W (t)∥ ≤ -λ t ∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 + λ 2 t L 2 µ 2 + CE t ∥W (t) -W (t)∥ ,where we assume that λ t ≤ 2/L 2 ν. Finally, by taking the expectation on both sides, we obtainE L(W (t))-L(W (t-1)) ≤ -λ t E ∥∇L(W (t-1))× j P Uj (t-1) ∥ 2 + λ 2 t L 2 µ 2 +CE ∥W (t)-W (t)∥By summing the last equation over t = 1, . . . , T we get-L(W (0)) ≤ E L(W (t)) -L(W (0)) ≤ - T t=1 λ t E ∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 + L 2 µ 2 T t=1 λ 2 t + C T t=1 E ∥W (t) -W (t)∥By rearranging the terms of the latter equation and by sending T → +∞, we finally obtain+∞ t=1 λ t E ∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 ≤ L(W (0)) + L 2 µ 2 +∞ t=1 λ 2 t + C +∞ t=1 E ∥W (t) -W (t)∥ < +∞The conclusion follows by the Robbins-Monro conditions, i.e.lim inf t→∞ E ∥∇L(W (t -1)) × j P Uj (t-1) ∥ 2 = 0NeurIPS Paper Checklist Justification: Limitations of the analysis is discussed after the corresponding theorem statements and in Section 5 Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. Justification: We provide an overview of the used compute resources in Section 4Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: The research is fundamental mathematical research. We do not see how it violates any aspect of the code of ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: Societal impact is discussed in Appendix A.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We use open datasets and very fundamental benchmarks, so we think our work does not has risk for missuses.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA]Justification: The paper does not release new assets Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix:  


Decision: As models continue to grow in size, it is crucial to explore more parameter-efficient alternatives. This paper examines a "parameter saving" technique, namely imposing low-rank constraints. Low-rank constraints have proven effective even for modern transformers. What distinguishes this paper is its combination of theoretical insights and practical aspects. On the theoretical side, the paper analyzes the local convergence of gradient-based optimization under specific low-rank constraints. On the practical side, it proposes adapting these constraints to address optimization challenges, particularly focusing on the geometric aspects of the decomposition. The paper targets an important practical topic and addresses it both theoretically and empirically, I therefore recommend accepting the paper.

There is a substantial body of literature on applying constraints to neural network weight matrices to enhance optimization performance, including constraints based on the Stiefel manifold. The paper could benefit from reviewing this existing research.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: -
questions: 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.
limitations: limitations: I don't see any direct potential negative societal impact.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Yes, the authors discuss about the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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).
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: N/A

Comments: title: Reply to authors' rebuttal
comment: We thank again the reviewer for the feedback.
We thank again the reviewer for the feedback.
We thank again the reviewer for the feedback.
We thank again the reviewer for their response and we are glad we were able to clarify all doubts.

Regarding the U-net, we applied LoRA to the attention and convolutional layers, using the official implementation of LoRA in the Huggingface PEFT package [1]. Specifically, the LoRA implementation for convolutions does not perform a proper tensor decomposition; instead, it corresponds to a low-rank factorization of a flattened version of the convolutional kernel (flattened to a matrix). Our proposed TDLRT has been applied to the same layers as those described above, maintaining their original tensor/matrix structure.

[1] S. Mangrulkar, S. Gugger, L. Debut, Y. Belkada, S. Paul and B. Bossan, "PEFT: State-of-the-art Parameter-Efficient Fine-Tuning methods", github 2022.

We thank the reviewer once again for their interest and remain available to clarify any further doubts.
Thank you for answering my questions. I read other reviews and the responses and have decided to keep my score at 7.
Thanks for the new experiment results. I am raising my score to 6.
Thanks for your response. I would keep my score.
Thank you for your thorough response! I was uncertain about the details of the one-step integration and the practicality of the proposed method, which were all addressed in the rebuttal. 

In particular, I thank the authors for providing additional experimental results on language and diffusion models. They clarify my doubts about the practicality. 

Regarding the diffusion model, which layers were updated with LoRA and TDLRT? Although it is common practice to apply LoRA to the attention layers of the U-Net, I suppose TDLRT can also update the convolution layers. I wonder if this was the case, and if so, it could be argued as an additional strength compared to LoRA.

Also, I retract Limitation 1 on QR decomposition overhead after rethinking it based on the authors' response. I don't think it will limit the usability of TDLRT as long as the rank is small.

Since I already gave 7, I keep my score the same. Instead, I am inclined to raise my confidence.