Title: Bottleneck Structure in Learned Features: Low-Dimension vs Regularity Tradeoff
Venue: NeurIPS 2023
Accepted: yes
Keywords: Feature Learning

Symmetry Learning

Theory of Deep Learning

Weight Decay
Abstract: Previous work has shown that DNNs with
large depth $L$ and $L_{2}$-regularization are biased towards learning
low-dimensional representations of the inputs, which can be interpreted
as minimizing a notion of rank $R^{(0)}(f)$ of the learned function
$f$, conjectured to be the Bottleneck rank. We compute finite depth
corrections to this result, revealing a measure $R^{(1)}$ of regularity
which bounds the pseudo-determinant of the Jacobian $\left\|Jf(x)\right\|\_\+$
and is subadditive under composition and addition. This formalizes
a balance between learning low-dimensional representations and minimizing
complexity/irregularity in the feature maps, allowing the network
to learn the `right' inner dimension. Finally, we prove the conjectured
bottleneck structure in the learned features as $L\to\infty$: for
large depths, almost all hidden representations are approximately
$R^{(0)}(f)$-dimensional, and almost all weight matrices $W_{\ell}$
have $R^{(0)}(f)$ singular values close to 1 while the others are
$O(L^{-\frac{1}{2}})$. Interestingly, the use of large learning rates
is required to guarantee an order $O(L)$ NTK which in turns guarantees
infinite depth convergence of the representations of almost all layers.

Bottleneck Structure in Learned Features: Low-Dimension vs Regularity Tradeoff

Abstract

Previous work [Jac23] has shown that DNNs with large depth L and L 2 -regularization are biased towards learning low-dimensional representations of the inputs, which can be interpreted as minimizing a notion of rank R (0) ( f ) of the learned function f , conjectured to be the Bottleneck rank. We compute finite depth corrections to this result, revealing a measure R (1) of regularity which bounds the pseudo-determinant of the Jacobian | Jf ( x ) | + and is subadditive under composition and addition. This formalizes a balance between learning low-dimensional representations and minimizing complexity/irregularity in the feature maps, allowing the network to learn the 'right' inner dimension. Finally, we prove the conjectured bottleneck structure in the learned features as L → ∞ : for large depths, almost all hidden representations are approximately R (0) ( f ) -dimensional, and almost all weight matrices W ℓ have R (0) ( f ) singular values close to 1 while the others are O ( L -1 2 ) . Interestingly, the use of large learning rates is required to guarantee an order O ( L ) NTK which in turns guarantees infinite depth convergence of the representations of almost all layers.

1 Introduction

The representation cost R ( f ; L ) = min θ : f θ = f ∥ θ ∥ 2 [DKS21] can be defined for any model and describes the bias in function space resulting from the minimization of the L 2 -norm of the parameters. While it can be computed explicitly for linear networks [DKS21] or shallow nonlinear ones [Bac17], the deep non-linear case remains ill-understood [JGHG22].

Previous work [Jac23] has shown that the representation cost of DNNs with homogeneous nonlinearity σ converges to a notion of rank over nonlinear functions lim L →∞ R ( f ; L ) L → R (0) ( f ) . Over a large set of functions f , the limiting representation cost R (0) ( f ) was proven the so-called Bottleneck (BN) rank Rank BN ( f ) which is the smallest integer k such that f : R d in → R d out can be factored f = R d in g -→ R k h -→ R d out with inner dimension k (it is conjectured to match it everywhere). This suggests that large depth L 2 -regularized DNNs are adapted for learning functions of the form f ∗ = g ◦ h with small inner dimension.

This can also be interpreted as DNNs learning symmetries, since a function f : Ω → R d out with symmetry group G (i.e. f ( g · x ) = f ( x ) ) can be defined as mapping the inputs Ω to an embedding of the modulo space Ω / G and then to the outputs R d out . Thus a function with a lot of symmetries will have a small BN-rank, since Rank BN ( f ) ≤ dim( Ω / G ) where dim( Ω / G ) is the smallest dimension Ω / G can be embedded into.

A problem is that this notion of rank does not control the regularity of f , but results of [Jac23] suggest that a measure of regularity might be recovered by studying finite depths corrections to the R (0) approximation. This formalizes the balance between minimizing the dimension of the learned features and their complexity.

Another problem is that minimizing the rank R (0) does not uniquely describe the learned function, as there are many fitting functions with the same rank. Corrections allow us to identify the learned function amongst those.

Finally, the theoretical results and numerical experiments of [Jac23] strongly suggest a bottleneck structure in the learned features for large depths, where the (possibly) high dimensional input data is mapped after a few layers to a low-dimensional hidden representation, and keeps the approximately same dimensionality until mapping back to the high dimensional outputs in the last few layers. We prove the existence of such a structure, but with potentially multiple bottlenecks.

1.1 Contributions

We analyze the Taylor approximation of the representation cost around infinite depth L = ∞ :

The first correction R (1) measures some notion of regularity of the function f that behaves subadditively under composition R (1) ( f ◦ g ) ≤ R (1) ( f ) + R (1) ( g ) and under addition R (1) ( f + g ) ≤ R (1) ( f ) + R (1) ( g ) (under some constraints), and controls the Jacobian of f : ∀ x, 2 log | Jf ( x ) | + ≤ R (1) ( f ) , where |·| + is the pseudo-determinant , the product of the non-zero singular values.

This formalizes the balance between the bias towards minimizing the inner dimension described by R (0) and a regularity measure R (1) . As the depth L grows, the low-rank bias R (0) dominates, but even in the infinite depth limit the regularity R (1) remains relevant since there are typically multiple fitting functions with matching R (0) which can be differentiated by their R (1) value.

For linear networks, the second correction R (2) guarantees infinite depth convergence of the representations of the network. We recover several properties of R (2) in the nonlinear case, but we also give a counterexample that shows that norm minimization does not guarantee convergence of the representation, forcing us to look at other sources of bias.

To solve this issue, we show that a Θ( L -1 ) learning rate forces the NTK to be O ( L ) which in turn guarantees the convergence as L →∞ of the representations at almost every layer of the network.

Finally we prove the Bottleneck structure that was only observed empricially in [Jac23]: we show that the weight matrices W ℓ are approximately rank R (0) ( f ) , more precisely W ℓ has R (0) ( f ) singular values that are O ( L -1 2 ) close to 1 and all the other are O ( L -1 2 ) . Together with the O ( L ) NTK assumption, this implies that the pre-activations α ℓ ( X ) of a general dataset at almost all layer is approximately R (0) ( f ) -dimensional, more precisely that the k +1 -th singular value of α ℓ ( X ) is O ( L -1 2 ) .

1.2 Related Works

The representation cost has mostly been studied in settings where an explicit formula can be obtained, such as in linear networks [DKS21], or shallow nonlinear networks [Bac17], or for deep networks with very specific structure [OW22, LJ22]. A low rank phenomenon in large depth L 2 -regularized DNNs has been observed in [TVS22].

Regarding deep fully-connected networks, two reformulations of the representation cost optimization have been given in [JGHG22], which also shows that the representation becomes independent of the width as long as the width is large enough.

The Bottleneck structure that we describe in this paper is similar to the Information Bottleneck theory [TZ15]. It is not unlikely that the bias towards dimension reduction in the middle layers of the network could explain the loss of information that was observed in the first layers of the network in [TZ15].

2 Setup

In this paper, we study fully connected DNNs with L +1 layers numbered from 0 (input layer) to L (output layer). The layer ℓ ∈ { 0 , . . . , L } has n ℓ neurons, with n 0 = d in the input dimension and n L = d out the output dimension. The pre-activations ˜ α ℓ ( x ) ∈ R n ℓ and activations α ℓ ( x ) ∈ R n ℓ are defined by

for the n ℓ × n ℓ -1 connection weight matrix W ℓ , the n ℓ -dim bias vector b ℓ and the nonlinearity σ : R → R applied entry-wise to the vector ˜ α ℓ ( x ) . The parameters of the network are the collection of all connection weights matrices and bias vectors θ = ( W 1 , b 1 , . . . , W L , b L ) . The network function f θ : R d in → R d out is the function that maps an input x to the pre-activations of the last layer ˜ α L ( x ) .

We assume that the nonlinearity is of the form σ a ( x ) = { x if x ≥ 0 ax otherwise for some α ∈ ( -1 , 1)

(yielding the ReLU for α = 0 ), as any homogeneous nonlinearity σ (that is not proportional to the identity function, the constant zero function or the absolute function) matches σ a up to scaling and inverting the inputs.

The functions that can be represented by networks with homogeneous nonlinearities and any finite depth/width are exactly the set of finite piecewise linear functions (FPLF) [ABMM18, HLXZ18].

Remark 1 . In most of our results, we assume that the width is sufficiently large so that the representation cost matches the infinite-width representation cost. For a dataset of size N , a width of N ( N +1) suffices, as shown in [JGHG22] (though a much smaller width is often sufficient).

2.1 Representation Cost

The representation cost R ( f ; Ω , L ) is the minimal squared parameter norm required to represent the function f over the input domain Ω :

where the minimum is taken over all weights θ of a depth L network (with some finite widths n 1 , . . . , n L -1 ) such that f θ ( x ) = f ( x ) for all x ∈ Ω . If no such weights exist, we define R ( f ; Ω , L ) = ∞ .

The representation cost describes the natural bias on the represented function f θ induced by adding L 2 regularization on the weights θ :

for any cost C (defined on functions f : Ω ↦→ R d out ) and where the minimization on the right is over all functions f that can be represented by a depth L network with nonlinearity σ .

For any two functions f, g , we denote f → g the function h such that g = h ◦ f , assuming it exists, and we write R ( f → g ; Ω , L ) = R ( h ; f (Ω) , L ) .

Remark 2 . The representation cost also describes the implicit bias of networks trained with the cross-entropy loss [SHN + 18, GLSS18, CB20].

3 Representation Cost Decomposition

Since there are no explicit formula for the representation cost of deep nonlinear networks, we propose to approximate it by a Taylor decomposition in 1 / L around L = ∞ . This is inspired by the behavior of the representation cost of deep linear networks (which represent a matrix as a product A θ = W L · · · W 1 ), for which an explicit formula exists [DKS21]:

where ∥·∥ p p is the L p -Schatten norm, the L p norm on the singular values s i ( A ) of A .

Approximating s 2 L = 1 + 2 L log s + 2 L 2 (log s ) 2 + O ( L -3 ) , we obtain

where log + is the pseudo-log, which replaces the non-zero eigenvalues of A by their log.

We know that gradient descent will converge to parameters θ representing a matrix A θ that locally minimize the loss C ( A ) + λR ( A ; L ) . The approximation R ( A ; L ) ≈ L Rank A fails to recover the local minima of this loss, because the rank has zero derivatives almost everywhere. But this problem is alleviated with the second order approximation R ( A ; L ) ≈ L Rank A +2log | A | + . The minima can then be interpreted as first minimizing the rank, and then choosing amongst same rank solutions the matrix with the smallest pseudo-determinant. Changing the depth allows one to tune the balance between minimizing the rank and the regularity of the matrix A .

3.1 First Correction: Regularity Control

As a reminder, the dominating term in the representation cost R (0) ( f ) is conjectured in [Jac23] to converge to the so-called Bottleneck rank Rank BN ( f ; Ω) which is the smallest integer k such that f can be decomposed as f = Ω g -→ R k h -→ R d out with inner dimension k , and where g and h are FPLF. A number of results supporting this conjecture are proven in [Jac23]: a sandwich bound

for the Jacobian rank Rank J ( f ; Ω) = max x Rank Jf ( x ) | T x Ω , and three natural properties of ranks that R (0) satisfies:

R (0) ( f ◦ g ; Ω) ≤ min { R (0) ( f ) , R (0) ( g ) } ,

R (0) ( f + g ; Ω) ≤ R (0) ( f ) + R (0) ( g ) ,

R (0) ( x ↦→ Ax ; Ω) = Rank A for any full dimensional and bounded Ω .

These results imply that for any function f = ϕ ◦ A ◦ ψ that is linear up to bijections ϕ, ψ , the conjecture is true R (0) ( f ; Ω) = Rank BN ( f ; Ω) = Rank A .

The proof of the aforementioned sandwich bound in [Jac23] actually prove an upper bound of the form L Rank BN ( f ; Ω) + O (1) thus proving that the R (1) term is upper bounded. The following theorem proves a lower bound on R (1) as well as some of its properties:

Theorem 3. For all inputs x where Rank Jf ( x ) = R (0) ( x ) , R (1) ( f ) ≥ 2 log | Jf ( x ) | + , furthermore:

If R (0) ( f ◦ g ) = R (0) ( f ) = R (0) ( g ) , then R (1) ( f ◦ g ) ≤ R (1) ( f ) + R (1) ( g ) .

If R (0) ( f + g ) = R (0) ( f ) + R (0) ( g ) , then R (1) ( f + g ) ≤ R (1) ( f ) + R (1) ( g ) .

If P Im A T Ω and A Ω are k = Rank A dimensional and completely positive (i.e. they can be embedded isometrically into R m + for some m ), then R (1) ( x ↦→ Ax ; Ω) = 2 log | A | + .

Notice how these properties clearly point to the first correction R (1) ( f ) measuring a notion of regularity of f instead of a notion of rank. One can think of L 2 -regularized deep nets as learning functions f that minimize

The depth determines the balance between the rank regularization and regularity regularization. Without the R (1) -term, the above minimization would never be unique since there can be multiple functions f with the same training outputs f ( X ) with the same rank R (0) ( f ) .

Under the assumption that R (0) only takes integer value the above optimization can be rewritten as

5

10

15

20

25

30

L

0

20

40

60

80

100

120

2

A

B

C

D

Rank

(

WL

/2) = 3

Rank

(

WL

/2) = 2

Rank

(

WL

/2) = 1

Rank

(

WL

/2) = 0

3

L

+32

2

L

+34

L

+25

(a) Parameter norm and depth

0.0

0.5

1.0

1.5

si

(

W

)

A

0.0

0.5

1.0

1.5

2.0

C

0

5

10

0.0

0.5

1.0

1.5

si

(

W

)

B

0

5

10

15

20

0.0

0.5

1.0

1.5

D

(b) Bottleneck structure at different depths.

Every inner minimization for a rank k that is attained inside the set { f : R (0) ( f ) = k } corresponds to a different local minimum. Note how these inner minimization do not depend on the depth, suggesting the existence of sequences of local minima for different depths that all represent approximately the same function f , as can be seen in Figure 1. We can classify these minima according to whether they recover the true BN-rank k ∗ of the task, underestimate or overestimate it.

In linear networks [WJ23], rank underestimating minima cannot fit the data, but it is always possible to fit any data with a BN-rank 1 function (by mapping injectively the datapoints to a line and then mapping them nonlinearly to any other configuration). We therefore need to also differentiate between rank-underestimating minima that fit or do not fit the data. The non-fitting minima can in theory be avoided by taking a small enough ridge (along the lines of [WJ23]), but we do observe them empricially for large depths in Figure 1.

In contrast, we have never observed fitting rank-underestimating minima, though their existence was proven for large enough depths in [Jac23]. A possible explanation for why GD avoids these minima is their R (1) value explodes with the number of datapoints N , since these network needs to learn a space filling surface (a surface of dimension k < k ∗ that visits random outputs y i that are sampled from a k ∗ -dimensional distribution). More precisely Theorem 2 of [Jac23] implies that the R (1) value of fitting BN-rank 1 minima explodes at a rate of 2(1 -1 k ∗ ) log N as the number of datapoints N grows, which could explain why we rarely observe such minima in practice, but another explanation could be that these minima are very narrow, as explained in Section 4.1.

In our experiments we often encountered rank overestimating minima and we are less sure about how to avoid them, though it seems that increasing the depth helps (see Figure 1), and that SGD might help too by analogy with the linear case [WJ23]. Thankfully overestimating the rank is less problematic for generalization, as supported by the fact that it is possible to approximate BN-rank k ∗ with a higher rank function with any accuracy, while doing so with a low rank function requires a pathological function.

3.2 Second Correction

We now identify a few properties of the second correction R (2) :

Proposition 4. If there is a limiting representation as L → 0 in the optimal representation of f , then R (2) ( f ) ≥ 0 . Furthermore:

If R (0) ( f ◦ g ) = R (0) ( f ) = R (0) ( g ) and R (1) ( f ◦ g ) = R (1) ( f ) + R (1) ( g ) , then √ R (2) ( f ◦ g ) ≤ √ R (2) ( f ) + √ R (2) ( g ) .

If R (0) ( f + g ) = R (0) ( f ) + R (0) ( g ) and R (1) ( f + g ) = R (1) ( f ) + R (1) ( g ) , then R (2) ( f + g ) ≤ R (2) ( f ) + R (2) ( g ) .

If A p Ω is k = Rank A -dimensional and completely positive for all p ∈ [0 , 1] , where A p has its non-zero singular taken to the p -th power, then R (2) ( x ↦→ Ax ; Ω) = 1 2 ∥ ∥ log + A T A ∥ ∥ 2 .

While the properties are very similar to those of R (1) , the necessary conditions necessary to apply them are more restrictive. There might be case where the first two terms R (0) and R (1) do not uniquely determine a minimum, in which case the second correction R (2) needs to be considered.

In linear networks the second correction R (2) ( A ) = 1 2 ∥ ∥ log + A T A ∥ ∥ 2 plays an important role, as it bounds the operator norm of A (which is not bounded by R (0) ( A ) = Rank A nor R (1) ( A ) = 2 log | A | + ), thus guaranteeing the convergence of the hidden representations in the middle of network. We hoped at first that R (2) would have similar properties in the nonlinear case, but we were not able to prove anything of the sort. Actually in contrast to the linear setting, the representations of the network can diverge as L →∞ , which explains why the R (2) does not give any similar control, which would guarantee convergence.

3.3 Representation geodesics

One can think of the sequence of hidden representations ˜ α 1 , . . . , ˜ α L as a path from the input representation to the output representation that minimizes the weight norm ∥ W ℓ ∥ 2 required to map from one representation to the next. As the depth L grows, we expect this sequence to converge to a form of geodesic in representation space. Such an analysis has been done in [Owh20] for ResNet where these limiting geodesics are continuous.

Two issues appear in the fully-connected case. First a representation ˜ α ℓ remains optimal after any swapping of its neurons or other symmetries, but this can easily be solved by considering representations ˜ α ℓ up to orthogonal transformation, i.e. to focus on the kernels K ℓ ( x, y ) = ˜ α ℓ ( x ) T ˜ α ℓ ( y ) . Second the limiting geodesics of fully-connected networks are not continuous, and as such they cannot be described by a local metric.

We therefore turn to the representation cost of DNNs to describe the hidden representations of the network, since the ℓ -th pre-activation function ˜ α ( ℓ ) : Ω → R n ℓ in a network which minimizes the parameter norm must satisfy

Thus the limiting representations ˜ α p = lim L →∞ ˜ α ℓ L (for a sequence of layers ℓ L such that lim L →∞ ℓ L / L = p ∈ (0 , 1) ) must satisfy

Let us now assume that the limiting geodesic is continuous at p (up to orthogonal transformation, which do not affect the representation cost), meaning that any other sequence of layers ℓ ′ L converging to the same ratio p ∈ (0 , 1) would converge to the same representation. The taking the limits with two sequences lim ℓ L L = p = lim ℓ ′ L L such that lim ℓ ′ L -ℓ L = + ∞ and and taking the limit of the equality

we obtain that R (0) ( σ (˜ α p ) → ˜ α p ; Ω) = R (0) ( f ) and R (1) ( σ (˜ α p ) → ˜ α p ; Ω) = 0 . This implies that σ (˜ α ( x )) = ˜ α ( x ) at any point x where Rank Jf ( x ) = R (0) ( f ; Ω) , thus R (0) ( id ; ˜ α p (Ω)) = R (0) ( f ; Ω) and R (1) ( id ; ˜ α p (Ω)) = 0 if Rank Jf ( x ) = R (0) ( f ; Ω) for all x ∈ Ω .

3.3.1 Identity

When evaluated on the identity, the first two terms R (0) ( id ; Ω) and R (1) ( id ; Ω) describe properties of the domain Ω .

For any notion of rank, Rank( id ; Ω) defines a notion of dimensionality of Ω . The Jacobian rank Rank J ( id ; Ω) = max x ∈ Ω dim T x Ω is the maximum tangent space dimension, while the Bottleneck rank Rank BN ( id ; Ω) is the smallest dimension Ω can be embedded into. For example, the circle Ω = S 2 -1 has Rank J ( id ; Ω) = 1 and Rank BN ( id ; Ω) = 2 .

On a domain Ω where the two notions of dimensionality match Rank J ( id ; Ω) = Rank BN ( id ; Ω) = k , the first correction R (1) ( id ; Ω) is non-negative since for any x with dim T x Ω = k , we have R (1) ( id ; Ω) ≥ log | P T x | + = 0 . The R (1) ( id ; Ω) value measures how non-planar the domain Ω is, being 0 only if Ω is k -planar, i.e. its linear span is k -dimensional:

Proposition 5. For a domain with Rank J ( id ; Ω) = Rank BN ( id ; Ω) = k , then R (1) ( id ; Ω) = 0 if and only if Ω is k -planar and completely positive.

This proposition shows that the R (1) term does not only bound the Jacobian of f as shown in Theorem 3, but also captures properties of the curvature of the domain/function.

Thus at ratios p where the representation geodesics converge continuously, the representations ˜ α p (Ω) are k = R (0) ( f ; Ω) -planar, proving the Bottleneck structure that was only observed empirically in [Jac23]. But the assumption of convergence over which we build this argument does not hold in general, actually we give in the appendix an example of a simple function f whose optimal representations diverges in the infinite depth limit. This is in stark contrast to the linear case, where the second correction R (2) ( A ) = 1 2 ∥ ∥ log + A T A ∥ ∥ 2 guarantees convergence, since it bounds the operator norm of A . To prove and describe the bottleneck structure in nonlinear DNNs, we therefore need to turn to another strategy.

4 Bottleneck Structure in Large Depth Networks

Up to now we have focused on one aspect of the Bottleneck structure observed in [Jac23]: that the representations α ℓ ( X ) inside the Bottleneck are approximately k -planar. But another characteristic of this phenomenon is that the weight matrices W ℓ in the bottleneck have k dominating singular values, all close to 1 . This property does not require the convergence of the geodesics and can be proven with finite depth rates:

Theorem 6. Given parameters θ of a depth L network, with ∥ θ ∥ 2 ≤ kL + c 1 and a point x such that Rank Jf θ ( x ) = k , then there are w ℓ × k (semi-)orthonormal V ℓ such that ∑ L ℓ =1 ∥ ∥ W ℓ -V ℓ V T ℓ -1 ∥ ∥ 2 F ≤ c 1 -2 log | Jf θ ( x ) | + thus for any p ∈ (0 , 1) there are at least (1 -p ) L layers ℓ with

Note how we not only obtain finite depth rates, but our result has the advantage of being applicable to any parameters with a sufficiently small parameter norm (close to the minimal norm solution). The bound is tighter at optimal parameters in which case c 1 = R (1) ( f θ ) , but the theorem shows that the Bottleneck structure generalizes to points that are only almost optimal.

To prove that the pre-activations ˜ α ℓ ( X ) are approximately k -dimensional for some dataset X (that may or may not be the training set) we simply need to show that the activations α ℓ -1 ( X ) do not diverge, since ˜ α ℓ ( X ) = W ℓ α ℓ -1 ( X ) + b ℓ (and one can show that the bias will be small at almost every layer too). By our counterexample we know that we cannot rule out such explosion in general, however if we assume that the NTK [JGH18] Θ ( L ) ( x, x ) is of order O ( L ) , then we can guarantee to convergence of the activations α ℓ -1 ( X ) at almost every layer:

Theorem 7. Given balanced parameters θ of a depth L network, with ∥ θ ∥ 2 ≤ kL + c 1 and a point x such that Rank Jf θ ( x ) = k then if Tr [ Θ ( L ) ( x, x ) ] ≤ cL , then ∑ L ℓ =1 ∥ α ℓ -1 ( x ) ∥ 2 2 ≤ c max { 1 ,e c 1 k } k | Jf θ ( x ) | 2 /k + L and thus for all p ∈ (0 , 1) there are at least (1 -p ) L layers such that

Note that the balancedness assumption is not strictly necessary and could easily be weakened to some form of approximate balancedness, since we only require the fact that the parameter norm ∥ W ℓ ∥ 2 F is well spread out throughout the layers, which follows from balancedness.

The NTK describes the narrowness of the minima [JGH20], and the assumption of bounded NTK is thus related to stability under large learning rates. There are multiple notions of narrowness that have been considered:

The operator norm of the Hessian H (which is closely related to the top eigenvalue of the NTKGrammatrix Θ ( L ) ( X,X ) especially in the MSE case where at any interpolating function ∥ H ∥ op = 1 N ∥ ∥ Θ ( L ) ( X,X ) ∥ ∥ op ) which needs to be bounded by 2 / η to have convergence when training with gradient descent with learning rate η .

The trace of the Hessian (in the MSE case Tr H = 1 N TrΘ ( L ) ( X,X ) ) which has been shown to describe the bias of stochastic gradient descent or approximation thereof [DML21, LWA21].

Thus boundedness of almost all activations as L → ∞ can be guaranteed by assuming either 1 N ∥ ∥ Θ ( L ) ( X,X ) ∥ ∥ op ≤ cL (which implies d out TrΘ ( L ) ( X,X ) ≤ cL ) or 1 N TrΘ ( L ) ( X,X ) ≤ cL directly, corresponding to either gradient descent with η = 2 / cL or stochastic gradient descent with a similar scaling of η ).

Note that one can find parameters that learn a function with a NTK that scales linearly in depth, but it is not possible to represent non-trivial functions with a smaller NTK Θ ( L ) ≪ L . This is why we consider a linear scaling in depth to be the 'normal' size of the NTK, and anything larger to be narrow.

Putting the two Theorems together, we can prove the Bottleneck structure for almost all representations ˜ α ℓ ( X ) :

Corollary 8. Given balanced parameters θ of a depth L network with ∥ θ ∥ 2 ≤ kL + c 1 and a set of points x 1 , . . . , x N such that Rank Jf θ ( x i ) = k and 1 N Tr [ Θ ( L ) ( X,X ) ] ≤ cL , then for all p ∈ (0 , 1) there are at least (1 -p ) L layers such that

4.1 Narrowness of rank-underestimating minima

We know that the large R (1) value of BN-rank 1 fitting functions is related to the explosion of its derivative, but a large Jacobian also leads to a blow up of the NTK:

Proposition 9. For any point x , we have

where ∂ 2 xy Θ( x, x ) is understood as a d in d out × d in d out matrix.

Furthermore, for any two points x, y such that the pre-activations of all neurons of the network remain constant on the segment [ x, y ] , then either ∥ Θ( x, x ) ∥ op or ∥ Θ( y, y ) ∥ op is lower bounded by

0

5

10

15

20

25

10 0

si

(

W

)

(a) Spectra of the

W

ℓ

s

0.04

0.03

0.02

0.01

0.00

0.12

0.10

0.08

0.06

0.04

0.02

0.00

0.02

(b) PCA of

α

6

(

X

)

.

(c) Dist.

∥

α

2

(

x

)

-

α

2

(

x

0

)

∥

.

With some additional work, we can show the the NTK of such rank-underestimating functions will blow up, suggesting a narrow minimum:

Theorem 10. Let f ∗ : Ω → R d out be a function with Jacobian rank k ∗ > 1 (i.e. there is a x ∈ Ω with Rank Jf ∗ ( x ) = k ∗ ), then with high probability over the sampling of a training set x 1 , . . . , x N (sampled from a distribution with support Ω ), we have that for any parameters θ of a deep enough network that represent a BN-rank 1 function f θ that fits the training set f θ ( x i ) = f ∗ ( x i ) with norm ∥ θ ∥ 2 = L + c 1 then there is a point x ∈ Ω where the NTK satisfies

One the one hand, we know that c 1 must satisfy c 1 ≥ R (1) ( f θ ) ≥ 2(1 -1 k ∗ ) log N but if c 1 is within a factor of 2 of this lower bound c 1 < 4(1 -1 k ∗ ) log N , then the above shows that the NTK will blow up a rate N α L for a positive α .

The previous explanation for why GD avoids BN-rank 1 fitting functions was that when N is much larger than the depth L (exponentially larger), there is a rank-recovering function with a lower parameter norm than any rank-underestimating functions. But this relies on the assumption that GD converges to the lower norm minima, and it is only true for sufficiently small depths. In contrast the narrowness argument applies for any large enough depth and does not assume global convergence.

Of course the complete explanation might be a mix of these two reasons and possbily some other phenomenon too. Proving why GD avoids minima that underestimate the rank with a rank 1 < k < k ∗ also remains an open question.

5 Numerical Experiment: Symmetry Learning

In general, functions with a lot of symmetries have low BN-rank since a function f with symmetry group G can be decomposed as mapping the inputs Ω to the inputs module symmetries Ω / G and then mapping it to the outputs, thus Rank BN ( f ; Ω) ≤ dim Ω / G where dim Ω / G is the smallest dimension Ω / G can be embedded into. Thus the bias of DNNs to learn function with a low BN-rank can be interpreted as the network learning symmetries of the task. With this interpretation, overestimating the rank corresponds to failing to learn all symmetries of the task, while underestimating the rank can be interpreted as the network learning spurious symmetries that are not actual symmetries of the task.

To test this idea, we train a network to predict high dimensional dynamics with high dimensional symmetries. Consider the loss C ( v ) = ∥ ∥ vv T -( ww T + E )∥ ∥ 2 F where w ∈ R d is a fixed unit vector and E is a small noise d × d matrix. We optimize v with gradient descent to try and fit the true vector w (up to a sign). One can think of these dynamics as learning a shallow linear network vv T with a single hidden neuron. We will train a network to predict the evolution of the cost in time C ( v ( t )) .

For small noise matrix E , the GD dynamics of v ( t ) are invariant under rotation around the vector w . As a result, the high-dimensional dynamics of v ( t ) can captured by only two summary statistics u ( v ) = (( w T v ) 2 , ∥ ∥ ( I -ww T ) v ∥ ∥ 2 ) : the first measures the position along the axis formed by w and the second the distance to this axis [AGJ22]. The evolution of the summary statistics is (approximately) self-consistent (using the fact that ∥ v ∥ 2 = ( w T v ) 2 + ∥ ∥ ( I -ww T ) v ∥ ∥ 2 ):

Our goal now is to see whether a DNN can learn these summary statistics, or equivalently learn the underlying rotation symmetry. To test this, we train a network on the following supervised learning problem: given the vector v (0) at initialization, predict the loss ( C ( v (1)) , . . . , C ( v ( T ))) over the next T GDsteps. For E = 0 , the function f ∗ : R d → R T that is to be learned has BN-rank 2, since one can first map v (0) to the corresponding summary statistics u ( v (0)) ∈ R 2 , and then solve the differential equation on the summary statistics ( u (1) , . . . , u ( T )) over the next T steps, and compute the cost C ( v ) = ∥ v ∥ 4 -2( w T v ) 2 +1+ O ( ∥ E ∥ ) from u .

We observe in Figure 2 that a large depth L 2 -regularized network trained on this task learns the rotation symmetry of the task and learns two dimensional hidden representations that are summary statistics (summary statistics are only defined up to bijections, so the learned representation match u ( v ) only up to bijection but they can be recognized from the fact that the GF paths do not cross on the 2D representation).

6 Conclusion

We have computed corrections to the infinite depth description of the representation cost of DNNs given in [Jac23], revealing two regularity R (1) , R (2) measures that balance against the dominating low rank/dimension bias R (0) . We have also partially described another regularity inducing bias that results from large learning rates. We argued that these regularity bias play a role in stopping the network from underestimating the 'true' BN-rank of the task (or equivalently overfitting symmetries).

We have also proven the existence of a bottleneck structure in the weight matrices and under the condition of a bounded NTK of the learned representations, where most hidden representations are approximately k = R (0) ( f θ ) -dimensional, with only a few high-dimensional representations.

References

[ABMM18] Raman Arora, Amitabh Basu, Poorya Mianjy, and Anirbit Mukherjee. Understanding deep neural networks with rectified linear units. In International Conference on Learning Representations , 2018.

[AGJ22] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for SGD: Effective dynamics and critical scaling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

[Bac17] Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research , 18(1):629-681, 2017.

[BHH59] Jillian Beardwood, J. H. Halton, and J. M. Hammersley. The shortest path through many points. Mathematical Proceedings of the Cambridge Philosophical Society , 55(4):299-327, 1959.

[CB20] Lénaïc Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory , volume 125 of Proceedings of Machine Learning Research , pages 1305-1338. PMLR, 09-12 Jul 2020.

[DKS21] Zhen Dai, Mina Karzand, and Nathan Srebro. Representation costs of linear neural networks: Analysis and design. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021.

[DML21] Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems , 34:2744927461, 2021.

[GLSS18] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 1832-1841. PMLR, 1015 Jul 2018.

[HLXZ18] Juncai He, Lin Li, Jinchao Xu, and Chunyue Zheng. Relu deep neural networks and linear finite elements. arXiv preprint arXiv:1807.03973 , 2018.

[Jac23] Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions. In The Eleventh International Conference on Learning Representations , 2023.

[JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems 31 , pages 8580-8589. Curran Associates, Inc., 2018.

[JGH20] Arthur Jacot, Franck Gabriel, and Clément Hongler. The asymptotic spectrum of the hessian of dnn throughout training. In International Conference on Learning Representations , 2020.

[JGHG22] Arthur Jacot, Eugene Golikov, Clément Hongler, and Franck Gabriel. Feature learning in l 2 -regularized dnns: Attraction/repulsion and sparsity. In Advances in Neural Information Processing Systems , volume 36, 2022.

[LJ22] Thien Le and Stefanie Jegelka. Training invariances and the low-rank phenomenon: beyond linear networks. In International Conference on Learning Representations , 2022.

[LWA21] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:2110.06914 , 2021.

[OW22] Greg Ongie and Rebecca Willett. The role of linear layers in nonlinear interpolating networks. arXiv preprint arXiv:2202.00856 , 2022.

[Owh20] Houman Owhadi. Do ideas have shape? plato's theory of forms as the continuous limit of artificial neural networks. arXiv preprint arXiv:2008.03920 , 2020.

[SHN + 18] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research , 19(1):2822-2878, 2018.

[TVS22] Nadav Timor, Gal Vardi, and Ohad Shamir. Implicit regularization towards rank minimization in relu networks. arXiv preprint arXiv:2201.12760 , 2022.

[TZ15] Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw) , pages 1-5. IEEE, 2015.

[WJ23] Zihan Wang and Arthur Jacot. Implicit bias of sgd in l 2 -regularized linear dnns: Oneway jumps from high to low rank, 2023.

A Properties of the Corrections

A.1 First Correction

Theorem 11 (Theorem 3 from the main) . For all inputs x where Rank Jf ( x ) = R (0) ( f ; Ω) , we have R (1) ( f ) ≥ 2 log | Jf ( x ) | + , furthermore:

If R (0) ( f ◦ g ) = R (0) ( f ) = R (0) ( g ) , then R (1) ( f ◦ g ) ≤ R (1) ( f ) + R (1) ( g ) .

If R (0) ( f + g ) = R (0) ( f ) + R (0) ( g ) , then R (1) ( f + g ) ≤ R (1) ( f ) + R (1) ( g ) .

If P Im A T Ω and A Ω are k = Rank A dimensional and completely positive (i.e. they can be embedded with an isometric linear map into R m + for some m ), then R (1) ( x ↦→ Ax ; Ω) = 2 log | A | + .

Proof. For the first bound, we remember that R ( f ; Ω , L ) ≥ L ∥ Jf ∥ 2 / L 2 / L , therefore

where we used s 2 L -1 = e 2 L log s -1 ≥ 2 L log s .

Since R ( f ◦ g ; Ω , L 1 + L 2 ) ≤ R ( f ; L 1 ) + R ( g ; L 2 ) , we have

Since R ( f + g ; Ω , L ) ≤ R ( f ; Ω , L ) + R ( g ; Ω , L ) , we have

By the first bound, we know that R (1) ( x ↦→ Ax ; Ω) ≥ 2 log | A | + , we now need to show R (1) ( x ↦→ Ax ; Ω) ≤ 2 log | A | + . Let us define the set of completely positive representations as the set of bilinear kernels K ( x, y ) = x T B T By such that Bx has non-negative entries for all x ∈ Ω (we say that a kernel K is completely positive over Ω if it can be represented in this way for some choice of B ). The set of completely positive representations is convex, since for K ( x, y ) = x T B T By and ˜ K ( x, y ) = x T ˜ B T ˜ By , we have

The conditions that there are O in and O out with O T in O in = P Im A T and O T out O out = P Im A such that O in Ω ∈ R k 1 + and O out A Ω ∈ R k 2 + is equivalent to saying that the kernels K in ( x, y ) = x T P Im A T y and K out ( x, y ) = x T A T Ax are completely positive over Ω .

By the convexity of completely positive representations, the interpolation K p = pK in +(1 -p ) K out is completely positive for all p ∈ [0 , 1] . Now choose for all depths L and all layers ℓ = 1 , . . . , L -1

a matrix B L,ℓ such that K p = ℓ L ( x, y ) = x T B T L,ℓ B L,ℓ y and then choose the weights W ℓ of the depth L network as

using the convention B L, 0 = I d in and B L,L = I out . By induction, we show that for any input x ∈ Ω the activation of the ℓ -th hidden layer is B L,ℓ x . This is true for ℓ = 1 , since W 1 = B L, 1 and therefore p (1) ( x ) = B L, 1 x which has positive entries so that q (1) ( x ) = σ ( p (1) ( x ) ) = B L, 1 x . Then by induction

which has positive entries, so that again q ( ℓ ) ( x ) = σ ( p ( ℓ ) ( x ) ) = B L,ℓ x . In the end, we get p ( L ) ( x ) = Ax as needed.

Let us now compute the Frobenius norms of the weight matrices ∥ W ℓ ∥ 2 F = Tr [ B T L,ℓ B L,ℓ ( B T L,ℓ -1 B L,ℓ -1 ) + ] as L →∞ , remember that B T L,ℓ B L,ℓ = ℓ L P Im A T +(1 -ℓ L ) A T A , therefore the matrices B T L,ℓ B L,ℓ and B T L,ℓ -1 B L,ℓ -1 converge to each other, so that at first order B T L,ℓ B L,ℓ ( B T L,ℓ -1 B L,ℓ -1 ) + converges to P Im A T , so that ∥ W ℓ ∥ 2 F → Rank A , so that ∑ L ℓ =1 ∥ W ℓ ∥ 2 F -L Rank A converges to a finite value as L →∞ . To obtain this finite limit, we need to study approximate the next order

But as we sum all these second order terms, they cancel out, and we are left with

We have therefore build parameters θ that represent the function x ↦→ Ax with parameter norm ∥ θ ∥ 2 = L Rank A +2log | A | + + O ( L -1 ) , which upper bounds the representation cost, thus implying that R (1) ( x ↦→ Ax ; Ω) ≤ 2 log | A | + as needed.

A.2 Identity

Proposition 12 (Proposition 5 from the main) . For a domain with Rank J ( id ; Ω) = Rank BN ( id ; Ω) = k , then R (1) ( id ; Ω) = 0 if and only if Ω is k -planar and completely positive.

Proof. First if Ω is completely positive and k -planar one can represent the identity with a depth L network of parameter norm Lk , by taking W 1 = O,W ℓ = P Im O , W L = O T where O is the m × d so that O Ω ⊂ R m + and O T O = P spanΩ . Thus R (1) ( id ; Ω) = 0 (and all other corrections as well).

We will show that for any two points x, y ∈ Ω with k -dim tangent spaces, their tangent spaces must match if R (1) ( id ; Ω) = 0 .

Let A = Jα ( L -1) ( x ) | T x Ω and B = Jα ( L -1) ( y ) | T y Ω be the be the Jacobian of the last hidden activations restricted to the tangent spaces, we know that

so that given any weight matrix W L whose image contains T x Ω and T y Ω , we can write

Without loss of generality, we may assume that the span of T x Ω and T y Ω is full output space, and therefore that W L W T L is invertible.

Now we now that any parameters that represent the identity on Ω and has A = Jα ( L -1) ( x ) | T x Ω and B = Jα ( L -1) ( y ) | T y Ω must have parameter norm at least

Subtracting kL and taking L →∞ , we obtain that

If we optimize W L only up to scaling (i.e. optimize aW L over a ) we see that at the optimum, we always have ∥ W L ∥ 2 F = k . This allows us to rewrite the optimization as

̸

The only way to put the first term inside the maximum to 0 is to have W L W T L = P T x Ω , but this leads to an exploding second term if P T x Ω = P T y Ω .

Under the assumption of C -uniform Lipschitzness of the representations (that for all ℓ , the functions ˜ α ℓ and ( α ℓ → f θ ) are C -Lipschitz), one can show a stronger version of the above:

Proposition 13. For a C -uniformly Lipschitz sequence of ReLU networks representing the function f , we have

Proof. The decomposition of the difference

for the w ℓ × w ℓ diagonal matrices D ℓ ( x ) = diag( ˙ σ (˜ α ℓ ( x ))) , implies the bound

since ∥ AB ∥ ∗ ≤ ∥ A ∥ 2 F + ∥ B ∥ 2 F 2 and ( D ℓ ( x ) -D ℓ ( y )) 2 = ( D ℓ ( x ) -D ℓ ( y )) . Now since

with the convention D 0 ( x ) = I d in and D L ( x ) = I d out . We obtain that

This implies the bound

and thus

A.3 Second Correction

Proposition 14 (Proposition 4 from the main) . If there is a limiting representation as L → 0 in the optimal representation of f , then R (2) ( f ) ≥ 0 . Furthermore:

If R (0) ( f ◦ g ) = R (0) ( f ) = R (0) ( g ) and R (1) ( f ◦ g ) = R (1) ( f ) + R (1) ( g ) , then √ R (2) ( f ◦ g ) ≤ √ R (2) ( f ) + √ R (2) ( g ) .

If A p Ω is k = Rank A -dimensional and completely positive for all p ∈ [0 , 1] , where A p has its non-zero singular taken to the p -th power, then R (2) ( x ↦→ Ax ; Ω) = 1 2 ∥ ∥ log + A T A ∥ ∥ 2 .

Proof. We start from the inequality

We subtract ( L f + L g ) R (0) ( f ◦ g ) + R (1) ( f ◦ g ) divide by L f + L g and take the limit of increasing depths L f , L g with lim L g ,L f →∞ L f L f + L g = p ∈ (0 , 1) to obtain

̸

If K p is the limiting representation at a ratio p ∈ (0 , 1) , we have R (2) ( f ; Ω) = 1 p R (2) ( K p ; Ω) + 1 1 -p R (2) ( K p → f ; Ω) and p must minimize the RHS since if it was instead minimized at a different ratio p ′ = p , one could find a lower norm representation by mapping to K p in the first p ′ L layers and then back to the outputs. Now there are two possiblities, either R (2) ( K p ; Ω) and R (2) ( K p → f ; Ω) are non-negative in which case the minimum is attained at some p ∈ (0 , 1) and R (2) ( f ; Ω) ≥ 0 , or one or both is negative in which case the above is minimized at p ∈ { 0 , 1 } and R (2) ( f ; Ω) = -∞ . Since we assumed p ∈ (0 , 1) , we are in the first case.

(1) To prove the first property, we optimize the RHS of 1 over all possible choices of p (and assuming that R (2) ( f ; g (Ω)) , R (2) ( g ; Ω) ≥ 0 ) we obtain

This follows from the inequality R ( f + g ; Ω , L ) ≤ R ( f ; g (Ω) , L ) + R ( g ; Ω , L ) after subtracting the R (0) and R (1) terms, dividing by L and taking L →∞ .

(3) If A = USV T , one chooses W ℓ = U ℓ S 1 L U T ℓ -1 with U 0 = V , U L = U and U ℓ chosen so that U ℓ S ℓ L V T Ω ∈ R n ℓ + , choosing large enough widths n ℓ . This choice of representation of A is optimal, i.e. its parameter norm matches the representation cost L Tr [ S 2 L ] = L Rank A + 2log | A | + + 1 2 L ∥ ∥ log + A T A ∥ ∥ 2 + O ( L -2 ) .

We know that

Proof. Since

If the preimage of A matches the image of B then | AB | + = | A | + | B | + . We therefore have

This implies that

with the convention D L ( x ) = I w out .

Our goal is to show that the LHS is a sum of positive value which sum up to a finite positive value, which will imply that most of the summands must be very small.

First observe that

which is positive since the eigenvalues of D ( x ) are ≤ 1 .

To show that the other part of the summands ∥ W ℓ ∥ 2 F -k -2 log ∣ ∣ P Im J (˜ α ℓ → f θ )( x ) T W ℓ P Im Jα ℓ -1 ( x ) ∣ ∣ + is positive, we give lower bound it. First, it can be rewritten as

B Bottleneck Structure

This first result shows the existence of a Bottleneck structure on the weight matrices:

Theorem 15 (Theorem 6 from the main) . Given parameters θ of a depth L network, with ∥ θ ∥ 2 ≤ kL + c 1 and a point x such that Rank Jf θ ( x ) = k , then there are w ℓ × k (semi-)orthonormal U ℓ , V ℓ such that

thus for any p ∈ (0 , 1) there are at least (1 -p ) L layers ℓ with

Now for a general matrix A , we have

for the SVD decomposition A = USV T . We can thus lower bound

where U ℓ S ℓ V T ℓ is the SVD decomposition of P Im J (˜ α ℓ → f θ )( x ) T W ℓ P Im Jα ℓ -1 ( x ) (which we know has rank k since it must match the rank of Jf θ ( x ) ).

Since ∥ θ ∥ 2 = ∑ ℓ ∥ W ℓ ∥ 2 F + ∥ b ℓ ∥ 2 ≤ kL + c 1 , we have

And for any p ∈ (0 , 1) there are at most pL layers ℓ with

The fact that almost all weight matrices W ℓ are approximately k -dim would imply that the preactivations ˜ α ℓ ( X ) = W ℓ α ℓ -1 ( X ) are k -dim too under the condition that the activations α ℓ -1 ( X ) do not diverge. Assuming a bounded NTK is sufficient to guarantee that these activations converge at almost every layer:

Theorem 16 (Theorem 7 from the main) . Given balanced parameters θ of a depth L network, with ∥ θ ∥ 2 ≤ kL + c 1 and a point x such that Rank Jf θ ( x ) = k then if 1 N Tr [ Θ ( L ) ( x, x ) ] ≤ cL , then ∑ L ℓ =1 ∥ α ℓ -1 ( x ) ∥ 2 2 ≤ c max { 1 ,e c 1 k } k | Jf θ ( x ) | 2 /k + L and thus for all p ∈ (0 , 1) there are at least (1 -p ) L layers such that

Proof. We have

we therefore need to lower bound ∥ J (˜ α ℓ → α L )( x ) ∥ 2 F to show that the activations ∥ α ℓ -1 ( x ) ∥ 2 2 must be bounded at almost every layer.

We will lower bound ∥ J (˜ α ℓ → α L )( x ) ∥ 2 F by ∥ J (˜ α ℓ → α L )( x ) P ℓ ∥ 2 F for P ℓ the orthogonal projection to the image Im J ˜ α ℓ ( x ) . Note J (˜ α ℓ → α L )( x ) P ℓ and Jf θ ( x ) have the same rank.

By the arithmetic-geometric mean inequality, we have ∥ A ∥ 2 F ≥ Rank A | A | 2 / k + , yielding

Now the balancedness of the parameters (i.e. ∥ W ℓ,i · ∥ 2 + b 2 ℓ,i = ∥ W ℓ +1 , · i ∥ 2 , ∀ ℓ, i ) implies that the parameter norms are increasing ∥ W ℓ +1 ∥ 2 F ≥ ∥ W ℓ ∥ 2 F and thus

Thus

and

and therefore

Thus

which implies that there are at most pL layers ℓ with

Note that for the MSE loss in the limit λ ↘ 0 , the Hessian at a global minimum (i.e. f θ interpolates the training set) equals Tr[ HL ( θ )] = 1 N Tr [ Θ ( L ) ( X,X ) ] . If we then assume that the trace of the Hessian is bounded by cL , we get that Tr [ Θ ( L ) ( X,X ) ] ≤ cNL and thus there are at least (1 -p ) L layers where

thus guaranteeing the infinite depth convergence of training set activations α ℓ -1 ( X ) on those layers. Putting the two above theorems together, we can prove that the pre-activations are k -dim at almost every layer:

Corollary 17 (Corollary 8 from the main) . Given balanced parameters θ of a depth L network with ∥ θ ∥ 2 ≤ kL + c 1 and a set of points x 1 , . . . , x N such that Rank Jf θ ( x i ) = k and 1 N Tr [ Θ ( L ) ( X,X ) ] ≤ cL , then for all p ∈ (0 , 1) there are at least (1 -p ) L layers such that

Proof. Since ˜ α ℓ ( X ) = W ℓ α ℓ -1 ( X ) + b ℓ 1 T N we know that

By Theorems 15 and 16, there are for any p ∈ (0 , 1 2 ) at least (1 -2 p ) L layers such that

so that

C Local Minima Stability

In this section we motivate the assumption that the Jacobian J ˜ α ℓ ( x ) is uniformly bounded in operator norm as L →∞ . The idea is that solutions with a blowing up Jacobian J ˜ α ℓ ( x ) correspond to very narrow local minima.

The narrowness of a local minimum is related to the Neural Tangent Kernel (or Fisher matrix). We have that

A large Jacobian Jf θ ( x ) leads to a blow up of the derivative of the NTK:

Proposition 18 (Proposition 9 from the main) . For any point x , we have

where ∂ 2 xy Θ( x, x ) is understood as a d in d out × d in d out matrix.

Furthermore, for any two points x, y such that the pre-activations of all neurons of the network remain constant on the segment [ x, y ] , then either ∥ Θ( x, x ) ∥ op or ∥ Θ( y, y ) ∥ op is lower bounded by

Proof. (1) For any point x , we have

On the other hand, we have

where we used the fact that D ℓ ( x ) D ℓ ( x ) = D ℓ ( x ) . This applies to the case ℓ = L and ℓ = 1 too, using the definition D L ( x ) = I d out and D 0 ( x ) = I d in . This implies

where we used the geometric/arithmetic mean inequality for the second inequality.

If u, v are right and left singular vectors of Jf θ ( x ) with singular value s , then the above bound equals Ls 2 -2 L .

(2) Now let us consider a segment γ ( t ) = (1 -t ) x + ty between two points x, y with no changes of activations on these paths (i.e. D ℓ ( γ ( t )) is constant for all t ∈ [0 , 1] ). Defining u = y -x ∥ y -x ∥ and v = Jf θ ( x ) u ∥ Jf θ ( x ) u ∥ , we have

∂ t v T Θ( γ ( t ) , γ ( t )) v = ∥ x -y ∥ ∂ x ( v T Θ( γ ( t ) , γ ( t )) v ) [ u ] + ∥ x -y ∥ ∂ y ( v T Θ( γ ( t ) , γ ( t )) v ) [ u ] and since ∂ xx Θ( γ ( t ) , γ ( t )) = 0 and ∂ yy Θ( γ ( t ) , γ ( t )) = 0 for all t ∈ [0 , 1] , we have ∂ 2 t ( v T Θ( γ ( t ) , γ ( t )) v ) = 2 ∥ x -y ∥ 2 ∂ 2 xy ( v T Θ( γ ( t ) , γ ( t )) v ) [ u, u ] ≥ 2 L ∥ x -y ∥ 2 ∥ Jf θ ( x ) u ∥ 2 -2 / L 2 Since v T Θ( γ ( t ) , γ ( t )) v ≥ 0 for all t ∈ [0 , 1] then either v T Θ( x, x ) v ≥ L 4 ∥ x -y ∥ 2 ∥ Jf θ ( x ) u ∥ 2 -2 / L 2 or v T Θ( y, y ) v ≥ L 4 ∥ x -y ∥ 2 ∥ Jf θ ( x ) u ∥ 2 -2 / L 2 .

Rank-underestimating fitting functions typically feature exploding derivatives, which was used to show in [Jac23] that BN-rank 1 fitting functions must have a R (1) term that blows up iwith the number of datapoints N . With some additional work, we can show that the NTK will blow up at some x :

Theorem 19 (Theorem 10 from the main) . Let f ∗ : Ω → R d out be a function with Jacobian rank k ∗ > 1 (i.e. there is a x ∈ Ω with Rank Jf ∗ ( x ) = k ∗ ), then with high probability over the sampling of a training set x 1 , . . . , x N (sampled from a distribution with support Ω ), we have that for any parameters θ of a deep enough network that represent a BN-rank 1 function f θ that fits the training set f θ ( x i ) = f ∗ ( x i ) with norm ∥ θ ∥ 2 = L + c 1 then there is a point x ∈ Ω where the NTK satisfies

Proof. For all i we define d 1 ,i and d 2 ,i to be the distance between y i and its closest and second closest point in y 1 , . . . , y N . Following the argument in [BHH59], the shortest path that goes through all points must be at least ∑ N i =1 d 1 ,i + d 2 ,i 2 (which would be tight if it is possible to always jump to the closest or second closest point along the path). Since the expected distances d 1 ,i and d 2 ,i are N -1 k ∗ since the y i are sampled from a k ∗ -dimensional distribution, the expected length of the shortest path is of order N 1 -1 k ∗ . Actually most of the distance d 1 ,i and d 2 ,i will be of order N -1 k ∗ with only a few outliers with larger or smaller distances, thus for any subset of indices I ⊂ [1 , . . . , N ] that contains a finite ratio of all indices, the sum ∑ i ∈ I d 1 ,i + d 2 ,i 2 will be of order N 1 -1 k ∗ too.

Following the argument in the proof Theorem 2 from [Jac23], we can reorder the indices so that the segment [ x 1 , x N ] will mapped f θ to a path that goes through y 1 , . . . , y N . We can therefore define the points ˜ x 1 , . . . , ˜ x N that are preimages of y 1 , . . . , y N on the segment.

.

On the interval [˜ x i , ˜ x i +1 ] there must a point x with ∥ Jf θ ( x ) ∥ op ≥ ∥ y i +1 -y i ∥ ∥ ˜ x i +1 -˜ x i ∥ . Now since ∑ ∥ ˜ x i +1 -˜ x i ∥ = ∥ x N -x 1 ∥ ≤ diagΩ there must be at least p ( N -1) intervals i with ∥ ˜ x i +1 -˜ x i ∥ ≤ diagΩ (1 -p )( N -1) , and amongst those i s they would all satisfy ∥ y i +1 -y i ∥ ≥ cN -1 k ∗ except for a few outliers. Thus we can for example guarantee that there are at least 5 6 ( N -1) intervals [˜ x i , ˜ x i +1 ] that contain a point x with ∥ Jf θ ( x ) ∥ op ≥ cN 1 -1 k ∗ .

First observe that by Theorem 2 of [Jac23], there must be a point x ∈ Ω such that ∥ Jf θ ( x ) ∥ op ≥ N 1 -1 k thus by Theorem 6, there are at least pL layers where

Consider one of those pL layers ℓ , with activatons z i = α ℓ -1 (˜ x i ) . Let i 1 , . . . , i N be the ordering of the indices so that u T ℓ z i m is increasing in m . Then the hidden representations must satisfy

and

For any two indices m 1 < m 2 separated by pN indices (where p > 0 remains finite), we have

where we used the fact that up to a few outliers d 1 ,i + d 2 ,i 2 = Ω( N 1 k ∗ ) .

Thus for L ≥ 4 c 2 e c 1 +1 N 2 k ∗ c 1 -2 log | Jf θ ( x ) | + p (diamΩ) 2 , we have u T ℓ ( z i m 2 -z i m 1 ) ≥ ( m 2 -m 1 ) c 2 e -∥ θ ( ℓ +1: L ) ∥ 2 -( L -ℓ ) 2 N -1 k ∗ . which implies that at least half of the activations z i have norm larger than c 8 e -∥ θ ( ℓ +1: L ) ∥ 2 -( L -ℓ ) 2 N 1 -1 k ∗ .

This implies that at least one fourth of the indices i satisfy for at least one fourth of the pL layers ℓ

Nowamongst these indices there are at least some such that there is a point x in the interval [˜ x i , ˜ x i +1 ] with ∥ Jf θ ( x ) ∥ op ≥ cN 1 -1 k ∗ . Since x is O ( N -1 ) -close to ˜ x i one can guarantee that ∥ α ℓ -1 ( x ) ∥ ≥ c ′ e -∥ θ ( ℓ +1: L ) ∥ 2 -( L -ℓ ) 2 N 1 -1 k ∗ for some constant c ′ .

But the Jacobian J (˜ α ℓ → f θ ) also explodes at x since

We can now lower bound the NTK

This suggests that such points are avoided not only because they have a large R (1) value, but also (if not mostly) because they lie at the bottom of a very narrow valley.

D Technical Results

D.1 Regularity Counterexample

We guve here an example of a simple function whose optimal representation geodesic does not converge, due to it being not uniformly Lipschitz:

Example 20. The function f : Ω → R 3 with Ω = [0 , 1] 3 defined by

satisfies R (0) ( f ; Ω) = 3 and R (1) ( f ; Ω) = 0 . The optimal representations of f are not uniformly Lipschitz as L →∞ .

Proof. While we are not able to identify exactly the optimal representation geodesic for the function f , we will first show that R (1) ( f ; Ω) = 0 , and then show that the uniform Lipschitzness of the optimal representations would contradict with Proposition 13.

know by Theorem 11 that R (1) ( f ; Ω) ≥ 2 log | I 3 | + = 0 . We therefore only need to construct a sequence of parameters of different depths that represent f with a squared parameter norm of order 3 L + o (1) . For simplicity, we only do this construction for even depths (the odd case can be constructed similarly). We define:

We have for all x ∈ R 3 +

and and

and

The norm of the parameters is

If we take ϵ = L -γ for γ ∈ ( 1 2 , 1) , then the terms 2 e -( L -2) ϵ and e -2( L -2) ϵ decay exponentially (at a rate of e L 1 -γ ), in addition the terms 2 ( e 2 ϵ -1 ) +( e -4 ϵ -1) and 2 ( e -2 ϵ -1 ) +( e 4 ϵ -1) are of order L -2 γ . This proves that R (1) ( f ; Ω) = 0 .

(2) Let us now assume that the optimal representation of f is C -uniform Lipschitz for some constant C , then by Proposition 13, we have that

which contradicts with the fact that R (1) ( f ; Ω) = 0 .

D.2 Extension outside FPLFs

Since all functions represented by finite depth and width networks are FPLFs, the representation cost of any such function is infinite. But we can define the representation cost of a function f that is the limit of a sequence of FPLF as the infimum over all sequences f i → f converging of lim i →∞ R ( f i ; Ω) (for some choice of convergence type that implies convergence of the Jacobians Jf i ( x ) → Jf ( x ) ). Note that since the representation cost R ( f ; Ω) is lower semi-continuous, i.e. lim inf f → f 0 R ( f ; Ω) ≥ R ( f 0 ; Ω) , this does not change the definition of the representation cost on the space of FPLFs.

E Numerical Experiments

For the first numerical experiment, the data pairs ( x, y ) were generated as follows. First we sample a 8 -dimensional 'latent vector' z , from which we define x = g ( z 1 , . . . , z 8 ) ∈ R 20 and y = h ( z 1 , z 2 ) ∈ R 20 for two random functions g : R 8 → R 20 and h : R 2 → R 20 given by two shallow networks with random parameters. Assuming that g is injective (which it is with high probability), the function f ∗ = h ◦ g -1 which maps x to y has BN-rank 2.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies the leading order expansions in $L$ of the representation cost for a large $L$, where $L$ is the depth of the model. The work concludes that there is a low-dimension and regularity tradeoff as one varies the depth\n\nSTRENGTHS:\nI think the main strength of the work is the novelty of the message it delivers: there is a low-dimension and regularity tradeoff as one varies the depth. The simplicity of the proofs is a plus.\n\nAnother important contribution is the technique to expand representation loss for a large depth, though I believe to require a little more technical soundness. See the weakness section\n\nAnother interesting insight of the paper is the concept of \"symmetry learning,\" and relating the \"spurious symmetries\" to the representation cost is also insightful\n\nAt this moment, I have some reservation about accepting this paper. If the authors answer my questions below in a satisfactory manner, I would be happy to accept it\n\nWEAKNESSES:\nThere are quite a few aspects that I think this work can be improved.\n\n1. Let $R(L)$ denote the representation cost as a function of $L$, the depth. The first step of the theory is to expand $R$ in $1/L$, which feels unjustified, why is this function Taylor-expandable? Why can we treat $L$ as a continuous variable? These points are not crucial problems in my opinion, but the authors do not explain these with sufficient emphasis, given how important they are for the results\n\n2. In my opinion, the most problematic aspect is the fact that it studies the representation cost of an almost infinite depth neural network at a finite weight decay (this also relates to point 1). See https://arxiv.org/abs/2202.04777. This reference (for example, see proposition 3 / theorem 3) implies that for a sufficiently deep neural network and a fixed weight decay, the global minimum is the origin ($\\theta=0$), independent of the depth $L$. As one decreases the depth, the global minimum jumps from the origin to a nonzero value. This means that if we restrict to the global minimum of the cost function $C$, the representation cost is not a differentiable function in $L$, and thus it is not Taylor-expandable. I think the authors need to explain why the expansion remains reasonable/correct in light of this result\n\n3. Given how strong the assumptions are (point 1), I think the authors really should have carried out some experiments to directly check the prediction. In particular, I believe the authors should present a numerical example where the author directly compute the first and second order terms of $R$ in $1/L$ and show how they change for real neural networks to make the arguments of the paper more convincing.\n\nQUESTIONS:\nSee weakness\n\nLIMITATIONS:\nSee weakness", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-02"}

comments: SUMMARY:
This paper looks at the representation cost $R(f)$ of networks as the $L_2$ norm of all of the parameters. The paper claims that for deep neural networks, computing the minimum cost representation of a function is not tractable. Hence, the look at the representation cost normalized by the number of layers $L$ as $L \to \infty$. The paper then computes the Taylor expansion with respect to $1/L$ at $L = \infty$. 

Using this expansion, they recover three terms, $R^{(0)}(f),R^{(1)}(f),R^{(2)}(f)$. They note that $R^{(0)}(f)$ is used as a regularize before and corresponded to regularizing this notion of rank called the bottleneck rank. However, they argue that this is not enough and present results that determine the type of regularity imposed by using $R^{(1)}(f), R^{(2)}(f)$ as regularizers as well.

STRENGTHS:
The paper is very interesting, as understanding the representation of functions via neural networks is very important. The description of the two new terms $R^{(1)}(f), R^{(2)}(f)$ are explored in depth from a variety of different angles. Hence I think it makes a good contribution to the understanding of neural representation learning.

WEAKNESSES:
My main issue with the paper is its presentation. The paper feels very disjointed. It would be helpful if the authors included more connecting discussion and discussed the high level picture in more detail.

QUESTIONS:
1. On line 91, what is the infinite width representation cost?

Typos

1. Line 156 instead of "of he task" -> "of the task"

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

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces two corrections to previous work that relates the representation cost of infinite-depth nonlinear networks to the bottleneck rank. These two corrections are obtained by Taylor expansion of the representation cost. In addition to the previously-known rank term $R^{(0)}(f)$ that is directly related to the bottleneck rank, the authors consider two corrections (1) $R^{(1)}(f)$, which bounds the pseudo-determinant of the Jacobian, and (2) $R^{(2)}(f)$, which plays a role in the intermediate representations in the hidden layer of the network. The authors specifically discuss several theoretical properties of $R^{(1)}(f)$. This term acts as a notion of regularity in the learning process of deep nonlinear networks, and it explains various interesting properties, such as how deep networks trained with conventional L2-regularized loss are uniquely determined (i.e., which $f$ among candidates with the same $R^{(0)}(f)$ is selected) and why the rank of the true function is not underestimated. The authors also briefly demonstrate the relationship between the curvature of the domain manifold and $R^{(1)}$ using the rank of the identity function, and discuss the connection between the learning rate and the Lipschitzness of learned deep networks with the results of the neural tangent kernel. Finally, the authors discuss the dynamics of hidden representations, which in the few initial layers the representations undergo changes to minimize $R^{(0)}$, then, as they move into the low-dimensional space, they change smoothly influenced by $R^{(1)}$ and $R^{(2)}$, i.e., exhibit the bottleneck structure of networks.\n\nSTRENGTHS:\nThis paper is a direct follow-up to [1], and since I am not very familiar with [1], my evaluation may not be very accurate.\n\nTheoretical understanding of deep nonlinear networks is a highly significant field and holds considerable interest among the NeurIPS audience. For me, the theoretical results of this paper are highly novel and intriguing. The paper successfully modifies the results of [1], and provides a rigorous yet intuitive explanation of why neural networks trained with our L2-regularized loss do not underestimate the rank of the true function. Additionally, it demonstrates that under some conditions, the representation dynamics of deep nonlinear dynamics exhibit a bottleneck structure, which consists of a sequence of long bottleneck dimensional representations. This result is also noteworthy.\n\n***\n\n[1] Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions. In The Eleventh International Conference on Learning Representations, 2023.\n\nWEAKNESSES:\nI couldn't find any major drawbacks in this paper.\n\nThe frequent mention of $R^{(2)}$ in Remark 10, before its properties were clearly stated, was somewhat confusing to me, although it is not a critical issue.\n\n\nThere are just a few very minor typos:\n\nIn line 32, “… not control the the …”\n\nIn line 34, “… formalizes the a …”\n\nIn line 118, “… because the the …”\n\nIn line 156, “… $k^{*}$ of he task ...\"\n\nQUESTIONS:\nI couldn't find any major drawbacks in this paper. One small question is:\n\nI would like to hear the authors’ thoughts on whether the optimal latent dimensionality of a well-trained auto-encoder $f$, which is likely to be similar to the identity $id$, might be associated with $R^{(0)}(f;\\Omega)$ and $R^{(1)}(f;\\Omega)$, for the case when $Rank_{J}(id;\\Omega) = Rank_{BN}(id;\\Omega)$, and if possible, when $Rank_{J}(id;\\Omega) < Rank_{BN}(id;\\Omega)$\n\nLIMITATIONS:\nThe authors do not explicitly address the limitations. However, in some of the theorems/propositions, they briefly mention the limitations of the proposed one, indicating what it can only explain or what constraints it has.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper extends the theoretical framework around the representation cost of DNNs as defined in previous work. The authors introduce two corrections to the existing infinite depth description, namely two regularity measures $R^{(1)}$ and $R^{(2)}$, which balance against the dominating rank bias $R^{(0)}$. The authors propose that these regularity measures prevent networks from underestimating the 'true' bottleneck rank (BN-rank) and argue that large learning rates can also induce a bias toward regularity. The paper provides a theoretical description of the limiting representation geodesics as depth approaches infinity, proving under certain conditions a bottleneck structure in the learned representations. Simple experiment is presented to test the network's ability to learn underlying symmetries.\n\nSTRENGTHS:\n* The paper makes a great analysis of the representation cost under the previous theoretical framework, better explaining some observed results and conjecture.\n* The theoretical arguments are well supported, with clear intuition and explanation.\n* Different aspects are considered in the paper such as the effect of learning rate and the representation geodesics which is particularly interesting.\n\nWEAKNESSES:\n* The definition of the term \"regularity\" requires a more specific definition for the problem. It is claimed in line-32 that \"this notion of rank does not control the regularity of f\". Intuitively, the rank $R^{(0)}$ still measures some notion of complexity of the function as a higher-rank matrix can represent a more complex linear transformation. The corrections are other levels of complexity measures. I guess the \"regularity\" is more like how simple a function is with the same rank?\n* Theorem 9 proves bottleneck structure for convergent function sequence with an infinite depth limit. But in practice, the network is trained with finite data points and depth. Could the authors comment on how the conclusions would change for practical networks?\n* Analyze of the learning rate appears somewhat disjointed from the rest of the paper. Could there be a more integrated discussion about its role?\n* I found the numerical experiment a bit hard to follow and would appreciate it if more explanation on the idea could be added. Additionally, is it applicable for some experiments on justifying the theoretical findings, such as the effect of learning rate?\n* Minor typos:\n    * Line32: \"...control the the regularity...\"\n    * Line34: \"This formalizes the a balance...\"\n    * Line156: \"...$k^*$ of he task...\"\n\nQUESTIONS:\n* Can these findings be used to improve the training of DNNs, for example, by adjusting the learning rate or regularization?\n* Considering that the argument in Proposition 7 is incomplete, how much further work would be needed to provide a complete argument?\n\nLIMITATIONS:\nIncluded in weakness.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-20"}