Shedding Light on the Blackbox: Understanding Deep Neural Network Training Dynamics
Authored by William Redman
Deep neural networks power everything from language models to image recognition, yet the process by which they learn remains one of modern science’s great mysteries. By viewing neural network training as a dynamical system, researchers are uncovering patterns and principles that were once hidden in complexity. Leveraging Koopman operator theory, AIMdyn scientists are shedding light on how networks evolve, learn, and even mirror each other’s behavior — offering a path toward more interpretable, efficient, and robust AI.
What is Deep Neural Network Training?
At the heart of the artificial intelligence (AI) revolution that the world has experienced over the past 10 years is the training of deep neural networks (DNN). This training is a process whereby a naive DNN has its parameters (e.g., “weights”) iteratively updated so that the DNN becomes highly performant on a specific task. The optimal way to train DNNs has been intensively investigated and there exist dozens of approaches that are standardly used, each of which strike different balances between efficiency in training and resulting accuracy. As DNNs have become increasingly large, going from models of ten of thousands of parameters to tens of billions of parameters, training approaches have become increasingly complex and multi-faceted. This growing complexity of DNNs and their training has made the need for a deeper understanding of how training shapes learning in DNNs an especially important question.
The Challenges We Face
On the face of it, understanding how DNN training shapes learning should be straightforward. DNN practitioners have access to every parameter of the model, every input and output of the model, and the exact equations governing the iterative optimization of the model parameters. However, the high-dimensional nature of DNNs and complex statistics present in the data used to train DNNs (e.g., images, text) makes analytic treatment of DNN learning incredibly difficult. Indeed, except in the simplest cases – where many simplifying assumptions are made – there exists little precise understanding of DNN training behavior. This lack of understanding has emphasized the importance of developing new approaches for studying DNN training.
A Dynamic Perspective
One view of the iterative optimization associated with DNN training is that it corresponds to a dynamical system, evolving along a trajectory in the high-dimensional space of parameters. This dynamical perspective has led researchers to identify several interesting phenomena. For instance, examining the behavior of the dynamics associated with training different DNNs led to the observation that training quickly becomes confined to a much lower-dimensional subspace [Guy-Ari 2018] [Brokman et al., 2024]. This suggests that there are a few key “directions” along which learning occurs. Additionally, the training different DNNs leads to trajectories that have been found to have similarities between them [Mao et al., 2024], despite the fact that the models themselves differ significantly in size and complexity. This suggests that learning occurs similarly on the same task, for different types of DNN models. While these studies demonstrate the potential insight that can be shed on DNN training by taking the dynamical perspective, traditional tools developed for studying dynamical systems are challenging to use in the context of DNNs and previous work has relied on coarse-level descriptions of the dynamics.
Advances in data-driven dynamical systems theory over the past two decades have enabled characterization of dynamics from complex systems at a level that has not been possible previously. In particular, Koopman operator theory [Mezic, 2005, Budisic et al., 2012] provides a framework for learning linear representations of non-linear dynamical systems from data. This has suggested that it might be possible to capture finer-scale dynamical information about DNN training by leveraging Koopman operator theory. The first evidence of this came from work that learned linear models that could mimic the dynamical behavior of standard DNN training [Dogra and Redman, 2020] [Tano et al., 2020]. Subsequent work has utilized Koopman representations to provide the first rigorous comparison of DNN training dynamics, identifying when the training of two different DNNs is “equivalent” and when it is not [Redman et al., 2024]. This offers the potential for constructing so-called equivalence classes, which can lead to better understanding of how different aspects of DNN models affect their training.
The success of Koopman operator theory in studying DNN training has led to a number of other applications. For instance, Koopman representations were leveraged to achieve better meta-learning (an approach for learning how to learn) [Simanek et al., 2022] [Hataya and Kawhara, 2024] and the dynamical information captured by the Koopman representations have been used to identify which DNN parameters can be removed (i.e., “pruned”) [Mohr et al., 2021] [Redman et al., 2022]. In addition, Koopman operator theory has been used to study the behavior of recurrent neural networks (RNNs), providing interpretable information about what they have learned [Naiman and Azencot, 2023] [Ostrow et al., 2023].
Understanding and Accelerating DNN learning
The ability to capture detailed representations of DNN training dynamics via Koopman operator theory opens up a wealth of exciting new possibilities for understanding and accelerating DNN learning. In particular, by advancing and expanding the developed methods to the ever more complex DNNs underlying the AI that has become ubiquitous in society, greater robustness of models and efficiency in training may be possible.