Identifying Equivalent Training Dynamics

We thank the reviewer for their time and detailed comments. We are encouraged that they found our work well written and our framework novel! Below, we respond to the specific questions the reviewer had.

"It would be impactful if the proposed method identifies conjugacy of seemingly non-equivalent training dynamics"

This is a great suggestion. To address this, and to put it in the context of Entezari et al. (2021), we have analyzed the Koopman eigenvalues associated with training the $h = 40$ FCN on MNIST, across 25 initializations. We computed the Wasserstein distance between all pairs of eigenvalues and found that the distribution has a similar mean as what is seen in comparing $h = 10$ and $h = 40$ (compare Fig. 2B in the global response pdf with Fig. 3F in our submission). In addition, we find examples where the eigenvalues are nearly exactly overlapping (Fig. 2A in the global response pdf). This provides evidence that our framework can identify similar dynamics for FCNs that converge to solutions previously shown to be separated by a loss barrier [Fig. 2 left of Entezari et a. (2021)]. We will add these results to the main text in our revised manuscript (Sec. 4.2). We appreciate this suggestion.

"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"

This is a great idea and we believe that the inclusion of a non-conjugate optimization algorithm will additionally further emphasize how training dynamics can be different. To make this point most clear, we have chosen to compare online mirror and online gradient descent with the bisection method (BM). Unlike OMD and OGD, which take small local steps, the BM takes large, in some cases global, steps. We illustrate one trajectory in Fig. 3A of the global response pdf where these global steps lead to parameters that "hop" between positive and negative values (light blue lines). As such, we can intuitively expect that there is not a conjugacy between BM and OMD/OGD. Comparing the Koopman eigenvalues associated with each optimizer, we indeed see a clear difference between BM and OMD/OGD. In particular, BM has eigenvalues that are complex and negative, which correspond to the "hopping" in the parameter space. In contrast, OMD/OGD have positive, real-only Koopman eigenvalues. This further emphasizes what differences in Koopman eigenvalues correspond to. We will add these results to the main text in our revised manuscript (Sec. 4.1). We appreciate this suggestion.

"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"

This is a great comment. Thank you for pointing this out. We were focused on the permutation symmetry, but you are correct that there are other symmetries that are relevant in the DNN training setting that the Koopman spectra are invariant to. We will discuss this in the Introduction and Sec. 3 of the revised manuscript. We believe that this further demonstrates the utility of our framework.

"There are some standard assumptions used in Koopman operator theory that may not strictly hold for mini-batch stochastic gradient descent"

This is a great point. We agree that more discussion on why we think that Koopman mode decomposition can be reasonably approximated in these settings, and what we believe can be done to improve the approximation, would strengthen our paper. Additionally, it would make it more clear what challenges might be expected in utilizing our framework to study other problems. We will add this in Sec. 5 of the revised manuscript.

We note that in the FCN setting, no learning rate scheduling or momentum-based optimization was used, thus making it most directly in-line with previous work that has used dynamic mode decomposition with time-delays to study noisy dynamical systems [Arbabi and Mezic (2017); Brunton et al. (2017)].

"The notations used in Equation (5) may be not exact"

Yes! Thank you for that correction. We will correct this in the revised version of the manuscript.

"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"

This is a good point and we agree that in general it is challenging to show conjugate dynamics with only the Wasserstein distance. For this reason, we explicitly examined the Koopman spectra of $h = 5$, $10$, and $40$ (Fig. 3D and E in submission). We find that the Koopman eigenvalues between $h = 10$ and $h = 40$ are not only close to each other, but have similar properties (i.e., if an eigenvalue for $h = 40$ is real only, the eigenvalue of $h = 10$ that is closest to it is also real only). This is not the case for $h = 5$, where there is a pair of complex conjugate eigenvalues that do not match the real only eigenvalues of $h = 10$/$40$. For this reason, we note that this “suggests conjugate dynamical behavior” (lines 187-188). In the revised manuscript, we will make it more clear why we are concluding this, and that this conclusion is a hypothesis, given the current tools we have.

"In the caption of Figure 6, should "... that are conjugate" be "... that are non-conjugate"? "

Oops yes! Thank you for pointing that out. We have corrected that typo.

References:

Arbabi and Mezic (2017) “Ergodic theory, dynamic mode decomposition, and computation of spectral properties of the Koopman operator”

Brunton et al., (2017) “Chaos as an intermittently forced linear system”

Entezari et al. (2021) "The role of permutation invariance in linear mode connectivity of neural networks"

Thank you for the comprehensive response. The added experiments and discussions would greatly improve the impact of the paper and I have raised my score accordingly, and my only remaining concern is on the proper establishment of the null for at least one experiment, as in line with the request of reviewer 42wo. This seems important as the added results also use Wasserstein distances to show the conjugacy of differently initialized FCNs. If this is addressed, I am willing to raise my score to 8.

Sorry for requesting this close to the end of the discussion period, but I think this may be feasible since the random permutation baseline (take two sets of eigenvalues, measure their Wasserstein distance, then permute them, then measure their Wasserstein distance, then compare these two distances) does not require running the training again.

We thank the reviewer for their time and helpful comments. We are encouraged that they found our work well written and our framework applicable. Below, we respond to the specific questions the reviewer had.

"There remains a concern about the computational efficiency of computing Koopman eigenvalues in the proposed framework, especially when the parameter space is high-dimensional."

We thank the reviewer for pointing out that this limitation was not addressed in our submission. We agree that computing the Koopman eigenvalues can be a computationally intensive procedure. However, there are several ways of mitigating this. First, for especially large parameter spaces, dimensionality reduction (e.g., via PCA) can be used to reduce the size of the matrices used to approximate the Koopman mode decomposition, without sacrificing much dynamical information. Second, instead of considering all the Koopman eigenvalues, the top-$n$ can be considered (where small $n$ can often be sufficient and can reduce the noise present in the dynamics). This enables a reduced SVD, which can aid in reducing the computational costs. Third, a subset of the total number of parameters can be chosen as observables. This is motivated by the fact that in many DNNs there are a small number of highly correlated sets of weights [Brokman et al. (2023)], and thus, only a small number of the total parameters is needed to get an accurate picture of the dynamics. And fourth, numerical implementations of Koopman mode decomposition have been integrated into LAPACK [Drmac (2024)] (one of the most widely integrated libraries for numerical linear algebra), and are being developed for high performance computing. Therefore, tools exist to examine even larger parameter spaces. We will discuss this limitation in the Sec. 5 of the revised manuscript.

"The framework may be not suitable to compare the dynamics of different dimensions"

This is a good point that we did not address in our submission. In dynamical systems theory, in addition to the notion of a topological conjugacy, there is the notion of a semi-conjugacy which is defined for systems of different dimensions. In this case, the semi-conjugacy $h$ is a smooth but non-invertible mapping satisfying Eq. 1. Semi-conjugacies can be identified by one dynamical system having Koopman eigenvalues that are a subset of another dynamical system. We will discuss semi-conjugacies in Sec. 3.1 and 3.3 of our revised manuscript.

References:

Brokman et al. (2023) "ENHANCING NEURAL TRAINING VIA A CORRELATED DYNAMICS MODEL"

Drmac (2024) "A LAPACK Implementation of the Dynamic Mode Decomposition"

We thank the reviewer for their time and detailed comments. We are encouraged that they found our work well written and our framework novel. Below, we respond to the specific questions the reviewer had.

"The main issue is what would constitute a "low" vs. "high" Wasserstein distance between the Koopman eigenvalues?"

We appreciate the reviewer pointing out this limitation, as it brings to our attention that we can strengthen our explanation of this in the manuscript. We agree that there is, at the moment, no single number that enables us to identify a high Wasserstein distance vs. a low Wasserstein distance. However, we believe that there are several features that make our framework novel and interpretable.

First, it is worth noting that there existed no principled way of studying training dynamics (indeed, this is what motivated us). This has hindered prior work attempting to study training dynamics. For instance, Frankle et al. (2020) plotted several different measurements of training dynamics (e.g., loss, weight trajectories, weight norms) to identify transitions. We believe this is less interpretable and rigorous than our method.

Second, the Koopman eigenvalues can be understood as corresponding to different time-scales of the dynamics. When different features of the Koopman eigenvalues are present (e.g., real-only Koopman eigenvalues, complex conjugate Koopman eigenvalue pairs, Koopman eigenvalues with magnitude > 1), this corresponds to distinct properties of the dynamics (e.g., exponential decay, oscillations, exponential growth). Comparison of the individual Koopman eigenvalues can aid in the interpretation of the results. For instance, in the case of the fully connected neural networks, the wider networks have a Koopman eigenvalue that is real only, while the narrowest network has a Koopman eigenvalue complex conjugate pair. This, in addition to the large Wasserstein distance, led us to conclude that there was a non-conjugacy. This is a point that we will make more clear in the revised version of the manuscript by adding text to Secs. 3.2 and 4.1-4.4.

"It is hard to conceptually understand what similar versus distinct training dynamics look like."

We agree that making the distinction between conjugate and non-conjugate training dynamics intuitive is important to effectively communicating our method. To address this, we considered the bisection method (BM), an optimizer that is not conjugate to online mirror and gradient descent. Unlike OMD and OGD, which take small local steps, the BM takes large, in some cases global, steps. We illustrate one trajectory in Fig. 3A of the global response pdf where these global steps lead to parameters that "hop" between positive and negative values (light blue lines). As such, we can intuitively expect that there is not a conjugacy between BM and OMD/OGD. Comparing the Koopman eigenvalues associated with each optimizer, we indeed see a clear difference between BM and OMD/OGD. In particular, BM has eigenvalues that are complex and negative, which correspond to the "hopping" in the parameter space. In contrast, OMD/OGD have positive, real-only Koopman eigenvalues. This further emphasizes what differences in Koopman eigenvalues correspond to. We will add these results to the main text in our revised manuscript (Sec. 4.1).

"The interpretations are all confirming previously known results, so it would be helpful to present a new result or refine a previously known result"

We would like to emphasize that, while all our experiments were motivated by previous work, our framework enabled us to obtain novel results. Namely, we showed that LeNet and ResNet-18 have different training dynamics (while having similar transitions in their dynamics) and that Transformers that do and do not undergo grokking exhibit distinct training dynamics. We see these as being illustrative of the power of our method. Additionally, inspired by the comments of reviewer 4cZ9, we examined whether our framework could identify conjugate training dynamics in FCNs that have different initializations, as hypothesized by Entezari et al. (2021). As shown in Fig. 2 of the global response pdf, we find evidence in support of this hypothesis. This is another novel result. We will add it to the main text of the revised manuscript (Sec. 4.2).

"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?"

This is a great idea! We will incorporate this into the revised version of the manuscript. We appreciate the reviewer suggesting this.

"How does this method allow grokking to be identified early?"

The hypothesis that our method may enable grokking to be identified early in training comes from the fact that, when there are Koopman eigenvalues with magnitude > 1, there is an instability in the dynamics of the system. To see this, note that in the Koopman mode decomposition, the Koopman eigenvalues are raised to powers of the number of time steps in the future (Eq. 4). Thus, if an eigenvalue has magnitude > 1, its mode will quickly dominate the decomposition, and then explode. In the context of DNN training, this could be indicative of either exploding gradients, or a short-term transition in the training dynamics that, over a brief time window, appears to be an instability. That we find Transformers that do grok to have more eigenvalues > 1 suggests a greater instability in the training dynamics. Therefore, we hypothesized that looking for these eigenvalues may enable the identification of Transformers that will undergo grokking ahead of time. We appreciate the reviewer pointing out that this was not sufficiently explained in our submission and we will add text to clarify this in Secs. 3.2 and 4.4.

References:

Frankle et al. (2020) "The Early Phase of Neural Network Training"

We thank the reviewer for their time and detailed comments. We are encouraged that they found our work enjoyable to read and insightful! Below, we respond to the specific questions the reviewer had.

"How can we investigate topological conjugacy if the training dynamics are more complex, such as in the case of chaotic dynamics with a mixed Koopman spectrum?"

This is a great question! Recent work by Avila and Mezic (2024) has hypothesized that conjugacies in the case of continuous or mixed spectra can be identified by studying extensions of Koopman operator theory. We will explicitly discuss this in the revised manuscript. Additionally, we will make clear that systems with continuous spectra do not have guaranteed conjugacies if their point spectra match. However, we are unaware of any DNN training algorithms that would lead to chaotic trajectories (if you are familiar with any, we would be very interested in hearing!).

"I am interested in knowing how choosing another activation function (any activation function except ReLU) can affect the results?"

This is a good question and we appreciate the suggestion. We performed a new set of experiments, examining the effect of FCN width when using GeLU activation functions. We find very similar results to FCNs using ReLU activation functions (Fig. 1 of the global response pdf). This illustrates that our framework is robust to changes in activation function. We will add this figure to the Appendix and reference the results in Sec. 4.2 of the revised manuscript.

"Why was the Wasserstein distance used to quantify differences between Koopman spectra?"

We appreciate the reviewer pointing out that this choice was not well motivated in our submission. We chose the Wasserstein distance as it provides a metric of how far apart individual Koopman eigenvalues are from each other. This is in contrast to the KL divergence, which quantifies how different the distributions of eigenvalues are (but not whether the distributions could be easily made to be similar). The need for a metric that provides the distance between individual eigenvalues comes from the fact that the eigenvalues correspond to time-scales of the dynamics. Therefore, we expect two dynamical systems with Koopman eigenvalues that are far apart to have more distinct dynamics than two dynamical systems with Koopman eigenvalues that are near each other. We will add this rationale in Sec. 3.3 to make our choice more transparent.

"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."

This is a good question that we hope future work will investigate in rigorous detail. However, we can provide some insight on this question.

There are two ways that numerical methods for approximating the Koopman mode decomposition might impact the ability to properly identify topological conjugacies between DNN training dynamics. The first is if one set of DNN models has more noise in its training dynamics than another model. In this case, the approximated Koopman eigenvalues associated with the noisier training may not match the approximated Koopman eigenvalues associated with the less noisy training, even if a topological conjugacy between the average dynamical behavior exists. To mitigate this, we chose to use time-delays, which have been found to be more robust to noise [Arbabi and Mezic (2017); Brunton et al. (2017)], and we used a reduced SVD, so as to not compare Koopman eigenvalues associated with modes that are weaker, and thus more likely to be noise. Additionally, by sampling more than one trajectory, we reduced the impact an individual noisy training trajectory has.

The second way numerical implementations of Koopman mode decomposition can impact the identification of topological conjugacies is if the numerical implementation is biased towards certain errors for specific dynamics. For instance, if the numerical implementation was biased towards generating complex conjugate pairs of eigenvalues when the dynamics have a fast exponential decay, but not if the dynamics have a slow exponential decay, this could lead to the conclusion of a greater difference in dynamics than what is actually present. However, we are unaware of any work showing that dynamic mode decomposition with time-delays (the numerical implementation we used) has this feature. Therefore, we do not believe it will have a major effect on our results.

We will add this discussion to the main text of our revised manuscript (Sec. 5).

"It would be helpful to understand how their approach could be applied to RNNs for time series forecasting."

We had not thought about using our method in the context of RNNs for time-series forecasting, but our initial thoughts are that it could be leveraged in two ways. First, identifying that two seemingly distinct time-series induce conjugate training dynamics in a given RNN model would suggest that the time-series have some fundamental similarities. This could provide a unique way of doing systems identification and coarse-graining. And second, training RNNs can be challenging due to vanishing and exploding gradients. By computing the Koopman spectra associated with the early training of models that have these issues, and comparing them to ongoing training runs, it may be possible to identify ahead of time whether an RNN model undergoing training is likely to experience these problems. This could allow for more efficient early stopping and architecture search.

References:

Avila and Mezic (2024) “Spectral Properties of Pullback Operators on Vector Bundles of a Dynamical System”

Arbabi and Mezic (2017) “Ergodic theory, dynamic mode decomposition, and computation of spectral properties of the Koopman operator”

Brunton et al., (2017) “Chaos as an intermittently forced linear system”