Identifying Equivalent Training Dynamics
This addresses a methodological gap for researchers studying neural network training dynamics, though it appears incremental as it applies existing dynamical systems concepts to a new context.
The paper tackled the problem of identifying when deep neural network models have equivalent training dynamics, which limits insights from prior work, by developing a framework using Koopman operator theory to detect conjugate and non-conjugate dynamics, and validated it by correctly identifying known equivalences and uncovering non-conjugate dynamics in various architectures like shallow vs. wide networks, CNNs, and Transformers.
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.