Paper: Dynamic Training Guided by Training Dynamics (ICLR-2024)

Reviewer comment (Reproducibility):
The methods and experiments do not seem to completely described. With no code provided, I had to guess and fill in potential details myself.

Author response:
**Weakness 1** - Regarding the code of proposed methods.
 
Pseudocodes for KGA and KWM have been provided in the Appendix A.1.

Reviewer comment (Experiments):
It is unclear whether the proposed methods indeed make consistent improvement compared with SOTA optimisers. The empirical comparisons were not carefully designed, with arbitrary and inconsistent selection of baselines. For example, Adam was compared, but it is also used as a fall-back when KGA (the proposed method) does not seem to help. In addition, the KGA seems to have a significant memory footprint and requires additional running time; there should be fair comparison with the SOTA optimisers, under the same budget of computing resources.

Author response:
**Weakness 2** - Regarding the tradeoff between cost vs. performance gain and the consistency of experiments:
 
We completely agree that the evaluation between cost and performance gain should be the center of experiments and results. During the rebuttal phase, we spent quite some effort to streamline the discussion of experiments and the way the results are presented. Here’s a summary of changes related to the cost vs. performance gain comparison.
1. We revised Fig. 4 by moving Table 1 (originally in Appendix A.3) to be a subfigure in Fig. 4, such that the left two subfigures show the performance gain of KGA as compared to SGD (in terms of training loss), and the right subfigure shows the runtime cost of KGA-Iter. Through this comparison, we conclude the better tradeoff achieved using the KGA-Epoch approach.
2. We revised Fig. 5 by selecting the most informative results from Figs. 8 and 11 (originally in Appendix A.3) , such that the left two subfigures show the performance gain of Efficient-KGA in terms of validation accuracy and the right subfigure shows the memory cost. Through this comparison, we conclude the better tradeoff achieved using the Efficient-KGA over the KGA-Epoch approach.
3. We revised Fig. 6 by moving Fig. 10 (originally in Appendix A.3) as a subfigure to the left. It shows the performance gain of Efficient-KGA over state-of-the-art prediction-based Koopman acceleration approach.
4. We conducted an additional experiment during the rebuttal phase to provide a comprehensive runtime comparison for the integrated strategy of masking+KGA. This is added to the main paper as Table 1. This complements the performance gain comparison in Fig. 7.

Reviewer comment (Experiments):
Section 3.3 seems to be there just to add more technical content. It is not clear to me whether it actually work significantly better than simpler baselines such as Delta masking or Norm masking.

Author response:
**Weakness 3** - Regarding Section 3.3, the comparison of KWM, Delta masking, and Norm masking:
 
Sec. 3.3 is there to describe a masking strategy based on the analysis of Koopman-based training dynamics. It is necessary to further reduce the cost of having to perform SVD on a large matrix. The KWM approach can be regarded as a generalized version of some existing ad-hoc masking approaches, like delta masking and norm masking. This is demonstrated in the Appendix A.5.

Reviewer comment (Presentation):
I do not understand Figure 3. Why does alpha=10 mean the vector is outside the unit circle? How does it related to Figure 12?

Author response:
**Question 1** - I do not understand Figure 3. Why does alpha=10 mean the vector is outside the unit circle? How does it related to Figure 12?
 
In the theoretical example shown in Figure 3, moderate values of $\alpha$ result in acceleration, the change in angle is a result of the underlying optimization algorithm (SGD/Adam). Choosing a large value of $\alpha$ will result in accelerating the eigenvalue/eigenmode pair to become unstable, resulting in a significant drop in performance. In Fig. 9 of the Appendix, we show empirical results after training with various values of $\alpha$ using an FCN model on Cifar 10. The figure shows that using higher values of $\alpha$ yields larger standard deviations in accuracy, indicating that the performance increase is not consistent with large values.

Reviewer comment (Writing):
In equation 14, how are lamdas indexed? What are "steady" modes?

Author response:
**Question 2** - In equation 14, how are lamdas indexed? What are "steady" modes?
 
To clarify and avoid any potential misunderstandings, we have decided to revise the terminology in our manuscript. The term 'steady modes' has been replaced with 'neutral mode' to better convey our intended meaning in the context. Regarding the indexing of lamdas, they are sorted based on their distance to the unit circle. 
We appreciate your attention to detail and have revised the paper to correct any ambiguities and typographical errors. Your feedback has been instrumental in enhancing the clarity and accuracy of our paper.

Reviewer comment (Evaluation):
In my opinion, by far the most significant weakness is how the cost vs performance gains comparison is handled. For ex, consider Fig. 7d, which is missing the comparison to the costs incurred in SGD (the blue curve is invisible). The proposed method outperforms in accuracy, but if the costs are astronomically high, are those gains worth it? Indeed, this should be tabulated and presented as a central focus of the work (whether in the limitations section or the advantages section I don't know, because I couldn't easily see what the overall cost versus performance trade-offs were). My score will change substantially depending on what the true results are (and my deepest apologies are extended to the authors if I have simply missed the comparison somehow in the main body - which is itself a weakness of the work as well, since this comparison should be unmissable). Table 1, A.3 seems to be the relevant comparison but it being in the appendix is a strong oversight. I am happy to be proven wrong (and change my score aptly) based on this point.

Author response:
> **Weakness A** - Regarding the tradeoff between cost vs. performance gain: 
> 
We completely agree that the evaluation between cost and performance gain should be the center of experiments and results. We also agree that this evaluation should not be pushed to the Appendix even with a page limit. During the rebuttal phase, we spent quite some effort to streamline the discussion of experiments and the way the results are presented. Here’s a summary of changes related to the cost vs. performance gain comparison.

1. We revised Fig. 4 by moving Table 1 (originally in Appendix A.3) to be a subfigure in Fig. 4, such that the left two subfigures show the performance gain of KGA as compared to SGD (in terms of training loss), and the right subfigure shows runtime cost of KGA-Iter. Through this comparison, we conclude the better tradeoff achieved using the KGA-Epoch approach.

2. We revised Fig. 5 by selecting the most informative results from Figs. 8 and 11 (originally in Appendix A.3) , such that the left two subfigures show the performance gain of Efficient-KGA in terms of validation accuracy and the right subfigure shows the memory cost. Through this comparison, we conclude the better tradeoff achieved using the Efficient-KGA over the KGA-Epoch approach.

3. We revised Fig. 6 by moving Fig. 10 (originally in Appendix A.3) as a subfigure to the left. It shows the performance gain of Efficient-KGA over state-of-the-art prediction-based Koopman acceleration approach.

4. About Fig. 7d: It shows the “extra” runtime for calculating DMD in each epoch, which is an additional cost. Since SGD does not need to calculate DMD, it does not have a corresponding value in this figure. As such, SGD would be at y=0 and was left out. We have updated Fig. 7 and Fig. 13 to draw the blue curve corresponding to SGD and Adam.

5. We conducted an additional experiment during the rebuttal phase to provide a comprehensive runtime comparison for the integrated strategy of masking+KGA. This is added to the main paper as Table 1. This complements the performance gain comparison in Fig. 7.

Reviewer comment (Novelty):
The literature review is missing two important early works in the modern ML/KOT intersection. To the best of my knowledge, the idea that KOT could be profitably used to study/boost optimization (especially in ML, but in general as well) was first mentioned in [1] and first properly investigated within the realm of ML in [2]:  

[1] Section 4.1, F. Dietrich, T. N. Thiem, I. G. Kevrekidis, On the Koopman Operator of Algorithms, arXiv:1907.10807 (2020)  

[2] Section 3, A. S. Dogra, Dynamical Systems and Neural Networks, arXiv:2004.11826 (2020)

Author response:
> **Weakness B**  - Regarding the two missing literature: 
>
Thank you for suggesting these two references! We were worried we might have missed some important literature, especially in the early stage of KOT and acceleration development. We have added both in Sec. 2 the literature review.

Reviewer comment (Presentation):
While not critical, there are presentation issues in the work. For ex, Fig. 3 is poorly presented and formatted (some symbols are almost impossible to read on the circle. In general the font sizes on all figures should be substantially larger).

Author response:
> **Weakness D** - Regarding Fig. 3: 
> 
We also realized the difficulty in reading and interpreting Fig. 3. We apologize for this inconvenience. We have redrawn Fig. 3 and modified the description. Basically, we wanted to provide a graphical illustration of how the different acceleration parameters ($\alpha$) might affect the acceleration of an eigenmode. In Sec. 4.1, we provided more comprehensive evaluation of the effect of $\alpha$ and the results are shown in Fig. 9 of Appendix A.3.

Reviewer comment (Theory):
My main concern is that the methods described are based on heuristics and no theoretical justification is presented. Hence, it is unclear whether this method will actually lead to convergence to an optimization dynamics (e.g. SGD) in a given non-convex optimization problem.

Author response:
We appreciate your concern regarding the convergence of our proposed method. To address this, we revisit the formulation of Koopman-based Gradient Acceleration (KGA).

The standard Stochastic Gradient Descent (SGD) can be simplified as follows:
$w_{i+1} = w_i - \eta \nabla$. 
In contrast, the proposed KGA is formulated as:
$w_{i+1} = w_i - \eta \nabla - \alpha \left( U Q_{stable} \Lambda_{stable} L_{stable} U^T \eta \nabla \right) $

From this comparison, it becomes evident that the convergence of our method is primarily guided by the original SGD dynamics. The additional component in KGA primarily contributes to accelerating specific directions. This type of formulation is widely used in optimization methods such as Momentum, Adam, etc.

It's important to note that these components, which do not strictly converge to a zero gradient, can sometimes offer advantages, such as avoiding becoming stuck in local minima and saddle points. This characteristic is one of the reasons why such methods have gained popularity despite the lack of strict convergence guarantees in every scenario. We believe this perspective aligns well with current practices in the field and contributes to the robustness and effectiveness of our proposed method.

Reviewer comment (Theory):
The terms stable manifold etc are used incorrectly -- these do not refer to manifolds tangent to the singular vectors of the Koopman operator approximations. The demarcation into "stable", "unstable" and "neutral" Koopman modes is also quite hand-wavy and it is not explained what projecting the weight updates to the basis of some singular vectors of a matrix approximation of the Koopman operator, actually does.

Author response:
We project the weight update onto the rank-r truncated left singular vectors of $W_i$. It’s in this space that we find a matrix approximation of the Koopman Operator ($\tilde{A}$). This follows the standard dynamic mode decomposition formulation (e.g., Schmid 2010). For the demarcation of ‘stable’, ‘unstable’, and ‘neutral’ Koopman modes, we follow the convention of Avila and Mezic from Nature Communications 2020 (https://doi.org/10.1038/s41467-020-15582-5). These are the terminologies defined to facilitate the reference and usage of the three categories of Koopman modes. We agree the usage of “stable”, “unstable”, and “neutral” might be unconventional, but it saves the trouble to have to use an entire sentence everytime we need to refer to these three groups of modes.

We have removed the term ‘stable manifold’ from the paper. In this case, the equilibrium that the stable manifolds approach is “convergence” of a network but we understand that this is less rigorous than the traditional notion of a stable manifold.

Reviewer comment (Theory):
First of all, the DMD matrix is only an approximation of the Koopman operator on L^2. It would change with the snapshots used. Projecting on the range of a particular approximation induces a representation error and further the ad hoc masking completely changes the optimization dynamics.

Author response:
The reviewer is absolutely right that DMD is the technique developed to “approximate” the infinite-dimensional Koopman operator with a finite-dimensional counterpart. There have indeed been many variants to standard DMD, including, for example, extended DMD (Williams et al. 2015)  and Hankel DMD (Arbabi & Mezic, 2017). The major challenge these variants attempt to resolve is to consider an extended basis of observables in order to provide better approximations of the Koopman operator. For example, Hankel DMD tries to lift the measurement data to a higher-dimensional space using delay embedding. In our study, however, since the observables are in the weight space, which is already very high dimensional (in the order of millions), the standard DMD is sufficient. In Sec. 2, we provided  a reference (Kutz et al., 2016) as a comprehensive discussion of DMD.

The reviewer is also correct that the proposed KGA and KWM would both change the dynamics of the training process which is why we re-calculate the dynamics in each training epoch. We did compare the performance of re-calculating the dynamics each iteration (KGA-Iter) and each epoch (KGA-Epoch), as shown in Fig. 4. KGA-Iter shows the best improvement in loss, but incurs a significant performance penalty. We choose to use KGA-Epoch for the tradeoff between computational efficiency and improvement in training loss. Please refer to the first set of comparisons in Sec. 4.1.

M. O. Williams, I. G. Kevrekidis, and C. W. Rowley. A data–driven approximation of the Koopman operator: Extending dynamic mode decomposition. *Journal of Nonlinear Science*, 25(6):1307–1346, 2015.

H. Arbabi and I. Mezic. Ergodic theory, dynamic mode decomposition, and computation of spectral properties of the Koopman operator. *SIAM Journal on Applied Dynamical Systems*, 16(4):2096–2126, 2017.

Reviewer comment (Theory):
Please check the equations for the Koopman mode decompositions (the DMD part seems correct to me). It should read as $$K g(x) = \sum_{k \in \mathbb{Z}^+} \langle g, \phi_k\rangle K \phi_k(x) =  \sum_{k \in \mathbb{Z}^+} \langle g, \phi_k \rangle \lambda_k \phi_k(x),$$ here, $\phi_k$ are the eigenfunctions of the Koopman operator on $L^2.$ In the text, are the projections (the inner products) $c_k$ and the eigenfunctions $\phi_k$? If yes, $\phi_k$ should be a function of $x.$

Author response:
We have double-checked the notations in Eqs. 2 and 3, and Fig. 1 and they are consistent with references cited (Mezic, 2021)(Mezic, 2005). In a sense, both the reviewer and the authors are right! $C_k(x)$ is indeed the projection of $g(x)$ onto the eigenspace $\phi_k$ and thus can be written as $<g(x), \phi_k>$. $\phi_k$, the eigenfunction of the Koopman operator, is indeed a function of $x$. However, here, it has been proven that $\phi$ is invariant under the dynamics of $T$ and its level sets are invariant sets. Please refer to Eqs. (51) and (60) in (Mezic, 2021).

Reviewer comment (Theory):
Are the different columns of the matrix $W_i$ time snapshots? In the DMD algorithm, they have to be.

Author response:
Yes. Different columns of the $W$ matrix are the snapshot (or weights) of each iteration or epoch. Please see response to #2 on why we choose KGA-Epoch over KGA-Iter.

Reviewer comment (Theory):
What is the justification for projection on to the “stable” singular vectors? It is true that the eigenvalues close to the unit circle, but within the unit disk, have a correspondence with almost invariant structures or "slowly decaying modes", but these are not precisely used (even if qualitatively) in designing KGA. The slow decay is in the correlations between observables. If KGA did involve identifying the observables that show slow decay of autocorrelations for instance, that might more sense qualitatively.

Author response:
We project onto the rank-r truncated singular vectors (the remaining singular vectors are assumed to explain noise). We then project to the stable left-eigenvectors of the linear operator in this rank-r subspace. By projecting onto the stable left-eigenvectors, we are getting the component of the weight that is slow-moving or “almost” invariant. The hypothesis is that the contribution from slowly decaying modes persists longer during standard SGD and that speeding up convergence of these slowly decaying modes will accelerate training.