Paper: Speed Limits for Deep Learning (ICLR-2024)

Reviewer comment (Theory):
I am not sure about the significance of this theoretical investigation. It seems that the paper only considers the lower bound of the time that goes from an initial parameter state to a target parameter state. It does not tell us information about e.g. the time lower bound to get to a near-stable parameter state, or the time lower bound to get to near-zero potential.

Author response:
The time bound is interesting because it sets a threshold of optimal learning. Unlike what the referee suggests, the speed limit bounds in the paper hold quite generally whether one is in an equilibrium state or not (see for instance Eq. 8 in the revised manuscript). Admittedly, computing it analytically during the dynamics and finite-width networks may be challenging. However, in our NTK section, we compute it analytically for infinite-width networks. Furthermore, in our numerical experiment, we compute it numerically for a finite-channel CNN working well outside the NTK regime and still find that the results are in qualitative agreement with NTK analytics (see also new App. Sec. D.)

Reviewer comment (Theory):
Characterizing the speed limit using the Wasserstein-2 distance is not easily interpretable since it is in general hard to compute the Wasserstein-2 distance. The paper seems only able to derive interpretable results under limiting conditions like no noise or infinite parameters.

Author response:
Based on your feedback, we've extended the discussion at the end of section 2.2 and the discussion section to emphasize that besides the analytical investigations like those we undertake in the paper, one could hope to study much more realistic scenarios (including ones with SGD noise) by numerically bounding the Wasserstein distance (using techniques like those in "2-Wasserstein Approximation via Restricted Convex Potentials with Application to Improved Training for GANs" by Taghvaei and Jalali, or "Wasserstein-2 Generative Networks" by Korotin et al.; references we've added) and by numerically bounding or calculating the entropy production directly, closer to Section 4. The latter is predominantly determined by the geometry of the loss landscape, in terms of the gradient and the Hessian along the training trajectory (Eq (5)). With help of contemporary finite-size network theory (e.g., Li \&,Sompolinsky 2020, Hanin\&Slopaka 2022, Zavatone-Veth\&Pehlevan (2021), Seroussi\&Ringel 2023), moreover, analytical estimates of the partition function required by our Eq. (4) for the entropy production, are within reach.  
Moreover, since all quantities involved in the time bound are readily measurable (see new paragraph "measurability" in the discussion), the inefficiency ratio may be used as an indicator to identify periods within the training that are furthest from optimality and hence would predominantly benefit from optimization, thus serving as a guide for the development of training algorithms.  
Finally, we comment that even using finite-batch SGD our bound still applies in the limit of very low learning rates, as the SGD noise averages out giving us again gradient flow type behavior. It seems plausible to us, that improving the architecture/training hyperparameters so that the dynamics travels closer to the time bound in this low learning rate region, would also improve the behavior at higher learning rates used in experiments.

Reviewer comment (Writing):
The writing of the paper can be improved. I hope to see formal theorems in the paper stating the main contribution, and notations need to be clearly stated. For instance, I am not sure what is $Z_T$ and $Z_0$ in Eq. (4). There is also no related works section.

Author response:
Thank you for pointing out this notation issue, which was indeed unclear. This has been fixed in the revised version. We added a subsection in the introduction describing related work, as well as related literature in the discussion.

Reviewer comment (Theory):
Given the conclusion in the paper that in the NTK regime the dynamic in the paper achieves a rate that is almost optimal, is this contradicting the fact that the continuous heavy-ball dynamic achieves a faster rate to converge to the stable point (see, e.g. https://link.springer.com/article/10.1007/s1022-01164-w)?

Author response:
Regarding the final question, the included link doesn't seem to work, so we're not sure which reference you meant, could you clarify?  In general, though, different dynamics will have different entropy production rates.  It would be interesting, and we hope to encourage others to measure entropy production rates for different optimization schemes.  It seems likely that differences in observed convergence times equate to measurable differences in the entropy production for different methods. Alternatively, we note that the notion of optimality we found in the NTK section is meant in the scaling sense - namely that $T_{SL}(t)/t$ does neither scale up with the number of data points nor with the training time. This still leaves an $O(1)$ room for improvement.

Reviewer comment (Experiments):
Only two setups have been studied: linear regression and NNs in the NTK regime. Moreover, the continuous-time SGD does not model faithfully the discrete SGD when training practical NNs on realistic data.

Author response:
This is an important insight, and we are grateful to the reviewer for
pointing it out. We agree that usual training algorithms employ
discrete updates for training, and we share the intuition by
the reviewer that this may cause important differences. In fact,
in the literature on stochastic thermodynamics, similar speed limits
also exist for discrete (Markov) systems; we are planning to address
this in future work. We will discuss this outlook in the revised manuscript.

Yet, we found that certain features seem to be more universal
and to carry over between discrete and continuous dynamics; also they
seem to generalize from the NTK regime (of infinite-sized networks)
to fully trained finite-sized networks:
The presented numerical experiments on Myrtle-5 networks in fact employ
a usual discrete training scheme. Also, we agree that training here is
not in the NTK regime proper (See also App. Sec. D.). Still, the initial transient of entropy
production, as explained by the NTK analysis, exists in a qualitatively
similar manner in the fully trained network. Moreover, the influence of
the spectrum of the residuals, as predicted by the NTK analysis, correctly
predicts the transition from the inefficient to the efficient training
regime.

Reviewer comment (Theory):
The authors do not discuss how $T_{SL}$ (lower bound on the training time) obtained in the NTK regime compares to a hypothetical $T_{SL}$ obtained in finite-width NNs. Would it be larger or smaller?

Author response:
We are thankful for this intriguing question. In fact, given the presented
framework in the manuscript, in particular our Eq. (4), and the various
existing finite-size theories for deep networks (e.g.,  Li \& Sompolinsky 2020, Zavatone-Veth \& Pehlevan 2021,
 Hanin \& Zlokapa 2022, Seroussi \& Ringel 2023), this question can be
addressed in a quantitative manner and be studied in relation to different choices of parameters, network architectures etc. Practically, this requires the computation of the irreversibility $R$, given by our Eq. (4) as well as
the Wasserstein distance.

Both of these quantities can of course readily be measured in numerical experiments;
For weak noise, $R$ as given by Eq. (24), is dominated by the gradient and
the Hessian of the loss. Likewise, the Wasserstein distance reduces to the
Euclidean length in weight space in the noiseless limit.

Also from the theoretical side, the reviewer's question is very interesting. We would like to share our preliminary thoughts on this point with the reviewer and are very much interested in
their thoughts on this. We feel, however, that a rigorous
answer deserves more attention, so we would like to defer to future work.

The finiteness of the network and hence the number of
parameters $N$ enters the difference of the free energies $F_0 = -\beta^{-1} \ln Z_0$
at initialization and $F_T = -\beta^{-1} \ln Z_T$ of the final posterior. These free energies (or partition functions $Z$) are central quantities of many contemporary finite-size network theories (see, e.g., Li \& Sompolinsky 2020, their Fig. (1)).
Such frameworks could yield quantitative answers.

We here would like to share a preliminary estimate with the reviewer. 

To compare the quantities for different network sizes, we need to know their scaling with the number of parameters $N$. First note that $W_2$ is extensive in $N$; this is most obvious from our Eq. (25), because the Euclidean squared length of of the thermodynamic velocity $\| v \|^2 = \mathcal{O}(N)$ is proportional to the number of parameters. By the same token, the irreversibility $R$ is extensive in the number of weights, as seen from our Eq. (23).
This overall scaling $\propto N$ of the two quantities thus drops out in ratio $W_2/R$ that determines the speed limit. What matters is the irreversibility per degree of freedom
and the Wasserstein distance per degree of freedom (per weight).

The irreversibility $R$, by our Eq. (4), can likewise be determined from the loss at
initialization, and the free energies at initialization and in the final state.
We, as well as the reviewer, would like to know if our estimates from kernel methods (such as NTK) provide a lower or an upper bound.

First, note that the free energy $F_0/N= -\beta^{-1} \ln Z_0 / N$ per degree of freedom of the random initialization is independent of $N$ and known exactly. The free energy of the posterior $F_T/N = -\beta^{-1} \ln Z_T / N$ depends on $N$ in a complicated manner. The question thus
reduces to knowing whether the true $F_T$ is larger or smaller than its mean-field estimate
by the NTK. To this end, remember that mean-field approximations can be derived in a
variational approach, where one seeks (within a given space of distributions) the distribution
that is closest to the true one in terms of the Kullback Leibler (KL) divergence. It is easy to show that from Gibb's inequality (KL divergence being positive semi-definite) it follows that at the minimum one has $0 \le \mathrm{KL} = F_{T, \mathrm{MF}} - F_{T, \mathrm{true}}$. So the mean-field free energy is always larger than or equal to the true one.
In terms of the speed limit, this implies that the irreversibility (Eq. (4)) $R = F_0 - F_{T} + \langle \mathcal{L}(0) \rangle$ in the denominator of Eq. (7), estimated by a mean-field method (such as NNGP or NTK), is smaller than the true value. Hence, the speed limit $T_{SL}$ estimated in the mean field is larger than the true value. 
The previous consideration holds for $W_2/N$, the traveled length in weight space, being the same. So in addition, one would need to analyze how $W_2/N$ changes between the
mean-field (NNGP or NTK) approximation and the full weight posterior at finite $N$. We deem
an analysis possible by employing contemporary approaches of feature learning (works cited above), but also acknowledge that this highly relevant question deserves more attention to
be answered in a satisfactory and rigorous manner. Any further thoughts or advice from the expert reviewer on this issue are of course highly welcome.

Reviewer comment (Theory):
Given the theoretical framework, it is not fully clear how this work is related to the usual challenges in deep learning. Can we use this work to improve optimization? To obtain theoretical guarantees?

Author response:
The reviewer here precisely asks a highly relevant question.
Yes, indeed, we see one merit of the presented approach in providing guarantees
for the achievable efficiency of the training procedure. But besides these
theoretical guarantees, all involved quantities are readily accessible numerically,
in particular, the irreversibility $R$ given by Eq. (24) being predominantly given
by the local slope and curvature of the loss function, this quantity may be used as
an indicator for periods within the training process that burn most of the "entropy budget";
in the presented numerical experiments and the NTK theory, we identified the initial training
phase as such a critical phase.
Efforts of optimizing the training procedure can hence be targeted to those
phases, reducing the training time most effectively.

Also, the finding that, depending on the observed power law exponents in the spectrum of
residuals, training is either in the suboptimal or in the optimal regime,
may be used to develop and investigate new forms of initialization that shape
these residuals precisely such that training starts closer to or within the optimal regime.

Reviewer comment (Evaluation):
Motivation: can we use this work to evaluate the quality of an optimizer?

Author response:
1. As stated above, we agree that the presented work may be used to evaluate
and assess the closeness to optimality of a given optimizer.

Reviewer comment (Theory):
Stronger results: how does the $T_{SL}$ obtained in the NTK limit relate to some "$T_{SL}$" in the finite width setting?

Author response:
We are thankful for this intriguing question. In fact, given the presented
framework in the manuscript, in particular our Eq. (4), and the various
existing finite-size theories for deep networks (e.g.,  Li \& Sompolinsky 2020, Zavatone-Veth \& Pehlevan 2021,
 Hanin \& Zlokapa 2022, Seroussi \& Ringel 2023), this question can be
addressed in a quantitative manner and be studied in relation to different choices of parameters, network architectures etc. Practically, this requires the computation of the irreversibility $R$, given by our Eq. (4) as well as
the Wasserstein distance.

Both of these quantities can of course readily be measured in numerical experiments;
For weak noise, $R$ as given by Eq. (24), is dominated by the gradient and
the Hessian of the loss. Likewise, the Wasserstein distance reduces to the
Euclidean length in weight space in the noiseless limit.

Also from the theoretical side, the reviewer's question is very interesting. We would like to share our preliminary thoughts on this point with the reviewer and are very much interested in
their thoughts on this. We feel, however, that a rigorous
answer deserves more attention, so we would like to defer to future work.

The finiteness of the network and hence the number of
parameters $N$ enters the difference of the free energies $F_0 = -\beta^{-1} \ln Z_0$
at initialization and $F_T = -\beta^{-1} \ln Z_T$ of the final posterior. These free energies (or partition functions $Z$) are central quantities of many contemporary finite-size network theories (see, e.g., Li \& Sompolinsky 2020, their Fig. (1)).
Such frameworks could yield quantitative answers.

We here would like to share a preliminary estimate with the reviewer. 

To compare the quantities for different network sizes, we need to know their scaling with the number of parameters $N$. First note that $W_2$ is extensive in $N$; this is most obvious from our Eq. (25), because the Euclidean squared length of of the thermodynamic velocity $\| v \|^2 = \mathcal{O}(N)$ is proportional to the number of parameters. By the same token, the irreversibility $R$ is extensive in the number of weights, as seen from our Eq. (23).
This overall scaling $\propto N$ of the two quantities thus drops out in ratio $W_2/R$ that determines the speed limit. What matters is the irreversibility per degree of freedom
and the Wasserstein distance per degree of freedom (per weight).

The irreversibility $R$, by our Eq. (4), can likewise be determined from the loss at
initialization, and the free energies at initialization and in the final state.
We, as well as the reviewer, would like to know if our estimates from kernel methods (such as NTK) provide a lower or an upper bound.

First, note that the free energy $F_0/N= -\beta^{-1} \ln Z_0 / N$ per degree of freedom of the random initialization is independent of $N$ and known exactly. The free energy of the posterior $F_T/N = -\beta^{-1} \ln Z_T / N$ depends on $N$ in a complicated manner. The question thus
reduces to knowing whether the true $F_T$ is larger or smaller than its mean-field estimate
by the NTK. To this end, remember that mean-field approximations can be derived in a
variational approach, where one seeks (within a given space of distributions) the distribution
that is closest to the true one in terms of the Kullback Leibler (KL) divergence. It is easy to show that from Gibb's inequality (KL divergence being positive semi-definite) it follows that at the minimum one has $0 \le \mathrm{KL} = F_{T, \mathrm{MF}} - F_{T, \mathrm{true}}$. So the mean-field free energy is always larger than or equal to the true one.
In terms of the speed limit, this implies that the irreversibility (Eq. (4)) $R = F_0 - F_{T} + \langle \mathcal{L}(0) \rangle$ in the denominator of Eq. (7), estimated by a mean-field method (such as NNGP or NTK), is smaller than the true value. Hence, the speed limit $T_{SL}$ estimated in the mean field is larger than the true value. 
The previous consideration holds for $W_2/N$, the traveled length in weight space, being the same. So in addition, one would need to analyze how $W_2/N$ changes between the
mean-field (NNGP or NTK) approximation and the full weight posterior at finite $N$. We deem
an analysis possible by employing contemporary approaches of feature learning (works cited above), but also acknowledge that this highly relevant question deserves more attention to
be answered in a satisfactory and rigorous manner. Any further thoughts or advice from the expert reviewer on this issue are of course highly welcome.

Reviewer comment (Theory):
Eqn (10) is difficult to interpret: in the proof, the effective learning rate is $\eta/n$; how such a proof can be interpreted in the limit $n \rightarrow \infty$?

Author response:
1. We thank the reviewer for pointing out this indeed unintuitive point,
which we will clarify in the revised version.
We are working with the squared error loss (not the mean squared error loss),
so that gradients are proportional to $n$, the number of training points.
Hence, for constant $\eta$,  the speed of changing weights would
naturally, scale with $n$ in a trivial manner (in practical terms,
training would turn unstable for fixed $\eta$ as $n \to \infty$).
To avoid this trivial "speedup" (and instability) and to instead
focus on the interesting scaling, we need to scale the learning rate
$\eta \propto n^{-1}$. We added an explanation about that in Appendix C.

Reviewer comment (Evaluation):
Figure 1.b seems to indicate that the training time is about $10$–$30$ times larger than $T_{SL}$, which is far from the "$O(1)$" factor written in the claims, while it is true that the ratio of the trajectory lengths (Fig. 1.e) is of order 1. How to solve this contradiction?

Author response:
The observation by the reviewer, moreover, is right to the point: While the training time in Fig 1b
is far from optimal (larger than $T_{SL}$ by a factor of about $30$), trajectories are still
somewhat straight (Fig 1d). Frankly, we were quite puzzled at first by this apparent contradiction, too.
The NTK analysis, however, provides an explanation: We find that $T_{SL} \propto T^{1 + \alpha^{-1} (1-\delta)}$. If $\delta>1$, the exponent is negative, so $T_{SL} \ll T$ for large T;
training is hence inefficient. This is what we observe in Fig 1b.
From the geometric point of view, in contrast, the geometric distance $l_{geo}$ and of the distance along the curve $l_\gamma$ both scale identically (Eq. (16)). This implies that the inefficiency we
observe is not due to the motion in a curved space, but rather due to the inhomogeneous speed.
The effect of speed on the entropy production can be appreciated by noting that the entropy
production rate is $ \propto \|v\|^2$, the thermodynamic speed in weight space (Eq. (23)).
We now see that we poorly phrased this point in the initial version, and thank the reviewer for
their helpful remark, which helped us to clarify this point.

Reviewer comment (Theory):
The result written in Eqn. (10) depends on the distribution of the data. To obtain such a result, the data are assumed to be Gaussian. It is then not surprising that Fig. 1.d contradicts Eqn. (10), since CIFAR-10 images are far from being Gaussian vectors.

Author response:
2-3. Thanks to the sharp observation by the reviewer, we realize that we failed to communicate this
quite an intricate point, clearly.
The reviewer is of course precisely right that the data distribution is crucial. While for the
analytical results for linear regression data is Gaussian, for the NTK results, which are
required to understand the optimality, the data distribution is more complicated. This distribution
enters only in the scaling form of the NTK kernel's eigenvalues $\lambda_k \propto k^{-\alpha}$ and
 the discrepancies $\Delta^2 \propto k^{-\delta}$.
Such power law distributions are observed throughout many settings (see, e.g., Maloney et al. 2022 https://arxiv.org/abs/2210.16859).