Title: Geometry-aware training of factorized layers in tensor Tucker format
Venue: NeurIPS 2024
Accepted: yes
Keywords: Tucker Tensors

Tensor Neural Networks

Low-Rank Compression

Optimization

Low-Rank Optimization

Dynamical Low-Rank Approximation
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.

Geometry-aware training of factorized layers in tensor Tucker format

Abstract

Reducing parameter redundancies in neural network architectures is crucial for achieving feasible computational and memory requirements during training and inference phases. Given its easy implementation and flexibility, one promising approach is layer factorization, which reshapes weight tensors into a matrix format and parameterizes them as the product of two small rank matrices. However, this approach typically 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 independently of 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.

1 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 rankr parametrization W = UV ⊤ for each layer W , with U, V rectangular matrices with only r columns. The network is then trained on the set of rankr matrices M r = { UV ⊤ : 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:

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];

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.

1.1 Related work

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

2 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 as

where C ∈ R r 1 ×···× r d is a core tensor of full Tucker rank ρ = ( r 1 , . . . , r d ) and the U i ∈ R n i × r i are matrices with orthonormal columns. From this representation, we immediately see that if W is represented in Tucker format, then the cost of storing W and of performing linear operations with W (e.g. matvecs or convolutions) is 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 . 5 r 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:

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.

We introduce a reparameterization in Theorem 2.1 that allows us to formulate robust dynamics for the reparametrized factors.

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.

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.

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

2.1 Dynamical low-rank approximation

For ρ = ( r 1 , . . . , r d ) , the set

is a manifold with the following tangent space at any point W = C × d i =1 U i ∈ M ρ [40]

̸

where 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 flow

where, 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 form W ( t ) = C ( t ) × d i =1 U i ( t ) ∈ M ρ . Then ˙ W ( t ) ∈ T W ( t ) M ρ , the tangent space of M ρ at W ( t ) . Thus, (2) boils down to

where 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 ) and U i ( t )

̸

where † 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 for K i .

Theorem 2.1. Let W = C × d i =1 U i ∈ M ρ be such that (3) 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,

̸

where 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 of W 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 :

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 ) ⊤

Form the new augmented matrix U new i by orthonormalizing the columns of [ U i , K new i ]

Compute K new i with one step integration of (5) starting from K i

Second, update the core tensor and truncate:

Lift the core tensor ˜ C = C × d i =1 ( U new i ) ⊤ U i using the new augmented basis matrices

Compute C new with one step integration of (5) starting from ˜ C using fixed basis matrices U new i

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:

where ρ = ( 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 tensor C ∈ M ρ ′ has the form C = C ′ × d i =1 U ′ i ∈ M ρ ′ and the computed U ′ i ∈ R 2 r i × r ′ i with r ′ i ≤ 2 r i are then pulled back to the initial dimension by the change of basis U i = U new i U ′ i ∈ R n i × 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,

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 ∇ U i 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 ∇ U i L in a single network evaluation and we use ∇ U i L to augment the basis. Note that stochastic gradient evaluations can be done in practice and that

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.

2.2 Computational Complexity

The computational costs for the full network training come from back and forward passes through each layer. For a layer with weight tensor W ∈ R n 1 ×···× 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 of O ( 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.

3 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 × d j =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.

60

65

70

75

80

85

90

95

100

compression rate[%]

55

60

65

70

75

80

85

test accuracy[%]

Baseline

TDLRT

Matrix

CP

(a) Alexnet Cifar10

60

65

70

75

80

85

90

95

100

compression rate [%]

80

82

84

86

88

90

92

test accuracy [%]

baseline

TDLRT

Matrix

CP

(b) VGG16 Cifar10

60

65

70

75

80

85

90

95

100

compression rate[%]

89

90

91

92

93

94

test accuracy[%]

Baseline

TDLRT

Matrix

CP

(c) ResNet18 Cifar10

60

65

70

75

80

85

90

95

100

compression rate[%]

46

47

48

49

50

51

test accuracy[%]

Baseline

TDLRT

(d)

ResNet18

Tiny-

Imagenet

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:

Suppose also that the spectral distribution stabilizes fast enough over time, i.e.,

and that the projected stochastic gradient has a controlled drift, namely

for some µ, ν > 0 , where P U is the orthogonal projection onto the range of U . Then, the following convergence condition holds

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 flow -∇L W ( Y ) is ' ε -close' to T Y M ρ . Then,

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

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

4.1 Compression Performance

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

10 5

10 6

10 7

convolution paramters

10 1

10 2

inference time [s]

c.r. 62.3%

c.r. 83.5%

c.r. 96.5%

c.r. 63.6%

c.r. 84.3%

c.r. 96.7%

Alexnet

VGG16

Alexnet - TDLRT

VGG16 - TDLRT

baseline

0

200

400

600

800

optimization step

0.0

0.2

0.4

0.6

0.8

1.0

TDLRT

CP

Method

time to

60%

acc.

2

.

35

s

4

.

70

s

TDLRT (fixed rank)

2

.

41

s

TDLRT

4

.

16

s

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

comparison), Tensor-Train (TT) and Tucker factorization, Riemmanian gradient descend (RGD) for Tucker decompositions, as well as the matrix-valued DLRT, due to the higher flexibility of the Tucker format where compression along each tensor mode individually is possible. The compression rate (c.r.) is computed as 1 -c/f , where c is the number of convolutional parameters in the compressed model after training and f is the number of convolutional parameters of the full model. While this is the compression rate after training, we emphasize that methods based on factorizations yield an analogous compression rate during the entire training process. We also remark that no DLRT version of CP decomposition is shown as CP is not suited for dynamical low-rank training due to its lack of a manifold structure. Similar results are obtained for the tiny-imagenet benchmark, see Fig. 1 (d) and Table 3. The rank evolution of the networks during training is discussed in Appendix C.2.

4.2 Robustness of the Optimization

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.

4.3 Computational Performance

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 120 K 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.

4.4 Fine-tuning with LoRA-like low-rank adapters

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 (( Cd 1 d 2 + F ) r ) against a O ( Cr 1 + Fr 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.

5 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].

Acknowledgments and Disclosure of Funding

This manuscript has been authored by UT-Battelle, LLC under Contract No. DE-AC05-00OR22725 with the U.S. Department of Energy. The United States Government retains and the publisher,by accepting the article for publication, acknowledges that the United States Government retains a non-exclusive, paid-up, irrevocable, world-wide license to publish or reproduce the published form of this manuscript, or allow others to do so, for United States Government purposes. The Department of Energy will provide public access to these results of federally sponsored research in accordance with the DOE Public Access Plan.

The work of E. Zangrando was funded by the MUR-PNRR project 'Low-parametric machine learning'. Francesco Tudisco is partially funded by the project FIN4GEO within the European Union's Next Generation EU framework, Mission 4, Component 2, CUP P2022BNB97.

References

P. Ablin and G. Peyré. Fast and accurate optimization on the orthogonal manifold without retraction. In G. Camps-Valls, F. J. R. Ruiz, and I. Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , volume 151 of Proceedings of Machine Learning Research , pages 5636-5657. PMLR, 28-30 Mar 2022.

P.-A. Absil and J. Malick. Projection-like retractions on matrix manifolds. SIAM Journal on Optimization , 22(1):135-158, 2012.

S. Arora, N. Cohen, W. Hu, and Y. Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems , 32, 2019.

M. Astrid and S.-I. Lee. Cp-decomposition with tensor power method for convolutional neural networks compression. In 2017 IEEE International Conference on Big Data and Smart Computing (BigComp) , pages 115-118, 2017.

B. Bah, H. Rauhut, U. Terstiege, and M. Westdickenberg. Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers. Inf. Inference , 11(1):307-353, 2022.

B. Bah, H. Rauhut, U. Terstiege, and M. Westdickenberg. Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers. Information and Inference: A Journal of the IMA , 11(1):307-353, 2022.

B. Baker, I. Akkaya, P. Zhokov, J. Huizinga, J. Tang, A. Ecoffet, B. Houghton, R. Sampedro, and J. Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems , 35:24639-24654, 2022.

N. Bansal, X. Chen, and Z. Wang. Can we gain more from orthogonality regularizations in training deep networks? In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

G. Ceruti, J. Kusch, and C. Lubich. A rank-adaptive robust integrator for dynamical low-rank approximation. BIT Numerical Mathematics , pages 1-26, 2022.

G. Ceruti and C. Lubich. An unconventional robust integrator for dynamical low-rank approximation. BIT Numerical Mathematics , 62(1):23-44, 2022.

G. Ceruti, C. Lubich, and D. Sulz. Rank-adaptive time integration of tree tensor networks. SIAM Journal on Numerical Analysis , 61(1):194-222, 2023.

G. Ceruti, C. Lubich, and H. Walach. Time integration of tree tensor networks. SIAM Journal on Numerical Analysis , 59(1):289-313, 2021.

M. Cisse, P. Bojanowski, E. Grave, Y. Dauphin, and N. Usunier. Parseval networks: Improving robustness to adversarial examples. In International Conference on Learning Representations (ICLR) , 2017.

M. Courbariaux, Y. Bengio, and J.-P. David. Training deep neural networks with low precision multiplications. In Workshop contribution at International Conference on Learning Representations , 2015.

M. Courbariaux, I. Hubara, D. Soudry, R. El-Yaniv, and Y. Bengio. Binarized neural networks: Training deep neural networks with weights and activations constrained to+ 1 or-1. arXiv:1602.02830 , 2016.

L. De Lathauwer, B. De Moor, and J. Vandewalle. A multilinear singular value decomposition. SIAM journal on Matrix Analysis and Applications , 21(4):1253-1278, 2000.

L. De Lathauwer, B. De Moor, and J. Vandewalle. On the best rank-1 and rank-(r 1, r 2,..., rn) approximation of higher-order tensors. SIAM journal on Matrix Analysis and Applications , 21(4):1324-1342, 2000.

M. Denil, B. Shakibi, L. Dinh, M. Ranzato, and N. de Freitas. Predicting parameters in deep learning, 2014.

F. Feppon and P. F. Lermusiaux. A geometric approach to dynamical model order reduction. SIAM Journal on Matrix Analysis and Applications , 39(1):510-538, 2018.

J. Frankle and M. Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations , 2019.

M. Gabor and R. Zdunek. Compressing convolutional neural networks with hierarchical tucker-2 decomposition. Applied Soft Computing , 132:109856, 2023.

T. Galanti, Z. S. Siegel, A. Gupte, and T. Poggio. Sgd and weight decay provably induce a low-rank bias in neural networks. 2022.

M. Ghadiri, M. Fahrbach, G. Fu, and V. Mirrokni. Approximately optimal core shapes for tensor decompositions. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 11237-11254. PMLR, 23-29 Jul 2023.

A. Gholami, S. Kim, Z. Dong, Z. Yao, M. W. Mahoney, and K. Keutzer. A survey of quantization methods for efficient neural network inference, 2021.

Y. Gong, L. Liu, M. Yang, and L. Bourdev. Compressing deep convolutional networks using vector quantization. In International Conference on Learning Representations (ICLR) , 2015.

Y. Guo, A. Yao, and Y. Chen. Dynamic network surgery for efficient dnns. Advances in neural information processing systems , 29, 2016.

S. Gupta, A. Agrawal, K. Gopalakrishnan, and P. Narayanan. Deep learning with limited numerical precision. In International conference on machine learning , pages 1737-1746. PMLR, 2015.

S. Han, J. Pool, J. Tran, and W. Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems , 28, 2015.

P. He, J. Gao, and W. Chen. Debertav3: Improving deberta using electra-style pre-training with gradient-disentangled embedding sharing, 2023.

Y. He, X. Zhang, and J. Sun. Channel pruning for accelerating very deep neural networks. In Proceedings of the IEEE international conference on computer vision , pages 1389-1397, 2017.

A. Hnatiuk, J. Kusch, L. Kusch, N. R. Gauger, and A. Walther. Stochastic aspects of dynamical low-rank approximation in the context of machine learning. Optimization Online , 2024.

E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models, 2021.

E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations , 2022.

Y. Idelbayev and M. A. Carreira-Perpinan. Low-rank compression of neural nets: Learning the rank of each layer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , June 2020.

Y. Ioannou, D. Robertson, J. Shotton, R. Cipolla, and A. Criminisi. Training cnns with low-rank filters for efficient image classification, 2016.

M. Khodak, N. Tenenholtz, L. Mackey, and N. Fusi. Initialization and regularization of factorized neural layers. In International Conference on Learning Representations , 2021.

E. Kieri, C. Lubich, and H. Walach. Discretized dynamical low-rank approximation in the presence of small singular values. SIAM Journal on Numerical Analysis , 54(2):1020-1038, 2016.

Y.-D. Kim, E. Park, S. Yoo, T. Choi, L. Yang, and D. Shin. Compression of deep convolutional neural networks for fast and low power mobile applications. In International Conference on Learning Representations (ICLR) , 2016.

O. Koch and C. Lubich. Dynamical low-rank approximation. SIAM Journal on Matrix Analysis and Applications , 29(2):434-454, 2007.

O. Koch and C. Lubich. Dynamical tensor approximation. SIAM , 31, 2010.

T. G. Kolda and B. W. Bader. Tensor decompositions and applications. SIAM review , 51(3):455500, 2009.

J. Kossaifi, A. Bulat, G. Tzimiropoulos, and M. Pantic. T-net: Parametrizing fully convolutional nets with a single high-order tensor. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 7822-7831, 2019.

V. Lebedev, Y. Ganin, M. Rakhuba, I. Oseledets, and V. Lempitsky. Speeding-up convolutional neural networks using fine-tuned CP-decomposition. In International Conference on Learning Representations , 2015.

V. Lebedev, Y. Ganin, M. Rakhuba, I. Oseledets, and V. Lempitsky. Speeding-up convolutional neural networks using fine-tuned CP-decomposition. In International Conference on Learning Representations (ICLR) , 2015.

N. Lee, T. Ajanthan, and P. H. S. Torr. Snip: Single-shot network pruning based on connection sensitivity, 2019.

M. Lezcano-Casado and D. Martínez-Rubio. Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 3794-3803. PMLR, 09-15 Jun 2019.

C. Li and C. J. R. Shi. Constrained optimization based low-rank approximation of deep neural networks. In Proceedings of the European Conference on Computer Vision (ECCV) , September 2018.

J. Li, F. Li, and S. Todorovic. Efficient riemannian optimization on the stiefel manifold via the cayley transform. In International Conference on Learning Representations , 2020.

Y. Li, S. Gu, L. V. Gool, and R. Timofte. Learning filter basis for convolutional neural network compression. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , October 2019.

V. Lialin, N. Shivagunde, S. Muckatira, and A. Rumshisky. Relora: High-rank training through low-rank updates, 2023.

B. Liu, M. Wang, H. Foroosh, M. Tappen, and M. Pensky. Sparse convolutional neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , June 2015.

C. Lubich. Time integration in the multiconfiguration time-dependent hartree method of molecular quantum dynamics. Applied Mathematics Research eXpress , 2015(2):311-328, 2015.

C. Lubich and I. V. Oseledets. A projector-splitting integrator for dynamical low-rank approximation. BIT Numerical Mathematics , 54(1):171-188, 2014.

C. Lubich, B. Vandereycken, and H. Walach. Time integration of rank-constrained tucker tensors. SIAM Journal on Numerical Analysis , 56(3):1273-1290, 2018.

S. Mangrulkar, S. Gugger, L. Debut, Y. Belkada, S. Paul, and B. Bossan. Peft: State-of-the-art parameter-efficient fine-tuning methods. https://github.com/huggingface/peft , 2022.

Z. Mariet and S. Sra. Diversity networks: Neural network compression using determinantal point processes. In International Conference on Learning Representations (ICLR) , 2016.

E. Massart. Orthogonal regularizers in deep learning: how to handle rectangular matrices? In 2022 26th International Conference on Pattern Recognition (ICPR) , pages 1294-1299, 2022.

B. Mishra, G. Meyer, S. Bonnabel, and R. Sepulchre. Fixed-rank matrix factorizations and riemannian low-rank optimization, 2013.

P. Molchanov, S. Tyree, T. Karras, T. Aila, and J. Kautz. Pruning convolutional neural networks for resource efficient inference. In International Conference on Learning Representations , 2017.

M. Nagel, M. v. Baalen, T. Blankevoort, and M. Welling. Data-free quantization through weight equalization and bias correction. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 1325-1334, 2019.

S. Narang, G. Diamos, S. Sengupta, and E. Elsen. Exploring sparsity in recurrent neural networks. In International Conference on Learning Representations (ICLR) , 2017.

A. Novikov, D. Podoprikhin, A. Osokin, and D. Vetrov. Tensorizing neural networks, 2015.

A. Phan, K. Sobolev, K. Sozykin, D. Ermilov, J. Gusak, P. Tichavský, V. Glukhov, I. Oseledets, and A. Cichocki. Stable low-rank tensor decomposition for compression of convolutional neural network. In European Conference on Computer Vision , 2020.

R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models, 2021.

N. Ruiz, Y. Li, V. Jampani, Y. Pritch, M. Rubinstein, and K. Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation, 2023.

T. N. Sainath, B. Kingsbury, V. Sindhwani, E. Arisoy, and B. Ramabhadran. Low-rank matrix factorization for deep neural network training with high-dimensional output targets. In 2013 IEEE International Conference on Acoustics, Speech and Signal Processing , pages 6655-6659, 2013.

D. Savostianova, E. Zangrando, G. Ceruti, and F. Tudisco. Robust low-rank training via approximate orthonormal constraints. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

S. Schotthöfer, E. Zangrando, J. Kusch, G. Ceruti, and F. Tudisco. Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations. In Advances in Neural Information Processing Systems , 2022.

D. Scieur, V. Roulet, F. Bach, and A. d'Aspremont. Integration methods and accelerated optimization algorithms. In Advances In Neural Information Processing Systems , 2017.

P. Sharma, J. T. Ash, and D. Misra. The truth is in there: Improving reasoning in language models with layer-selective rank reduction. In International Conference on Learning Representations (ICLR) , 2024.

D. Song, P. Zhang, and F. Li. Speeding up deep convolutional neural networks based on tucker-cp decomposition. In Proceedings of the 2020 5th International Conference on Machine Learning Technologies , ICMLT 2020, page 56-61, New York, NY, USA, 2020. Association for Computing Machinery.

E. Stoudenmire and D. J. Schwab. Supervised learning with tensor networks. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016.

K. Ullrich, E. Meeds, and M. Welling. Soft weight-sharing for neural network compression. In International Conference on Learning Representations (ICLR) , 2017.

B. Vandereycken. Low-rank matrix completion by riemannian optimization. SIAM Journal on Optimization , 23(2):1214-1236, 2013.

V. Vanhoucke, A. Senior, and M. Z. Mao. Improving the speed of neural networks on cpus. In Deep Learning and Unsupervised Feature Learning Workshop, NIPS 2011 , 2011.

G. Venkatesh, E. Nurvitadhi, and D. Marr. Accelerating deep convolutional networks using low-precision and sparsity. In 2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 2861-2865. IEEE, 2017.

A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding, 2019.

C. Wang, G. Zhang, and R. Grosse. Picking winning tickets before training by preserving gradient flow. In International Conference on Learning Representations , 2020.

H. Wang, S. Agarwal, and D. Papailiopoulos. Pufferfish: Communication-efficient models at no extra cost. Proceedings of Machine Learning and Systems , 3:365-386, 2021.

Y. Wang, Y. Lin, X. Zeng, and G. Zhang. Multilora: Democratizing lora for better multi-task learning, 2023.

Z. Wang, C. Li, and X. Wang. Convolutional neural network pruning with structural redundancy reduction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 14913-14922, 2021.

G. Wanner and E. Hairer. Solving ordinary differential equations II , volume 375. Springer Berlin Heidelberg New York, 1996.

J. Wu, C. Leng, Y. Wang, Q. Hu, and J. Cheng. Quantized convolutional neural networks for mobile devices. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 4820-4828, 2016.

H. Yang, M. Tang, W. Wen, F. Yan, D. Hu, A. Li, H. Li, and Y. Chen. Learning low-rank deep neural networks via singular vector orthogonality regularization and singular value sparsification. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops , pages 678-679, 2020.

Q. Zhang, M. Chen, A. Bukharin, N. Karampatziakis, P. He, Y. Cheng, W. Chen, and T. Zhao. Adalora: Adaptive budget allocation for parameter-efficient fine-tuning, 2023.

J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian. Galore: Memory-efficient llm training by gradient low-rank projection, 2024.

A Impact statement

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

For 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 n i × r i , 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 n i × 2 r i for ˜ K i and we form the augmented 2 r 1 ×··· × 2 r 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:

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 factorization ̂ C = 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 2 r i × r ′ i with r ′ i ≤ 2 r i are then pulled back to the initial dimension of the filter by setting U i = U new i U ′ i ∈ R n i × 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 Additional experiments

C.1 Additional experiments for ResNet18 on Tiny-Imagenet

Table 3 displays the compression to test accuracy results for ResNet18 on Tiny ImageNet as a supplement to the results in section 4.1.

C.2 Additional experiments VGG16 on Cifar10

The 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. The results of Fig 3 indicate that Alg. 1 can identify this low-rank manifold during training.

C.3 Additional experiments for AlexNet on Cifar 10

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 .

̸

0

100

200

300

optimization iter

0

166

333

500

basis ranks

(a)

τ

= 0

.

1

0

100

200

300

optimization iter

0

166

333

500

basis ranks

(b)

τ

= 0

.

08

0

100

200

300

optimization iter

0

166

333

500

basis ranks

(c)

τ

= 0

.

05

For more complicated datasets and architectures like Cifar10 on VGG16 and Alexnet, results in Tab.1 and Fig.1 show that our proposal consistently gets better results than all the methods under comparison at parity of compression.

D Proof of Theorem 2.1

Theorem. Let W = C × d i =1 U i ∈ M ρ be such that (3) 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,

̸

where Ten i denotes 'tensorization along mode i ', i.e. the inverse reshaping operation of Mat i .

Proof. The proof follows the path first suggested in [52, §4], i.e., the quantity V i defined next is frozen in time: ˙ V i = 0 . A detailed derivation of the matrix-tensor equations for K i and C is beyond the scope of the present work and it is provided in [40, 53, 54, 12, 11].

We begin recalling the evolution equations for the factors of the projected gradient flow (3):

̸

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

For sake of brevity, the parameter is going to be omitted. Next, we perform a QR-factorization

Thus, we observe that the matrix Mat i ( W ) admits an SVD-like decomposition as follows

̸

where the matrix Q i possess othornormal columns, i.e., Q ⊤ i Q i = I . We introduce the quantity

By construction, we observe that

The tiny matrix S i satisfies the following differential equation

The first null identity follows from the gauge condition on U i . The second null identity follows by the initial assumption ˙ V i = 0 . Therefore, the matrix K i satisfies the differential equation

̸

̸

To conclude, we set Z = K i V ⊤ i and let W = Ten i ( K i V ⊤ i ) = Ten i ( Z ) . We obtain

where we remind that V ⊤ i V i = I . Hence

The first K i -differential equation is then obtained combining (11) and (10)

The right-hand side can be further reduced using standard tensorization formulas [41]

̸

The second differential equations follows by observing that

̸

̸

The tensor C satisfies the differential equation

where the extra terms disappear due to the imposed gauge conditions U ⊤ i ˙ U i = 0 .

E Proof of Corollary 2.2

Corollary E.1. 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,

Proof. From Eq. (9) of the proof in Theorem 2.1 it is apparent, that

Moreover, we observe that using the chain rule,

Then, with (13) we have

Using full-rankness of S i concludes the proof.

F Proofs of descent and approximation theorems

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 that d dt L ( ̂ W ( t )) = ⟨∇L ( ̂ W ( t )) , ˙ ̂ W ( t ) ⟩ = ⟨∇L ( ̂ W ( t )) , ˙ ̂ C ( t ) × i ̂ U 1 i ⟩ = ⟨∇L ( ̂ W ( t )) × i ̂ U 1 , ⊤ i , ˙ ̂ C ( t ) ⟩ = ⟨∇L ( ̂ W ( t )) × i ̂ U 1 , ⊤ i , -∇L ( ̂ W ( t )) × i ̂ U 1 , ⊤ i ⟩ = -∥∇L ( ̂ W ( t )) × i ̂ U 1 , ⊤ i ∥ 2 . If we define α := min 0 ≤ τ ≤ 1 ∥∇L ( ̂ W ( τλ ) ) × i ̂ U 1 , ⊤ i ∥ 2 , it follows that for t ∈ [0 , λ ]

Integrating (14) from t = 0 until t = λ , we obtain

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

.

Lemma F.1. The following estimate holds

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

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

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,

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

It suffices to study the latter term. For t ∈ [0 , λ ] , we define the quantity

It satisfies the differential initial value problem

The term W ( t ) can be written as a perturbation of ˜ C

where

Then, we observe that

The remainder can be estimated as follows

Furthermore, the full gradient can be re-written as

where the defect D ( t ) is defined as

Because of the Lipschitiz assumption, we have that

Next, we compare the two differential equations

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 is ˜ C ( λ ) = W ( λ ) × j U ⊤ j . Therefore, the study of the local error follows by a direct application of the Gronwall inequality

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

G Proof of stochastic convergence

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.

The spectral distribution stabilizes fast enough over time, i.e.

The projected stochastic gradient has a controlled drift, namely

where P U = UU T is the orthogonal projection onto the range of U , and E [ ·| t ] = E [ ·| W ( t ) , { U i ( t ) } d i =1 ] denotes the conditional expectation. Then the following convergence condition holds

The 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 holds

Proof. By using Cauchy-Schwartz inequality, we have that

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 that

̸

Proof. We introduce first the function

̸

Let m be the the infimum on the set U as defined below

̸

Because the term ∇L ( W ) can be decomposed into a sum of ∇L B ( W ) × j = j ∗ P U j ∗ , 0 × j ∗ P U j, 1 and its orthogonal component, we observe that the infimum satisfies

̸

Hence, by taking the expectation on both sides, we can conclude that

̸

̸

The 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 by E 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 that

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 have

̸

Hence

where P U = UU T is the orthogonal projector onto the range of arbitrary matrix U . Applying G.2 we have that

Wenotice that for the augmented basis ˜ U j ( t ) = [ U j ( t -1) | U j ( t )] , it holds P ˜ U j ( t ) = P U j ( t -1) + P U j ( t ) . When we expand L B ( W ( t -1)) × j P ˜ U j ( t ) using its sum representation and apply Lemma G.3 to the mixed terms, we obtain

The locally bounded loss computed on the truncated approximation W ( t ) is bounded via

By combining equations (17) and (18), we arrive at the following bound

Following assumption (16), we have

where we assume that λ t ≤ 2 /L 2 ν . Finally, by taking the expectation on both sides, we obtain

By summing the last equation over t = 1 , . . . , T we get

By rearranging the terms of the latter equation and by sending T → + ∞ , we finally obtain

The conclusion follows by the Robbins-Monro conditions, i.e.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: See statements in abstract and introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All assumptions are stated in the corresponding theorems, and their proofs are linked and stated in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The numerical examples are fully described in Section 4.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code is given in the supplementary material.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The numerical examples, the neural network training parameters for all computer vision tests, and the data-augmentation for the used datasets are described in are described in Section 4.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The test case statistics are explained in the corresponding sections in Section 4.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide an overview of the used compute resources in Section 4

Guidelines:

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

9. 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 ?

Answer: [Yes]

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

10. 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).

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

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We neither crowdsource experiments nor perform researchw ith human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

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

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

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

QUESTIONS:
What is a standard Tucker decomposition like? Could you provide more details for the comparison of the proposed method and the vanilla standard Tucker decomposition, especially over the space/time complexity comparison, and their convergency analysis?

What are the compressed parameters in the neural networks? Are the parameters merely the convolutional kernel? Can the method be applied on the weights of linear layers?

What does ``geometry-aware’’ indicates? What kind of geometry is the method aware of? Could you give more explanation?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The paper is clearly written. Most parts are supported by sufficient technical details.\n- The proposed method is an extended version of DLRT. Yet, some unique challenges appear in the Tucker tensor format have addressed with interesting approaches, e.g., the reparameterization method and Corollary 2.2.\n\nWEAKNESSES:\nI have some minor concerns on writing and the experimental result\n\n- More detailed information about one-step integration methods could be provided for the readers who are not familiar with the concept.\n- Some sections assume a certain degree of background knowledge on DLRT, e.g., gauge conditions in Line 612.\n- 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.\n\nQUESTIONS:\n- How does the one-step integration work? Does it not add significant computational overhead?\n- 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?\n- 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?\n- 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.\n\n[1] Hu, Edward J., et al. \"Lora: Low-rank adaptation of large language models.\" arXiv preprint arXiv:2106.09685 (2021).\n\nLIMITATIONS:\n- 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.\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- A strong motivation as there is a clear need for further efficiency improvements\n- Introduces a novel rank-adaptive geometry-aware training method that dynamically updates ranks during training\n- Proposed to overcome the sensitivity to parameter initialisation and the need for a full-model warm-up phase\n- Thorough theoretical analysis, including guarantees on loss descent, convergence, and approximation\n\nWEAKNESSES:\n- 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.\n- Not very clear in many parts. In particular, Section 2.1 is difficult to follow.\n- 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. \n- 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.\n- 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.\n\nQUESTIONS:\n- 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?\n- Can the authors include data points for the proposed method within the full compression range for a direct comparison with other methods?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

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

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

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

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

LIMITATIONS:
Yes, the authors discuss about the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nReducing 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.\n\nIt 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.\n\nThe 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.\n\nThe paper is generally well-written and provides an appropriate method for a very important problem.\n\nI 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.\n\n[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.\n\nSTRENGTHS:\nThe paper is generally well-written and provides an appropriate method for a very important problem.\n\nWEAKNESSES:\n-\n\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nI don't see any direct potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-08-01"}