Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations
Anonymous Author(s)
AffiliationAddress email
Abstract
Neural networks have achieved tremendous success in a large variety of applications.
However, their memory footprint and computational demand can renderthem impractical in application settings with limited hardware or energy resources.
In this work, we propose a novel algorithm to find efficient low-rank subnetworks.
Remarkably, these subnetworks are determined and adapted already during thetraining phase and the overall time and memory resources required by both trainingand evaluating them are significantly reduced.
The main idea is to restrict theweight matrices to a low-rank manifold and to update the low-rank factors ratherthan the full matrix during training.
To derive training updates that are restrictedto the prescribed manifold, we employ techniques from dynamic model order reduction for matrix differential equations.
Moreover, our method automatically anddynamically adapts the ranks during training to achieve the desired approximationaccuracy.
The efficiency of the proposed method is demonstrated through a varietyof numerical experiments on fully-connected and convolutional networks.
Introduction
While showing great performance in terms of classification records, most state-of-the-art neuralnetworks require an enormous amount of computation and memory storage both for the training andthe evaluation phases [26].
These requirements not only increase infrastructure costs and energyconsumption, but also make the deployment of artificial neural networks to infrastructures withlimited resources such as mobile phones or smart devices prohibitive.
On the other hand, it is wellknown that networks’ weights contain structures and redundancies that can be exploited for reducingthe parameter space dimension without significantly affecting the overall accuracy [4, 10, 18].
Network pruning is a popular line of research that addresses this problem by removing redundantparameters from pre-trained models.
Typically, the initial network is large and accurate, and thegoal is to produce a smaller network with similar accuracy.
Methods within this area include weightsparsification [21, 24, 46] and quantization [11, 57], with different pruning techniques, includingsearch-based heuristics [24], reinforcement learning [3, 23] and genetic algorithms [41].
More recentwork has considered pruning during training, by formulating pruning as a data-driven optimizationproblem [21, 25, 26].
The resulting “dynamical pruning” boils down to a parameter-constrainedtraining phase which, however, has been mostly focused on requiring sparse or binary weights so far.
Rather than enforcing sparsity or binary variables, in this work we constrain the parameter space tothe manifold of low-rank matrices.
Neural networks’ parameter matrices and large data matricesin general are seldom full rank [16, 45, 50, 53].
Constraining these parameters to lie on a manifolddefined by low-rank matrices is thus a quite natural approach.
By interpreting the training problem asa continuous-time gradient flow, we propose a training algorithm based on the extension of recent
Dynamical Low-Rank Approximation (DLRA) algorithms [5, 6, 7].
This approach allows us to use low-rank numerical integrators for matrix Ordinary Differential Equations (ODEs) to obtainmodified forward and backward training phases that only use the small-rank factors in the low-rankrepresentation of the parameter matrices.
We perform several experimental evaluations on fully-connected and convolutional networks showingthat the resulting dynamical low-rank training paradigm yields low-parametric neural networkarchitectures which compared to their full-rank counterparts are both remarkably less-demandingin terms of memory storage and require much less computational cost to be trained.
Moreover, thetrained low-rank neural networks achieve comparable accuracy to the original full architecture.
Thisobservation is reminiscent of the so-called lottery tickets hypothesis — dense neural networks containsparse subnetworks that achieve high accuracy [18] — and suggests the presence of low-rank winningtickets : highly-performing low-rank subnetworks of dense networks.
Remarkably, our dynamicallow-rank training strategy seems to be able to find the low-rank winning tickets directly during thetraining phase independent of initialization.
Related work on low-rank methods
Low-rank factorization using the SVD and other matrix decomposition techniques have been extensively studied in the scientific computing and machine learning communities.
The challengeof compressing and speeding up large-scale neural networks using low-rank methods has sparkedwide-spread research interest in recent years and significant effort has been put towards developinglow-rank factorization strategies for deep neural networks.
Previous works can roughly be categorized in approaches with fixed low rank and variable low rankduring training time.
Fixed rank approaches decompose weight matrices using SVD or tensor decompositions of pre-trained networks and fine-tune the factorized network [13, 37, 48, 52], constrainweight matrices to have a fixed low rank during training [29], or create layers as a linear combinationof layers of different rank [28].
Hence, these methods introduce the rank of the matrix decompositionas another hyperparameter to be fine-tuned.
Rank-adaptive methods mitigate this issue by automaticdetermination and adaption of the low-rank structure after training.
In particular, [32, 33] applyheuristics to determine the rank of the matrix decomposition ahead of time, whereas [56] encourageslow-rank weight matrices by means of a penalized loss that depends on approximated matrix ranks.
Few methods have been proposed recently that adapt the ranks of the weight matrix alongside themain network training phase.
In [39], the authors set up the neural network training as a constrainedoptimization problem with an upper bound on the ranks of the weights, which is solved in analternating approach resulting in an NP-hard mixed integer program.
The authors of [27] formulatea similar constrained optimization problem resulting in a mixed discrete-continuous optimizationscheme which jointly addresses the ranks and the elements of the matrices.
However, both theseapproaches require knowledge of the full weight matrix (and of its singular value decomposition)during training and overall are more computational demanding than standard training.
In this work we overcome the above issues and propose a training algorithm with reduced memoryand computational requirements.
To this end, we reinterpret the optimization problem of a neuralnetwork as a gradient flow of the network weight matrices and thus as a matrix ODE.
This continuousformulation allows us to use recent advances in DLRA methods for matrix ODEs which aim atevolving the solution of the differential equation on a low-rank manifold.
The main idea of DLRA[34], which originates from the Dirac-Frenkel variational principle [15, 19], is to approximate thesolution through a low-rank factorization and derive evolution equations for the individual factors.
Thereby, the full solution does not need to be stored and the computational costs can be significantlyreduced.
To ensure robustness of the method, stable integrators have been proposed in [42] and [7].
Instead of evolving individual low-rank factors in time, these methods evolve products of low-rankfactors, which yields remarkable stability and exactness properties [31], both in the matrix and thetensor settings [9, 35, 43, 44].
In this work, we employ the “unconventional” basis update & Galerkinstep integrator [7] as well as its rank-adaptive extension [5], see also [8, 36].
The rank-adaptiveunconventional integrator chooses the approximation ranks according to the continuous-time trainingdynamics and allows us to find highly
Low-rank training via gradient flow
Consider a feed-forward fully-connected neural network N ( x ) =
z M , with z 0 = x ∈ R n 0 ,z k
=
σ k
( W k z k − 1 + b k ) ∈ R n k , k = 1 , . .
, M (the convolutional setting is discussed in thesupplementary material §6.8).
We consider the training of N based on the optimization of a lossfunction L ( W 1 , . . .
, W M ; N ( x ) , y ) by means of a gradient-based descent algorithm.
For example,when using gradient descent, the weights of N at iteration t ∈ N are updated via
[Equation]
with a learning rate λ .
When the weight matrices W k are full, both the forward and gradient evaluations of the network require a large number of full matrix multiplications, with high computationalexpense and large memory footprint.
This renders the training and the use of large-scale neuralnetworks a difficult challenge on limited-resource devices.
At the same time, a wealth of evidenceshows that dense networks are typically overparameterized and that most of the weights learned thisway are unnecessary [16, 45].
In order to reduce the memory and computation costs of training, wepropose a method that performs the minimization over the manifold of low-rank matrices.
To this end, we assume that the ideal W k can be well-approximated by a matrix of rank r k ≪ n k , n k +1of the form U k S k V ⊤ k ∈ R n k × n k − 1 , where U k ∈ R n k × r k , V k ∈ R n k − 1 ×
r k are thin and tall matriceshaving orthonormal columns spanning optimal subspaces which capture essential properties ofparameters, and S k ∈ R
r k × r k is a tiny full-rank matrix that allows us to extrapolate the usefulinformation from the learned subspaces U k and V k .
Traditional descent algorithms such as (1) do not guarantee the preservation of the low-rank structure
U k S k V ⊤ k when updating the weights during training and require knowledge of the whole W k ratherthan the factors U k , S k , V k .
To overcome these limitations, we move to a gradient flow formulationof the loss-minimization problem, and we derive a robust and efficient training method that preservesthe low-rank structure and works only with the factors U k , S k , V k , rather than their product.
Minimizing the loss function with respect to the k -th weight matrix is equivalent to evaluating thelong time behaviour of the matrix ODE
[Equation]
where the “dot” denotes the time-derivative.
Thus, using (2) we can interpret the training phase asa continuous evolution of the weights, rather than a discrete update.
Let M r k denote the manifoldof matrices with rank r k and assume at a certain time t 0 the weights are in the manifold, i.e.,
W k ( t 0 ) ∈
M r k .
Using this continuous-time interpretation allows us to derive a strategy to evolve theweights according to the dynamics in (2) so that W k ( t ) ∈ M r k for all t ≥ t 0 .
To this end, in §3.1 weexploit the fact that W k ( t ) admits a time-dependent factorization [14] W k ( t ) =
U k
( t ) S k ( t ) V k ( t ) ⊤to rewrite (2) as a system of matrix ODEs for each of the individual factors U k , S k and V k .
Then, in§4 we propose an algorithm to efficiently integrate the system of ODEs.
We show in §4.1 that suchalgorithm reduces the loss monotonically and is accurate, in the sense that U k S k V ⊤ k ≈ W k , i.e. thelearned low-dimensional subspaces U k
and V k well-match the behaviour of the full-rank network
W k solution of (2), through the action of the learned S k .
Remarkably, using the dynamics of theindividual factors will also allow us to adaptively adjust the rank r k throughout the continuous-timetraining process.
3.1 Coupled dynamics of the low-rank factors via DLRA
We consider the dynamical system of a single weight matrix W k , while the remaining weight matricesare fixed in time and are treated as parameters for the gradient.
In the following, we omit writing thesesparameters down for efficiency of exposition.
Assuming
W k ( t ) ∈
M r k , we can formulate (2) as
[Equation]
where T W k ( t ) M r k is the tangent space of M r k at position W k ( t ) .
In order to solve (3), we furtherobserve that (3) can be equivalently formulated as the following Galerkin condition [34]:
[Equation]
From W k
= U k S k V ⊤ k , a generic element δW k of the tangent space T W k
( t ) M r k can be written as
[Equation]
where δU k and δV k are generic elements of the tangent space of the Stiefel manifold with r korthonormal columns at the points U k and V k , respectively, and δS k is a generic r k × r k matrix, seee.g.
[34, §2] for details.
Additionally, the Gauge conditions U ⊤ k δU k = 0 and V ⊤ k δV k = 0 must beimposed to ensure orthogonality of the basis matrices, and the uniqueness of the representation of thetangent space elements.
Similarly, by the chain rule applied several times we have
[Equation]
Now, the Galerkin condition (4) becomes
[Equation]
with U ⊤ k ˙
U k
=
and V ⊤ k ˙ V k = 0 .
If we choose δW k = U k δS k V ⊤ k in (5), we obtain
[Equation]
Thus, using the Gauge conditions, we obtain ⟨ ˙ S k + U ⊤ k ∇ W k L ( W k ( t ))
V k , δS k ⟩ = 0 , which has tohold for a generic r k × r k matrix δS k .
We obtain this way an evolution equation for the S k ( t ) factor.
Similarly, specifying (5) for the two choices δW k = δU k S k V ⊤ k and δW k = U k S k δV ⊤ k , allows us toobtain the following system of differential equations for the individual factors of W k :
[Equation]
KLS-based training algorithm
In order to perform an efficient and robust rank-constrained training step, we numerically integrate thesystem of ODEs (6).
Our approach is based on the “unconventional KLS integrator” [7] and its rankadaptive version [5].
The pseudocode of the proposed training strategy is presented in Algorithm 1.
The main idea of the KLS algorithm is to alternately represent the product W k
= U k S k V ⊤ k as
W k = K k V ⊤ k and W k
= U k L ⊤ k , consider the corresponding coupled ODEs from (6), and thenperform three main steps:1,2.
K&L-steps (in parallel).
Update the current K k and L k by integrating the differential equations
[Equation]
from t = 0 to t =
η ; then form new orthonormal basis matrices (cid:101)
U k
and (cid:101)
V k spanning the range ofthe computed K k ( η ) and L k ( η ) .
S-step .
Update the current S k by integrating the differential equation
[Equation]
from t = 0 to t = η , with initial value condition S k (0) = (cid:101)
U ⊤ k
U k S k V ⊤ k (cid:101)
V k .
An important feature of this algorithm is that it can be extended to rank adaptivity in a relativelystraightforward manner [5], letting us dynamically evolve the rank of S k (and thus the rank of W k )during the computation.
This is particularly useful, as we may expect the weight matrices to have lowranks but we may not know what the “best” ranks for each layer are.
Typically, dynamically adaptingthe ranks of a low-rank optimization scheme is a challenging problem as moving from the manifold
M r k to M r k ± 1 introduces singular points [2, 20].
Instead, treating the training problem as a systemof matrix differential equations allows us to overcome this issue with a simple trick: at each step ofthe KLS integrator we double the dimension of the basis matrices (cid:101) U k and (cid:101) V k computed in the K- and
L-steps by computing orthonormal bases spanning [ K k ( η ) | U k ] and [ L k ( η ) | V k ] , respectively, i.e.by augmenting the current basis with the basis computed in the previous time step.
Then, after thenew S k matrix is computed via the S-step, a truncation step is performed, removing from the newlycomputed S k matrix all the singular values that are under a certain threshold ϑ .
Of course, adding the rank-adaptivity to the integrator comes at a cost.
In that case, each step requiresto perform an SVD decomposition of twice the size of the current rank of S k in order to be able tothreshold the singular values.
Moreover, the dimension of the bases U k and V k may grow, whichalso may require additional computational effort.
However, if the ranks remain small throughout thedynamics, this computational overhead is negligible, as we will further discuss in §4.3 and §5.
4.1 Error analysis and convergence
In this section we present our main theoretical results, showing that (a) the low-rank matrices
U k S k V ⊤ k formed by the weights’ factors computed with Alg.
1 are close to the true solution of (2),and (b) that the loss function decreases during DLRT, provided the singular value threshold ϑ is nottoo large, i.e., is bounded by a constant times the square of the time-step size η
(see Theorem 1).
Inthe version we present here, part of the statements are presented in an informal way for the sake ofbrevity.
We refer to the supplementary material §6.1 for details and for the proofs.
Assume the gradient flow F k ( Z ) =
−∇ W k L ( W 1 , . . .
, Z, . .
, W M , N ( x ) , y ) in (2) is locallybounded and locally Lipschitz continuous, with constants C 1 and C 2 , respectively.
Then,
Theorem 1.
Fixed x and y , let W k ( t ) be the (full-rank) continuous-time solution of (2) and let
U k , S k , V k be the factors computed with Algorithm 1 after t steps.
Assume that the K,L,S steps (7)and (8) are integrated exactly from 0 to η .
Assume moreover that, for any Z ∈ M r k sufficiently closeto W k ( tη ) , the whole gradient flow F k ( Z ) is “ ε -close” to M r k .
Then,
[Equation]
where the constants c 1 , c 2 and c 3 depend only on C 1 and C 2 .
In particular, the approximation bounddoes not depend on the singular values of the exact nor the approximate solution.
Observe that, while the loss function L decreases monotonically along any continuous-time solution
W k ( t ) of (2), it is not obvious that the loss decreases when the integration is done onto the low-rankmanifold via Algorithm 1.
The next result shows that this is indeed the case, up to terms of the orderof the truncation tolerance ϑ .
More precisely, we have
Theorem 2.
Let W tk =
U tk S tk
( V tk ) ⊤ be the low rank weight matrix computed at step t of Algorithm 1and let L ( t ) =
L ( W t 1 , . .
, W tM , N ( x ) , y ) .
Then, for a small enough time-step η we have
[Equation]
where α and β are positive constants that do not depend on t , η and ϑ .
4.2 Efficient implementation of the gradients
All the three K,L,S-steps require the evaluation of the gradient flow of the loss function with respectto the whole weight matrix W k .
Different approaches to compute this gradient in an efficient way maybe used.
The strategy we discuss below aims at reducing memory requirements and computationalcosts by avoiding the computation of the full gradient ∇ W k L and instead working with the gradientwith respect to the low-rank factors.
To this end, we note that for the K-step it holds ∇ W k L ( K k ( t ) V ⊤ k ) V k = ∇ K k L ( K k ( t ) V ⊤ k ) .
Hence,the whole gradient can be computed through a forward run of the network with respect to K k
[Equation]
and taping the gradient with respect to K k .
In this way, the full gradient does not need to be computedand the overall computational costs are comprised of running a forward evaluation while tapinggradients with respect to K k , analogously to the traditional back-propagation algorithm.
The L- and
S-steps can be evaluated efficiently in the same manner, by evaluating the network while taping thegradients with respect to L k and S k , respectively.
Hence, instead of a single gradient tape (or chainrule evaluation) of the full weight matrix network, we have three gradient tapes, one for each lowrank step, whose combined computational footprint is less than the full matrix tape.
We providedetailed formulas for all the three gradient tapes in the supplementary material §6.7.
K tk ← U tk S tk /* K-step */
K t +1 k ← one-step-integrate (cid:8) ˙
K ( t ) =
−∇ K L ( K ( t )( V tk ) ⊤ z k − 1 + b tk ) , K (0)
=
K tk (cid:9)
L tk ← V tk ( S tk ) ⊤ /
* L-step */
L t +1 k ← one-step-integrate (cid:8) ˙
L ( t ) =
−∇ L L ( U tk L ( t ) ⊤ z k − 1 + b tk ) , L (0)
=
L tk (cid:9) adaptive then /
* Basis augmentation step */
K t +1 k ←
[ K t +1 k | U tk ]
L t +1 k ←
[ L t +1 k | V tk ]
U t +1 k ← the
K
+1 k /* S-step */
M k ← ( U t +1 k )
⊤ U tk
V t +1 k ← basis for the range of L t +1 k
N k ← ( t
1 k )
V tk (cid:101) S
tk ← M k S tk N k
U t +1 k S ( t )(
[Table]
t
[Table]
⊤
z k − 1 + b tk (cid:1) , S (0)= (cid:101) S tk (cid:9)
[Table]
if adaptive
then /
* Rank compression step */
[Table] [Table]
truncate
[Table] [Table]
columns of P ]
[Table]
columns of Q ] /* Bias update step */
[Table]
L ( U t +1 k
S t +1 k
( V t +1 k )
⊤ z k − 1 + b ( t )) , b (0)= b tk (cid:9)
4.3 Implementation details, computational costs and limitations
Each step of Alg.
1 requires the computation of two orthonormal bases for the ranges of K t +1 kand
L t +1 k .
There are of course different techniques to compute such orthonormal matrices.
In ourimplementation we use the QR algorithm, which is known to be one of the most efficient and stableapproaches for this purpose.
In the adaptive strategy the singular values of S t +1 k are truncatedaccording to a parameter ϑ .
To this end, in our implementation, we use the Frobenius norm of Σ .
Precisely, we truncate Σ = diag( σ i ) at step 19 of Alg. 1 by selecting the smallest principal r × rsubmatrix such that ( (cid:80) i ≥ r +1 σ 2 i ) 1 / 2 ≤ ϑ .
Finally, one-step-integrate denotes a numerical procedurethat integrates the corresponding ODE from time t = 0 to t =
η .
In practice one can employ differentnumerical integrators, without affecting the ability of the algorithm to reduce the loss function (see[5, Thm. 5]) while maintaining the low-rank structure.
In our implementation we used two methods:
Explicit Euler.
This method applied to the gradient flow coincides with one step of StochasticGradient Descent (SGD), applied to the three factors K k , L k , S k independently.
Adam.
Here we formally compute the new factors by modifying the explicit Euler step as in the
Adam optimization method.
Note that Nesterov accelerated SGD is known to coincide with aparticular linear multistep ODE integrator [49].
While Adam does not directly correspond to a
numerical integrator to our knowledge, in our tests it resulted in a faster decrease of the loss than
both Euler (SGD) and Nesterov accelerated SGD.
For both choices, the target time step η corresponds to the value of the learning rate, which we set to0 .
2 for Euler.
For Adam, we use the default dynamical update, setting 0 .
001 as starting value.
Computational cost.
To obtain minimal computational costs and memory requirements for the K-step,the ordering of evaluating K k V ⊤ k z k − 1 in (9) is important.
First, we compute (cid:101) z :=
V ⊤ k z k − 1 ∈ R r kwhich requires O ( r k n k − 1 ) operations.
Second, we compute K k (cid:101) z which requires O ( r k n k ) operations.
Adding the bias term and evaluating the activation function requires O ( n k ) operations.
Hence,combined over all layers we have asymptotic cost of O ( (cid:80) k r k ( n k + n k +1 )) .
Taping the forwardevaluation to compute the gradient with respect to K k as discussed in §4.2 does not affect theasymptotic costs, i.e. the costs of computing the K-step at layer k assuming a single data point xrequires C K ≲ (cid:80)
k r k ( n k + n k +1 ) operations.
In a similar manner, we obtain the computational costsof the L- and S-steps, which are again C L,S ≲ (cid:80) k r k ( n k + n k +1 ) .
Moreover, the QR decompositionsused in the K- and L-step require O (cid:0)(cid:80) k r 2 k ( n k + n k − 1 ) (cid:1) operations and computing the SVD inthe truncation step has worst-case cost of O (cid:0)(cid:80) k r 3 k (cid:1) .
Hence, assuming r k ≪ n k , n k +1 , the costper step of our low-rank method is C DLRA ≲ (cid:80) k r 2 k ( n k + n k − 1 ) , opposed to the dense networktraining, which requires C dense ≲ (cid:80) k n k n k +1 operations.
In terms of memory cost, note that weonly need to store r k (1 + n k + n k +1 ) parameters per layer during the algorithm, corresponding tothe matrices S tk , U tk , V tk .
Moreover, at the end of the training we can further compress memory bystoring the product of the trained weight factors U k S k , rather than the individual matrices.
Limitations.
A requirement for DLRT’s efficiency is that r k ≪ n k , n k +1 .
When the truncationthreshold ϑ is too small, Alg.
1 does not provide advantages with respect to standard training.
This isalso shown by Fig.
Moreover, in the terminology of [51], DLRT is designed to reduce trainingcosts corresponding to model parameters and to the optimizer.
To additionally decrease activationcosts, DLRT can be combined with micro-batching or checkpointing approaches.
Finally, the choiceof ϑ introduces one additional hyperparameter which at the moment requires external knowledgefor tuning.
However, our experiments in §5 show that relatively large values of ϑ yield competingperformance as compared to a number of baselines, including standard training.
Numerical Results
We illustrate the performance of DLRT Algorithm 1 on several test cases.
The code is implemented inboth Tensorflow [1] and PyTorch [47], and it is available in the supplementary material.
The networksare trained on an AMD Ryzen 9 3950 X CPU and a Nvidia RTX 3090 GPU.
Timings are measuredon pure CPU execution.
We inspect and compare the performance of the dynamical low-rank trainingapproach on the MNIST data-set [12], which we partition in randomly sampled train-validation-testsets of size 50 K - 10 K - 10 K : The images are pixelwise normalized and no further data augmentationor weight regularization has been used.
The MNIST dataset is made available under the terms of the
Creative Commons Attribution-Share Alike 3.0 license.
5.1 Fixed-rank fully-connected feed-forward network timings
First, we compare the training time of the adaptive DLRT Algorithm 1 on a 5 -layer fully-connected[5120 , 5120 , 5120 , 5120 , 10] network fixing the ranks of layers 1-4, i.e. choosing a specific startingrank r 0 k for the input weight factors and truncating Σ at line 19 of Alg.
1 to the principal r 0 k × r 0 ksubmatrix, rather than via a threshold.
Next, we measure the average prediction time of the whole
MNIST dataset over 1000 runs.
Figures 1(a) and 1(b) show that both timings scale linearly with therank of the factorizations, and that for sufficiently small ranks the dynamical low-rank approach isfaster than the full baseline both in terms of training and of prediction.
5.2 Rank evolution and performance of DLRT for different singular value thresholds
Next, we demonstrate the capabilities of DLRT to determine the rank of the network’s weight matricesautomatically during the network training using Algorithm 1.
The Adam optimizer with defaultlearning rate is used for the gradient update.
We train fully connected 5 -layer networks, of which thefirst 4 are replaced by low-rank layers in the subsequent tests.
The activation function is chosen to be
ReLU for the hidden layers, and softmax for the output layer.
The training loss is sparse categoricalcross entropy and we additionally measure the model’s accuracy.
We use batch size 256 and trainfor 250 epochs.
We choose ϑ = τ ∥ Σ ∥ , thus we truncate the singular values of the current S tk by afraction τ of the total Frobenius norm.
The smaller τ , the more singular values are kept.
[Figure]
[Figure]
Figures 2 (a) and (b) show the evolution of the rank
adaptive layers of a 5 -layer 500 -neuron networkin a long time case study for τ = 0 .
05 and τ = 0 .
We can see that within the first epoch the initialfull matrix ranks are reduced significantly, to 27 for τ = 0 . 15 , and to ∼ 85 for τ = 0 .
respectively.
Within the first 50 epochs, the layer ranks are already close to their final ranks.
This indicates that therank adaptive algorithm is only needed for the first few training epochs, and can then be replaced bythe computationally cheaper fixed-low-rank training (by setting the Boolean variable adaptive to
False in Algorithm 1).
Figure 3 compares the mean test accuracy of 5 -layer networks with 500 and784 neurons with different levels of low-rank compression, over five independent runs with randomlysampled train-test-val sets.
The networks can be compressed via dynamical low-rank training bymore than 95% , while only losing little more than 1% test accuracy compared to the dense referencenetwork marked in red.
Remark that restricting the space of possible networks to a given rankregularizes the problem, since such a restriction can be understood as adding a PCR regularizationterm to the loss function.
This can be seen from the tendency of not overfitting and reaching improvedtest accuracies compared to the corresponding dense network for moderate compression ratios.
Alsonote that adaptive-low rank training eliminates the need for hyperparameter grid search in terms oflayer-weights, due to automatic rank adaptation.
The rank dynamics for all configurations can beseen in the supplementary material §6.3.
Finally, in the supplementary material §6.6 we compare theuse of DLRT with the vanilla approach which simply thresholds the singular values of the full-ranknetwork.
Our results show that advantageous low-rank winning tickets exist, but are not easy to find.
In fact, the vanilla low-rank subnetworks perform very poorly.
From this point of view, our approachcan be seen as an efficient dynamical pruning technique, able to determine high-performing low-ranksubnetworks in a given dense network.
Remarkably, our numerical experiments suggest that low-rankwinning tickets can be trained from the start and do not to heavily depend on the initial weight guess.
5.3 Convolutional layers: LeNet
Here we compare the proposed dynamical low-rank training scheme on LeNet5 [38] on MNIST,against the full-rank reference and several baselines.
SVD prune [58] and LRNN [27] are the closestapproaches to our DLRT: they dynamically train low-rank layers by adding a rank-penalty to the loss
[Table]
2% < 0% function, and by complementing the standard training step via an SVD projection step in the latter anda pruning step in the former.
While computing low-rank factors for each layer, thus reducing memorystorage of the network, this training approach is more expensive than training the full network.
GAL[40], SSL [59], and NISP [55] are pruning methods which aim at learning optimal sparse
weights(rather than low-rank) by adding sparsity-promoting regularization terms to the training loss.
As for
LRNN, these methods do not reduce the computational cost of the training phase (as indicated withthe < 0% in Table 1).
Analogously to [27], our adaptive low-rank training technique is applied tothe convolutional layers by flattening the tensor representing the convolutional kernel into a matrix.
Details are provided in the supplementary material §6.8.
All the models are trained for 120 epochsusing SGD with a fixed learning rate of 0 . 2 .
Results in Table 1 show that the DLRT algorithm is ableto find low-rank subnetworks with up to 96 .
4% less parameters than the full reference, while keepingthe test accuracy above 95% .
Compared to the baseline methods, we achieve better compression ratesbut observe lower accuracy.
However, unlike the baseline references, DLRT automatically prunes thesingular values during training, without requirement to solve any additional optimization problem,thus significantly improving time and memory efficiency of both forward and backward phase, withrespect to the full reference.
5.4 Results on Cifar10 and Cifar100 with AlexNet and VGG16
We compare DLRT with standard full-rank training and a number of low-rank alternative baselines onthe VGG16 and AlexNet models on the Cifar10 and Cifar100 datasets using the SGD optimizer withmomentum factor 0 .
The last layers of the networks have been adapted to match the correspondingclassification tasks.
Detailed results are reported in Tables 7 and 8 in the supplementary material.
In
Cifar10, see Table 7, we archive a train compression of 77 .
5% with an accuracy loss of just 1 .
1% for
VGG16 and 84 .
2% compression at 1 .
79% accuracy loss.
For the Cifar100 test case, see Table 8, wehave 72% train compression and a 2% gain of accuracy for VGG16, and 85% compression at 3 .
88%accuracy loss over the corresponding baseline for AlexNet.
Appendix
6.1 Proofs of the main results
We provide here a proof of Theorems 1 and 2.
The proof is based on a number of classical results aswell as recent advances in DLRA theory, including [5, 7, 31, 34, 42].
Recall that, for a fixed layer k , we reinterpret the training phase as a continuous-time evolution of theweights on the manifold of low-rank matrices, as illustrated in Fig.
4(a-b).
This boils down to solvingthe manifold-constrained matrix differential equation
[Equation]
where ∥ · ∥ is the Frobenius norm and F k denotes the gradient flow of the loss with respect to thek -th matrix variable, namely
[Equation]
For the sake of simplicity and for a cleaner notation, as all the results we will present hold for ageneric k , we drop the subscript k from now on.
In particular, we assume W is the weight matrix ofa generic hidden layer with n input and m output neurons.
In order for our derivation to hold, we require the following two properties:
P1.
The gradient flow F is locally bounded and locally Lipschitz continuous, with constants C 1 andC 2 ,
respectively.
Namely, we assume there exist C 1 , C 2 > 0 (independent of k ) such that
[Equation]
for all Z, ˜ Z ∈ R m × n .
P2.
The whole gradient flow is “not too far” from the rank- r manifold M r .
Precisely, we assume that for any Z ∈ M r arbitrary close to W ( t ) , the whole gradient flow F ( Z ) near t is such that
[Equation]
where P( Z ) denotes the orthogonal projection onto T Z M r .
Note that both assumptions are valid for low-rank neural network training.
In particular, Lipschitzcontinuity and boundedness of the gradient are standard assumptions in optimization and are satisfiedby the gradient of commonly used neural networks’ losses.
Moreover, assuming the gradient flow tobe close to the low-rank manifold is an often encountered empirical observation in neural networks[16, 45, 50].
In order to derive the proof of Theorems 1 and 2 we first present a number of relevant backgroundlemmas.
The first lemma shows that the subspace generated by the K-step in Algorithm 1 after the
QR-decomposition is O ( η ( η + ε )) close to the range of the exact solution.
Lemma 1 ([7, Lemma 2]) .
Let W 1 be the solution at time t
= η of the full problem (2) with initial condition W 0 .
Let U 1 be the matrix obtained with the K-step of the fixed-rank Algorithm 1, after one step.
Under the assumptions P1 and P2 above, we have
[Equation]
where
[Equation]
Proof.
The local error analysis of [31] shows that there exists L 1 such that
[Equation]
It follows that,
[Equation]
Therefore,
[Equation]
Hence, since both terms must be bounded by θ 2 individually, we obtain the stated result.
In the next lemma we show that also the space generated by the L step is close by the exact solution.
Namely, combined with the previous result, we have
Lemma 2 ([7, Lemma 3]) .
Let W 1 , U 1 be defined as above.
Let V 1 be the matrix obtained from the
L-step of the fixed-rank Algorithm 1, after one step.
The following estimate holds:
[Equation]
Proof.
The L-step is obtained as the K-step applied to the transposed function G ( Y )
=
F ( Y ⊤ )
⊤ .
Due to the invariance of the Frobenius norm under transposition, property P1 holds.
Similarly,property P2 continues to be satisfied because
[Equation]
where the equality P( Y ) Z ⊤ = (cid:2) P( Y ⊤ ) Z (cid:3) ⊤ has been used [34, §4].
It follows from Lemma 1 that
[Equation]
This implies that
[Equation]
Because ∥ V 1 V 1 , ⊤ ∥ 2 = 1 , the stated result follows from (11).
With the previous lemmas, we are in the position to derive the local error bound for the fixed-rank
KLS integrator of Section 4.
Lemma 3 (Local Error, [7, Lemma 4]) .
Let W 1 , U 1 , V 1 be defined as above and let S 1 be the matrix obtained with the S-step of Algorithm 1 after one step.
The following local error bound holds:
[Equation]
where the constants
[Equation]
Proof.
From Lemma 2 and the equality Y 1 = U 1 S 1 V 1 , ⊤ , we have that
[Equation]
The local error’s analysis is reduced to bound the term ∥ S 1 − U 1 , ⊤ W 1 V 1 ∥ .
For 0 ≤ t ≤ η , we thus introduce the following auxiliary quantity:
[Equation]
We observe that the term W ( t ) can be re-written as
[Equation]
where R ( t ) denotes the term in big brackets.
For 0 ≤ t ≤ η , it follows from Lemma 2 and the bound
C 1 of the function F that
[Equation]
Therefore, the term R ( t ) is bounded by
[Equation]
We re-cast the function F (cid:0) W ( t ) (cid:1) as
[Equation]
where the defect D ( t ) is given by
[Equation]
Via the Lipschitz constant C 2 of the function F , the defect is bounded by
[Equation]
Now, we compare the two differential equations
[Equation]
The solution S 1 obtained in the second differential equation is the same as given by the S-step of the
KLS integrator of Section 4.
By construction, the solution obtained in the first differential equation at time t
= η is (cid:101) S ( η )
=
U 1 , ⊤ W 1 V 1 .
With the Gronwall inequality we obtain
[Equation]
The result yields the statement of the theorem using the definition of θ .
We are now in the position to conclude the proof of Theorem 1.
Proof of Theorem 1.
In Lemma 3, the local error for the fixed-rank integrator of §4 has been provided.
The local error in time of the rank-adaptive version is directly obtained via a triangle inequality:
[Equation]
where ϑ is the tolerance parameter chosen for the truncation procedure.
Here, we abuse the notationand we maintain the same nomenclature U 1 , S 1 , and V 1 also for the novel low-rank approximationobtained via the truncation procedure.
Thus, we conclude the proof using the Lipschitz continuity of the function F .
We move from thelocal error in time to the global error in time by a standard argument of Lady Windermere’s fan [22,
Section II.3].
Therefore, the error after t steps of the rank-adaptive Algorithm 1 is given by
[Equation]
To conclude with, we prove that after one step the proposed rank-adaptive DLRT algorithm decreasesalong the low-rank approximations.
We remind that only property P1 needs to be assumed here.
Proof of Theorem 2.
Let (cid:98) Y ( t )
=
U 1 S ( t ) V 1 , ⊤ .
Here, S ( t ) denotes the solution for t ∈ [0 , η ] of the
S-step of the rank-adaptive integrator .
It follows that
[Equation]
The last identities follow by definition of the S-step.
For t ∈ [0 , η ] we have
[Equation] [Equation]
) where α = min 0 ≤ τ ≤ 1 ∥ U 1 , ⊤ ∇L (cid:0)(cid:98) Y ( τη
[Equation]
(cid:1) V 1 ∥ .
Integrating (12) from t = 0 until t =
η , we obtain
[Equation]
Because the subspace U 1 and V 1 contain by construction the range and co-range of the initial value, we have that (cid:98) Y
= U 0 S 0 V 0 , ⊤ [5, Lemma 1].
The truncation is such that ∥ Y 1 − (cid:98) Y 1 ∥ ≤ ϑ .
Therefore,
L
[Equation]
≤ L [Equation] Y 1 ) + βϑ where β
[Equation]
∥ .
Hence, the stated result is obtained.
6.2 Detailed timing measurements
Table 2 displays the average batch training times of a 5 -layer, 5120 -neuron dense network on the
MNIST dataset, with a batch size of 500 samples.
We average the timings over 200 batches andadditionally display the standard deviation of the timings corresponding to the layer ranks.
The batchtiming measures the full K,L and S steps, including back-propagation and gradient updates, as wellas the loss and metric evaluations.
[Table]
Table 3 shows the average test time of a 5 -layer, 5120 -neuron dense network, for different low-rankfactorizations and the full rank reference network.
The timings are averaged over 100 evaluations ofthe
60 K sample MNIST training data set.
We measure the K step forward evaluation of the low-ranknetworks as well as the loss and prediction accuracy evaluations.
6.3 Detailed training performance of adaptive low-rank
[Table]
Tables 4 and 5 display a detailed overview of the adaptive low-rank results of §5.2.
The displayedranks are the ranks of the converged algorithm.
The rank evolution of the 5 -Layer, 500 -Neuron test
[Table]
case can be seen in Fig.
The Evaluation parameter count corresponds to the parameters of the
K step of the dynamical low-rank algorithm, since all other matrices are no longer needed in the
[Table]
evaluation phase.
The training parameter count is evaluated as the number of parameters of the Sstep of the adaptive dynamical low rank training, with maximal basis expansion by 2 r , where r isthe current rank of the network.
We use the converged ranks of the adaptive low-rank training tocompute the training parameters.
Note that during the very first training epochs, the parameter countis typically higher until the rank reduction has reached a sufficiently low level.
[Table]
NN metrics
[Table]
[Table]
6.3.1 Lenet5 experiment
[Table]
In Table 6 we report the results of five independent runs of the dynamic low-rank training scheme on
Lenet5; we refer to §5.3 for further details.
For each column of the table
[Table]
we report the mean value
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
together with its relative standard deviations.
No seed has been applied for splitting the dataset andgenerating the initial weights configuration.
[Table]
with training of UV
[Table]
factorization
[Table]
In Figure 6, we study the influence of the SVD-like decomposition used for the derivation of the
DLRT algorithm on the test accuracy.
Contrary to the proposed strategy used for the DLRT algorithm,performing a low-rank training on a factorization of the form W k
= U k V ⊤ k as used in [30, 54] isnot robust with respect to the presence of small singular values.
The problem is peculiar to thelow-rank manifold itself, as discussed in [17, 34].
In fact, the local curvature of the manifold offixed-rank matrices is proportional to the inverse of the smallest singular value of the solution (orweight matrices in the neural network methodology).
Therefore, a naive approach for training theweights on the low-rank manifold might translate in an ill-conditioned optimization which requiresreduced learning rates.
In contrast to this approach, the DLRT algorithm takes advantage of theruled structure of the manifold, and it moves along flat sub-spaces mitigating problems arising fromthe high-curvature of the manifold [31].
This can be seen from the bound of Theorem 1, where theconstants are independent of the singular values of the weight-matrices.
In the following, we providean ad-hoc example to illustrate these situations in the context of low-rank neural network’s training.
Therefore, we compare the DLRT algorithm with a training performed directly on the factors of an
UV ⊤ decomposition[30], where the initial value for the training exhibits exponentially decayingsingular values.
The resulting test accuracy over the epoch number can be found in Figure 6.
Oneobserves that for the chosen learning rate, the UV ⊤ based training converges significantly slowerthan DLRT.
[Figure]
6.5 Cifar10 and Cifar100 experiments with AlexNet and VGG16
Table 7 shows the low-rank compression results of AlexNet and VGG16 for the Cifar10 dataset.
Baseline and DLRT training was computed with SGD and momentum factor of 0 .
We compareto literature result, which compress the same baseline-networks.
Table 8 shows the correspondingresults for the Cifar100 test case.
The same optimizer settings as in Cifar10 are used.
[Table] [Table]
6.6 Detailed training performance of low-rank pruning
The proposed low-rank training algorithm does not need to be applied to train a network from randominitial weight guesses.
When an already trained network is available, the proposed method can beemployed as a memory-efficient pruning strategy.
A straightforward approach to reduce a trainedfully-connected network to a rank r network is to compute an SVD for all weight matrices and totruncate those decompositions at rank r .
However, while this choice is optimal to present weightmatrices, it might significantly reduce the accuracy of the network.
Hence, retraining the determinedlow-rank subnetwork is commonly necessary to obtain desirable accuracy properties.
Three keyaspects are important to obtain an efficient pruning method for low-rank methods:
Retraining preserves the low-rank structure of the subnetwork.
Retraining does not exhibit the memory footprint of the fully connected network.
Retraining finds the optimal network among possible low rank networks.
Let us note that the attractor of the proposed dynamical low-rank evolution equations fulfills thesethree requirements.
Recall that for the evolution equations we have (3):
[Equation]
The condition ˙
W k ( t ) ∈ T W k
( t ) M r k ensures that the weight matrices remain of low-rank.
Moreover,as previously discussed, the training method only requires memory capacities to store low-rankfactors.
At the attractor, i.e., when ˙ W k = 0 , the last condition ensures that the attractor minimizes∥∇ W k L ( W k ( t ))
∥ F .
That is, the attractor is the optimal low-rank subnetwork in the sense that it picksthe network with minimal gradient.
To underline the effectiveness of our low-rank method as a pruningtechnique, we take the fully connected network from Table 5.
To demonstrate the poor validationaccuracy when simply doing an SVD on the full 784 by 784 weight matrices and truncating at a givensmaller rank, we perform this experiment for ranks r ∈ { 10 , 20 , 30 , 40 , 50 , 60 , 70 , 80 , 90 , 100 } .
Itturns out that though reducing memory requirements, this strategy leads to unsatisfactory accuracy ofabout 10% , see the first column of Table 9.
Then, we use the proposed low-rank training methodswith fixed rank r to retrain the network.
As starting points, we use the low-rank networks whichhave been determined by the truncated SVD.
Retraining then reaches desired accuracies that arecomparable to the previously determined low-rank networks in Table 5.
[Table]
6.7 Detailed derivation of the gradient
In this section, we derive the computation of the gradients in the K, L and S steps in detail.
For this,let us start with the full gradient,
i
[Table]
e., the gradient of the loss with respect to the weight matrix W k .
We have
[Equation]
For a general α
[Equation]
let us define
[Equation]
and note that for α ̸ = ℓ
[Equation]
whereas for α = ℓ we have
[Equation]
Therefore, recursively plugging (15), (16) and (17) into (14) yields
[Equation]
Written in matrix notation and making use of the Hadamard product defined as y ◦ A ◦ x
=
( y i A
ij x j ) ij ,for A ∈ R m × n , x ∈ R n and y ∈ R m , we have:
[Equation]
Now, let us move to deriving the K, L and S-steps for the dynamical low-rank training.
For the K-step,we represent the weight matrix W ℓ as W ℓi ℓ i ℓ − 1 = (cid:80) m K ℓi ℓ m V ℓi ℓ − 1 m .
Hence, reusing the intermediateresult (18) yields
[Equation]
In matrix notation we obtain
[Equation]
which is exactly the right-hand side of the K-step.
Hence, the K-step can be computed by a forwardevaluation of L and recording the gradient tape with respect to K ℓ .
Similarly, for the L-step, werepresent W ℓ as W ℓi ℓ i ℓ − 1 =
[Equation]
m U ℓi ℓ m L ℓi ℓ − 1 m .
Hence,
[Equation]
In matrix notation, we obtain
[Equation]
Lastly, for the S-step we write W ℓi ℓ i ℓ − 1 = (cid:80) n,m U ℓi
[Equation]
m S mn V ℓi ℓ − 1 n . Then,
[Equation]
=
[Equation]
In matrix notation, we have
[Equation]
6.8 Low-rank matrix representation and implementation of convolutional layers
A generalized convolutional filter is a four-mode tensor W ∈ R F × C × J × K consiting of F filters ofshape C ×
J × K , which is applied to a batch of N input C − channels image signals Z of spatialdimensions U × V as the linear mapping,
[Equation]
In order to train the convolutional filter on the low-rank matrix manifold, we reshape the tensor Winto a rectangular matrix W resh ∈ R F × CJK .
This reshaping is also considered in e.g. [27].
An option is, to see the convolution as the contraction between an three-mode tensor Z unfolded ofpatches and the reshaped kernel matrix W resh using Pytorch’s fold-unfold function.
We can constructthe unfold by stacking the vectorized version of sliding patterns of the kernel on the original input,obtaining in this way a tensor Z unfolded ∈ R N × CJK × L , where L denotes the dimension of flattenversion of the output of the 2-D convolution.
Thus, equation 19 can be rewritten as a tensor modeproduct:
[Equation]
As it is shown in (20), we can a decompose the starting weight W resh =
USV ⊤ and then do all thetraining procedure as a function of the factors ( U, S, V ) , without never reconstructing the kernel.
Then we can apply the considerations of fully connected layers.