Identifying Equivalent Training Dynamics

TL;DR

This paper introduces a novel framework using Koopman operator theory to identify when deep neural network training dynamics are equivalent, enabling better understanding of different training regimes and architectures.

Contribution

We develop a method leveraging Koopman eigenvalues to determine dynamical equivalence of neural network training processes, addressing a key challenge in analyzing training dynamics.

Findings

Successfully identified equivalence between online mirror descent and online gradient descent.

Detected non-conjugate dynamics between shallow and wide neural networks.

Characterized early training phases and grokking phenomena in CNNs and Transformers.

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.