The Low-Rank Simplicity Bias in Deep Networks

1 INTRODUCTION . It has become conventional wisdom that the more layers one adds , the better a deep neural network ( DNN ) performs . This guideline is supported , in part , by theoretical results showing that deeper networks can require far fewer parameters than shallower networks to obtain the same modeling “ capacity ” ( Eldan & Shamir ( 2016 ) ) . While it is not surprising that deeper networks are more expressive than shallower networks , the fact that state-of-the-art deep networks do not overfit , despite being heavily over-parameterized1 , defies classical statistical theory ( Geman et al . ( 1992 ) ; Zhang et al . ( 2017 ) ; Belkin et al . ( 2019 ) ) . The belief that over-parameterization via depth improves generalization is used axiomatically in the design of neural networks . Unlike conventional regularization methods that penalize model complexity ( e.g. , ` 1/ ` 2 penalty ) , over-parameterization does not . Yet , like explicit regularization , over-parameterization appears to prevent the model from over-fitting ( Belkin et al . ( 2018 ) ; Nakkiran et al . ( 2019a ) ) . Why this implicit regularization works is still an ongoing area of research . Even in the zero-training error regime , it is commonly observed that increasing the number of parameters improves generalization performance . Currently , the status quo explanation for this phenomenon is that gradient descent in over-parameterized models acts as a nuclear norm regularizer ( Gunasekar et al . ( 2017 ) ; Arora et al . ( 2019a ) ; Li et al . ( 2020 ) ) . This work provides a new set of observations that expand the growing body of work on overparameterization and highlights the central role of depth in finding solutions that map to low effective rank embeddings for both linear and non-linear networks . Mainly , we make series of empirical observations that indicate deep networks have an inductive bias to find lower effective rank embeddings . First , we observe that random deep networks are biased to map data to a feature space whose Gram matrix has a low effective rank . Next , we find that this low effective rank phenomenon exists even after training with gradient descent . We further observe that the bias towards low effetive rank embeddings exists in a wide variety of common optimizers — even optimizers that do not use gradient descent . Moreover , we find that regardless of the initialization , the effective rank of the 1e.g. , Dosovitskiy et al . ( 2020 ) trains a 632 million parameter , 200+ layer model , on 1.3 million images . converged solution is largely dependent on the depth of the model . This set of observations leads us to conjecture that deeper networks are implicitly biased to find lower effective rank embeddings because the volume of functions that map to low effective rank embeddings increases with depth . We then leverage our observations to demonstrate how one could use “ depth '' as a practical regularizer to achieve better generalization performance on standard benchmarks such as CIFAR ( Krizhevsky et al . ( 2009 ) ) and ImageNet ( Russakovsky et al . ( 2015 ) ) . 2 PRELIMINARIES . 2.1 NEURAL NETWORKS AND OVER-PARAMETERIZATION . Simple linear network A simple linear neural network transforms input x ∈ Rn×1 to output ŷ ∈ Rm×1 , with a learnable parameter matrix W ∈ Rm×n , ŷ = Wx . ( 1 ) For notational convenience , we omit the bias term . Over-parameterized linear networks One can over-parameterize a linear neural network by defining d matrices { Wi } di=1 and multiplying them successively with input x : ŷ = WdWd−1 · · ·W1x = Wex , ( 2 ) whereWe = ∏d i=1Wi . As long as the matrices are of the correct dimensionality — matrixWd hasm columns , W1 has n rows , and all intermediate dimensions { dim ( Wi ) } d−1i=2 ≥ min ( m , n ) — then this over-parameterization expresses the same set of functions as a single-layer network . We disambiguate between the collapsed and expanded set of weights by referring to { Wi } as the over-parameterized weights and We as the end-to-end or the effective weights . Non-linear networks For non-linear network , activation function ψ ( e.g . ReLU ) is interleaved in between the weights matrices : ŷ = Wdψ ( Wd−1 . . . ψ ( W1 ( x ) ) ) ( 3 ) In contrast to linear networks , non-linear models become more expressive as more layers are added . 2.2 EFFECTIVE RANK . We characterize the rank of a matrix using a continuous measure known as the effective rank : Definition 1 ( Effective rank ; Roy & Vetterli ( 2007 ) ) . For any matrix A ∈ Rm×n , the effective rank ρ is defined as the Shannon entropy of the normalized singular values : ρ ( A ) = − min ( n , m ) ∑ i=1 σ̄i log ( σ̄i ) , where σ̄i = σi/ ∑ j σj are normalized singular values , such that ∑ i σ̄i = 1 . Also referred to as the spectral entropy . Without loss of generality , we drop the exponentiation for convenience . This measure gives us a meaningful representation of “ continuous rank ” , which is maximized when the magnitude of the singular values are all equal and minimized when a single singular value dominates relative to others . The effective rank provides us with a metric that summarizes the distribution envelope . Effective rank has been used in prior works ( Arora et al . ( 2019a ) ; Razin & Cohen ( 2020 ) ; Baratin et al . ( 2021 ) ) and we use this measure extensively throughout our work . We have also found that our observations are consistent with closest definition of rank in which we threshold the smallest singular values after normalization ( Appendix C ) . Similar to prior works , we refer to low effective rank as low rank here on out . 2.3 EMBEDDING MAPS . A parameteric function f { W } ∈ FW is a neural network parameterized with a set weights { W } = { W1 , . . . , Wd } that maps the input space to the output space X → Y . For a training dataset of size q , the input and output data is X ∈ Rn×q and Y ∈ Rm×q . Then , the predicted output is Ŷ = Wdψ ( Φ ) = f { W } ( X ) , where Φ ∈ Rn ′×q is the last-layer embedding and Wd ∈ Rm×n ′ is the last layer of the network . We analyze the embedding space by computing the effective rank on the Gram/kernel matrix K ∈ Rp×p where p is the size of the test set . The ij-th entry of the Gram matrix corresponds to a distance kernel Kij = κ ( φi , φj ) where φi corresponds to the i-th column of Φ . We use the model ’ s intermediate features before the linear classifier and use cosine distance kernel 2 , a common method for measuring distances in feature space ( Kiros et al . ( 2015 ) ; Zhang et al . ( 2018 ) ) . Since the dimensionality of the Gram-matrix does not depend on the model parameters , we can compare neural networks with different modeling capacities in the zero training error regime . Gram matrices are often used to analyze optimization and generalization properties of neural networks ( Zhang et al . ( 2019 ) ; Du et al . ( 2018 ; 2019 ) ; Wu et al . ( 2019 ) ; Arora et al . ( 2019b ) ) . In natural data , it is often assumed that we are trying to discover a low-rank relationship between the input and the label . For example , a model that overfits to every training sample without inferring any structure on the data will generally have a test gram-matrix that is higher rank than that of a model that has learned parsimonious representations . Lower rank on held-out data indicates less excess variability and is indicative for studying generalization and robustness . The intuition becomes clearer in linear networks , since the rank of Gram matrix depends on the rank of the linear transformation computed by the network . We illustrate this empirically in Appendix L , where we see that there is a tight relationship between the rank of the linear weight matrix and the resulting Gram matrix . 2.4 LEAST SQUARES . Given a dataset X , Y generated from W ∗ , the goal is to regress a parameterized function f { W } ( · ) to minimize the squared-distance ‖f { W } ( X ) − Y ‖22 . The rank ( W ∗ ) is a measure of the “ intrinsic dimensionality ” of the data and we refer to it as the task rank . In this work , we exclusively operate in the under-determined regime where we have fewer training examples than model parameters . This ensures that there is more than one minimizing solution . 3 THE INDUCTIVE PARAMETERIZATION BIAS OF DEPTH . Given that our models can always fit the data , what are the implications of searching for the solution in the over-parameterized model ? In linear models , this is equivalent to searching for solutions in { Wi } versus directly in We . One difference is that the gradient direction∇ { Wi } L ( { Wi } ) is usually different than ∇WeL ( We ) for a typical loss function L ( see Appendix I ) . The consequences of this difference have been previously studied in linear models by Arora et al . ( 2018 ; 2019a ) , where the 2cosine kernel : κ ( φi , φj ) = φiφ T j ‖φi‖‖φj‖ over-parameterized update rule has been shown to accelerate training and encourage singular values to decay faster , resulting in low nuclear-norm solution . Here we motivate a result from the perspective of parameter volume space . Conjecture 1 . Deeper networks have a greater proportion of parameter space that maps the input data to lower-rank embeddings ; hence , deeper models are more likely to converge to functions that learn simpler embeddings . We now provide a set of empirical observations that supports our conjecture . Our work and existing theoretical works on gradient descent biases are not mutually exclusive and are a likely compliment . We emphasize that we do not make any claims on the simplicity of the function , but only on the simplicity – lower effective rank – of the embeddings . 3.1 LOW-RANK SIMPLICITY BIAS OF DEEP NETWORKS . Observation 1 . Randomly initialized deep neural networks are biased to correspond to Gram matrices with a low effective rank . When sampling random neural networks , both linear and non-linear , we observed that the Gram matrices computed from deeper networks have lower effective rank . We quantify this observation by computing the distribution over the effective rank of the Gram matrix in Figure 1 . Here , the weights of the neural networks are initialized using uniform Wi ∼ U ( · , · ) or Normal distributions Wi ∼ N ( · , · ) . The input , output , and intermediate dimensions are 32 , giving parameters { Wi } ∈ Rd×32×32 for a network with d layers . We draw 4096 random parameter samples and compute the effective rank on the resulting Gram matrix . We see that the distribution density shifts towards the left ( lower effective rank ) when increasing the number of layers . These distributions have small overlap and smoothen out with increased depth . This observation shows that depth correlates with lower effective rank embeddings . The low-rank bias becomes more intuitive in linear models as there is a simple way to relate the Gram matrix to the weights of the model K ≈ ( Wd−1:1X ) T ( Wd−1:1X ) . Intuitively , if any constituent matrices are low-rank , then the product of matrices will also be low-rank – the product of matrices can only decrease the rank of the resulting matrix : rank ( AB ) ≤ min ( rank ( A ) , rank ( B ) ) ( Friedberg et al . ( 2003 ) ) . In Appendix L , we show that as the depth of the model increases , both the effective rank of the Gram matrix and the weights decrease together . Another way to interpret our observation is that for linear models , over-parameterization does not increase the expressivity of the function but re-weights the likelihood of a subset of parameters – the hypothesis class . For non-linear models , we can not make the same claims . Although uniformly sampling under the parameter distribution is an unbiased estimator of the volume of the parameter space , it is certainly possible that a sub-space of the parameters is more likely to be observed under gradient descent . Hence , by naively sampling networks , we may never encounter model parameters that gradient descent explores . In light of this , we repeat our experiment above by computing the PDF on randomly sampled parameters after taking n gradient descent steps . Observation 2 . Deep neural networks trained with gradient descent also learns to map data to simple embedding with low effective rank . Figure 2 illustrates the change in distribution as we train our model to convergence using gradient descent . Each randomly drawn network sample is trained to minimize the least-squares error . The initial distribution is plotted with dotted lines , and the converged distribution is plotted with solid lines . As the model is trained , the distribution of the rank shifts towards the ground-truth rank ( green line ) . Training the model with gradient descent results in a distribution that is still largely dependent on depth ; this reaffirms that the role of optimization does not remove the parameterization bias in deep models . In fact , if the bias stems from the model ’ s parameterization , the same bias must also exist under other common and natural choices of optimizers . We investigate this claim in the next section .

Summary: The paper explores the effect of depth on a continuous estimate of rank (effective rank) of the data embeddings computed by feed-forward neural networks.  On artificial datasets, the paper shows that at initialization (with common uniform or normal weight distributions), the effective rank of both linear and non-linear MLPs decreases with depth, and this is preserved after training. The results does not extend to network architectures with skip-connections such as ResNets.  The authors claim that this low-rank bias of the embeddings help generalization on natural datasets, and show that a reparametrization of common architectures where the linear layers are replaced by products of linear layers improves test accuracy on CIFAR-10 and CIFAR-100.