On Alignment in Deep Linear Neural Networks

1 INTRODUCTION . Although overparameterized deep networks can interpolate randomly labeled training data ( Du et al. , 2019 ; Wu et al. , 2019 ) , training overparameterized networks with modern optimizers often leads to solutions that generalize well . This suggests that there is a form of implicit regularization occurring through training ( Zhang et al. , 2017 ) . As an example of implicit regularization , the authors in Ji & Telgarsky ( 2018 ) proved that the layers of linear neural networks used for binary classification on linearly separable datasets become aligned in the limit of training . That is , for a linear network parameterized by the matrix product WdWd−1 . . .W1 , the top left/right singular vectors ui and vi of layer Wi satisfy |vTi+1ui| → 1 as the number of gradient descent steps goes to infinity . Alignment of singular vector spaces between adjacent layers allows for the network representation to be drastically simplified ( see Equation 3 ) ; namely , the product of all layers becomes a product of diagonal matrices with the exception of the outermost unitary matrices . If alignment is an invariant of training , then optimization over the set of weight matrices reduces to optimization over the set of singular values of weight matrices . Thus , importantly , alignment of singular vector spaces allows for the gradient descent update rule to be simplified significantly , which was used in Ji & Telgarsky ( 2018 ) to show convergence to a max-margin solution . In this work , we generalize the definition of alignment to the multidimensional setting . We study when alignment can occur and moreover , under which conditions it is an invariant of training in linear neural networks under gradient descent . Prior works ( Gidel et al. , 2019 ; Saxe et al. , 2014 ; 2019 ) have implicitly relied on invariance of alignment as an assumption on initialization to simplify training dynamics for 2 layer networks . In this work , we provide necessary and sufficient conditions for when alignment is an invariant for networks of arbitrary depth . Our main contributions are as follows : 1 . We extend the definition of alignment from the 1-dimensional classification setting to the multi-dimensional setting ( Definition 2 ) and characterize when alignment is an invariant of training in linear fully connected networks with multi-dimensional outputs ( Theorem 1 ) . 2 . We demonstrate that alignment is an invariant for fully connected networks with multidimensional outputs only in special problem classes including autoencoding , matrix factorization and matrix sensing . This is in contrast to networks with 1-dimensional outputs , where there exists an initialization such that adjacent layers remain aligned throughout training under any real-valued loss function and any training dataset . 3 . Alignment largely simplifies the analysis of training linear networks : We provide an explicit learning rate under which gradient descent converges linearly to a global minimum under alignment in the squared loss setting ( Proposition 1 ) . 4 . We prove that alignment can not occur , let alone be invariant , in networks with constrained layer structure ( such as convolutional networks ) , when the amount of training data dominates the dimension of the layer structure ( Theorem 3 ) . 5 . We support our theoretical findings via experiments in Section 6 . As a consequence , our characterization of the invariance properties of alignment provides settings under which the gradient descent dynamics can be simplified and the implicit regularization properties can be fully understood , yet also shows that further results are required to explain implicit regularization in linear neural networks more generally . 2 RELATED WORK . Implicit regularization in overparameterized networks has become a subject of significant interest ( Gunasekar et al. , 2018a ; b ; Martin & Mahoney , 2018 ; Neyshabur et al. , 2014 ) . In order to characterize the specific form of implicit regularization , several works have focused on analyzing deep linear networks ( Arora et al. , 2019b ; Gunasekar et al. , 2018b ; 2017 ; Soudry et al. , 2018 ) . Even though such networks can only express linear maps , parameter optimization in linear networks is non-convex and is studied in order to obtain intuition about optimization of deep networks more generally . One such form of implicit regularization is alignment , identified by Ji & Telgarsky ( 2018 ) to analyze linear fully connected networks with 1-dimensional outputs trained on linearly separable data . They proved that in the limit of training , each layer , after normalization , approaches a rank 1 matrix , i.e . lim t→∞ W ( t ) i ‖W ( t ) i ‖F = uiv T i and that adjacent layers , Wi+1 and Wi become aligned , i.e . |vTi+1ui| → 1 . In addition , Ji & Telgarsky ( 2018 ) proved that alignment in this setting occurs concurrently with convergence to the max-margin solution . Follow-up work mainly focused on this convergence phenomenon and gave explicit convergence rates for overparameterized networks trained with gradient descent ( Arora et al. , 2019c ; Zou et al. , 2018 ) . Our definition of invariance of alignment extends assumptions on initialization appearing in various prior works ( Gidel et al. , 2019 ; Saxe et al. , 2014 ; 2019 ) . While the connection to alignment was not mentioned in their work , the authors in Gidel et al . ( 2019 ) begin to generalize alignment to multidimensional outputs by considering two-layer networks initialized so that layers are aligned with each other and to the data . We generalize this to networks of any depth , showing that our definition of alignment corresponds to the initialization considered in Gidel et al . ( 2019 ) . Moreover , we establish necessary and sufficient conditions for when alignment is an invariant of training in Theorem 1 instead of assuming these conditions . Furthermore , their result on sequential learning of components can be derived via our singular value update rule in Corollary 1 . Balancedness is another closely related form of implicit regularization in linear neural networks . It was introduced in Arora et al . ( 2018 ) and defined as the property that if WTi Wi = Wi+1W T i+1 for all i at initialization , then this property is invariant under gradient flow . Du et al . ( 2018 ) present a more general form , that WTi Wi −Wi+1WTi+1 is constant under gradient flow . In practice , analyses rely on this quantity being close to or exactly zero . In this exact setting , balancedness indeed implies alignment of singular vector spaces between consecutive layers . To study gradient descent , slightly more general notions such as approximate balancedness ( Arora et al. , 2019a ) and -balancedness have been introduced . Du et al . ( 2018 ) also defined balancedness with respect to convolutional networks , showing that under gradient flow , the difference in the norm of the weights of consecutive layers is an invariant . Generally , the goal of identifying invariants of training such as balancedness or alignment is to help understand both the dynamics of training and properties of solutions at the end of training . 3 DEFINITION OF ALIGNMENT IN THE MULTI-DIMENSIONAL SETTING . In this section , we first define alignment for linear neural networks with multi-dimensional outputs . We then define when alignment is an invariant of training . We consider linear neural networks . Let f : Rk0 → Rkd denote such a d-layer network , i.e . f ( x ) = WdWd−1 . . .W1x , ( 1 ) where Wi ∈ Rki×ki−1 for i ∈ [ d ] , where we follow the convention that [ d ] = { 1 , 2 , . . . d } . Let ( X , Y ) ∈ Rk0×n × Rkd×n denote the set of training data pairs { ( x ( i ) , y ( i ) ) } for i ∈ [ n ] . Gradient descent with learning rate γ is used to find a solution to the following optimization problem : arg min f∈F 1 2n n∑ i=1 ` ( f ( x ( i ) ) , y ( i ) ) , ( 2 ) where F is the set of linear functions represented by f and ` is a real-valued loss function . When not stated otherwise , we assume ` ( f ( x ( i ) ) , y ( i ) ) = ‖y ( i ) − f ( x ( i ) ) ‖22 , which is the squared loss ( MSE ) . In addition , we denote by W ( t ) i for t ∈ Z≥0 the weight matrix Wi after t steps of gradient descent . When there are no additional constraints on the matrices Wi , then f is a fully connected network . We next introduce a generalized form of the singular value decomposition : Definition 1 . An unsorted , signed singular value decomposition ( usSVD ) of a matrix A ∈ Rm×n is a triple U ∈ Rm×m , Σ ∈ Rm×n , V ∈ Rn×n such that U , V are orthonormal matrices , Σ is diagonal , and A = UΣV T . In contrast to the usual definition of singular value decomposition ( SVD ) of a matrix , the diagonal entries of Σ may be in any order and take negative values . Throughout , we will refer to the entries of Σ in a usSVD as singular values and the vectors in U , V as singular vectors . Using the usSVD , we now generalize the notion of alignment from Ji & Telgarsky ( 2018 ) to the multi-dimensional setting . Definition 2 . Let f = WdWd−1 . . .W1 be a linear network . We say that f is aligned if there exists a usSVD Wi = UiΣiV Ti with Ui = Vi+1 for all i ∈ [ d− 1 ] . ( We also say that a matrix A is aligned with another matrix B if there exist usSVD ’ s A = UAΣAV TA , B = UBΣBV T B such that VA = UB . ) Note that if Wi and Wi+1 are rank 1 matrices in an aligned network f , then the inner product of the first columns of Vi+1 and Ui is 1 in absolute value . Hence Definition 2 is consistent with alignment in the 1-dimensional setting from Ji & Telgarsky ( 2018 ) . We next define when alignment is an invariant of training for deep linear networks . Again , such invariants are of interest since they may provide insights into properties of trained networks and significantly simplify the dynamics of gradient descent . Definition 3 . Alignment is an invariant of training for a linear neural network f if there exists an initialization { W ( 0 ) j } dj=1 such that W ( ∞ ) 1 , W ( ∞ ) 2 , . . . , W ( ∞ ) d achieves zero training error 1 and for all gradient descent steps t ∈ Z≥0 ( a ) the network f is aligned ; ( b ) W ( t ) i = UiΣ ( t ) i V T i for all i ∈ { 2 , . . . d− 1 } , that is , Ui , Vi are not updated ; ( c ) W ( t ) 1 = U1Σ ( t ) 1 V ( t ) 1 T and W ( t ) d = U ( t ) d Σ ( t ) d V T d , that is , U1 and Vd are not updated . If additionally , V1 and Ud are not updated for any t ∈ Z≥0 , then we say that strong alignment is an invariant of training . When alignment is an invariant of training , there are important consequences for training . In particular , note that when the network f is aligned with usSVDs Wi = UiΣiV Ti for all 1 ≤ i ≤ d , then f ( x ) = Wd · · ·W1x = Ud ( d−1∏ i=0 Σd−i ) V T1 x . ( 3 ) 1The interpolation condition in this definition ( i.e. , achieving zero training error ) is important in ruling out several architectures where the layers are trivially aligned . For example , if all layers are constrained to be diagonal matrices throughout training , then the layers are all trivially aligned , but can not interpolate datasets where the target is not the product of a diagonal matrix with the input . Hence if alignment is an invariant of training , then the singular vectors of layers 2 through d− 1 are never updated and the analysis of gradient descent can be limited to the singular values of the layers and the matrices V1 and Ud . Remarks . For the remainder of the paper , we assume that the gradient of the loss function at initialization { W ( 0 ) i } di=1 is non-zero . Otherwise , training with gradient descent would not proceed . We also only consider datasets ( X , Y ) for which there is a linear network that achieves loss zero . This is consistent with the assumptions in Ji & Telgarsky ( 2018 ) .

Summary: This paper considers deep linear networks trained by gradient descent with the squared loss, and characterizes when alignment happens, meaning that for any pair of adjacent weight matrices W_{i+1} and W_i, the (unsorted, signed) right singular vectors of W_{i+1} are identical to the (unsorted, signed) left singular vectors of W_i. This paper further gives a few examples where the conditions of alignment hold, and proves a convergence rate for aligned networks. For networks with constrained layer structure, such as convolutional networks, this paper shows that aligned networks in general cannot achieve zero training error with the squared loss. Finally, empirical support of the theory is provided.