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: The paper presents an extension of the idea of alignment in linear neural networks, that can help in providing convergence analysis of such networks. Such a notion was previously studied for networks with a single output. The current paper extends it to networks with multi-dimensional outputs. The paper offers multiple interesting results: a) conditions on the datasets where alignment can remain invariant b) lack of alignment or invariance for networks with constrained layers