Identifying Equivalent Training Dynamics
Authors: William T Redman
Juan M. Bello-Rivas
Maria Fonoberova
Ryan Mohr
Yannis Kevrekidis
Igor Mezic
Venue: NeurIPS 2024 spotlight (2024)
Keywords: training dynamics
dynamical systems theory
Koopman operator theory
grokking
learning phases
Primary area: other
TL;DR: We make use of Koopman operator theory to identify when the dynamics associated with training two different neural network models are equivalent.

Abstract: Study of the nonlinear evolution deep neural network (DNN) parameters undergo during training has uncovered regimes of distinct dynamical behavior. While a detailed understanding of these phenomena has the potential to advance improvements in training efficiency and robustness, the lack of methods for identifying when DNN models have equivalent dynamics limits the insight that can be gained from prior work. Topological conjugacy, a notion from dynamical systems theory, provides a precise definition of dynamical equivalence, offering a possible route to address this need. However, topological conjugacies have historically been challenging to compute. By leveraging advances in Koopman operator theory, we develop a framework for identifying conjugate and non-conjugate training dynamics. To validate our approach, we demonstrate that comparing Koopman eigenvalues can correctly identify a known equivalence between online mirror descent and online gradient descent. We then utilize our approach to: (a) identify non-conjugate training dynamics between shallow and wide fully connected neural networks; (b) characterize the early phase of training dynamics in convolutional neural networks; (c) uncover non-conjugate training dynamics in Transformers that do and do not undergo grokking. Our results, across a range of DNN architectures, illustrate the flexibility of our framework and highlight its potential for shedding new light on training dynamics.

Identifying Equivalent Training Dynamics
William T. Redman"
Juan Bello-Rivas Johns Hopkins University
Maria Fonoberova AIMdyn Inc.
Ryan Mohr AlMdyn Inc.
AIMdyn Inc. UC Santa Barbara
Yannis G. Kevrekidis
Igor Mezié! AIMdyn Inc.
Johns Hopkins University
UC Santa Barbara
Abstract
Study of the nonlinear evolution deep neural network (DNN) parameters undergo during training has uncovered regimes of distinct dynamical behavior. While a detailed understanding of these phenomena has the potential to advance improve- ments in training efficiency and robustness, the lack of methods for identifying when DNN models have equivalent dynamics limits the insight that can be gained from prior work. Topological conjugacy, a notion from dynamical systems theory, provides precise definition of dynamical equivalence, offering a possible route to address this need. However, topological conjugacies have historically been challenging to compute. By leveraging advances in Koopman operator theory, we develop a framework for identifying conjugate and non-conjugate training approach, demonstrate that comparing Koopman eigenvalues can correctly identify a known equivalence between online mirror descent and online gradient descent. We then utilize our approach to: (a) identify non-conjugate training dynamics between shallow and wide fully connected neural networks; (b) characterize the early phase of training dynamics in convolutional neural networks; (c) uncover non-conjugate training dynamics in Transformers that do and do not undergo grokking. Our results, across range DNN architectures, illustrate the flexibility of our framework and highlight its potential for shedding new light on training dynamics.
dynamics. To validate
1 Introduction
The analysis and experimentation of deep neural network (DNN) training continues to uncover new and in some cases, surprising phenomena. By changing the architecture, optimization hyper-parameters, and/or initialization, it possible to identify regimes which DNN parameters evolve along trajectories (in parameter space) with linear dynamics I 21. low-dimensional dynamics 33 correlated dynamics [4], lazy/rich dynamics [5] [ and oscillatory dynamics 7 g] In some cases, the training dynamics have been linked with the performance of the trained model 7 9 [T0), providing new insight DNN generalization. Additionally, detailed understanding of the dynamics has led to improvements in training efficiency work can provide.
demonstrating the practical implications such
To obtain more complete picture of DNN training, it is necessary to have method by which equivalent dynamics can identified and distinguished from other, non-equivalent dynamies. The construction of equivalence classes, which has fundamentally shaped the study of complex systems
'redmanv@aimdnn.com com
mezi ci Gaimdyn.com
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

in other domains (e.g., phase transitions [[2], bifurcations [13], defects in materials [14), would advance the understanding of how architecture, optimization parameeer,,,,ess, and initialization shape DNN training and could be leveraged to search for new phenomena. However, identifying equivalent and non-equivalent training dynamics is challenged by the need for methods that:
Go beyond the coarseness of loss. While useful metrics, training and test loss can be shaped by non-dynamical features the training (e.g., different initializations, different number of hidden units). Thus, different losses is neither necessary sufficient to conclude non-equivalent training dynamics.
Respect permutation symmetry. DNNs are invariant to the re-ordering, within layers, of their hidden units [15]. Thus, identifying that DNN parameters evolve along trajectories that occupy distinct parts of parameter space is not sufficient to conclude non-equivalent dynamics 16 17
We propose to use topological conjugacy [8].
notion of dynamical equivalence developed in the field of dynamical systems theory (Sec. 3.1. to address these limitations. Historically, topological
conjugacy has been difficult to compute 19I. especially when the equations governing the dynamical systems under study are not known. However, recent advances in Koopman operator theory [20 21
[22) (Sec. 3.2p have enabled the identification of topological conjugacies from data [23 (Sec. 3.3
We explore the potential of this Koopman- based approach for identifying topological conjugacies in the domain of DNN training, finding that it is able to:
Recover a known nonlinear topological conjugacy between the training dynamics of online
mirror descent and online gradient descent (24 25 26 (Sec. 4.1):
Identify non-conjugate training dynamics between narrow and wide fully connected neural networks (FCNs) (Sec. 4.2);
Demonstrate the existence of conjugate training dynamics across different random initializa- tions of FCNs 16 (Appendix C.4):
Characterize the early phase of training dynamics 27
in convolutional neural networks
(CNNs) (Sec. 4.3
Uncover non-conjugate training dynamics across Transformers that do, and that do not undergo delayed generalization (i.e., "grokking") 28
29
(Sec.
That the same framework can be used across number of DNN architectures study variety of dynamical phenomena during training demonstrates the generality the approach. We conclude by discussing how it can be further improved to enable greater resolution of equivalent dynamics, and how can be used to shed greater light on DNN training (Sec. 5)
2
Related work
2.1 Identification of DNN training dynamics phenomena
Analytical results have been obtained for the DNN training dynamics of shallow student-teacher
30 31 and infinite width 1 networks. For modern architectures (e.g. CNNs, Transformers), the training dynamics have been probed via analysis of computational experiments. Application of dimensionality reduction has led to the observation that parameters are quickly constrained to being optimized along low-dimensional subspaces of the high-dimensional parameter space [3] 4) Inspection of losses, parameter and gradient magnitudes, etc. led to the identification of several transitions in the training dynamics of CNNs during the initial few epochs
(27). While insightful, this prior work cannot except at coarse-graing level - be used to determine whether the dynamics associated with training different DNN models (or training the same DNN model with different choices in paramtteas or initialization) are equivalent.
2.2 Koopman operator theory applied to DNN training
Data-driven implementations of Koopman operator theory have been used to model the dynamics of DNN training [32 33 34). Because of the linearity of the learned Koopman models (Sec. 3.2).

---

using them in place of standard gradient- gradient-based based methods has led to reductions in computational costs associated with DNN training Koopman-bon based methods have additionally been used meta-learn optimizers for DNNs 35
36 The ability of Koopman models to capture features of training dynamies has been leveraged to develop new methods for pruning DNN parameters [37 38 and new adaptive training methods (39]. While this prior work has demonstrated that accurate Koopman operator representations of the nonlinear training dynamics can extracted, none have utilized the theory to identify topological conjugacies.
3 Identifying equivalent training dynamics
3.1 Topological conjugacy
Given discrete-time dynamical maps Ti X X and T2: Y, natural question to ask whether they induce equivalent dynamical behavior. There are various possibilities for defining equivalence, but dynamical systems theory has made use of the notion of topological conjugacy [18] to identify when smooth invertible mapping can be used to transform trajectories of Ti to those of T2 (and vice versa). Formally, Ti and T2 said be topologically conjugate if there exists homeomorphism, h X -
+ Y, such that
ho Ti = 0 h.
It is straightforward to identify and construct conjugacies for linear systems. Let X = R" and let Ti = A and T2 = B, where A, B E R=x=" These describe linear dynamical systems, as at + 1) At(() and y(t + 1) By((). In this setting. A and B conjugate if there exists an H E TR such that y(t = Hx(t) and A = H-1BH. BH. This can happen if and only if the eigenvalues of A are the same the eigenvalues of B. Thus, for linear systems, topological conjugacy can be used to construct equivalence classes, partitioning the space of dynamical systems into families of matrices that have the same spectra. However, for nonlinear systems, it is challenging prove the existence or non-existence of conjugacies [19], limiting its use as a tool. In addition, historically it has been possible compute topological conjugacies for systems where the underlying dynamics are not analytically known.
3.2 Koopman mode decomposition
Over the past two decades, Koopman operator theory has emerged as powerful framework for study- ing nonlinear dynamical systems [20 21 (2). The Koopman operator, U, is an infinite dimensional linear operator that describes the time evolution of observables (i.e. functions of the underlying state-space variables, æ E X) that live in an appropriately defined function space, F (Fig. TA). That is, the observable g E evolves as
Ug(z(t)] [()x)()),
where t E N and T: + X is the underlying dynamical map on the state-space X.
The linearity of U enables mode decomposition [termed "Koopman Mode Decomposition" (KMD)).
The KMD is similar to the mode decomposition used for linear systems analysis, except that is defined in J, instead of X. In particular, the KMD is defined
U' U'g(x)
Σ Mor(z)); 1=1
(3)
where the triplet (Ai, 9i, v.) describes the Koopman eigenvalues. eigenfunctions, and modes, respec- tively. there exists a subspace F finite dimension, E N, that invariant to the action of the Koopman operator, then finite dimensional representation of the KMD can constructed,
U'9(x)= Σ (x)()))
4
FFor the sake of space, we describe only discrete- siscrete-time dynamical systems, but the theory extends to continuous time dynamical systems.
3

---

Function space
Koopman mode decompcation
Linear dynamics [(x(++1]] - Ugfat0]
U2((0) A 2
MVM
Lining
9,000
9,00
space
Norlinear dynamics x(++() 1[x6]
Equivalent
Conjugacy
XTU
h-T, TTy
Figure 1: Schematic of Koopman operator theory-based identification of conjugate dynamical systems. (A) By lifting nonlinear dynamics from finite dimensional state-space to an infinite dimensional function space, a linear representation can be achieved (from which finite dimensional approximation can be obtained). (B) The lincarity of the Koopman operator enables mode decomposition, which includes Koopman eigenvalues (orange), cigenfunctions (green), and modes (blue). (C) Dynamical systems with the same Koopman cigenvalues are topologically conjugate.
In cases of chaotic dynamics, representation by a finite number of Koopman modes is not achievable. Such systems are said to have continuous spectra. order for DNN training algorithm to be useful, it likely must avoid chaotic behavior. Therefore, we focus on training dynamics where Eq. is assumed to be valid.
From Eq. 4 it can be seen that the evolution of observable functions is described as sum of Koopman modes, each evolving at a specific time-scale (which is determined by the Koopman eigenvalues) (Fig. [B). The Koopman eigenvalues and their associated Koopman modes and eigenfunctions can be connected to the state-space geometry the underlying dynamical system 40
An important feature of the Koopman eigenvalues is that they are invariant to permutations of the ordering of state-space variables. Let x = [T., znJ and
(",( {1,..., n - {1, 3 is permutation and Pa : - is the permutation mapping. That is, a equivalent to x via a re-ordering of its labels. In this case, the action of the Koopman operator is
where
U'g(T) Σ (o)
(5)
where ī(Ø) (ā)] and ((Ø) o (Z)] Thus, the Koopman cigenvalues the same as they were for the non non-permuted system. We note that the Koopman spectrum is the same for other invariances that are known to exist in DNNs, such rescaling (of cascaded linear layers) and rotations (of query and key projections used in attention in Transformers) [41] This makes it a generally powerful approach for studying DNN training dynamics.
While Eq. iis true for deterministic dynamical systems and does hold for training via stochastic gradient descent (SGD), we believe it is still to appropriate to compute the KMD from weight trajectories for two reasons. First, theoretical work has expanded the notion of Koopman operator theory to stochastic dynamical systems [22 and defined Eq. 4 in terms of the expectation of the dynamics. This suggests that the KMD associated with SGD training will be able to inform us of the "average" dynamics during training. This will be useful to comparing different network behaviors. And second, prior work computing KMD on DNN training has found able to sufficiently approximate the training dynamics so to allow for the Koopman operator to be used to optimize and sparsify 38 DNNs. This suggests KMD can capture important aspects of the training.
32 34
Many numerical methods have been developed for approximating the KMD from data. This has enabled its successful application as a tool for spatio-temporal decomposition in providing insight into complex, real-world dynamical systems 43 44
45
46
47
Dynamic mode decomposition
(DMD) 48 449, the most popular of these methods, has spawned many variants
50 51 52
54
In general, DMD-basea approaches collect T + snapshots of data e R construct data matrices

---

Z = [=(0), x - 1)] and Z' = [x(1), x(T),, where Z, Z' E R and then approximate the Koopman operator by
U Z'z'z
where i denotes the pseudo-inverse. Utilizing dictionaries with nonlinear functions [50] has led to improved results, demonstrating how usage of the underlying Koopman operator theory can enhance the capture of complex dynamics. In addition, leveraging Takens' Embedding Theorem [55] and using time-delayed observables has proved to be a generally powerful approach for approximating Koopman eigenvalues
44
an approach we make use of (Sec.
3.3 Equivalent Koopman spectra implies topological conjugacy
Given that KMD provides a linear representation of nonlinear dynamical systems, identifying topological conjugacies through matching eigenvalues (Sec. [.1 again becomes viable. Indeed, has been proven that two discrete-time dynamical maps Ii and 12, each in the basin of attraction of a stable fixed point, are topologically conjugate if and only if the Koopman eigenvalues of the associated Koopman operators, U1 and U2 are the same [3 (Fig. (C). That is, topological conjugacy exists and only if
1 x
Vi 1, N
where A(1 and (2) correspond to the eigenvalues associated with U, and Uz. respectively, and N is the number of Koopman modes. When the dynamical systems under study have continuous spectra, Eq. does not imply topological conjugacy. As noted earlier, we do not believe this to be a major limitation when studying meaningful training dynamics. However, recent work has suggested that topological conjugacies may still be identifiable in the case of continuous spectra by using extensions of Koopman operator theory [56). We believe this will be fruitful direction for future work.
When the number of Koopman eigenvalues of U1 is larger than the number of Koopman eigenvalues of U2. the strict equivalence of Eq. cannot be satisfied. However, there may exist smooth, but non-invertible mapping h from X onto Y. In such case, T1 and T- are said to be semi-conjugate,
and this can be identified when {A No eigenvalues of U1 and U2, respectively.
1 c (2
Jj=1, where N₂ < Ni are the number of Koopman ( 3
Computing the KMD from data is unlikely to yield the same exact Koopman eigenvalues for conjugate dynamical systems, due to the presence of noise and finite sampling. Therefore, a method for computing the distance between eigenvalues is necessary when making comparisons. Here, we make use of the Wasserstein distance (5), a metric developed in the context of optimal transport that quantifies how much one distribution must be changed to match another. This notion of "distance" is important as the Koopman cigenvalues correspond to time-scales and we expect dynamical systems with increasingly large differences between their cigenvalues will have increasingly large differences in their dynamical behavioo") In the case where a small, finite number of Koopman eigenvalues computed (which can be achieved, even for systems with large number of observables, by performing dimensionality reduction or residual based pruning of modes (53)), ), the Wasserstein distance can be efficiently computed by using linear sum assignment.
4 Results
4.1
Identifying conjugate optimizers
We begin by validating that numerical approximations of KMD can indeed correctly identify conjuga- cies in settings relevant to DNN training. To do this, we consider a recently discovered nonlinear topological conjugacy between the optimization dynamics of Online Mirror Descent (OMD) and Online Gradient Descent (OGD) [24] 25 26 (Appendix A. This work has been of particular interest as OMD occurs on a convex loss landscape and OGD occurs on a non-convex loss landscape,
suggesting a potential route for studying behavior of OGD in simpler setting.
The conjugacy between OMD and OGD relies on a reparametrization of the loss function. Without prior knowledge of this reparametrization, it is challenging to identify the conjugacy by looking at only
Note that other metrics, such as the KL-divergence, may capture this distinction.
5

---

A
OMD
OCD
BM
0.5
1.-
OMD OGD
1.0
BM
0.0
Loss
imag(2)
0.4
-0.5-
-1.5
-1.0
20
40
80
80
100
20
40
60
80
100
lterations
Iterations
-1.0
-0.5
0.0
0.5
real(A)
Figure 2: Conjugacy between online mirror descent online gradient descent identifiable from Koopman spectra. (A) Comparing example trajectories of variables optimized via OMD (x1, (x1, OGD (u1, U2), and BM (21, 22). the existence of a conjugacy between OMD and OGD is not obvious. (B) Similarly, the existence of conjugacy is apparent when looking at the loss incurred using OMD and OGD. (C) Comparing the Koopman eigenvalues associated with optimizing using OMD, OGD, and BM correctly identifies the existence conjugacy between OMD and OGD, and lack of conjugacy between OMD/OGD and BM. The function optimized is in all subfigures is fxx Σ tan(z).
the training trajectories the losses (Fig. ZA, B Appendix A.3 for details implementation of OGD and OMD). This highlights some of the current challenges present in identifying dynamical equivalence from data.
We compute the KMD associated with optimization using OMD and OGD by considering trajectories of both, from many initial conditions, and compare the resulting Koopman eigenvalues (Appendix A..4.. We find high overlap between the two spectra (Fig. 2C, red and black dots). Additionally, the two sets of eigenvalues have the same structure. Namely, they consist only of real, positive eigenvalues. In contrast, the bisection method (BM), another optimization algorithm that is not conjugate to OMD OGD (Appendix A.2), has associated spectra that complex (Fig. c light blue dots). Performing a randomized shuffle of the eigenvalues between algorithms (Appendix [B). we find that 25% of the shuffles between OMD and OGD eigenvalues result in Wasserstein distance greater than or equal the true Wasserstein distance. This suggests the distributions are not statistically significantly distinct. However, 0% of the shuffles have Wasserstein distance greater than or equal to the true Wasserstein distance for OMD and BM, and OGD and BM, respectively. This provides evidence that the spectra of OMD/OGD and BM are statistically significantly distinct.
Similar results are obtained when applying KMD to OMD and OGD optimization a different func- tion (Fig. (T). Collectively, these results demonstrate that the Koopman-basee spectral identification of topological conjugacies can successfully recover a known equivalence and provide support that can be used more broadly in uncovering equivalences in DNN training dynamics.
4.2 Identifying the effect of width on fully connected neural network training
To start exploring the potential of our framework for identifying topological conjugacies in DNN training, we begin with a small-scale example. Namely, we consider fully connected neural network (FCN) with only a single hidden layer, trained on MNIST (Appendix C.I). Consistent with other architectures, we find that the wider the FCN (i.e., the more units in the hidden layer), the better the performance and the lower the loss (Fig. BA). Whether this is due to an increase in capacity, with more hidden units enabling dynamics, leading to a better traversal the loss landscape,
refined solution, whether this is due to change in the training at this point unclear.
Computing the Koopman eigenvalues associated with training FCNs of varying width (Fig. see Appendix C.2 for details), we find that narrow (h = 5) and wide (h 40) FCNs have training dynamics that are non-conjugate, as their Koopman spectra are non-overlapping (Fig. BE F). This suggests that the training dynamics undergo fundamental change as width increases. However, for FCNs with intermediate width (h- 10), the training dynamics are more aligned with the wide FCNs (Fig. E, F), suggesting conjugate dynamical behavior. The dynamical difference in training narrow and wide FCNs is also supported by performing the cigenvalue shuffle analysis (Appendix B).

---

8
1.-
1.0
Training loss
0.5
3 0.0
3
10'
-26
-1.0
0
10
20
30
50
60
Training iteration
X10P
Training iterations
x10
Training images
X1D
x10
0.04
imsg(AA imag(A)
0.04
Wasserstain distance
8.0-
0.00
imag(A)
0.02
0.04
8.0-
0.04
-0.0--
h
0.00
2.0
0.04
-0.04-
0.0
0.58
real(-)
1,00
1.02
0.5
8.00
0.85
real(A)
0.55
h-
h-5-h 40
h-
Figure 3: Narrow and wide fully connected neural networks have non-conjugate training dynamics. (A) for FCNs with hidden layer widths h 10, and 40. Solid line is mean and shaded
Training loss area is È standard deviation across 25 independently trained networks. (B). (C) Example weight trajectories. across training iterations, for narrow, intermediate, and wide FCNs. (D) Koopman eigenvalues associated with training FCNs of varying width. (E) Same as (D), but zoomed out and with the eigenvalues associated with h = 5 and h 10 compared to those associated with h 40. Dashed line in (D) and (E) denotes unit circle. (F) Wasserstein distance between Koopman cigenvalues associated with training FCNs of varying width. Error bars are + standard deviation across 25 independently trained FCNs. Kolmogorov-Smingg
(KS) tests were performed to assess statistical significance of distance: denotes p < 0.01 and denotes < 0.0001.
In particular, a much larger number of the shuffles between / 10 and h 40 eigenvalues have Wasserstein distance greater than or equal to the true Wasserstein distance than between /t = 5 and h = 40 (81% % vs. 55%), although significance is not reached. Similar results were found when using the GeLU instead of ReLU activations [59] (Fig. S3 demonstrating that our results are consistent across FCNs with similar activation functions. Thus, conclude that the additional improvement in performance observed when increasing the network width from / = 10 to h = 40 comes more from increase in capacity, than from change in training dynamics. Identifying this not possible by solely comparing the loss or weights (Fig. 1g..31--5 C-C), demonstrating the advantage of the Koopman- based approach for identifying equivalent and non-equivalent DNN training dynamics.
To further study the behavior of FCN training dynamics, we also compared the computed Koopman spectra of & 40 networks trained from different random initial conditions (Appendix C.4). Prior work has proven that different random initializations of sufficiently wide FCNs converge to local minima that have loss barrier along the linear interpolation between them, when taking into account permutation symmetry [16]. This suggests conjugate training dynamics, although this has not been explicitly shown. In support of this hypothesis, we find examples of FCNs, trained from different random initializations, with nearly identical Koopman spectra (Fig. S4A).
4.3 Identifying dynamical transitions convolutional neural network training
Prior work has argued that CNNs undergo transitions in their training dynamics during the early part of training (i.c. the first several epochs), and that these transitions are similar across different CNN architectures [27]. However, dynamical systems based methods were not used for analysis. Instead, this observation relied on coarse- coarse-grained grained observables (e.g., training loss, magnitude of gradients) to define the transitions and to determine when they occur.
To understand whether such results hold when considering the training dynamics finer-scale, we utilize our Koopman- based framework. To do this, we split the first epoch of training into windows of 100 training iterations. We compute the Koopman eigenvalues associated with dynamics that oceur in each window and denote them by Atytt. where ( < la are the first and last training iteration in the window. We then measure the Wasserstein distance between all combinations of pairs of eigenvalues. This enables to guantitatively assess transient dynamical behavior and identify when
7

---

LeNet (MNIST): Epoch
ReeNet-20 (CIFAR-10): Epoch
Joietp uww.ustooaaa B807
100-
-2.
Training lttrations
300
200
Training
-3.0
-3.0
datance
700
700
100
300
500
700
100
200
500
700
Training Iterations
Training ieeration
Training
Treining
Treining flerations
0.15
Lehee Reshin:23
0.010
R0
LeHlet
0.004
Reshes120
Imag/A]
0.05-
imaglt)
0.006
0.002-
0.00
0.000
0.000
0005
0.002-
0.00
1.10
0.965
1,010
0.905
0.988
1.002
1,004
(oaNA)
real(A)
real(.)
Figure 4: Koopman-based framework enables identification of transitions in dynamics during the early phase training for LeNet and ResNet-20. associated with LeNet training over windows of 100 training iterations during epoch 1. (B) Same as (A), but for ResNet-20 training. (C) Koopman eigenvalues associated with the dynamics that occur during training iterations intervals 0-99, 400-499, and 600-699. Dashed line denotes the unit circle.
(A) Logio Wasserstein distance between Koopman eigenvalues
in the early phase of training the dynamics transition from one equivalence class to another. We apply our approach to LeNet [60), simple CNN trained on MNIST, and ResNet-20 [6], trained on CIFAR-10 10 (see Appendix D.I for details).
We find that, for both LeNet and ResNet-20, the first 100 training iterations have the most distinct Koopman eigenvalues, as the Wasserstein distance between A0:99 and all other cigenvalues is large (Fig. AA, brrighht yellow first column and row). In addition, for both LeNet and ResNet-20, the training dynamics become similar after 700 training iterations, as the Wasserstein distance between A600:699 and A700:799 is small (Fig. 4A, dark blue square around diagonal in lower righthand corner). This is in agreement with the timeline found by Frankle et al. (2020) (27). However, we additionally find that the dynamics that occur between greater change for ResNet-20 than for LeNet, as there is larger Wasserstein distance between Koopman eigenvalues. This suggests a difference dynamical behavior between the architectures. By examining the Koopman eigenvalues associated with different training iteration windows, find non-overlapping spectra (Fig. 4C). Performing the eigenvalue shuffle analysis (Appendix B), we find evidence that the first 4 splits of 100 training steps have statistically significant different associated Koopman eigenvalues, as the 2%, 2%, 0%, and 4% of the shuffles had Koopman eigenvalues greater than or equal to the true Wasserstein distance. This suggests lack of topological conjugacy between the earliest training dynamics of Re and LeNet, despite the fact that the general timeline in transitions in dynamics is similar between the architectures.
and 700 training iterations exhibit
To understand how the training dynamics change over larger span of training time, we perform the same analysis, but computing the Koopman eigenvalues from the dynamics that occur during each epoch (Appendix D.3). We find that, at this coarser grain scale, both architectures see a similar evolution of their training dynamics. In particular, dynamics (Fig. AS5,, yellow first column and row), and the subsequent epochs have dynamics that become increasingly more similar (Fig. S5A. B increasing size of dark blue blocks centered on the diagonal).
find that the first epoch has the most distinct
Taken together, our Koopman-basse based analysis supports prior decomposition of the early phase of CNN training dynamics into regimes separated by transitions that occur after a similar number of training iterations across architectures (27]. However, with a finer-scale resolution of the dynamics, we additionally find that LeNet and ResNet-20 have non-conjugate training, demonstrating that the exact training dynamics are architecture-specffi.

---

Unconstrained: Grokking
B
Constrained: grokking
Toain
Tren
0.8-
0.8-
Accuracy
Accuracy
0.6-
0.8-
0.2-
0.2-
0.0
0.0
10a
10
10
10P
Training iterations
Training iterations
C
0.10
Train: Unconstrained
Unconstrained Constrained
0.5
Tram: Constraine
Accuracy
0.4-
Test Unsonstraieed
lest
imag(*)
0.05-
0.3-
0.00
0.24
0.14
-0.05-
0.0
-0.10
0
20
40
60
80
100
0.96
0.58
1.00
1.02
1.04
Training iterations
real()))
Figure 5: Transformers that do, and that do not undergo grokking have early training dynamics that are not conjugate. (A) Train and test loss, as a function of training steps, for a Transformer model that undergoes grokking. (B) Same as (A), but for Transformer whose training is constrained to have constant weight norm [22] In this case, grokking observed. (C) In the first 100 training steps, little difference loss of Transformers with and without constrained training. Lines are mean and shaded area is standard deviation across 20 independently trained networks. (D) Koopman eigenvalues associated with the dynamics that occur over the first 100 training iterations for Transformers that do, and that not undergo grokking.
between the
4.4 Identifying non-conjugate training dynamics for Transformers that do and do not grok
Since the discovery that Transformers trained on algorithmic data (c.g.. modular addition) undergo de- layed generalization ("grokking" - Fig. 5AA [28], considerable effort has been invested to understand how this arises. One particularly influential theory is that the norm of the weights at initialization plays a large role. In particular, it was shown that single layer Transformers, initialized at weights with sufficiently small norm, have training and test loss landscapes that are "aligned", while the same single layer Transformers, initialized at weights with sufficiently large norm, have training and test loss landscapes that "mis-aligned" (62). Constraining the norm of the weights to be small prevents grokking, with train and test accuracy increasing at similar training iterations (Fig. 5 [62]
What role the training dynamics play in grokking remains less understood. In particular, the extent to which constraining the weight norm changes the training dynamics (which could shed additional light on grokking) has yet be explored. We therefore compute the Koopman eigenvalues associated with
the training of constrained and unconstrained Transformers on modular addition (Appendix E We the dynamics from the earliest part of training, namely the first 100 training iterations (Fig. 5C).
We do this to avoid trivially seeing a difference, given the small weight changes that Transformers which undergo grokking make when the training accuracy is high.
We find that the Koopman eigenvalues are distinet (Fig. 5P). In addition a gap between the computed eigenvalues, we find that the dynamics associated with training the constrained Transformer has pair of complex conjugate cigenvalues that lie along the unit circle, whereas the dynamics associated with training the unconstrained Transformer has pair of complex conjugate eigenvalues outside of the unit circle. This suggests difference in stability properties, as Koopman eigenvalues with magnitude greater than 1 (i.e. those that lie outside the unit circle) correspond to unstable dynamics. Similar results were found when computing the Koopman eigenvalues associated with the training of the unconstrained Transformer over longer training time window (Fig. [S6).
These results suggest non-conjugacy the training dynamics of Transformers that do, and those that do not undergo grokking. In particular, constraining the weight norm appears to lead to more stable training dynamics, which may be due to the selection of a better subnetwork train (3 Additionally, these results suggest that it may be possible to identify grokking before it happens [64

---

5
Discussion
Motivated by the need for quantitative methods that can determine when the training dynamics of DNN parameter trajectories are equivalent, utilized advances in Koopman operator theory to develop a framework that can identify topological conjugacies. This Koopman based identification of
conjugate training dynamies is invariant to permutations of DNN parameters (Eq. 5 necessary
feature for methods used to study DNN training dynamics [15 16 17I. By applying our approach to the dynamics associated with optimization using OMD and OGD, we were able validate that
numerical implementations of KMD can identify conjugacies which are known to exist [24] 25 26
(Fig. O Additionally, this example demonstrates challenges existing approaches for comparing DNN training dynamics face. as comparing the losses and the parameter evolutions of OMD and OGD does not lead to a clear indication of the underlying equivalence (Fig. ( B).
Leveraging our Koopman-ben based approach on the training dynamics of DNNs of varying architectures led to several insights. First, we found evidence that shallow and wide FCNs have non-conjugate training dynamics (Fig. 3). This is consistent with theoretical and experimental work showing that FCN width can lead to lazy and rich training regimes [6] This provides further evidence that our Koopman- based approach can correctly identify equivalent and non-equivalent training dynamics. In addition, we find that FCNs of intermediate width have Koopman eigenvalues that are more similar to those of wide FCNs (Fig. 3) demonstrating that our approach can provide insight beyond the wide and shallow regimes. Second, applying our framework to the dynamics of CNNs, we found transitions in the dynamics during the early phase of training, consistent with prior work [27] (Fig. However, by closely examining the Koopman eigenvalues, we found non-conjugate dynamics between different CNN architectures, suggesting fundamental differences in training. These distinct dynamical features are aligned with previous observations of different behaviors when training sparse CNNs [1]3]] And third, we found that Transformers that do, and that do not undergo grokking have non-conjugate training dynamies (Fig. 5). By focusing on the early phase of Transformer training, we avoid trivially finding this due to differences in the training loss. Additionally, this provides evidence for the ability to anticipate grokking before it happens (64)
Our framework is similar in spirit to an approach that categorizes iterative algorithms from a controlled dynamical systems perspective 665 66). However, such an approach requires access to the underlying equations to identify equivalence classes, which our data-driven, Koopman based framework does not 67 68). Work concurrent ours has leveraged a similar approach study the dynamics of recurrent neural networks [69]. However, Ostrow et al. (2023) [69] studied the dynamics of the activations and not the dynamics of network parameters, which the focus of this paper.
Limitations. Numerical implementations that compute the KMD are only approximations to the action of the true Koopman operator. As such, they are subject to the same difficulties as other data- driven approaches. These include the selection of hyper-parameters associated with the construction of the Koopman operator, the choice of observable functions, and the number of modes considered. To mitigate the effect these limitations might have on our analysis, we used DMD-RRR. state-of- the-art numerical approach for KMD (53). and time-delayed observables, which have been found to provide robust results across range of domains 44 45 5T]. Determining the existence of a topological conjugacy between two dynamical systems requires assessing whether their associated Koopman eigenvalues are sufficiently similar. While in some cases this is clear (e.g. identical Koopman eigenvalues associated with optimization using OMD and OGD- - Fig. P distinct Koopman eigenvalues associated with training LeNet and ResNet-20- Fig. in other cases it is less apparent. To quantify these differences, we made use of the Wasserstein distance and attempted to compute significance with a randomized shuffle control. While a natural choice, additional work remains to connect the magnitude of the Wasserstein distance to the divergence of the dynamical propertics associated with training DNN models.
Future directions. The ability of our Koopman-based approach to identify conjugacies between iterative optimization algorithms suggests its potential for data-driven discovery and generation of new classes of algorithms 770 7 (2). By identifying equivalent training dynamics of DNNs, may be possible to use our approach for learning mappings that transform one DNN model to another (73. Finally, the characterization of the Koopman eigenvalues associated with training a wide range of DNN models varying in architecture, optimization hyper-parameters, and initialization will enable a detailed understanding of how these properties shape DNN training. Leveraging this understanding may lead to improved methods for DNN training and model development.
10

---

Acknowledgments
We thank members of AlMdyn Inc., Mitchell Ostrow, Adam Eisen, lla Fiete, and members of the Fiete Group for helpful discussion surrounding this work. We thank the anonymous NeurIPS reviewers for their thorough and detailed feedback, which strengthened or work. This material is based upon work supported by the Air Force Office of Scientific Research under award number PA9550-22-1-0531.
References
[1] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[2] Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks any depth evolve linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
[3] Guy Gur-Ari, Daniel Roberts, and Ethan Dyer. Gradient descent happens a tiny subspace. arXiv preprint (o2.:v:1112204770 2018.
[4] Jonathan Brokman, Roy Betser, Rotem Turjeman, Tom Berkov, Ido Cohen, and Guy Gilboa. Enhancing neural training via a correlated dynamics model. In The Twelfth International Conference on Learning Representations, 2024.
[5] Lenaic Chizat, Edouard Oyallon, and Francis Bach, On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
[6] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673, PMLR, 2020.
[7] Nisha Chandramoorthy, Andreas Loukas, Khashayar Gatmiry, and Stefanie Jegelka. On the generalization learning algorithms that not converge. Advances in Neural Information Processing Systems, 35:34241- 34257, 2022.
[8] Daniel Kunin, Javier Sagastuy- Lauren Gillespie, Eshed Margalit, Hidenori Tanaka, Surya Ganguli, and Daniel LK Yamins. The limiting dynamics of sgd: Modified loss, phasc-space oscillations, and anomalous diffusion. Neural Computation, 36(1): 36(1):151-17) 174, 2023.
[9] Alessandro Achille, Matteo Rovere, and Stefano Soatto, Critical learning periods in deep networks. In International Conference on Learning Representations. 2019.
[10] Julian Zilly, Alessandro Achille, Andrea Censi, and Emilio Frazzoli. On plasticity, invariance, and mutually frozen weights in sequential task learning. Advances in Neural Information Processing Systems, 14:1666--1239 2021.
[11] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pages 3259-3269. PMLR, 2020.
[12] Kenneth Wilson. The renormalization group and critical phenomena. Reviews of Modern Physics. 55(3):583, 1983.
[13] John Guckenheimer and Philip Holmes. Nonlinear oscillations, dynamical systems, and bifurcations of vector fields, volume 42. Springer Science & Business Media, 2013.
[14] N David Mermin. The topological theory of defects ordered media. Reviews of Modern Physics. 51(3):591, 1979.
[15] Robert Hecht-Nielsen. the algebraic structure feedforward network weight spaces. Advanced Neural Computers, pages 129-135. Elsevier, 1990.
[16] Rahim Entezari, Hanic Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. In International Conference on Learning Representations, 2022.
[17] Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re- ra-insin: Merging models modulo permutation symmetries. In The Eleventh International Conference on Learning Representations, 2023.

---

[18] Stephen Wiggins. Introduction to applied nonlineur dynamical systems and Chaus. Springer, 1996.
[19] Joseph D Skufca and Erik M Bollt. A concept of homeomorphic defect for defining mostly conjugate dynamical systems. Chaos: An Intendisciplinary Journal of Nonlinear Science, 18(1), 2008.
[20] Bernard O Koopman. Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 17(5):31--318, 1931.
[21] Igor Mezić. Spectral properties of dynamical systems, model reduction and decompositions. Nonlinear Dynamics, 41:309-325, 2005.
[22] Marko Budišić, Ryan Mohr, and Igor Mezić. Applied koopmanism. Chaos: An Interdisciplinary Journal of Nonlinear Science, 22(4), 2012.
[23] Igor Mezić. Spectrum of the koopman operator, spectral expansions in functional spaces, and state-space geometry. Journal of Nonlinear Science, 30(5):2091-2155, 2020.
[24] Ehsan Amid and Manfred KK Warmuth. Reparameterizing mirror descent as gradient descent. Advances in Neural Information Processing Systems, 33:8430-8439, 2020.
[25] Ehsan Amid and Manfred K Warmuth. Winnowing with gradient descent. In Conference on Learning Theory, pages 163-182, 182. PMLR, 2020.
[26] Udaya Ghai, Zhou Lu, and Elad Hazan. Non-convex online learning via algorithmic equivalence. Advances in Neural Information Processing Systems, 35:22161-223 22172. 2022.
[27] Jonathan Frankle, David Schwab, and Ari S. Morcos. The early phase of neural network training. In International Conference on Learning Representations, 2020.
[28] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:200.00217,, 2022.
[29] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. The Eleventh International Conference on Learning Representations, 2022.
[30] David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters, 74(21):433) 1995.
[31] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information procesing systems, 32, 2019.
[32] Akshunna S Dogra and William T Redman. Optimizing neural networks via koopman operator theory. Advances in Neural Information Processing Systems, 33:2087-200 2020.
[33] Mauricio E Tano, Gavin D Portwood, and Jean C Ragusa. Accelerating training in artificial neural networks with dynamic mode decomposition. arXiv preprint arXiv:2006. 14371, 2020.
[34] Di Luo, Jiayu Shen, Rumen Dangovski, and Marin Soljacic. Quack: Accelerating gradient-based quantum optimization with koopman operator learning. Advances in Neural Information Processing Systems, 36, 2024.
[35] Petr Šimánek, Daniel Vašata, and Pavel Kordík. Learning to optimize with dynamic mode decomposition. In 2022 International Joint Conference on Neural Networks (IJCNN), pages 1-8. IEEE, 2022.
[36] Ryuichiro Hataya and Yoshinobu Kawahara. Glocal hypergradient estimation with koopman operator. arXiv preprint a022 2024.
[37] Ryan Mohr, Maria Fonoberova, Iva Manojlovic, Alcksandr Andrejeuk, Zlatko Drmac, Yannis G Kevrekidis, and Igor Mezic. Applications of koopman mode analysis to neural networks. In AAAI Spring Symposium: MLPS, 2021.
[38] William T Redman, Maria Fonoberova, Ryan Mohr, Yannis Kevrekidis, and Igor Mezic. An operator theoretic view on pruning deep neural networks. In International Conference on Learning Representations, 2022.
[39] Fanqi Wang, Landon Harris, Weisheng Tang, Hairong Qi, Dan Wilson, and Igor Mezic. Dynamic training guided by training dynamics. 2023.
12

---

[40] Alexandre Mauroy. Igor Mezić, and Jeff' Moehlis. Isostables, isochrous, and koopman spectrum for the action-angle representation of stable fixed point dynamics. Physica D: Nonlinear Phenomena, 261:19-30. 2013.
[41] Liu Ziyin. Symmetry leads structured constraint of learning. arXiv preprint orXiv:2309. 16932, 2023.
[42] Nelida Cmjarii-7ie,, Senka Maćcšić, and Igor Mezić. Koopman operator spectrum for random dynamical systems. Journal of Nonlinear Science, 30:2007-2056, 2020.
[43] Bingni W Brunton, Lise A Johnson, Jeffrey G Ojemann, and Nathan Kutz. Extracting spatial-temporal coherent patterns in large-scale neural recordings using dynamic mode decomposition. Journal of neuro- science methods, 258:1 1- 15, 2016.
[44] Steven L Brunton, Bingni W Brunton, Joshua L Proctor, Eurika Kaiser, and Nathan Kutz. Chaos an intermittently forced linear system. Nature communications, 8(1):19, 2017.
[45] Allan M Avila and Igar Mezić. Data-driven analysis and forecasting of highway traffic dynamics. Nature Communications, 11(1):1-16, 2020.
[46] James Hogg, Maria Fonoberova, and Igor Mezić. Exponentially decaying modes and long-term prediction of sea ice concentration using koopman mode decomposition. Scientific reports, 10(1):16313, 2020.
[47] Igor Mezić, Zlatko Drmač, Nelida Čmjaarć, Senka MaćeSić, Maria Fonoberova, Ryan Mohr, Allan M Avila, Iva Manojlović, and Aleksandr Andrejčuk. A koopman operator- prediction algorithm and its application covid-19 pandemic and influenza cases, Scientific reports, 14(1):5788, 2024.
[48] Peter Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of fluid mechanicy, 656:5-28, 2010.
[49] Clarence Rowley, Igor Mezić, Shervin Bagheri, Philipp Schlatter, and Dan Henningson. Spectral analysis of nonlinear flows. Journal of fluid mechanics, 641:115 127, 2009.
[50] Matthew O Williams, Joannis Kevrekidis, and Clarence W Rowley. data driven approximation of the koopman operator: Extending dynamic mode decomposition. Journal of Nonlinear Science, 15:1307-1346, 2015.
[51] Hassan Arbabi and Igor Mezic. Ergodic theory, dynamic mode decomposition, and computation of spectral properties of the koopman operator. SIAM Journal Applied Dynamical Systems, 16(4):2096-2) 2126. 2017.
[52] Travis Askham and J Nathan Kutz. Variable projection methods for an optimized dynamic mode decompo- sition. SIAM Journal on Applied Dynamical Systems, 17(1):300-416, 2018
[53] Zlatko Drmac, Igor Mezic, and Ryan Mohr. Data driven modal decompositions: analysis and enhancements. SIAM Journal on Scientific Computing, (44::22223322233530 2018.
[54] Matthew Colbrook and Alex Townsend. Rigorous data-driven computation of spectral propertics of koopman operators for dynamical systems. Communications on Pure and Applied Mathematics, (71122 283, 2024.
[55] Floris Takens. Detecting strange attractors in turbulence. In Dynamical Systems and Turbulence, Warwick 1980: proceedings of a symposium held at the University Warwick 1979/80, pages 366-381. Springer. 2006.
[56] Allan M Avila and Igor Mezić. Spectral properties of pullback operators on vector bundles of dynamical system. SIAM Journal on Applied Dynamical Systems, 22(4):3059-309,, 2023.
[57] Leonid V Kantorovich. Mathematical methods of organizing and planning production. Management science, 6(4):366-422, 1960.
[58] Leonid Nisonovich Vaserstein. Markov processes over denumerable products of spaces, describing large systems of automata. Problemy Peredachi Informatsii, 5(3):64-72, 1969.
[59] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv: 180.04160..08455. 2016.
[60] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- learning applied to document recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998.
13

---

[61] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and puttern recognition, pages 770-778. 2016.
[62] Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In The Eleventh International Conference on Learning Representations, 2022.
[63] Gouki Minegishi, Yusuke Iwasawa, and Yutaka Matsuo. Grokking tickets: Lottery tickets accelerate grokking. arXiv preprint arXiv:2310.19477 2023.
[64] Pascal Notsawo Jr, Hattie Zhou, Mohammad Pezeshki, Irina Rish, Guillaume Dumas, et al. Predicting grokking long before it happens: A look into the loss landscape of models which grok. arXiv preprint arXiv:2306. 13253, 2023.
[65] Laurent Lessard, Benjamin Recht, and Andrew Packard. Analysis and design of optimization algorithms via integral quadratic constraints. SIAM Journal on Optimization, 26(1):57-995 95, 2016.
[66] Shipu Zhao, Laurent Lessard, and Madeleine Udell. An automatic system detect equivalence between iterative algorithms. arXiv preprint arXiv:2105.04660 2021.
[67] Felix Dietrich, Thomas N Thiem, and Joannis G Kevrekidis. On the koopman operator of algorithms. SIAM Journal on Applied Dynamical Systems, 19(2):60-8883 2020.
[68]
William T Redman, Maria Fonoberova, Ryan Mohr, Ioannis G Kevrekidis, and Igor Mezić. Algorithmic (semi-) conjugacy via koopman operator theory. 2022 IEEE 61 Conference on Decision and Control (CDC), pages 6006-6011. 1. IEEF, 2022.
[69] Mitchell Ostrow, Adam Eisen, Leo Kozachkov, and Fiete. Beyond geometry: Comparing the temporal structure of computation in neural circuits with dynamical similarity analysis. Advances in Neurul Information Processing Syxtems, 36, 2024.
[70] Alexander Mitsos, Jaromit Najman, and Ioannis G Kevrekidis. Optimal deterministic algorithm generation. Journal of Global Optimization, 71 71:891-913, 2018
[71] Yue Guo, Felix Dietrich, Tom Bertalan, Danimir T Doncevic, Manuel Dahmen, Ioannis G Kevrekidis, and Qianxiao Li. Personalized algorithm generation: A case study in learning ode integrators. SIAM Journal Scientific Computing, 44(4):A111-AA99) 2022.
[72] Danimir T Doncevic, Alexander Mitsos, Yue Guo, Qianxiao Li, Felix Dietrich, Manuel Dahmen, and Ioannis G Kevrekidis. A recursively recurrent neural network (r2n2) architecture for learning iterative algorithms. SIAM Journal on Scientific Computing, 46(2)A719- 2024.
[73] Tom Bertalan, Felix Dietrich, and Ioannis G Kevrekidis. Transformations between deep neural networks. arXiv preprint NXX::2000.0064466 2020.
[74] Davis Blalock, Jose Javier Gonzalez Ortiz, Jonathan Frankle, and John Guttag. What is the state of neural network pruning? Proceedings of machine learning and syxtems, 2:129-146, 2020.
14

---

A Online mirror and online gradient descent
A.1 Conjugacy between OMD and OGD
Here we outline the conjugacy between Online Mirror Descent (OMD) and Online Gradient Descent (OGD). We follow the notation and framing presented in Ghai al. (2022) 26
Let K. be convex set. OMD is applied find a minimum of the function ff on K, subject to convex regularizer R. For each iteration of OMD, the state of the algorithm (initialized z<0) E K.). updated by performing ((x()) where is the learning rate. Because this step may not K., the Bregman projection operator, [ (x) arg gin Pn(vl Ir), is used. OGD, on the other hand, is applied to (possibly)
(VR) [V][a[O)] 7/ V
non-convex set K"' and a (possibly) non-convex function f. As with OMD, each iteration of OGD involves updating the state (initialized at u(0) E K'') and projecting the update back into the set K'. In this case, the update is u(() v/ [[u(t]. where again 1y is the learning rate and the projection is done via the Euclidean projection operator. Ih (ar) (see Algorithms and for pseudocode implementations of both algorithms).. V (q (u((] and KC'
When V/[[((]]
(KC). the outputs of the two algorithms equivalent via the mapping 4 (i.e.. q is the copological conjugacy or reparameterizaaionn). key theorem of [25 showed that, in continuous-tin time, if (t) = q( (()(t))), ) then ou
dt [(((((()) (26).
Algorithm 1 Online Mirror Descent [26]
Algorithm 2 Online Gradient Descent [26]
O: Input: z00) K. R. f for t 0...., T'- 1 do
Input: u(0) E K'' KT, J 0: for t = 0,.... T-
O: 1(++1)= = (VHlc())))))
0: u(l + 1) =((1= V fu(() u(t + 1) = -I [u( + 1)]
axt + 1) n + 1)]
A.2 Bisection method
Let K. [c, dI where f(c) < 0, fd) 0, and there exists only one x" [c, dJ s.t f(=") = 0 Let b(0)))(0) E K" s.t. [a((0)] 0 and [[]( O. Define z(1) = [a(I) b(()] /2. For each iteration of the Bisection Method (BM), if [[=[]] < then [a(t), b(4)] is updated to [=(1). 6(().. Otherwise [a(t), 6(t) is updated to (0().=()), (see Algorithm 3 for pseudocode implementation).
Because not one but three variables are being updated at each iteration of the BM, [a(t), b(t), (t) the BM can exhibit several distinct properties from OMD and OGD. First, how much the updated =(t + 1) differs from 2(t) depends on a(t) and b(t). Thus, || + 1) =t)]|22 can be much larger than the steps OMD and OGD takes. This global property enables to escape local minima that OMD/OGD get stuck in, but can also lead to increases in loss. This behavior can be seen in Fig. B. Second, because f[((t) < 0 < {b((t) and because =(t + 1) is either aff) or b(f), the outputs of the BM [i.e.. =ff] can flip sign. This kind of oscillatory behavior can be seen in Fig. A. This would not happen with OMD or OGD, assuming sufficiently small learning rate. For these reasons, expect that the bisection method is not conjugate to OMD OGD. The distinct Koopman spectra (Fig. [C) demonstrate that our framework properly identifies this non-conjugacy.
Algorithm 3 Bisection Method
Input: f,a(0) E K", b(0) E K" s.t. fla(0)) < 0 < fla(b)]
0: for f = 0,...,
do
O: =(t) = [a(t) + b(t))/2
0: if f[x(()) < 0 then
aft + 1) x(t)
b(t + 1) = b(t)
0:
else
0:
alt + 1) a())
0: b(t + 1) z(()
A.3 Numerical experiments
To validate that the numerically computed Koopman eigenvalues, corresponding to OMD and OGD applied to specific problems, encode sufficient information to correctly identify the conjugacy, we chose test them on the example of log barrier regularization with exponential reparameterization. as presented in Ghai et al. (2022) 26]
In particular. we set R = -> -' log(2.). KC- = j0.01, 1.0j". K = [[4466, 0.0]*. and K. [-4/3, 133.8/7 8/7] (with d = 2), and q(u) exp(u). This is an example of nonlinear conjugacy.
15

---

Non-equivalent trajectories
Non-equivalent loss
C
Equivalent spectra
0.04
0.0
Loss
imag(A)
0.00
0.15-
4.00
-1.0
0.05
-0.04
20
60
100
20
60
100
0.8
0.7
088
Iterations
Iterations
real(-)
Figure S1: Conjugacy between online mirror descent and online gradient descent is identifiable from Koopman spectra. (A) Comparing example trajectories of variables optimized via OMD (2)) and OGD (w1, u2), the existence of conjugacy is not obvious.
22)
(B) Similarly, the existence when looking at the loss incurred by using OMD OGD. (C) Comparing the Koopman eigenvalues associated with optimizing using OMD and OGD correctly identifies the existence of a conjugacy. The function optimized is in all subfigures is f(z) Σ x". Note that the BM was not used in this figure fxx is symmetric around its minimum, which does satisfy the assumption that (a) < 0.
conjugacy not apparent
We applicd OMD and OGD
Σd Li=1 2 (Fig.
S1 and f Σ tan(x.) (Fig. 2 for
100 time steps, with a learning rate of rj 0.01. evolved 25 initial conditions, with a())
(..1, 0.5, 0.7,0.7.... 0.9) x (0.1, 0.5, 0.7, 0.9) and u(0) E 2-330 -1.75, -1.20,100000 -0.65, -0.6,-00.0001100 -0.10) x
(-2.30, -175, -.200 0....-00.010000 00.10). When we used the BM (Fig. 2) we sampled 25 initial conditions,
with a(0) E 1-16/12,-1/112, -10/12,-7/11 -1/12_ x (1/7,0.393, 0.643,, 0.893, 8/7). Note that, for simplic-
(111)]((11611111111111111111 (-16/12 13/12, 10/112 -7/12, -4/12]
and b(0) ity, we consider each a(O) and b(0) only once (leaving initial conditions). Using the resulting tra- jectories, we computed the KMD (see Appendix (.4). All experiments were run on a MacBook Air with an Apple MI chip. 1 CPU, and no GPUs. Code implementing our experiments can be found at https://github www.woooml.cccccccccoc..o.wc.wwe wallllamm-redmmm.. woddmmnIddeniffinng Eew.ttiannnguialletn....mm Terierrrriiinnnnnieeinnin Dynamics.
0.393, 0.643, 0.893,8/77
A.4 Computing the Koopman mode decomposition
To compute the KMD associated with optimization using OMD and OGD on f(x) Et tan(x;) (Fig.
2 and f(x) Ei xf (Fig. S1 we saved the values of xff) and u(t) across the 100 training steps.
These were concatenated into tensors X,U U E 22x100225 Four time-delays 44 45 51) were applied, and
the resulting tensors were flattened along the last dimension. This led to matrices Zx Zu B 230x2876 We
applied DMD-RRR (53) on these matrices to compute Koopman eigenvalues. The same approach was used to compute the KMD associated with optimization via BMM
B Randomized shuffle control
While the Wasserstein distance gives natural to quantify how similar the Koopman spectra associated
with dynamical systems are, it remains open question how to best interpret the magnitude of the computed distance. In particular, if A (L) {x(") ...., x and A (2) [x(x) .....
(2) }. what value of the
Wasserstein distance w Wa [ ()) (A is "sufficiently small" enough we can confidently conclude
two dynamical systems are topologically conjugate?
To help increase the transparency and interpretability of our results, we develop a randomized shuffle control baseline assess how significantly distinct (1) and (?) are. In particular, by comparing a distribution of Wasserstein distances computed from randomly shuffling (2) and A (2)) we may estimate whether the true Koopman eigenvalues more less "distant" than would be expected by chance.
One challenge that emerges in creating good randomized shuffie control is the fact that the Koopman cigenvalues correspond dynamical propertics the underlying system (c.g., decay rates, oscillations, fixed points). Therefore, if we take the naive approach and randomly assign half of A (4) (A4A, (22) y to A'()) and the other half to A (2) then A (1) and A (2) could be "non- having only slow or fast decays modes, no fixed
points, etc. In such case, we expect < w o== W (A()) (3),,A) AA()) in general. This makes the naive approach poor baseline.
16

---

To attempt to construct stronger and more informative baseline, we randomly assign eigenvalues to A and A in the following manner. First, by computing W22 (A(), A we identify the "assignment"
{1, {L, N) + {I,..,, N that maps the order of A (2) to be as close A A possible. In particular, is defined
N min EIA
(8)
To generate the shuffled cigenvalues, we then place x! in A with 50% probability and in A with 50% probability. The eigenvalue in (22 that is closest to X! via the assignment a [i.e., X is then placed in
whichever A'O) or A that ( is not. Performing TetuiT = 100 shuffles, we get distribution of Wasserstein distances between the shuffled eigenvalues, (whsat.. J. We report the number of shuffles that have Wasserstein distance greater than equal us. In the case where the two are highly distinct, we expect there be few shuffles that satisty w' 2 w. In contrast, when the Koopman eigenvalues are highly overlapping, we expect there to be more shuffles that satisfy w' 2 w. For this reason, we report the percent of shuffles with w' 2 w in the main text.
of Koopman eigenvalues
This approach defining the randomized shuffle control has two useful properties. First, A and A²2) have similar magnitudes of eigenvalues to A (1) and A possible. And second, if x and X complex conjugate pair of eigenvalues, whose closest matches are x and (2) ((++1) that are real only, then the complex conjugate pair can be "split" when creating A ((1)) and A A This can lead to wي < w. Thus, there is "penalty" when the number complex conjugate eigenvalues does not match between A (1) and A (2))
(s)
We note that this randomized shuffle control is not perfect construct and represents one possible way of gener- ating significance. We hope that future work will enable rigorous comparison between Koopman eigenvalues by developing computation schemes for estimating how much two dynamical systems diverge given the differences in their approximate spectra.
C Fully connected neural networks
C.1 Training experiment details
FCNs with only a single hidden layer are trained on MNIST, for one epoch, using SGD. Training was performed using PyTorch. All hyper-parameters used for training are presented in Tab All experiments were run on a MacBook Air with an Apple M1 chip. CPU, and no GPUs. Code implementing our experiments can be found thubb thub. http//wlliammmmeee
httpp::/ //gi
Equivalent Crtraiiing_DDnammc..
Hyper-parameters
Values
Learning rate (1))
0.1
Batch size (b)
60
Optimizer
SGD
Epochs
Activation function
ReLU
Table S1: Hyper-parameters used for FCN training in Sec. 4.2
C.2 Computing the Koopman mode decomposition
A challenge in computing the KMD associated with training FCNs is that, in order to accurately approximate the Koopman eigenvalues, multiple trajectories must be sampled. However, given that FCNs (and other DNN models) have loss landscapes with multiple local minima, training from different random initializations lead to different trajectories with different dynamical properties. To address this, we perform the following three steps:
We randomly sample initialization for the network input and output weights. We use the standard PyTorch initialization scheme, with weights being uniformly sampled in [-VK, vTe), where is the number of input features. We denote this initialization by Wo. We then train the FCN, from Wa. for full epoch.
We sample another n. 1 initializations of the FCN, for n, E N. Instead of randomly sampling a new set of parameters, we consider a perturbation of Wg. Namely, Wol1 + eN 1)]. where (0, 1) is Gaussian distribution with zero mean and unit variance. The FCNs were then trained
17

---

from these initializations, using the same batch order as the one used to train the FCN from Wo. In experiments, we nt 10 and 0.001. To investigate whether training from the perturbed initialization indeed lead to dynamics that were the same W6. we computed the ratio of their end test loss with the end test loss of the network initialized at Wo- We find that the ratio is centered around 1 (Fig. S22. providing evidence for that the training dynamics restricted to the basin of attraction of the same local minimum.
attraction as training from
We repeated the steps above The 1 times, for na E N. The weight evolutions, from each set of networks, were saved separately and the KMD was computed each one independently. In our experiments, Th 25.
h
40
48-
48
32-
Percent
Percent
-2-
52-
Percent
16-
0-
01
0.95
1.00
1.02
1.04
0.94
0.95
0.90
1.00
1.02
1.04
0.98
0.98
1.00
1.02
04
Relative test loss
Relative test loss
Relative test loss
Figure S2: Perturbing network weights leads similar loss relative original model. To evaluate whether training the FCN from the perturbed initialization led to trajectories that lay within the same hasin of attraction the unperturbed initialization, compute the relative test loss (lest loss for perturbed initialization divided by test loss for unperturbed initialization.. For all FCN widths, loss is peaked at 1, demonstrating the perturbed FCNs converge to a similar test loss as unperturbed initialization. This provides evidence that the trajectories lie in the same local minimum basin.
see that the distribution of relative test
To compute the KMD, considered observables the weights from the hidden layer to the output layer. This choice made because: the values of these weights determine the weight evolution of the earlier layers (due backprop);
there are fewer weights from the hidden layer to the output layer, than there are from the input layer to the hidden layer, enabling our approach to be more computationally tractable. To enrich our observables, we considered time-delays 44 45 51.. Because we considered FCNs of differing width, using the same number of time-delays leads to matrices of different dimension. Therefore, fixed the ratio of d (number of time-delays) to (the number of units the hidden layer), setting 64 FCNs with h 5, 32 for FCNs with 10, and 8 for FCNs with h
We applied DMD-RRR 5 on these time-delayed observables. When principal component analysis (PCA) was performed on these observables, it was found that lower-dimensional subspace contained a large percentage of the variance. This is in-line with previous work finding the weights of DNNs traverse low dimensional subspaces Bl Therefore, considered only the top 10 Koopman modes, using reduced singular value decomposition (SVD).
C.3 GeLU FCNs
understand how robust our conclusion that shallow wide FCNs have non-conjugate training dynamics is, we perform the same Koopman-bras based analysis on FCNs with GeLU activation functions (59). Given the similarity between ReLU and GeLU, we expect that they should have similar training dynamics behavior. Computing the Koopman eigenvalues and comparing across FCNs of varying widths, find very similar results (compare Fig. [S3 with Fig. (3). This demonstrates that Koopman based framework is robust (minor) hyper-parameter differences. A full examination of other choices activation function (especially those that introduce "squashing" ee.g. sigmoid) impact the training dynamics would be an interesting future direction to pursue.
C.4 Conjugate training dynamics across random initializations of FCNs
Motivated by recent work arguing that different random initializations of DNNs converge solutions with no loss barrier the linear interpolation between them [66.. when taking into account the inherent permutation symmetry the parameters, we asked whether our Koopman- dynamics for different initialized FCNs. We examine how similar the Koopman spectra for all 25 independently trained FCNs of width h = 40 are. We chose this because Entezari et al. (2022) [16) proved their result for sufficiently wide FCNs. but sufficiently narrow FCNs exhibited loss barrier along the linear interpolation, when
framework identified conjugate training
18

---

A
x10
0.04
imag(A) imag(?))
0.04
Wasserstein distance
100-
N.S.
0.00
0.02
imag(A)
0.04
0.00
0.04
5
0.00
3.04
40
0.96
0.88
1.00
1.02
-0.04-
0.0
h 5-
h- 10 20-h-m h
real(-)
real(7.)
h=5-h=40
Figure S3: Narrow and wide fully connected neural networks have non-conjugate training dynamics when using GeLU activation functions. Same as Fig. 3 in the main text, but for FCNs using GeLU activation functions. Error bars are + standard deviation across 10 independently trained FCNs.
0.04
50-
40-
Count
Seed 4
10
0
x10*
real(?)
Wasserstein distance
Figure S4: Conjugate training dynamics across random initializations FCNs. (A) Example Koopman spectra associated with training FCNs, = 40, from two different random seeds. (B) Distribution of Wasserstein distance between the Koopman spectra of all possible pairs of 25 independently trained FCNs.
permutation symmetry was not taken into account [66. Our h = 40 FCNs strike a balance between these two points, making it good point for analysis.
We find examples of different random initializations with nearly identical Koopman spectra (Fig. S4 A). Across all possible pairs the 25 independently trained FCNs, we find that the Wasserstein distance between Koopman spectra centered around values similar to what was seen when comparing h = = 40 FCNs. These results support the hypothesis that, when taking into account permutation symmetry (which computation of the Koopman eigenvalues naturally does), at least some random initializations have conjugate training dynamics.
D Convolutional neural network training phases
D.1 Training experiment details
LeNet [60) and ResNet-20 [6]] models were trained on MNIST and CIFAR-10 (respectively), for 20 epochs. Training was performed using the Serch framework 74, which makes of PyTorch. The open source code can be found here: https://giis
used training are presented Table [S] They were chosen heweeebennch/teeeeeeeee cree/master. All hyper-parameter: same between two architectures to make the comparison between the two fair and were selected to match the hyper-parameten that were previously used to study LeNet's training dynamics [II]. Three independent seeds were trained (n., 3). each of which MacBook with an Apple MI chip, CPU, and GPUs. Code implementing our experiments can found at https::/// //ghthub//btt// com/si /w/willl willis fillliammmmm wawaam-edaaa/aaaaaaaaaaaa Identi yingi hariinnnngiilleeeeneeeee Twwiirrrriinnnnaaaaaaaain
com/ htt./////hinbbeeeeeeeeeee
initialized from 10 perturbed initializations (n,, 10). All experiments
D.2 Computing Koopman mode decomposition
As with the FCNs (Sec. the observables used to construct the Koopman operator were time-delays of weights going from the last hidden layer to the output layer. Eight time-delass, 8, were used and the top 10 Koopman modes were considered, based on a reduced SVD. The appropriateness of this was again verified by examining the amount of variance captured by the first 10 principal components.
19

---

Hyper- Sypee--prameters
Value
Learning rate (7)
0.0012
Batch size (b)
60
Optimizer
Adam
Epochs
20
Activation function
ReLU
Table S2: Hyper-parameters used for CNN training in Sec.
LeNet (MNIST)
B
ResNet-20 (CIFAR-10)
-30
6980
Training eprochs
Booggii weetenenoeoooooo
Training spochs
woweeeeeoooo destance
4.0
4.0
4.5
20
Training epochs
Training opochs
Figure S5: Koopman-based framework enables identification of transitions in dynamics across the training of LeNet and ResNet-20. (A)-BB Same Fig. 4 but for dynamics computed over individual epochs.
D.3 Evaluating CNN training phases at a coarser timescale
addition studying earliest part CNN training, which was focus of previous work [27]. we examined how CNN training dynamics evolved across the first 20 epochs of training. To reduce the computational cost associated with computing and comparing Koopman spectra, we computed the Koopman mode decomposition across all training iterations within single epoch.
At this coarser timescale, we find that LeNet, trained MNIST, and ResNet-20, trained CIFAR-10 10, exhibit very similar evolutions training dynamics (Fig. (Fig.S5).. Indeed, both see continual reduction Wasserstein distance between neighboring epochs as training time goes on. Interestingly, in both cases, there is growing block diagonal of dark blue, that becomes especially strong at training epoch 13. This supports the general similarity in the evolution of training dynamics between different CNN architectures, was previously observed 27
E
Transformer grokking
E.1
Training experiment details
Single hidden layer Transformers. with four attention heads, were trained on modular arithmetic us- ing open source code [62]: https://github com/KindXiaoming wowno//nn//mmigg//rr////////////
As noted in the repository, this is modified version of previously developed code 29 We keep all hyper- parameters the same. Therefore, refer the interested reader the details presented in the repos- itory and related papers.
independent seeds were trained (n,, = 20), each of which was ini- tialized from 10 perturbed initializations (n, = 10). All experiments were run on a MacBook Air with an Apple MI chip, 1 CPU, and no GPUs. Code implementing our experiments can be found at com/willl wwwlliillmmmmdmmaae aawiaaamannanaaeeeee-eeeee-aaeae Saaiaaa//ddnntfiing.eeee.. Eeuivaleet_TTaiineee
TraininggDynamicca Dynamics.
E.2 Computing Koopman mode decomposition
As with the FCNs (Appendix 2 and CNNs (Appendix D.2). we use as observables time-delays of the weights from the last hidden layer to the output. d = 32 time-delays were used. Because the number of weights was 65000, before performing the time-delays and flattening the tensor that stored all the weights, we performed dimensionality reduction. This was achieved by applying PCA the flattened tensor that contained all weights (without tine-delays) and then projecting the weight trajectories corresponding to the training of each perturbed initialization onto the top 10 principal components. These dimensionally reduced weights were then time-de deleyee and used to construct the KMD. The top 10 Koopman modes considered, based reduced SVD. The
20

---

B
0.10
Teak: Unson.
Unconstrained
Test: Uroon.
0.05
Accuracy
Teat: Con.
imag(2)
0.
0.14
-0.05-
-0.10
50 Training iterations
0.96
1.04
real().)
Figure S6: Distinct Koopman eigenvalues between Transformers that do and do not undergo grokking consistent when considering windows of training time that have more similar training loss. (A) As Fig. [C. we train Transformers with unconstrained weight norm for I 100 (thin vertical black line) iterations and use the associated weight trajectories to approximate the KMD. In contrast, here we train Transformers with constrained weight norm for I' 200 iterations (thin vertical red line) and use the associated weight trajectories approximate the KMD. This enables the two networks reach more comparable (although not perfectly matching) training losses. Lines are mean across 20 indepedently trained Transformers and shaded area + one standard deviation. (B) Same as Fig. 5P, but with the Koopman cigenvalues associated with training the constrained Transformer for T 200 iterations.
appropriateness of this was again verified by examining the amount of variance captured by the first 10 principal components.
We found in Sec. 4.4 that the Transformers trained with constrained weight norm (that do not undergo grokking) have non-conjugate dynamics with the Transformers trained with unconstrained weight (that do undergo grokking). To ensure that this was not due to for computing the Koopman eigenvalues (the first 7 100 training iterations) was distinct hetween two Transformers, we performed the following control experiment. Namely, we performed the same analysis, hut we considered the weight trajectories of Transformers with constrained weight norm over twice long training time interval (I' 200). In this the unconstrained Transformer (Fig, S compare solid red black lines), although there more variability between independent seeds (Fig. S6A shaded red area). However. this case we again find that the Koopman cigenvalues associated with unconstrained and constrained training are distinct (Fig. B) In particular, the constrained Transformer again has pair of complex conjugate Koopman eigenvalues that lie along the unit cirele, while the unconstrained Transformer has pair of complex conjugate Koopman eigenvalues outside of the unit circle. This suggests distinct stability properties, and further emphasizes the absence conjugacy.
used training loss over window training
the constrained Transformer reaches training loss that is closer that of
21

---

NeurIPS Paper Checklist
1. Claims
Question: the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims in the abstract and introduction are grounded in experimental results discussed in detail in Sec. F Hnnd Figs. 25
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The ahstract 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 he expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of the developed framework at discussed in the Limitations subsection in Sec.
Guidelines:
The answer NA means that paper has no limitation while the answer means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust 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 the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-ter system might not be used reliably to provide closed captions for online lectures because 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, worse outcome might that reviewers discover limitations that aren't I 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.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: No new theoretical results were achieved. References are made to relevant work that proved theorems that are used.
Guidelines:
The answer NA means that the paper does not include theoretical results.
22

---

All the theorems, formulas, and proofs in the paper should be numbered and cross cross-reference..
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can cither appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Details all the numerical experiments, DNN training, and computation of KMD are reported the corresponding sections of the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments,
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (c.g., in the case of a large language model), releasing model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility. which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g.. a large language model), then there should either be a way to access this model for reproducing the results a way to reproduce the model (e.g.. with open- dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed source models, it may be that access to the model is limited in some way (e.g.. registered users), hut it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: The Githuh repository where our code is available referenced in the Appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurTPS code and data submission guidelines hatzpe://nips.cc/pinl for more details.
guides/Cool htte://ooeeuaaaiiiooooooooo
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).
23

---

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines htttps://nips.cc/ppue for more details.
Suadde//DodeSubmisssooooo
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 new proposed method and baselines. If only a subset experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. 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: In Appendix A.3 the parameters used for numerically evaluating online mirror descent and online gradient are reported. In Tables (SblessST/(S2 the hyper-parameters used training FCNs and CNNs are reported. In Appendix[E.] [E.]] the code used for training is linked. As no changes were made to the code, this serves as a direct way to identify all training and test details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that necessary appreciate the results and make sense them.
The full details can be provided either with the code, appendix. or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [Yes]
Justification; All error bars figures are described and statistical test performed Fig. reported.
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, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors 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 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 the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error har than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
crror bars are reported tables or plots, The authors should explain 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 execution) needed reproduce the experiments?
Answer: [Yes]
24

---

The computing resource used for all experiments is reported in the corresponding Appendix subsections.
Justification:
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 well 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 make into the paper).
9. Code Of Ethics
Question: Does research conducted the paper conform, every respect, with NeurIPS Code of Fthics httpp:///eeuiippp.ccpppld///iiii
Answer: [Yes]
Justification: Potential harm and negative societal impact were mitigated.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require 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).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: There is no significant negative societal impact of this work. The positive socictal impact
of this work was discussed in Secs. [
Guidelines:
The answer NA means that there is societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has societal impact or why the paper does address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (c.g.. disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deploy- ment technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimatc point out that improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, it not needed point out a generic algorithm optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, scraped datascts)?
Question:
Answer: [NA]
Justification: This paper poses such risks.
25

---

Guidelines:
answer means the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary
safeguards allow for controlled use of model, for example by requiring users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We reongnize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
Open source code that was used for training and developing upon is cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
authors should state which version
asset is used and, possible, include a URL.
The name of the license (e.g.. c. CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g.. website), the copyright and terms of service of that source should provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasetssetssllle/s has curated licenses for some datasets. Their licensing guide can help determine the license a dataset.
existing datasets that are re-packaged, both original license and license of derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged reach out to asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The Github repository where general code for identifying conjugate training dynamics is available is referenced in the Appendix.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the dctails of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations,
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No crowdsourcing nor research with human subjects was performed.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
26

---

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: No human or animal work was performed in this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval. you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
initial submissions, ble), such as institution conducting the review.
include any information that would break anonymity applica-
27


Reviews:
[{"review_id": "CuRiT9il5J", "reviewer": "Reviewer_C2X3", "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:\n\n- **Shallow vs. Wide Fully Connected Networks**: The framework reveals non-conjugate training dynamics between these types of networks.\n- **CNNs**: The framework helps characterize the early training phase dynamics in CNNs.\n- **Transformers**: The framework identifies non-conjugate training dynamics in Transformer models, including instances of \"grokking\".\n\nOverall, the results demonstrate the framework's versatility and potential to illuminate the complex training dynamics of different DNN architectures.", "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?\n\nII) 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?", "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.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "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.\n\n- 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.\n\n- They investigated the training dynamics across various DNN architectures.", "weaknesses": "There are only a few concerns:\n\n1- 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?\n\n2- 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.\n\n3- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**\"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?\"**\n\nThis 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!). \n\n**\"I am interested in knowing how choosing another activation function (any activation function except ReLU) can affect the results?\"**\n\nThis 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. \n\n**\"Why was the Wasserstein distance used to quantify differences between Koopman spectra?\"**\n\nWe 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.\n\n**\"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.\"** \n\nThis is a good question that we hope future work will investigate in rigorous detail. However, we can provide some insight on this question.\n\nThere 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. \n\nThe 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. \n\nWe will add this discussion to the main text of our revised manuscript (Sec. 5). \n\n**\"It would be helpful to understand how their approach could be applied to RNNs for time series forecasting.\"**\n\nWe 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.\n\n**References:**\n\nAvila and Mezic (2024) “Spectral Properties of Pullback Operators on Vector Bundles of a Dynamical System”\n\nArbabi and Mezic (2017) “Ergodic theory, dynamic mode decomposition, and computation of spectral properties of the Koopman operator”\n\nBrunton et al., (2017) “Chaos as an intermittently forced linear system”", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uHViDrRJMv", "reviewer": "Reviewer_42wo", "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).", "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?\n\nWhat 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)?\n\nWould sampling Bayesian neural networks allow for arbitrarily many weights from the same training dynamic?\n\nHow does this method allow grokking to be identified early?", "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.\n\nThe 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "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.\n\nThe 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.\n\nIt 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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**\"The main issue is what would constitute a \"low\" vs. \"high\" Wasserstein distance between the Koopman eigenvalues?\"**\n\nWe 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. \n\nFirst, 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. \n\nSecond, 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.  \n\n**\"It is hard to conceptually understand what similar versus distinct training dynamics look like.\"**\n\nWe 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). \n\n**\"The interpretations are all confirming previously known results, so it would be helpful to present a new result or refine a previously known result\"**\n\nWe 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). \n\n**\"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?\"**\n\nThis is a great idea! We will incorporate this into the revised version of the manuscript. We appreciate the reviewer suggesting this. \n\n**\"How does this method allow grokking to be identified early?\"**\n\nThe 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.\n\n**References:**\n\nFrankle et al. (2020) \"The Early Phase of Neural Network Training\"", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qhqRGT4Scl", "reviewer": "Reviewer_ZFpv", "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.", "questions": "See weaknesses.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written, and well-structured.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**\"There remains a concern about the computational efficiency of computing Koopman eigenvalues in the proposed framework, especially when the parameter space is high-dimensional.\"**\n\nWe 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. \n\n**\"The framework may be not suitable to compare the dynamics of different dimensions\"**\n\nThis 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. \n\n**References:**\n\nBrokman et al. (2023) \"ENHANCING NEURAL TRAINING VIA A CORRELATED DYNAMICS MODEL\" \n\nDrmac (2024) \"A LAPACK Implementation of the Dynamic Mode Decomposition\"", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9e5HVNLBta", "reviewer": "Reviewer_4cZ9", "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:\n\n- 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,\n- 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.\n- 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.\n- 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.", "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.\n\nLiu Ziyin, Symmetry induces structure and constraint of learning (2023)", "limitations": "The authors have discussed the limitations in Section 5.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- S1. The paper is very well written and easy to follow.\n- 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.\n- 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.\n- 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).\n- 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).\n\nBielecki, Topological conjugacy of discrete time-map and Euler discrete dynamical systems generated by a gradient flow on a two-dimensional compact manifold (2002)\n\nEntezari et al. The role of permutation invariance in linear mode connectivity of neural networks (2021)\n\nBokman 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.\n- 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).\n- 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.)\n- 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.\n- 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).\n\nMinor comments and typos\n- In the caption of Figure 6, should \"... that are conjugate\" be \"... that are non-conjugate\"?\n\nEntezari et al. The role of permutation invariance in linear mode connectivity of neural networks (2021)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**\"It would be impactful if the proposed method identifies conjugacy of seemingly non-equivalent training dynamics\"**\n\nThis 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.\n\n**\"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\"** \n\nThis 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.\n\n**\"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\"**\n\nThis 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.\n\n**\"There are some standard assumptions used in Koopman operator theory that may not strictly hold for mini-batch stochastic gradient descent\"**\n\nThis 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. \n\nWe 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)]. \n\n**\"The notations used in Equation (5) may be not exact\"** \n\nYes! Thank you for that correction. We will correct this in the revised version of the manuscript. \n\n**\"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\"**\n\nThis 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.\n\n**\"In the caption of Figure 6, should \"... that are conjugate\" be \"... that are non-conjugate\"? \"**\n\nOops yes! Thank you for pointing that out. We have corrected that typo.\n\n**References:**\n\nArbabi and Mezic (2017) “Ergodic theory, dynamic mode decomposition, and computation of spectral properties of the Koopman operator”\n\nBrunton et al., (2017)  “Chaos as an intermittently forced linear system”\n\nEntezari et al. (2021) \"The role of permutation invariance in linear mode connectivity of neural networks\"", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time and thoughtful comments. A response to each individual reviewer’s comments is provided in the thread of the associated review. We believe that addressing these questions has greatly improved the quality of our work. Here, we highlight three new results that we obtained (see attached global response pdf): 

**Robustness to different activation functions:** Motivated by reviewer C2X3, we trained fully connected networks of various widths on MNIST, using the GeLU activation function, instead of the ReLU activation function. We find very similar results when computing and comparing the Koopman eigenvalues (Fig. 1, global response pdf), illustrating that our framework is robust to changes in the activation function. 

**Similar dynamics for FCNs with different initializations:** Motivated by reviewer 4cZ9, we examined whether our widest FCNs had similar Koopman eigenvalues, across different initializations and SGD seeds. This directly tests a hypothesis by Entezari et al. (2021) on different initializations having equivalent training dynamics, when taking into account the permutation symmetry (which the Koopman eigenvalues are invariant to). We find evidence supporting this hypothesis (Fig. 2, global response pdf), illustrating another example of when and how our framework can be used to address phenomena that have been challenging to address with existing methods. 

**Distinct Koopman eigenvalues for non-conjugate optimizer:** Motivated by reviewer 4cZ9, we compared the Koopman spectra of online mirror and online gradient descent to a non-conjugate optimizer, the bisection method (BM). We find that OMD/OGD have significantly different Koopman eigenvalues than those associated with BM. In particular, because BM takes large, global steps that can lead to the changing of signs of the parameters being optimized, the Koopman eigenvalues are complex. This is in contrast with OMD/OGD, which takes local steps and has positive, real-only Koopman eigenvalues. This illustrates the significance of similarity in eigenvalues associated with OMD and OGD, and provides additional insight into what makes two training trajectories non-conjugate.

**References:**

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

Meta-review: The paper proposes a new method for identifying equivalent training dynamics of neural networks using Koopman operator theory and the notion of topological conjugacy. The paper received positive feedback overall: reviewers recognized the proposed analyzes of the training dynamics as novel, addressing important questions about neural network similarity and offering a general method that accounts for various symmetries; the potential impact of the method in understanding and improving neural network training dynamics was also acknowledged. The reviewers appreciated the clear presentation and found the method to be natural and applicable to various training dynamics. As per reviewers’ requests, the authors improved the paper by adding additional examples and experiments, and clarifying some of the assumptions.
Decision: Accept (spotlight)