Hamiltonian Mechanics of Feature Learning: Bottleneck Structure in Leaky ResNets

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:28:35

Abstract: We study Leaky ResNets, which interpolate between ResNets ($\tilde{L}=0$)
and Fully-Connected nets ($\tilde{L}\to\infty$) depending on an 'effective
depth' hyper-parameter $\tilde{L}$. In the infinite depth limit,
we study 'representation geodesics' $A_{p}$: continuous paths in
representation space (similar to NeuralODEs) from input $p=0$ to
output $p=1$ that minimize the parameter norm of the network. We
give a Lagrangian and Hamiltonian reformulation, which highlight the
importance of two terms: a kinetic energy which favors small layer
derivatives $\partial_{p}A_{p}$ and a potential energy that favors
low-dimensional representations, as measured by the 'Cost of Identity'.
The balance between these two forces offers an intuitive understanding
of feature learning in ResNets. We leverage this intuition to explain
the emergence of a bottleneck structure, as observed in previous work:
for large $\tilde{L}$ the potential energy dominates and leads to
a separation of timescales, where the representation jumps rapidly
from the high dimensional inputs to a low-dimensional representation,
move slowly inside the space of low-dimensional representations, before
jumping back to the potentially high-dimensional outputs. Inspired
by this phenomenon, we train with an adaptive layer step-size
to adapt to the separation of timescales.

Main paper: 
 introduction: 
 Feature learning is generally considered to be at the center of the recent successes of deep neural networks (DNNs), but it also remains one of the least understood aspects of DNN training.There is a rich history of empirical analysis of the features learned by DNNs, for example the appearance of local edge detections in CNNs with a striking similarity to the biological visual cortex [19] , feature arithmetic properties of word embeddings [22] , similarities between representations at different layers [18, 20] , or properties such as Neural Collapse [24] to name a few. While some of these phenomenon have been studied theoretically [3, 8, 27] , a more general theory of feature learning in DNNs is still lacking.For shallow networks, there is now strong evidence that the first weight matrix is able to recognize a low-dimensional projection of the inputs that determines the output (assuming this structure is present) [4, 2, 1] . A similar phenomenon appears in linear networks, where the network is biased towards learning low-rank functions and low-dimensional representations in its hidden layers [13, 21, 29] .But in both cases the learned features are restricted to depend linearly on the inputs, and the feature learning happens in the very first weight matrix, whereas it has been observed that features increase in complexity throughout the layers [31] .The linear feature learning ability of shallow networks has inspired a line of work that postulates that the weight matrices learn to align themselves with the backward gradients and that by optimizing for this alignment directly, one can achieve similar feature learning abilities even in deep nets [5, 25] .Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.For deep nonlinear networks, a theory that has garnered a lot of interest is the Information Bottleneck [28] , which observed amongst other things that the inner representations appear to maximize their mutual information with the outputs, while minimizing the mutual information with the inputs. A limitation of this theory is its reliance on the notion of mutual information which has no obvious definition for empirical distributions, which lead to some criticism [26] .A recent theory that is similar to the Information Bottleneck but with a focus on the dimensionality/rank of the representations and weight matrices rather than the mutual information is the Bottleneck rank/Bottleneck structure [16, 15, 30] : which describes how, for large depths, most of the representations will have approximately the same low dimension, which equals the Bottleneck rank of the task (the minimal dimension that the inputs can be projected to while still allowing for fitting the outputs). The intuitive explanation for this bias is that a smaller parameter norm is required to (approximately) represent the identity on low-dimensional representations rather than high dimensional ones. Some other types of low-rank bias have been observed in recent work [9, 14] .In this paper we will focus on describing the Bottleneck structure in ResNets, and formalize the notion of 'cost of identity' as a driving force for the bias towards low dimensional representation.The ResNet setup allows us to consider the continuous paths in representation space from input to output, similar to the NeuralODE [6] , and by adding weight decay, we can analyze representation geodesics, which are paths that minimize parameter norm, as already studied in [23] .Our goal is to study a variant of the NeuralODE [6, 23] approximation of ResNet with leaky skip connections and with L 2 -regularization. The classical NeuralODE describes the continuous evolution of the activations α p (x) ∈ R w starting from α 0 (x) = x at the input layer p = 0 and then follows∂ p α p (x) = W p σ(α p (x))for the w × (w + 1) matrices W p and the nonlinearity σ : R w → R w+1 which maps a vector z to σ(z) = ( [z 1 ] + . . . [z w ] + 1 ) , applying the ReLU nonlinearity entrywise and appending a new entry with value 1. Thanks to the appended 1 we do not need any explicit bias, since the last column W p,•w+1 of the weights replaces the bias.This can be thought of as a continuous version of the traditional ResNet with activations α ℓ (x) for ℓ = 1, . . . , L: α ℓ+1 (x) = α ℓ (x) + W ℓ σ(α ℓ (x)).We will focus on Leaky ResNets, a variant of ResNets that interpolate between ResNets and FCNNs, by tuning the strength of the skip connections leading to the following ODE with parameter L:∂ p α p (x) = -Lα p (x) + W p σ(α p (x)).This can be thought of as the continuous version of α ℓ+1 (x) = (1 -L)α ℓ (x) + W ℓ σ(α ℓ (x)). As we will see, the parameter L plays a similar role as the depth in a FCNN. Finally we will be interested describing the paths that minimize a cost with L 2 -regularizationmin Wp 1 N N i=1 ∥f * (x i ) -α 1 (x i )∥ 2 + λ 2 L 1 0 ∥W p ∥ 2 F dp.The scaling of λ L for the regularization term will be motivated in Section 1.2.This type of optimization has been studied in [23] without leaky connections, but we will describe in this paper large L behavior which leads to a so-called Bottleneck structure [16, 15] as a result of a separation of time scales in p.Changing the leakage parameter L is equivalent (up to constants) to changing the integration range [0, 1] or to scaling the outputs.Integration range: Consider the weights W p on the range [0, 1] and leakage parameter L, leading to activations α p . Then stretching the weights to a new range [0, c], by defining W ′ q = 1 c W q/c for q ∈ [0, c], and dividing the leakage parameter by c, stretches the activations α ′ q = α p/c :∂ q α ′ q (x) = - L c α ′ q (x) + 1 c W q/c σ(α ′ q (x)) = 1 c ∂ p α q/2 (x),and the parameter norm is simply divided by c:c 0 W ′ q 2 dq = 1 c 1 0 ∥W p ∥ 2 dp.This implies that a path on the range [0, c] with leakage parameter L = 1 is equivalent to a path on the range [0, 1] with leakage parameter L = c up to a factor of c in front of the parameter weights.For this reason, instead of modeling different depths as changing the integration range, we will keep the integration range to [0, 1] for convenience but change the leakage parameter L instead. To get rid of the factor in front of the integral, we choose a regularization term of the form λ L . From now on, we call L the (effective) depth of the network.Note that this also suggests that in the absence of leakage ( L = 0), changing the range of integration has no effect on the effective depth, since 2 L = 0 too. Instead, in the absence of leakage, the effective depth can be increased by scaling the outputs as we now show.Output scaling: Given a path W p on the [0, 1] (for simplicity, we assume that there are no bias, i.e.W p,•w+1 = 0), then increasing the leakage by a constant L → L + c leads to a scaled down path α ′ p = e -cp α p . Indeed we have α ′ 0 (x) = α 0 (x) and∂ p α ′ p (x) = -( L + c)α ′ p (x) + W p σ(α ′ p (x)) = e -cp (∂ p α p (x) -cα p (x)) = ∂ p (e -cp α p (x)).Thus a nonleaky ResNet L = 0 with very large outputs α 1 (x) is equivalent to a leaky ResNet L > 0 with scaled down outputs e -Lα 1 (x). Such large outputs are common when training on cross-entropy loss, and other similar losses that are only minimized at infinitely large outputs. When trained on such losses, it has been shown that the outputs of neural nets will keep on growing during training [12, 7] , suggesting that when training ResNets on such a loss, the effective depth increases during training (though quite slowly).The optimization of Leaky ResNets can be reformulated, leading to a Lagrangian form.First observe that the weights W p at any minimizer can be expressed in terms of the matrix of activations A p = α p (X) ∈ R w×N over the whole training set X ∈ R w×N (similar to [17] ):W p = ( LA p + ∂ p A p )σ(A p ) +where (•) + is the pseudo-inverse.We therefore consider the equivalent optimization over the activations A p :min Ap:A0=X 1 N ∥f * (X) -A 1 ∥ 2 + λ 2 L 1 0 LA p + ∂ p A p 2 Kp dp.This is our first encounter with the norm ∥M ∥ Kp = ∥M σ(A p ) + ∥ F corresponding to the scalar product ⟨A, B⟩ Kp = Tr AK + p B for K p = σ(A p ) T σ(A p ) that will play a central role in our upcoming analysis. By convention, we say that ∥M ∥ Kp = ∞ if M does not lie in the image of K p , i.e. ImM T ⊈ ImK p .It can be helpful to decompose this loss along the different neuronsmin Ap:A0=X w i=1 1 N ∥f * i (X) -A 1,i ∥ 2 + λ 2 L 1 0 LA p,i• + ∂ p A p,i• 2 Kp dp,Leading to a particle flow behavior, where the neurons A p,i• ∈ R N are the particles. At first glance, it appears that there is no interaction between the particles, but remember that the norm ∥•∥ Kp depends on the covarianceK p = w i=1 σ(A i• )σ(A i• )T , leading to a global interaction between the neurons.If we assume that ImA T p ⊂ Imσ(A p ) T , we can decompose the inside of the integral as three terms:1 2 L LA p + ∂ p A p 2 K + p = L 2 ∥A p ∥ 2 Kp + L ⟨∂ p A p , A p ⟩ K + p + 1 2 L ∥∂ p A p ∥ 2 Kp .The middle term ⟨∂ p A p , A p ⟩ K + p plays a relatively minor role in our analysis 1 , so we focus more on the two other terms:Cost of identity ∥A p ∥ This suggests that the local optimal paths must balance two objectives that are sometimes opposed:the kinetic energy favors going from input representation to output representation in a 'straight line' that minimizes the path length, the COI on the other hand favors paths that spends most of the path in low-dimensional representations that have a low COI. The balance between these two goals shifts as the depth L grows, and for large depths it becomes optimal for the network to rapidly move to a representation of smallest possible dimension (not too small that it becomes impossible to map back to the outputs), remain for most of the layers inside the space of low-dimensional representations, and finally move rapidly to the output representation; even if this means doing a large 'detour' and having a large kinetic energy. The main goal of this paper is to describe this general behavior.Note that one could imagine that as L → ∞ it would always be optimal to first go to the minimal COI representation which is the zero representation A p = 0, but once the network reaches a zero representation, it can only learn constant representations afterwards (the matrix K p = 11 T is then rank 1 and its image is the space of constant vectors). So the network must find a representation that minimizes the COI under the condition that there is a path from this representation to the outputs.Remark. While this interpretation and decomposition is a pleasant and helpful intuition, it is rather difficult to leverage for theoretical proofs directly. The problem is that we will focus on regimes where the representations A p and σ(A p ) are approximately low-dimensional (since those are the representations that locally minimize the COI), leading to an unbounded pseudo-inverse σ(A p ) + . This is balanced by the fact that ( LA p + ∂ p A p ) is small along the directions where σ(A p ) + explodes, ensuring a finite weight matrix normLA p + ∂ p A p 2 K + p . But the suppression of ( LA p + ∂ p A p )along these bad directions usually comes from cancellations, i.e. ∂ p A p ≈ -LA p . In such cases, the decomposition in three terms of the Lagrangian is ill adapted since all three terms are infinite and cancel each other to yield a finite sumEQUATION. One of our goal is to save this intuition and prove a similar decomposition with stable equivalent to the cost of identity and kinetic energy where K + p is replaced by the bounded (K p + γI) + for the right choice of γ.1 In linear networks σ = id it can actually be discarded, since it is integrable1 0 Tr ∂pApσ(Ap) + σ(Ap) +T A T p dp = log |A1| + -log |A0| + , where |•| + is pseudo-determinant, the product of the non-zero singular values. Since its integral only depends on the endpoints, it has no impact on the representation path in between, which is the focus of this paper. In nonlinear networks, we are not able to discard in such a manner, but we will see that in the rest of analysis the two other terms play a central role, while the second term plays less role.The cost of identity can be thought of as a measure of dimensionality of the representation. It is obvious for non-negative representations because ∥A p ∥2 K + p = ∥A p A p + ∥ 2 F = RankA p , but in general,it can be shown to upper bound a notion of 'stable rank':Proposition 1. ∥Aσ(A) + ∥ 2 F ≥ ∥A∥ 2 * ∥A∥ 2 F for the nuclear norm ∥A∥ * = RankA i=1 s i (A).Proof. We know that ∥σ(A)∥ F ≤ ∥A∥ F , therefore ∥Aσ(A)+ ∥ 2 F ≥ min ∥B∥ F ≤∥A∥ F ∥AB + ∥ 2 F which is minimized when B = ∥A∥ F √ ∥A∥ * √ A, yielding the result.The stable rank∥A∥ 2 * ∥A∥ 2 Fis upper bounded by RankA, with equality if all non-zero singular values of A are equal, and it is lower bound by the more common notion of stable rank∥A∥ 2 F ∥A∥ 2 op, becauses i max s i ≥ s 2 i for the singular values s i .Note that in contrast to the COI which is a very unstable quantity because of the pseudo-inverse, the ratio∥A∥ 2 * ∥A∥ 2 Fis continuous except at A = 0. This also makes it much easier to compute empirically than the COI itself.We know that the COI matches the dimension or rank for positive representations, but it turns out that the local minima of the COI that are stable under the addition of a new neuron are all positive:Proposition 2. A local minimum of A → ∥Aσ(A) + ∥ 2F is said to be stable if it remains a local minimum after concatenating a zero vectorA ′ = A 0 ∈ R (w+1)×N .All stable minima are non-negative, and satisfy ∥Aσ(A)+ ∥ 2 F = RankA.Proof. The COI of the nearby point A ϵz for z ∈ Imσ(A) T equalsTr (A T A + ϵ 2 zz T ) (σ(A) T σ(A) + ϵ 2 σ(z)σ(z) T + = Aσ(A) + 2 + ϵ 2 z T σ(A) + 2 -ϵ 2 σ(z) T σ(A) + σ(A) +T A T 2 + O(ϵ 4 ).Assume by contradiction that there is a i = 1, . . . , N such that σ(A•i ) ̸ = A •i , then choosing z = σ(A) T σ(A •i )we have σ(z) = z and the two ϵ 2 terms are negative:ϵ 2 ∥σ(A i )∥ 2 -ϵ 2 ∥A i ∥ 2 < 0, which implies that A ′ it is not a local minimum.These stable minima will play a significant role in the rest of our analysis, as we will see that for large L the representations A p of most layers will be close to one such local minimum. Now we are not able to rule out the existence of non-stable local minima (nor guarantee that they are avoided with high probability), but one can show that all strict local minima of wide enough networks are stable.Actually we can show something stronger, starting from any non-stable local minimum there is a constant loss path that connects it to a saddle:Proposition 3. If w > N (N + 1) then if Â ∈ R w×N is local minimum of A → ∥Aσ(A) + ∥ 2 F that is not non-negative, then there is a continuous path A t of constant COI such that A 0 = Â and A 1 is a saddle.This could explain why a noisy GD would avoid such negative/non-stable minima, since there is no 'barrier' between the minima and a lower one, one could diffuse along the path described in Proposition 3 until reaching a saddle and going towards a lower COI minima. But there seems to be something else that pushes away from such non-negative minima, as in our experiments with full population GD we have only observed stable/non-negative local minimas. We can further reformulate the evolution of the optimal representations A p in terms of a Hamiltonian, 188 similar to Pontryagin's maximum principle.Let us define the backward pass variablesB p = -1 λ ∂ Ap C(A 1 ) for the cost C(A) = 1 2 ∥f * (X)-A∥ 2 F , 190which play the role of the 'momenta' of A p in this Hamiltonian interpretation, which follows the191 backward differential equation 192 B 1 = - 1 λ ∂ A1 C(A 1 ) = 2 λN (f * (X) -A 1 ) -∂ p B p = σ(A p ) ⊙ W T p B p -LB p .Now at any critical point, we have that∂ Wp C(A 1 ) + λ L W p = 0 and thus W p = 193 -L λ ∂ Ap C(A 1 )σ(A p ) T = LB p σ(A p )T , leading to joint dynamics for A p and B p :194 ∂ p A p = L(B p σ(A p ) T σ(A p ) -A p ) -∂ p B p = L σ(A p ) ⊙ σ(A p )B T p B p -B p . These are Hamiltonian dynamics ∂ p A p = ∂ Bp H and -∂ p B p = ∂ Ap H w.r.t. the Hamiltonian 195 H(A p , B p ) = L 2 B p σ(A p ) T 2 -LTr B p A T p .The Hamiltonian is a conserved quantity, i.e. it is constant in p. It will play a significant role in The Lagrangian and Hamiltonian are also very similar to the ones in [10, 11] , and the separation of 206 timescales and rapid jumps that we will describe also bear a strong similarity. Though a difference 207 with our work is that the norm ∥•∥ Kp depends on A p and can be degenerate.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The contribution section accurately describes our contributions, and all theorems/propositions are proven in the main or the appendix.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
separation of timescales: 
 If ImA T p ⊂ Imσ(A p ) T , then the Hamiltonian equals the sum of the kinetic and potential energies:H = 1 2 L ∥∂ p A p ∥ 2 Kp - L 2 ∥A p ∥ 2 Kp .This implies that ∥∂ p A p ∥ Kp = L ∥A p ∥ ≤ c < ∞, and any γ > 0, we have- 1 L ℓ γ, L + √ γc 2 ≤ - 2 L H -min p A L p 2 (Kp+γI)≤ γc,for the path length ℓ γ, L = 1 0 ∂ p A L p (Kp+γI) dp. Finally - L√ γc ≤ ∥∂ p A p ∥ (Kp+γi) -L ∥A p ∥ 2 (Kp+γI) + 2 L H ≤ 2 L√ γc.Note that the size of ∥B L p ∥ 2 can vary a lot throughout the layers, we therefore suggest choosing a p-dependent γ:γ p = γ 0 ∥σ(A L p )∥ 2 op = γ 0 ∥K p ∥ 2 op .There are two motivations for this: first it is natural to have γ scale with K p , ; and second, since W p = LB p σ(A p ) T is of approximately constant size (thanks to balancedness, see Appendix A.3), we typically have that the size of B p is inversely proportional to that of σ(A p ), so that γ p ∥B p ∥ 2 should keep roughly the same size for all p.Theorem 4 shows that for large L (and choosing e.g. γ = L-1 ), the Hamiltonian is close to the minimal COI along the path. Second, the norm of the derivative∥∂ p A p ∥ (Kp+γi) is close to L times the 'extra-COI' ∥A p ∥ 2 (Kp+γI) + 2 L H ≈ ∥A p ∥ 2 (Kp+γI) -min q ∥A q ∥ 2 (Kq+γI), which describes the separation of timescales, with slow (∼ 1) dynamics at layers p where the COI is almost optimal and fast (∼ L) dynamics everywhere the COI is far from optimal. Assuming a finite length ℓ γ, L < ∞, the norm of the derivative must be finite at almost all layers, meaning that the COI/dimensionality is optimal in almost all layers, with only a countable number of short high COI/dimension jumps. These jumps typically appear at the beginning and end of the network, because the input and output dimensionality and COI are (mostly) fixed, so it will typically be non-optimal, and so there will often be fast regions close to the beginning and end of the network.We have actually never observed any jump in the middle of the network, though we are not able to rule them out theoretically.If we assume that the paths A L p are stable under adding a neuron, then we can additionally guarantee that the representations in the slow layers ('inside the Bottleneck') will be non-negative: Proposition 5. Let A L p be a uniformly bounded sequence of local minima for increasing L, at any p 0 ∈ (0, 1) such that ∥∂ p A p ∥ is uniformly bounded in a neighborhood of p 0 for all L, thenA ∞ p0 = lim L A L p0 is non-negative.We therefore know that the optimal COI min q ∥A q ∥ 2 (Kq+γI) is close to the dimension of the limiting representations A ∞ p0 , i.e. it must be an integer k * which we call the Bottleneck rank of the sequence of minima since it is closely related to the Bottleneck rank introduced in [16] . The Hamiltonian H is then close to -L 2 k * .Figure 1 illustrates these phenomena: the Hamiltonian (and the stable HamiltoniansH γ = 1 2 L ∥∂ p A p ∥ 2 (Kp+γI) -L 2 ∥A p ∥ 2 (Kp+γI)) approach the rank k * = 3 from below, while the minimal COI approaches it from above; The kinetic energy is proportional to the extra COI, and they are both large towards the beginning and end of the network where the weights W p are higher dimensional.We see in Figure 1c that the (stable) Hamiltonian are not exactly constant, but it still varies much less than its components, the kinetic and potential energies.Because of the non-convexity of the loss we are considering, one can imagine that there could exist distinct sequences of local minima as L → ∞, which could have different rank, depending on what low-dimension they reach inside their bottleneck. Indeed in our experiments we have seen that the number of dimensions that are kept inside the bottleneck can vary by 1 or 2, and in FCNN distinct sequences of depth increasing minima with different ranks have been observed in [15] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss limitations of our results and approach after we state them.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community.Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
discretization scheme: 
 To use such Leaky ResNets in practice, we need to discretize over the range [0, 1]. For this we choose a set of layer-steps ρ 1 , . . . , ρ L with ρ ℓ = 1, and define the activations at the locationsp ℓ = ρ 1 + • • • + ρ ℓ ∈ [0, 1] recursively as α p0 (x) = x α p ℓ (x) = (1 -ρ ℓ L)α p ℓ-1 (x) + ρ ℓ W p ℓ σ α p ℓ-1 (x)and the regularized costL(θ) = C(α 1 (X)) + λ 2 L L ℓ=1 ρ ℓ ∥W p ℓ ∥ 2 ,for the parameters θ = (W p1 , . . . , W p L ). Note that it is best to ensure that ρ ℓ L remains smaller than 1 so that the prefactor (1 -ρ ℓ L) does not become negative, though we will also discuss certain setups where it might be okay to take larger layer-steps. Now comes the question of how to choose the ρ ℓ s. We consider three options:Equidistant: The simplest choice is to choose equidistant points ρ ℓ = 1 L . Note that the conditionρ ℓ L < 1 then becomes L > L.But this choice might be ill adapted in the presence of a Bottleneck structure due to the separation of timescales.Irregular: Since we typically observe that the fast layers appear close to the inputs and outputs with a slow bottleneck in the middle, one could simply choose the ρ ℓ to be go from small to large and back to small as ℓ ranges from 1 to L. This way there are many discretized layers in the fast regions close to the input and output and not too many layers inside the Bottleneck where the representations are changing less. More concretely one can chooseEQUATION2 ) for a ∈ [0, 1), the choice a = 0 leads to an equidistant mesh, but increasing a will lead to more points close to the inputs and outputs. To guarantee ρ ℓ L < 1, we need L > (1 + a 1 4 ) L.Adaptive: But this can be further improved by choosing the ρ ℓ to guarantee that the distances ∥A ℓ -A ℓ-1 ∥ /∥A p ∥ are approximately the same for all ℓ (we divide by the size of A p since it can vary a lot throughout the layers). Since the rate of change of A p is proportional to ρ ℓ(∥A ℓ -A ℓ-1 ∥ /∥A p ∥ = ρ ℓ c ℓ ), it is optimal to choose ρ ℓ = c -1 ℓ c -1 ℓ for c ℓ = ∥A ℓ -A ℓ-1 ∥ /ρ ℓ ∥Ap∥. The update ρ ℓ ← c -1 i c -1 ican be done at every training step or every few training steps.Note that the condition ρ ℓ L < 1 might not be necessary inside the bottleneck since we have the approximation W p σ(A p ℓ-1 ) ≈ LA p ℓ-1 , canceling out the negative direction. In particular with the adaptive layer-steps that we propose, a large ρ ℓ is only possible for layers where c ℓ is small, which is only possible when W p σ(A p ℓ-1 ) ≈ LA p ℓ-1 .Figure 2 illustrates the effect of the choice of ρ ℓ for different depths L, we see a small but consistent advantage in the test error when using adaptive or irregular ρ ℓ s. Looking at the resulting Bottleneck structure, we see that the adaptive ρ ℓ s result in more steps especially in the beginning of the network, but also at the end. This because the 'true function' f * : R 30 → R 30 we are fitting in these experiments is of the form f * = g 3 • g 2 • g 1 where the first inner dimension is 6 and the second is 3, thus resulting in a rank of k * = 3. But before reaching this minimal dimension, the network needs to represent g 2 • g 1 , which requires more layers, and one can almost see that the weight matrices are roughly 6-dimensional around p = 0.3. The adaptivity to this structure could explain the advantage in the test error.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Answer: [No] Justification: We use synthetic data, with a description of how to build this synthetic data.The code is not the main contribution of the paper, so there is little reason to publish it.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 We have given a description of the representation geodesics A p of Leaky ResNets. We have identified an invariant, the Hamiltonian, which is the sum of the kinetic and potential energy, where the kinetic energy measures the size of the derivative ∂ p A p , while the potential energy is inversely proportional to the cost of identity, which is a measure of dimensionality of the representations. As the effective depth of the network grows, the potential energy dominates and we observe a separation of timescales.At layers with minimal dimensionality over the path, the kinetic energy (and thus the derivative ∂ p A p ) is finite. Conversely, at layers where the representation is higher-dimensional, the kinetic energy must scale with L. This leads to a Bottleneck structure, with a short, high-dimensional jump from the input representation to a low dimensional representation, followed by slow dynamics inside the space of low-dimensional representations followed by a final high-dimensional jump to the high dimensional outputs. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Most details are given in the experimental setup section in the Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: The numerical experiments are mostly there as a visualization of the theoretical results, our main goal is therefore clarity, which would be hurt by putting error bars everywhere.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: In the experimental setup section of the Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]Justification: The paper is theoretical in nature, so it has no direct societal impact that can be meaningfully discussed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. 
2 bottleneck structure in representation geodesics: 
 A recent line of work [16, 15] studies the appearance of a so-called Bottleneck structure in large depth fully-connected networks, where the weight matrices and representations of 'almost all' layers of the layers are approximately low-rank/low-dimensional as the depth grows. This dimension k is consistent across layers, and can be interpreted as being equal to the so-called Bottleneck rank of the learned function. This structure has been shown to extend to CNNs in [30] , and we will observe a similar structure in our leaky ResNets, further showcasing its generality.More generally, our goal is to describe the 'representation geodesics' of DNNs: the paths in representation space from input to output representation. The advantage of ResNets (leaky or not) over FCNNs is that these geodesics can be approximated by continuous paths and are described by differential equations (as described by the Hamiltonian reformulation).This section provides an approximation of the Hamiltonian that illustrates the separation of timescales that appears for large depths, with slow layers with low COI/dimension, and fast layers with high COI/dimension. 
 
 Appendix: A.1 Cost of Identity Proposition 6 (Proposition 3 in the main.).If w > N (N + 1) then if Â ∈ R w×N is local minimum of A → ∥Aσ(A) + ∥ 2F that is not non-negative, then there is a continuous path A t of constant COI such that A 0 = Â and A 1 is a saddle.Proof. The local minimum Â leads to a pair of N × N covariance matrices K = ÂT Â and Kσ = σ( Â) T σ( Â).The pair ( K, Kσ ) belongs to the conical hullCone ( Âi• ÂT i• , σ( Âi• )σ( Âi• ) T ) : i = 1, .. . , w . Since this cones lies in a N (N + 1)-dimensional space (the space of pairs of symmetric N × N matrices), we know by Caratheodory's theorem (for convex cones) that there is a conical combination ( K, Kσ -β 2 1 N ×N ) = w i=1 a i ( Âi• ÂT i• , σ( Âi• )σ( Âi• ) T) such that no more than N (N + 1) of the coefficients are nonzero. We now define A t to have lines A t,i• = (1 -t) + ta i Âi• , so that A t=0 = Â and at t = 1 at least one line of A t=1 is zero (since at least one of the a i s is zero). First note that the covariance pairs remain constant over the path:K t = A T t A t = w i=1 ((1 -t) + ta i ) Âi• ÂT i• = (1 -t) K + t K = Kand similarly K σ t = Kσ , which implies that the cost∥A t σ(A t ) + ∥ 2 F = Tr K t K σ+ t is constanttoo. Second, since a representation A is non-negative iff the covariances satisfy K = K σ , the representation path A t cannot be non-negative either since it has the same kernel pairs ( K, Kσ ) withK ̸ = Kσ .Now (the converse of) Proposition 2 tells us that if A t=1 is not non-negative and has a zero line, then it is not a local minimum, which implies that it is a saddle. and any γ > 0, we have- 1 L ℓ γ, L + √ γc 2 ≤ - 2 L H -min p A L p 2 (Kp+γI) ≤ γc, for the path length ℓ γ, L = 1 0 ∂ p A L p (Kp+γI) dp. Finally - L√ γc ≤ ∥∂ p A p ∥ (Kp+γi) -L ∥A p ∥ 2 (Kp+γI) + 2 L H ≤ 2 L√ γc.Proof. First observe that1 L ∂ p A p + γB p 2 (Kp+γI) = ∥B p (K p + γ) -A p ∥ 2 (Kp+γI) = B p σ(A p ) T 2 + γ ∥B p ∥ 2 -2Tr B p A T p + ∥A p ∥ 2 (Kp+γI) = 2 L H + γ ∥B p ∥ 2 + ∥A p ∥ 2 (Kp+γI)and thus we have- 2 L H = ∥A p ∥ 2 (Kp+γI) - 1 L ∂ p A p + γB p 2 (Kp+γI) + γ ∥B p ∥ 2 .(1) The upper bound-2 L H -min p A L p 2 (Kp+γI)≤ γc then follows from the fact that ∥B p ∥ 2 ≤ c.For the lower bound, first observe thatEQUATIONand therefore1 L ℓ γ, L = 1 L 1 0 ∥∂ p A p ∥ (Kp+γI) dp ≥ 1 0 ∥A p ∥ 2 (Kp+γI) + 2 L H - √ γcdp ≥ min p ∥A p ∥ 2 (Kp+γI) + 2 L H - √ γcwhich implies the lower bound.(2) The lower bound follows from equation 1. The upper bound follows from1 L ∥∂ p A p ∥ (Kp+γI) ≤ 1 L ∂ p A p + γB p (Kp+γI) + ∥γB p ∥ (Kp+γI) ≤ ∥A p ∥ 2 (Kp+γI) + 2 L H + γ ∥B p ∥ 2 + √ γc ≤ ∥A p ∥ 2 (Kp+γI) + 2 L H + √ γ ∥B p ∥ + √ γc ≤ ∥A p ∥ 2 (Kp+γI) + 2 L H + 2 √ γc.Proposition 8 (Proposition 5 in the main.). Let A L p be a uniformly bounded sequence of local minima for increasing L, at any p 0 ∈ (0, 1) such that ∥∂ p A p ∥ is uniformly bounded in a neighborhood of p 0 for all L, then A ∞ p0 = lim L A L p0 is non-negative.Proof. Given a path A p with corresponding weight matrices W p corresponding to a width w, then A 0 is a path with weight matrix W p 0 0 0 . Our goal is to show that for sufficiently large depths, one can under certain assumptions slightly change the weights to obtain a new path with the same endpoints but a slightly lower loss, thus ensuring that if certain assumptions are not satisfied then the path cannot be locally optimal.Let us assume that ∥∂ p A p ∥ ≤ c 1 in a neighborhood of a p 0 ∈ (0, 1), and assume by contradiction that there is an input index i = 1, . . . , N such that A p0,•i has at least one negative entry, and therefore∥A p0,•i ∥ 2 -∥σ(A p0,•i )∥ 2 = c 0 > 0 for all L.We now consider the new weightsW p -Lϵ 2 t(p)A p,•i σ(A p,•i ) T ϵ Lt(p)A p,•i ϵ Lt(p)σ(A p,•i ) 0 for t(p) = max{0, 1 -|p-p0|r } a triangular function centered in p 0 and for an ϵ > 0.For ϵ and rsmall enough, the parameter norm will decrease:1 0 W p -Lϵ 2 t(p)A p,•i σ(A p,•i ) T ϵ Lt(p)A p,•i ϵ Lt(p)σ(A p,•i ) 0 2 dp = 1 0 ∥W p ∥ 2 + L2 ϵ 2 t(p) 2 - 2 L A T p,•i W p σ(A p,•i ) + ∥A p,•i ∥ 2 + ∥σ(A p,•i )∥ 2 dp. Now since W p σ(A p,•i ) = ∂ p A p,•i + LA p,•i , this simplifies to 1 0 ∥W p ∥ 2 + L2 ϵ 2 t(p) 2 -∥A p,•i ∥ 2 + ∥σ(A p,•i )∥ 2 - 1 L A T p,•i ∂ p A p,•i dp + O(ϵ 4 ).By taking r small enough, we can guarantee thatEQUATIONfor all p such that t(p) > 0, and for L large enough we can guarantee that1 L A T p,•i ∂ p A p,•i is smaller then c0 4 , so thatwe can guarantee that the parameter norm will be strictly smaller for ϵ small enough.We will now show that with these new weights the path becomes approximately A p ϵa p wherea p = L p 0 t(q)K p,i• e L(q-p) dq.Note that a p is positive for all p since K p has only positive entries. Also note that as L → ∞, a p → t(p)K p,i• and so that a 0 → 0 and a 1 → 1.On one hand, we have the time derivative∂ p A p ϵa p = W p σ(A p ) -LA p ϵ L (t(p)K p,i• -a p ).On the other hand the actual derivative as determined by the new weights:W p -Lϵ 2 t(p)A p,•i σ(A p,•i ) T ϵ Lt(p)A p,•i ϵ Lt(p)σ(A p,•i ) 0 σ(A p ) ϵσ(a p ) -L A p ϵa p = W p σ(A p ) -LA p -Lϵ 2 t(p) 2 A p,•i K p,i• + Lϵ 2 t(p)A p,•i a p ϵ Lt(p)K p,i• -ϵ La(p) .The only difference is the two terms-Lϵ 2 t(p) 2 A p,•i K i• + Lϵ 2 t(p)A p,•i a p = Lϵ 2 t(p)A p,•i (t(p)K i• -a p ) .One can guarantee with a Grönwall type of argument that the representation path resulting from the new weights must be very close to the path A p ϵa p .This paper will heavily focus on the Hamiltonian H p that is constant throughout the layers p ∈ [0, 1],and how it can be interpreted. Note that the Hamiltonian we introduce is distinct from an already known invariant, which arises as the result of so-called balancedness, which we introduce now.Though this balancedness also appears in ResNets, it is easiest to understand in fullyconnected networks. First observe that for any neuron i ∈ 1, . . . , w at a layer ℓ one can multiply the incoming weights (W ℓ,i• , b ℓ,i ) by a scalar α and divide the outcoming weights W ℓ+1,•i by the same scalar α without changing the subsequent layers. One can easily see that the scaling that minimize the contribution to the parameter norm is such that the norm of incoming weights equals the norm of the outcoming weights∥W ℓ,i• ∥ 2 + ∥b ℓ,i ∥ 2 = ∥W ℓ+1,•i ∥ 2 .Summing over the is we obtain∥W ℓ ∥ 2 F + ∥b ℓ ∥ 2 = ∥W ℓ+1 ∥ 2 F and thus ∥W ℓ ∥ 2 F = ∥W 1 ∥ 2 F + ℓ-1 k=1 ∥b k ∥ 2 F, which means that the norm of the weights is increasing throughout the layers, and in the absence of bias, it is even constant.Leaky ResNet exhibit the same symmetry:Proposition 9. At any critical W p , we have ∥W p ∥ 2 = ∥W 0 ∥ 2 + L p 0 ∥W p,•w+1 ∥ 2 dq.Proof. This proofs handles the bias W p,•(w+1) differently to the rest of the weights W p,•(1:w) , to simplify notations, we write V p = W p,•(1:w) and b p = W p,•(w+1) for the bias.First let us show that choosing the weight matrices Ṽq = r ′ (q)V r(q) and bias bq = r ′ (q)e L(r(q)-q) b r(q) leads to the path Ãq = e L(r(q)-q) A r(q) . Indeed the path Ãq = e L(r(q)-q) A r(q) has the right value when p = 0 and it then satisfies the right differential equation:∂ q Ãq = L(r ′ (q) -1) Ãq + e L(r(q)-q) r ′ (q)∂ p A r(q)= L(r ′ (q) -1) Ãq + e L(r(q)-q) r ′ (q) -LA r(q) + V r(q) σ(A r(q) ) + b r(q)= -L Zq + r ′ (q)A r(q) σ Zq + e L(r(q)-q) r ′ (q)b r(q)= Ṽq σ Ãq + bq -L ÃqThe optimal reparametrization r(q) is therefore the one that minimizes1 0 Wq 2 + bq 2 dq = 1 0 r ′ (q) 2 W r(q)2 + e 2 L(r(q)-q) b r(q) 2 dqFor the identity reparametrization r(q) = q to be optimal, we need1 0 2dr ′ (p) ∥W p ∥ 2 + ∥b p ∥ 2 + 2 Ldr(p) ∥b p ∥ 2 dp = 0for all dr(q) with dr(0) = dr(1) = 0. Since1 0 dr ′ (p) ∥W p ∥ 2 + ∥b p ∥ 2 dp = - 1 0 dr(p)∂ p ∥W p ∥ 2 + ∥b p ∥ 2 dq,we need1 0 dr(p) -∂ p ∥W p ∥ 2 + ∥b p ∥ 2 + L ∥b p ∥ 2 dp = 0and thus for all p∂ p ∥W p ∥ 2 + ∥b p ∥ 2 = L ∥b p ∥ 2 .Integrating, we obtain as needed∥W p ∥ 2 + ∥b p ∥ 2 = ∥W 0 ∥ 2 + ∥b 0 ∥ 2 + L p 0 ∥b q ∥ 2 dq.Our experiments make use of synthetic data to train leaky ResNets so that the Bottleneck rank k * is known for our experiments. The synthetic data is generated by teacher networks for a given true rank k * . To construct a bottleneck, the teacher network is a composition of networks for which the the inner-dimension is k * . Our experiments used an input and output dimension of 30, and a bottleneck of k * = 3. For data, we sampled a thousand data points for training, and another thousand for testing which are collectively augmented by demeaning and normalization.To train the leaky ResNets, it is important for them to be wide, usually wider than the input or output dimension, we opted for a width of 100. However, the width of the representation must be constant to implement leaky residual connections, so we introduce a single linear mapping at the start, and another at the end, of the forward pass to project the representations into a higher dimension for the paths. These linear mappings can be either learned or fixed. Figure 3 shows how their respective magnitudes remains relatively constant as the effective depth L grows.For compute resources, these small networks are not particularly resource intensive. Even on a CPU, it only takes a couple minutes to fully train a leaky ResNet.NeurIPS Paper Checklist 


Decision: This paper studies the dynamics of feature learning in leaky ResNets using tools from Hamiltonian dynamics. While the reviewers all agree that the results in the paper are interesting and relevant to the NeurIPS community, various reviewers also raised concerns about the clarity of writing and the presentation in the paper. I believe a careful update of the manuscript to address these concerns should be done before the paper is accepted, and I thus recommend rejection.

Reviews: review: paper_summary: paper_summary: This paper explores the feature learning dynamics in Leaky ResNets using Hamiltonian mechanics. By introducing the concept of 'representation geodesics',  the authors analyze continuous paths in representation space that minimize the parameter norm of the network. The study provides a Lagrangian and Hamiltonian reformulation that highlights the importance of balancing kinetic and potential energy during feature learning. The potential energy, which favors low-dimensional representations, becomes dominant as the network depth increases, leading to a bottleneck structure. This structure involves rapid transitions from high-dimensional inputs to low-dimensional representations, slow evolution within the low-dimensional space, and eventual rapid transitions back to high-dimensional outputs. The paper also proposes training with adaptive layer step-size to better handle the separation of timescales observed in the representation dynamics. These insights offer a novel perspective on the mechanisms underlying feature learning in deep neural networks and suggest potential improvements in network training methodologies.
strength_weakness: strength_weakness: 1. There are multiple typos in the article, which affect readability. Below are several obvious typos, and it is recommended that the authors carefully polish the language of the article.
	-  The third word in line 24, "phenomenon"$\rightarrow$ "phenomena".
	-   In line 27, "determines" $\rightarrow$ "determine".
	-   In line 40, "lead" $\rightarrow$ "leads".
	-   In line 68, the preposition "in" should be added after "interested".
	-  The formula at the end of line 78  should be$$\alpha_{q}^{'}=\alpha_{q/c}.$$
	-  The formula between lines 78 and 79 is also incorrect. The last term should be $$\frac{1}{c}\partial_{p}\alpha_{q/c}(x).$$
	-  The expression of $K_p$  in line 111 is incorrect, because it should depend on $p$.
	-  The formula between lines 112 and 113 is incorrect. The coefficient of the middle term on the right side of the equation should be 1 instead of $\tilde{L}$.
	- In line 145, " $||\tilde{L}A_{p}+\partial_{p}A_{p}||_{K_{p}^{+}}^{2}$ "  $\rightarrow$  "$||\tilde{L}A_{p}+\partial_{p}A_{p}||_{K_{p}}^{2}$".
	-  In line 159, "bound" $\rightarrow$ "bounded".
	-  The expression for $C(A)$ in line 190 is incorrect, which should be $$C(A)=\frac{1}{N}||f^{*}(X)-A||_{F}^{2} .$$
	- In line 282, "$\rho_{l}L <1$ " $\rightarrow$ "$\rho_{l}\tilde{L}<1$".
	- In line 415, "cones"  $\rightarrow$  "cone".
	- In Theorem 7 and its proof, it seems that all instances of $\sqrt{\gamma c}$ should be changed to $\gamma \sqrt{c}$.
	- In proposition 9, $\tilde{Z}_{q}$ should be $\tilde{A}_{q}$ in the formula above line 485.
2. The formula between lines 101 and 102 is a crucial one, so it is recommended that the authors provide the derivation process for this formula.
3. In line 59 of the paper, the definition of $\sigma$ is given. First, the "+" in this formula should be replaced with a comma. Secondly, my question is about the last component "1" in the definition of $\sigma$. In the proofs of some propositions, is it necessary that $\sigma$ does not include this last component "1"?
	- In line 153, why $||A_{p}||_{K_{p}}^{2}=||A_{p}A_{p}^{+}||_{F}^{2}$ for non-negative $A_{p}$ ?
	- In line 156, why $||\sigma(A)||_{F}\le ||A||_{F}$ ?
	- In line 450, why $||A_{p_{0}, \cdot i}||^{2}-||\sigma(A_{p_{0}, \cdot i})||^{2}=c_{0}>0$ for all $\tilde{L}$ ?
4. In the proof of Proposition 8, the authors did not explain why the limit exists. Secondly, the formula above line 454 is missing $O(\epsilon^4)$. Additionally, the formula above line 464 is incorrect.
questions: questions: 1. In practical applications, the training of neural networks is often affected by random initialization and noise. Could the authors elaborate on how the Hamiltonian maintains its stability and robustness in the presence of such randomness and noise? Are there theoretical or experimental supports demonstrating the effectiveness of the Hamiltonian in describing the system dynamics under different training conditions?
2. The authors propose an adaptive layer step-size to adapt to the separation of timescales observed in Leaky ResNets. How does this adaptive training approach impact the computational complexity and training time compared to standard training methods? Can the authors provide benchmarks or case studies demonstrating the trade-offs between performance gains and computational costs?
limitations: limitations: The theory they proposed aims to provide guidance for better training of networks, and their experimental and implementation details have been documented.
paper_summary: paper_summary: This paper studies feature learning in Leaky ResNets and shows the emergence of the previously studied Bottleneck structure under certain assumptions. In particular the paper provides a Hamiltonian formulation of the features and their dynamics to show that the ResNet will prefer low dimensional features (low potential energy) when the effective depth of the ResNet is large, which gives the Bottleneck structure. The paper also has a final section on choosing the scales of the residual layers across depths motivated by their theory.
strength_weakness: strength_weakness: 1. The paper is unclear in several important moments which compromises readability. For example, is the leakage parameter $\tilde{L}$ suppose to lie in [0,1] (as suggested by line 80) or in [0,\infty] (as is necessary for the “separation of timescales” arguments in section 2.1). Moreover, in line 224 the authors write closed forms for the Hamiltonian but it is not clear how they obtain this object, from the previously stated Hamiltonian on linear 195.
2. Theory seems tied to several simplifying assumptions which reduces its generality in describing/understanding feature learning in standard NNs, e.g. the reliance on ReLU activation, the need for weight decay to minimise parameter norms (though it has been called into question if the role of weight decay is actually to find minimal parameter norms in practice (https://arxiv.org/abs/2310.04415), or also the omission of normalisation layers.
3. On a related note, the theory in the paper seems to have a limited relevance for practice. The one glimmer of this is that the paper suggests changing the weighting on the residual branches in order to evenly balance the difference layers in terms of how much the representations are changing, but this seems underdeveloped at present. It would be worth investigating if this can improve training in practice. Moreover, the paper (and the works on the Bottleneck structure in general) seem to argue that the representations should be changing a lot at late layers because the representation shifts back from being low rank. But this is counter to existing practical works that suggest one can prune late residual layers for improved efficiency https://arxiv.org/abs/2207.07061. This represents a gap between this theory (of Bottleneck structures) and practice to me.
4. The paper studies properties of optimal solutions (e.g. geodesics with minimal parameter norm) in terms of Hamiltonian energies etc (Theorem 4) but does not seem to discuss whether training dynamics will lead to such solutions in practice.
questions: questions: 1. The motivation to study the Bottleneck structure seems not fully convincing to me. I appreciate that the authors spend quite a lot of the introduction to justify the (information) Bottleneck setting from various angles, but has it been shown that practical networks display such a bottleneck structure? Especially in the modern age of large pre-training datasets when the models are underparameterised. It seems like the theory relies on assuming that a low rank structure exists in the function being learnt?
2. What is the effect of $\gamma$ in the inverse in line 257 in terms of meaning that the Bottleneck rank is an integer? It seems that will only hold if $\gamma=0$, in which case is it necessary for the Bottleneck rank to be an integer?


Typos:

1. q/2 in equation below line 78
2. B^T in line 105.
3. no $\tilde{L}$ in middle term of equation below line 112.
4. "of the layers" repeated in line 212
limitations: limitations: There is no limitations section.
paper_summary: paper_summary: The paper maps the dynamics of representations across layers of leaky ResNets to a Largrangian and Hamiltonian formulation, giving an intuitive picture of a balance between two terms: a kinetic energy term which favors small layer derivatives and a potential energy
that favors low-dimensional representations. This intuition to explain the emergence of a bottleneck structure, as observed in previous work.
strength_weakness: strength_weakness: 1. Numerical experiments:
a. Many of the figures are poorly explained and have missing labels etc, e.g. in Figs 1b, 2b what is the color code? 
b. I failed to find a mention of what data the models were trained / evaluated on. 
c. Fig 2c - what is the projection on to? 

2. It is sometimes hard to follow the rationale and motivation for the "storyline" of the paper and its different sections could be better connected to each other. 


3. Novelty wrt previous works - in lines 206-208 a difference from similar works is mentioned but this seems is very brief and looks not very significant on the face of it. It would be better to highlight and elaborate on what is new here relative to these previous works. 



Technical comments: 
a. Eqs are not numbered. 
b. I found the notations to be confusing and non standard, e.g. $\alpha_p \in \mathbb{R}^w$ which make it harder to follow the derivations.



Typos:
a. Lines 211-212: "layers of the layers..."
questions: questions: 1. Does $K_p$ bare some analogy with the mass matrix from classical mechanics? 

2. The backward pass variables $B_p$ play the role of momenta but why are they defined as in line 190? why not via a Legendre transformation of the velocities? or are these equivalent? 

3. Line 224: the analogous condition in classical mechanics is that the potential energy depends only on the coordinates and not on the momenta, and that the kinetic energy is a quadratic form in the velocities. Is there some intuition that can be imported here? 

4. In Fig 1c - why are the Hamiltonians not exactly conserved? these are quite substantial deviations.
limitations: limitations: The authors have adequately addressed the limitations of their results.
paper_summary: paper_summary: The paper introduces the so-called "Leaky ResNet" ordinary differential equation. Leaky ResNets are a variant of the NeuralODE with an additional vector field that attracts trajectories to the origin, the strength of which is governed by a parameter $\tilde{L}$ that is later shown to correspond to a separation of timescales. 

The authors consider a system of ODEs resulting from the optimization of the empirical risk with a regularization term inversely proportional to $\tilde{L}$. 

They define the "Cost of Identity" (COI), a quantity that expresses the coupling of the ODEs and evolves along the trajectories of the ODEs with initial values depending on the training dataset. They then proceed to show that the solutions spend most time in regions where the COI is close to an optimal value.

The final part of the paper proposes three different discretization schemes for the ODE and provides numerical results.
strength_weakness: strength_weakness: * The main results are not clearly stated in the abstract or introduction. The author's stated goal is to study Leaky ResNets, it would be nice to have a rigorous statement of the results of that study at the beginning of the paper.

* In the abstract, the authors state that the paper explains the emergence of a bottleneck structure in ResNets. It is not clear how this claim can be derived from the results in the paper.

* There is no rigorous justification that the results from the study apply to neural networks with a finite number of layers.

* One of the main ingredients used in the derivations is quoted from another paper in (l.101). It would be nice to have a short proof in the appendix or some explanation to make the paper more self-contained.

* In general the paper lacks rigor, as the authors themselves note in many positions, it is not guaranteed that all the quantities are finite and that the decompositions are justified. It would be welcome if some of the informal discussions could be replaced by rigorous statements and proofs.

* The propositions/theorems in general lack clear statements of which assumptions are used. For example, at the top of page 4 it is stated that the decomposition only holds under a certain assumption (which is rarely satisfied in practice as noted later in the paper). Most of the discussion in the paper relies on this decomposition, but it is not immediately clear whether the formal results require the assumption.

* In general there is a lot of discussion mixing formal definitions and informal arguments, for example reasoning in terms of quantities that don't exist in most cases of interest. This makes the paper somewhat hard to read. I can understand the author's intention of providing some intuition, which might perhaps be better served by a combination of rigorous definitions together with a toy example where all the rigorous quantities simplify.
questions: questions: * What are the precise assumptions required in each of the Propositions/Theorems?
* How can this framework be used to analyze neural networks with a finite number of layers?
limitations: limitations: The authors appropriately discuss the limitations of their results next to their statements. However, the extrapolation to feature learning in ResNets made in the abstract is not justified by the results of the paper.

Comments: title: Thanks for the response
comment: Thanks for the response. I will keep my score.
fair enough, I consider this aspect sufficiently addressed. I will increase my score to weak accept.
Yes that would be the setup, we can then show that at any critical point of the loss with bounded $A_p,B_p$, the representation path will be close to the solution of a continuous solution of the Hamiltonian dynamics. Our results then apply to this continuous path under the required assumptions. 

Also note that our main result Theorem 4 applies for finite $\tilde{L}$, so another way to use our results would be to train a leaky ResNet for a fixed $\tilde{L}$ and $L = e^{c^{-1}\tilde{L}}$ until convergence to a critical point, then this solution is close to a continuous Leaky ResNet, and if the $B_p$ and the length of the path $\ell_{\gamma,\tilde{L}}$ are not too large (for the right choice of $\gamma$, e.g. $\gamma=c_2 \tilde{L}^{-1}$), then the LHS and RHS of both equations are small, which implies that the Hamiltonian will be close to the minimal cost of identity over the path, and that the kinetic energy is approximately proportional to the squared root of the extra-COI. Thus distinguishing slow layers, where the COI is close to minimal over the path and so the kinetic energy is small, from fast layers, where the extra-COI is sufficiently large, so that the kinetic energy can be shown to be large too.

In our numerical experiments we are in some sense implementing this strategy, and we do observe that the Hamiltonian of the discrete path is approximately constant (which suggests that we are close to a continuous solution), and it is close to the $\gamma$-Hamiltonians for the right choices of $\gamma$, and also to the minimal $\gamma$-COI over the path. We then observe that the $\gamma$-COI is close to minimal inside the bottleneck and becomes large close to the beginning and end.
I am willing to believe that the phenomenon studied appears for finite $\tilde{L}$. 
However, in the context of the present manuscript, the claimed contribution is about neural networks (discrete) and the presented "evidence" / theorems use $\tilde{L} \to \infty$ in a continuous-time setting. Thus the presented evidence only supports the claim if it is shown that the results transfer from the continuous to the discrete-time setting.

To clarify, the scheme / network architecture you propose is
$a_0 = x$, $a_l = (1-\frac{c \log L}{L}) a_{l-1} + \frac{1}{L} W_l \sigma(a_{l-1})$ with loss $L(\theta) = C(a_L) + \frac{\lambda}{2 c L \log L} \sum |W_l|^2$ with $c$ sufficiently small so that $\frac{c \log L}{L} < 1$ (which is clearly not a simple Euler scheme since the discretization step $1/L$ appears in some extra places)?

And you show that this is not converging to a dynamics with $a_L \approx a_0$ for $L$ sufficiently large, but to the leaky ResNet ODE instead, and that the bottleneck structure appears in this architecture as $L \to \infty$?
We have now worked a bit more on the proof, and $L$ actually needs to grow exponentially in $\tilde{L}$ for the naive Euler method type of bound to work. We will add this result to the appendix. This would answer your question on how to build a practical network that would converge to the continuous limit we study: take the discretization already introduced in the paper (with equidistant steps for simplicity), and let $L\to \infty$ with $\tilde{L} = c \log L$ for a sufficiently small constant $c$.

However, we strongly believe that this exponential dependence is not tight in practice, because inside the bottleneck the dynamics seem to be much slower, and one should be able to leverage this to improve the bound. More precisely, if operator norm of the weights $W_p$ inside the bottleneck are bounded by $\tilde{L} + cst$, then only $L \gg \tilde{L}$ should be sufficient. The reason why this could be the case is that inside the bottleneck the weight matrix should approach $\tilde{L}$ times the projection along the $k^*$-dimentional span of the data, but we are not yet able to capture how fast it converges to this rescaled projection to control its operator norm. 

Also we want to add that we study the infinite $\tilde{L}$ limit because it makes it easier to prove and illustrate the separation of time-scales in and out of the bottleneck. This is mainly a theoretical endeavor, whose goal is to identify quantities/phenomena of interest. In practice it is more optimal to choose $\tilde{L}$ just large enough to start seing a bottleneck (say such that the dynamics inside the bottleneck are roughly 3 times slower than outside) but not too large. In our experiments, this seems to happen for $\tilde{L}$ as small as 3 or 4, and for such $\tilde{L}$ a discretization is reasonable. We have plans to capture at which $\tilde{L}$ the bottleneck starts, but it will probably require some new concepts, and would thus best be done in a follow-up paper.

Though there is still a long way to go to rigorously prove what features are learned in a ResNet, our work identifies a very similar model (up to continuous approximation + leakage), for which a non-trivial phenomenon in the learned features - the bottleneck structure - can be interpreted as a rather simple separation of time-scales coming from the changing balance of two energies/quantities. There exists few simple models to study feature learning in deep networks, and most models that can be analyzed theoretically are either shallow networks or linear networks, and for those models it is impossible to describe depth related phenomenons such as the bottleneck structure.
Thank you for your response, and I am looking forward to the convergence analysis. 

If the final manuscript included a numerical scheme (equivalently, sequence of neural network architectures) that correctly captures the interaction between $\tilde{L} \to \infty$ and the convergence of the numerical scheme I would consider my concern regarding relevance to neural networks addressed. Whether it explains anything about ResNets in particular would of course depend on the concrete expression of the resulting scheme, you might also want to comment on this aspect.

My main concern is about the interaction between the two limits, especially given all the rescaling of time done in Section 1.2. An example that keeps $\tilde{L}$ fixed seems inconsistent with the stated main contribution of the paper (l. 202) and Proposition 5. I am willing to accept taking Euler scheme first and then limit with $\tilde{L}$ if the Euler scheme makes clear higher-order dependence on $\tilde{L}$ using for example little-o notation, so that it would apply to a real-life setting where $L$ is merely big and not infinite. I don't see any problem with using a smooth approximation to ReLU or similar simplifications unrelated to questions of timescale change and discretization.
Thanks for the response and clarifications. It is good to hear that the authors plan in future work to address concerns regarding more general settings, practical relevance and training dynamics. I will keep my score for this submission.
That's a good point, however:
- We are always allowed to choose $L$ as large as necessary to get a good match between the discrete and continuous versions. We are still writing the exact bound but something like $L \gg \tilde{L}$ should be sufficient, or at worst a polynomial of $ \tilde{L}$.
- It appears that the Bottleneck structure we study appears for pretty small $\tilde{L}\approx 5$, we can therefore easily choose an $L$ which is significantly larger than $\tilde{L}$.

Now it would be very intersting to closely capture how large $L$ needs to be to reach the continuous geodesic regime, and whether adaptive learning rates allows one to reach it earlier (our rough conjecture is that as $\tilde{L}\to\infty$ one needs $L \gg \tilde{L}$  for equidistant steps but $L \gg 1$ might be sufficient for adaptive steps). We have plans to takle this question, but it appears that a number of new concepts  are required to answer it in a satisfying manner. The aspect that we are not yet able to capture is how `deep' the first and last paths (going from high- to low- and finally back to high-dimension) need to be, to understand how many layers are needed to discretize them well. However we believe that large depths correspond to the presence of hierarchies of symmetries where the dimension is reduced sequentially in multiple steps, and we are still looking for the right concepts and tools to capture such structure.
I assumed L-tilde would have to depend on L since we want L-tilde to go to infinity in the theorem. And the claim is about bottleneck structure as the number of layers L goes to infinity if I understand correctly.
We will add a result in the Appendix to describe the convergence of the discrete version to the continuous one. We need to assume that the derivative of the nonlinearity $\dot{\sigma}$  is Lipschitz, so it would not apply to the ReLU, but it would apply to smooth approximations of the ReLU. The argument is a simple adaptation of the convergence of Euler methods (explicit for $A_p$ and implicit for $B_p$):

For a fixed $\tilde{L}$ (by the way $\tilde{L}$ does not depend on $L$, we are not sure where you got this impression?) consider parameters $W_{p_1},\dots,W_{p_L}$ that are critical points of the regularized cost with uniformly bounded $\tilde{A}\_{p\_\\ell},\tilde{B}\_{p\_\\ell}$ ( $\\|\tilde{A}\_{p\_\\ell}\\|\_F,\\|\tilde{B}\_{p\_\\ell}\\|\_F \leq C\_1$ for all $\ell$) then there is a continuous solution of the Hamiltonian dynamics $A_p,B_p$ such that for all $\ell$:
$$
\\|\tilde{A}\_{p\_\ell}-A\_{p\_\ell} \\|\_F,\\|\tilde{B}\_{p\_\ell}-B\_{p\_\ell} \\|\_F \leq ...
$$
where the RHS goes to zero as the $\rho_\ell$ go to zero, more precisely the dominant term should be proportional to $\sum_\ell \rho_\ell^2$. There is also a dependence on $\tilde{L}$, but since we consider it fixed, it is fine.

This implies that if one fixes $\tilde{L}$ and considers a uniformly bounded convergent sequence of critical points of the discrete regularized loss with increasing depth $L$, then as $L\to\infty$ it will converge to a solution of the Hamiltonian dynamics as long as $\sum_\ell \rho_\ell^2\searrow 0$. This is obviously the case for equidistant steps, but even adaptive steps will respect this condition under the same assumptions, so everything is fine.

With some more work and a few extra assumptions, it might be possible to extend this result to the ReLU. Intuitively, we need to assume that at almost all layers, the activations are away from the discontinuity of $\dot{\sigma}$, or something of this type. This assumption is quite reasonable for finite width $w$ and finite $N$ assuming that each neuron,training input only crosses zero activation a finite number of times throughout the layers. Hopefully, we can make this work for the final deadline.
Thank you for your response. As you point out, already the question of whether the NeuralODE models ResNets is delicate. In your discretization scheme in l. 275 you have both $\tilde{L}$ and $\rho\_l$ that depend on $L$ in non-trivial ways, present in both the network architecture (as $\rho \tilde{L}$) and the regularized cost (as $\rho / \tilde{L}$). Assuming $L$ is the quantity we want to take to infinity, it is not immediately obvious that the discretization scheme converges to the expected continuous-time process.

Can you give a formula for a discrete leaky ResNet (and associated loss) only in terms of $L$ which converges to a version of the ODE with bottleneck structure, as $L$ goes to infinity, even with some simplifying assumptions? And compare the resulting Leaky ResNet architecture to a standard ResNet?
Re item 2 - I can't point to any specific section, this was a general sense I had when reading the paper. 

yes, using $p$ as a continuous layer index and $w$ as the dimension is non-standard for me, but of course it's a matter of taste. 

I read the author's response and comments by other reviewers and am currently inclined to keep my score of 6, although I should say that my level of confidence is such that I would not have a solid objection if this paper were to be rejected.
Thanks for the thoughtful review. Regarding the weaknesses (and also
the questions) you raise:

- We agree with the sentiment and we usually try to follow this approach for most paper, but for this paper most results require several definitions to be stated so it is difficult to summarize all results at the beginning. But we will clarify the contribution section.

- "Explaining" is subjective, we then precise what we show after the
colon: that for large depth the dynamics switch between fast and slow
dynamics coinciding with high and low-dimensional representations.
It is true that our phrasing suggest that there will always be a fast
then slow then fast sequence (corresponding to high to low to high
dimensions), when in practice one can also observe half-bottlenecks
where the sequence is only fast then slow, or the dynamics could remains low throughout all the layers, but this only happens when the dimension remains constant throughout all the layers. We simply view these as edge cases of the Bottleneck structure, and focused on the main `full'
bottleneck for simplicity.

- The NeuralODE approximation of DNNs as continuous path has been
used in many previous paper, and our numerical experiments on finite
layers network seem to agree with our theoretical results. The match
between NeuralODEs and ResNets has been studied in previous papers
(https://arxiv.org/pdf/2205.14612), which offers a mixed conclusion:
basically that the match depends on the parameters of the network.
We do think that investigating whether the Bottleneck structure helps
or hinders this match could be an interesting follow up work.

- The formula of line 101 follows from the fact that $W_{p}$ is the
minimal Frobenius norm solution of $\partial_{p}A_{p}=-\tilde{L}A_{p}+W_{p}\sigma(A_{p})$.
We will add this derivation.

- The statement and proofs are rigorous, we do not rely on the assumption
$\mathrm{Im}A_{p}^{T}\subset\mathrm{Im}\sigma(A_{p})^{T}$, we simply
mention that under this assumption the separation of timescales and
Bottleneck structure would follow directly from the decomposition
of the Hamiltonian into the two energies. Since it is a useful intuition,
we describe this argument in the first paragraph of section 2.1, but
no proposition or theorem relies on it. Actually the whole point of
Theorem 4 is to show an approximate version of the same argument that
does not require this assumption.

- Again, while in the discussions we offer some intuition that rely
on specific assumptions, when it comes to the propositions and theorems
we state all of our assumptions explicitly (for the results about
the COI we will mention that the COI can be infinite, but since we
only prove lower bounds it is not an issue).

Your questions are answered in our responses to the weaknesses you raised.
Thanks for your thoughtful review. Regarding the weaknesses you point
to:

1. For Figures 1b and 2b the color have no meaning, we just assign
different colors to different singular values for aesthetic purpose.
The experimental setup and the synthetic data is described in the
Appendix B, we will add more details and refer to the appendix in
the main.

2. What part(s) did you think was not fitting the structure? 

3. This previous work defines the same Hamiltonian for non-leaky ResNets (as a minor point inside a more general paper), but there is no reformulation as the sum of two energies, nor any mention of a separation of timescales / Bottleneck structure. We will better discuss the difference in the main.

We will also connect the different formulas more and add numbers
when needed.

What part of $\alpha_{p}$ is confusing? Is it the use of $p$ to
represent the layer instead of $\ell$? We switched to $p$ because
at the end we do have discrete layers indexed by $\ell$.

Thanks, we'll fix the typos you pointed.

Regarding your questions:

1. Very nice observation, it seems to indeed match quite directly
the inverse of the mass matrix. We'll mention this connection.

2. We followed Pontryagin's maximum principle to derive the momenta/dual variables and the Hamiltonian, but it should also be possible to recover it with the Legendre transform.

3. Yes that's exactly the intuition.

4. We do not have a definite answer to this, but we believe it comes
from the layer discretization:

- The bumbs become smaller as the number of "layer
steps" ($L$) is increased, also switching to the adaptive layer step size reduced this error.

- The bumbs appear at $p$ with large kinetic energies, which should
exactly be the points were short layer steps are needed.

- We simply use the forward Euler method to discretize, since it is
what traditional ResNets can be interpreted as doing, but we are
interested in studying better discretization methods, or even methods
that guarantee a constant Hamiltonian in the future.
Thanks for the thoughtful review. To answer to the weaknesses you
mention:

1. Readability should be improved thanks to your and the other reviewers'
remarks. Regarding $\tilde{L}$, our proofs only need $\tilde{L}\geq0$,
and you seem to have misunderstood line 80, the range $[0,1]$ is
the integration range for $p$, not for $\tilde{L}$, we will clarify
this. The Hamiltonian was derived using Pontryagin's maximum principle
(which we will mention), but the derivation is technical and not very
enlightening. But note that it is easy to check that it is the right
Hamiltonian by computing the derivatives.

2. We have plans to handle more general nonlinearities and normalizations,
but we focused here on the simplest setting that exhibits a Bottleneck
structure, to make the analysis as clean as possible. We also have
plans for the training dynamics too and how weight decay affects earlier
training times, see our answer to Question 1 of Reviewer JYMC for
more details.

3. We also plan to more thoroughly analyse the effect of adaptive
layer steps on real world data, but this paper is mostly theoretical,
the experiments are mostly there to help visualize the theory. You
are right that in practice, and especially for classification tasks
with the Neural Collapse, the dimension does not increase back in
the last few layers, but this can simply be understood as a `half-bottleneck' (as observed in https://arxiv.org/abs/2402.08010) and our theory still applies.

4. One big improvement of this paper over previous Bottleneck structure papers is that it applies at any (stable) local minima and not the global minimizer, assuming convergence to a local minimum is much more reasonable.

Regarding your questions:

1. The Bottleneck structure has up to now been mostly studied theoretically, but it has the potential to explain multiple empirical observations
such as Neural Collapse, Information Bottleneck and other low-rank
bias observations with a unified theoretical framework of feature
learning. We think that theoretical analysis is not only useful for
describing empirical observations, but can also be used to identify
structures that can be then confirmed empirically afterwards.

2. Thanks for pointing that error, it will indeed not exactly be an integer but should approach an integer if one takes $\tilde{L}\to\infty$ and $\gamma\searrow0$. 

Thanks for finding these typos, we will fix them.
Thanks for the thorough and thoughtful review. Responding to the weaknesses
you point:

1. We will fix all the typos you have identified.

2. This line follows from the fact that $W_{p}$ is the minimal Frobenius
norm solution of $\partial_{p}A_{p}=-\tilde{L}A_{p}+W_{p}\sigma(A_{p})$.
We will add this derivation.

3. The pluses correspond to the notation $[x]_{+}=\max\\{0,x\\}$ but we forgot to define it. Also, thanks for pointing out to this error in the proof, the proofs of the results of Section 1.4 were done without bias. Thankfully we can show at each stable local minima of the COI with bias, the COI with and without bias are the same and all results for the no-bias COI then apply. We will add this result and a few other basic relations between the bias COI and no-bias COI in  this section.

4. You are correct that it may not converge, we will replace the statement by `the limit is non-negative if it exists'.
Note that the argument would also work for any convergent subsequence.

Regarding your questions:

1. Our long term goal is to prove how training dynamics converge to
such Bottleneck structure, and we believe that in practice the BN
structure is hidden behind the noise of the initialization, but remains
relevant. This is supported by the fact that when we train these networks
we observe an earlier time, where the train and test stop changing
much and we start to see $k^{*}$ singular values coming out of the
bulk of the weight spectra, but we need to train longer for the weight
decay to `get rid' of this noise. It seems that the Hamiltonian analysis
only works at the end of training, but we are confident that is might
be possible to extend it to earlier time with the right modifications.
Again, this is our end goal, but our strategy is to first understand
the BN structure in its cleanest form at the end of training and later
describe corrections to it.

2. The adaptive layer steps has no cost at inference time. During training, the computational cost is negligible (the training is approx. 2\% longer with adaptive learning), especially since
it is sufficient to update the $\rho$ every few steps (every 30 steps in our case). We will add some training time information in the appendix.