Identifying Equivalent Training Dynamics

Summary

This paper proposes a method for identifying equivalent training dynamics of deep neural networks from the perspective of dynamical systems theory, in particular the spectral analysis of Koopman operators. The authors propose to use the notion of topological conjugacy, i.e. two dynamical systems are considered equivalent if a smooth invertible map can transform the trajectories of one system into trajectories of the other. Since in case of linear systems topological conjugacy holds if and only if the eigenvalues of the two (linear) systems are identical, the authors propose to utilize Koopman operator theory which precisely concerns linearization and eigendecomposition of non-linear dynamical systems, along with an advantage in the context of analyzing training dynamics of neural networks as Koopman eigenvalues (eigenvalues of a linearized dynamical system given by eigendecomposition of Koopman operator) are invariant up to permutation of state variables, i.e. they automatically account for the permutation symmetries of neurons (weights). Using a variant of dynamic mode decomposition (DMD) algorithm for approximately obtaining Koopman eigenvalues, and by comparing the sets of eigenvalues obtained from different training trajectories using Wasserstein distance, the authors report the following results: - The proposed method recovers a recently discovered nonlinear topological conjugacy between the optimization dynamics of online mirror descent and online gradient descent on two synthetic regression problems, - The proposed method identifies a fundamental change of training dynamics as the width of a MLP with a single hidden layer and ReLU activation function is increased from 5 to 10, in comparison to an increase from 10 to 40 which is less significant in terms of the difference in the eigenvalue spectrum. - The proposed method identifies transitions in dynamics during early training of LeNet and ResNet-20 convolutional neural networks, and non-conjugacy of training dynamics between LeNet and ResNet-20 overall. - The proposed method identifies non-conjugacy of training dynamics for transformers that do and do not exhibit grokking (delayed perfect generalization on small algorithmic datasets) controlled by constraining weight norm.

Strengths

- S1. The paper is very well written and easy to follow. - S2. The use of topological conjugacy to define equivalence of training dynamics rather than trajectories of parameters or loss is convincing, and original as far as I am aware; the most related work I could find was Bielecki (2002) which studies topological conjugacy of a system under a gradient flow and its Euler discretization. - S3. The proposed method of estimating the eigenvalues of the Koopman operator, using a variant of DMD algorithm, and comparing them with Wasserstein distance, are natural and seems technically sound in implementing the idea of detecting topological conjugacy. - S4. The fact that Koopman eigenvalues are invariant to permutation of neurons (weights) is quite useful, especially given the role of permutation symmetries in training dynamics of neural networks (Entezari et al. 2021; Bokman and Kahl, 2023). - S5. The experimental results support the proposed idea and method, recovering a number of known results on equivalence of dynamics and falsifying some of the previously argued equivalences (e.g. of convolutional networks after the early transition of transition dynamics). Bielecki, Topological conjugacy of discrete time-map and Euler discrete dynamical systems generated by a gradient flow on a two-dimensional compact manifold (2002) Entezari et al. The role of permutation invariance in linear mode connectivity of neural networks (2021) Bokman and Kahl, Investigating how ReLU-networks encode symmetries (2023)

Weaknesses

- W1. There are some standard assumptions used in Koopman operator theory that may not strictly hold for mini-batch stochastic gradient descent, which includes non-stochasticity, forward-completeness, and time-invariance (violated e.g. due to learning rate scheduling and the use of momentum-based optimizers). Addressing these points may strengthen the paper. - W2. While most of the presented experiments on neural networks find non-conjugacy of training dynamics, it would be impactful if the proposed method identifies *conjugacy* of seemingly non-equivalent training dynamics, e.g. of differently initialized neural networks as conjectured by Entezari et al. (2021). - W3. The notations used in Equation (5) may be not exact; if we let $\rho\_\sigma:x\mapsto\tilde{x}$, it seems we might need to have $U^t \tilde{g}(\tilde{x})=\sum\_{i=1}^N\lambda\_i^t\tilde{\phi}(\tilde{x})v_i$ where $\tilde{g}\coloneqq g\circ \rho\_\sigma^{-1}$ and $\tilde{\phi}\coloneqq \phi\circ \rho\_\sigma^{-1}$. This will still leave Koopman eigenvalues invariant. (Please correct me if I'm wrong.) - W4. It may be beneficial to add a non-conjugate optimization algorithm to Figure 2 to show that the similarity of spectra of online mirror descent and online gradient descent is indeed significant. - W5. The use of Kolmogorov-Smirnov test in the panel F of Figure 3 seemed a bit weird since it shows that the difference of spectra of h=10 and h=40 is less than the difference of h=5 and h=10/40, but it does not precisely show the conjugacy of h=10 and h=40 as argued in Line 187-188; what would precisely show the conjugacy is a very low Wasserstein distance close to 0 (or a Wasserstein distance similar to random chancel level e.g. obtained from randomized partitioning of the eigenvalues of h=10 and h=40). Minor comments and typos - In the caption of Figure 6, should "... that are conjugate" be "... that are non-conjugate"? Entezari et al. The role of permutation invariance in linear mode connectivity of neural networks (2021)

Questions

- Q1. While the authors point out invariance to permutation symmetries of neurons (weights) in Equation (5), it seems that Koopman eigenvalues should be invariant for arbitrary invertible transformation symmetries on the parameters, which includes but not limited to permutations. Namely, rescaling symmetries (of cascaded linear layers) and rotation symmetries (of query and key projections used in attention in transformers) are other examples of invertible symmetries on the parameters (Ziyin, 2023). Please correct me if I am wrong. Liu Ziyin, Symmetry induces structure and constraint of learning (2023)

Rating

8

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have discussed the limitations in Section 5.

Summary

To compare two training dynamics, the paper proposes to compare Koopman operators for them, especially their eigenvalues, based on the previous result showing that consensus in the Koopman eigenvalues implies the equivalence of two dynamics up to some homeomorphism. Their framework is applied to compare various neural network architectures, whose results provide fine-grained insights to the previous (ad-hoc) findings.

Strengths

- The paper is clearly written, and well-structured. - The proposed framework seems natural, and can be applicable to various training dynamics.

Weaknesses

- There remain a concern about the computational efficiency of computing Koopman eigenvalues in the proposed framework, especially when the parameter space is high-dimensional. - The framework may be not suitable to compare the dynamics of different dimensions because the equality of the dimensions is required for the homeomorphic equivalence, according to basic results in topology.

Questions

See weaknesses.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

This paper applies Koopman mode decomposition (KMD) to examine whether the training dynamics of different model architectures/optimizers are equivalent or not. In particular the authors examine 4 cases: 1) the same objective learned by online mirror descent vs online gradient descent (which are equivalent), 2) the effect of width on similarity between FCNs with one hidden layer trained on MNIST, 3) training dynamics and especially the early phase of CNNs (known to rapidly evolve and then settle into a consistent phase), and 4) grokking in Transformers (known to have delayed test improvement if weight norms are unconstrained).

Strengths

The use of Koopman mode decomposition for analyzing training dynamics is novel as far as I can tell, and addresses important questions about neural network similarity, such as whether networks are equivalent after eliminating the effect of certain symmetries. These types of questions are difficult to answer and currently require specific methods to target specific symmetries. Having a general method that accounts for many symmetries is obviously advantageous, as is having a method that looks directly at the time evolution of models rather than their various outputs. The paper succinctly presents an easy to follow overview of Koopman operators and their applications with the necessary citations.

Weaknesses

Unfortunately, the results lack points of comparison and do not provide sufficiently novel or precise interpretations. The main issue is what would constitute a "low" vs. "high" Wasserstein distance between the Koopman eigenvalues? The first task examines similarity and the remainder dissimilarity, but all tasks need baselines for both known similarities and known dissimilarities between models. Otherwise, it is impossible to decide on a threshold between the two interpretations. This is similar to the issues with interpreting representational similarity metrics (e.g. Kornblith et al. Similarity of Neural Network Representations Revisited) and one may look at works in that field for an idea of how to make experiments more interpretable. It is hard to conceptually understand what similar versus distinct training dynamics look like. Some more toy examples like task 1, perhaps with distinct minima, could be helpful in this regard. The interpretations are all confirming previously known results, so it would be helpful to present a new result or refine a previously know result, that can also be confirmed by other means. For example, the method could be used to predict linear mode connectivity onset on a class of models/optimizers that has not been previously examined by Frankle et al. (using their empirical methods to confirm the onset time). This would establish that the method has real interpretive power.

Questions

Why not examine loss/error barriers instead of loss (as in Frankle et al.) to determine if the perturbed networks in appendix B.2 have similar trajectories? What is the significance of keeping the product of time delay and width constant? How does the frequency of time samples affect the estimated eigenvalues? How about noise (e.g. from SGD)? Would sampling Bayesian neural networks allow for arbitrarily many weights from the same training dynamic? How does this method allow grokking to be identified early?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

A major limitation of this method is that it requires many runs of a model to give good estimates of the Koopman eigenvalues. This makes designing experiments both more tricky and computationally expensive. Specifically, the difficulty is in having different runs that are known to have similar training dynamics, yet different weights. I'm this work, random multiplicative Gaussian perturbations are used. But there seems to be no way to directly check that these runs are actually similar in training dynamics. One baseline might be to evaluate KMD on independent sets of the same kind of run, and see how much the eigenvalues vary. The authors mention that eigenvalues greater than 1 correspond to unstable dynamics. It would be interesting to elaborate on what this means and how it relates to other observations or settings.

Summary

The authors utilize topological conjugacy and Koopman operator theory to create a framework for identifying between conjugate and non-conjugate training dynamics in DNNs. To validate their approach, they first show that their framework can accurately identify the known equivalence between online mirror descent and online gradient descent optimization methods. They then use the framework to gain new insights into training dynamics across various DNN architectures: - **Shallow vs. Wide Fully Connected Networks**: The framework reveals non-conjugate training dynamics between these types of networks. - **CNNs**: The framework helps characterize the early training phase dynamics in CNNs. - **Transformers**: The framework identifies non-conjugate training dynamics in Transformer models, including instances of "grokking". Overall, the results demonstrate the framework's versatility and potential to illuminate the complex training dynamics of different DNN architectures.

Strengths

- I very much enjoyed the paper. It tackles a significant and challenging problem in deep learning practice using a novel and precise analytical approach. Understanding training dynamics can enhance the efficiency and robustness of DNN models. Any effort to advance the understanding of training dynamics is highly valuable, and this paper provides excellent and helpful insights in this regard. - Topological conjugacies have traditionally been difficult to compute. Using Koopman operator theory, the authors developed a novel framework for identifying conjugate and non-conjugate training dynamics. - They investigated the training dynamics across various DNN architectures.

Weaknesses

There are only a few concerns: 1- Regarding the identification of topological conjugacies using Koopman Operator Theory, what happens if the training dynamics are more complex, such as in the case of chaotic dynamics with a **mixed Koopman spectrum**? How can we investigate topological conjugacy in such situations? 2- It is still unclear to me whether and how a numerical method for approximating the KMD might affect the results of their method for identifying topological conjugacies. 3- The Wasserstein distance is used to quantify the differences between Koopman eigenvalues. Why was only this metric chosen to compare distributions? Could other metrics, such as Kullback-Leibler Divergence (KL Divergence) or Hellinger Distance, be used for comparison? Specifically, it would be nice to better understand more advantages of the Wasserstein distance over these other metrics. Additionally, could using multiple metrics (at least two different metrics) help assess the robustness of the results?

Questions

I) On page 2, lines 59-60, the authors have mentioned that the same framework can be used across a number of DNN architectures to study a variety of dynamical phenomena during training demonstrates the generality of this approach. It would be helpful to understand how their approach could be applied to **RNNs for time series forecasting**. Could you please elaborate on this? II) According to Table S1, the activation function for FCN is ReLU. I am interested in knowing how choosing another activation function (any activation function except ReLU) can affect the results?

Rating

8

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

The authors have adequately addressed most of the limitations. The other limitation that occurs to me is the limitations of their method for training dynamics that are more complex, such as in the case of chaotic dynamics with a mixed Koopman spectrum.