Collective variables of neural networks: empirical time evolution and scaling laws

## Abstract

Abstract This work presents a novel means for understanding learning dynamics and scaling relations in neural networks.
We show that certain measures on the spectrum of the empirical neural tangent kernel, specifically entropy and trace, yield insight into the representations learned by a neural network and how these can be improved through architecture scaling.
These results are demonstrated first on test cases before being shown on more complex networks, including transformers, auto-encoders, graph neural networks, and reinforcement learning studies.
In testing on a wide range of architectures, we highlight the universal nature of training dynamics and further discuss how it can be used to understand the mechanisms behind learning in neural networks.
We identify two such dominant mechanisms present throughout machine learning training.
The first, information compression, is seen through a reduction in the entropy of the NTK spectrum during training, and occurs predominantly in small neural networks.
The second, coined structure formation, is seen through an increasing entropy and thus, the creation of structure in the neural network representations beyond the prior established by the network at initialization.
Due to the ubiquity of the latter in deep neural network architectures and its flexibility in the creation of feature-rich representations, we argue that this form of evolution of the network’s entropy be considered the onset of a deep learning regime.

## 1 Introduction

Scaling behaviour in neural networks has become a pivotal area of investigation in modern deep-learning research.
Traditional scaling laws, which correlate the number of network parameters with performance metrics, provide a foundational understanding but often fall short when comparing different architectures or capturing the intricacies of neural network dynamics arising from the architecture of the network itself.
To bridge this gap, our research leverages the neural tangent kernel (NTK) to elucidate neural networks’ learning dynamics and scaling laws, offering a more comprehensive framework for understanding how network representations evolve and how architecture scaling impacts these processes.
Specifically, this investigation utilizes a set of collective variables, the entropy and trace of the empirical NTK, to identify and explain learning regimes during the training process.
The NTK, introduced by Jacot et al. (), provides a powerful tool for analyzing the infinite-width limits of neural networks.
By focusing on the empirical NTK and its evolution during training, we can gain insights into the internal representations learned by finite-size neural networks and their dependence on architecture.
Our approach is validated across various tasks and network types, including transformers, auto-encoders, graph neural networks, and reinforcement learning models, demonstrating its broad applicability and robustness.

### 1.1 State of the Art

Previous work has established various scaling laws for neural networks, often relating to the number of parameters and the expected performance (Bahri et al., ).
However, these descriptions are predominantly phenomenological and do not fully capture the dynamics of different architectures or the nuances of the learning process.
In Bahri et al. (), neural scaling laws are studied in a so-called resolution-limited and variance-limited regime, highlighting changes in the network performance within each.
While this work highlights a useful relationship between the data manifold and neural scaling, it does not address the fundamental question of whether learning mechanisms within each regime change at all.
Leclerc and Madry () identify regimes during training and relate them to a learning-rate dependent mechanism, namely, large-step and small-step regimes.
They highlight that the process by which the network learns information differs between these two regimes and can be leveraged more effectively when considered as two separate training phases.
In their 2021 paper, Geiger et al. () explored the role of initialization in over-parameterized neural networks and how it changes the training process of the neural network.
The problem has also been explored with more practical uses in finding, for example, in Mirzadeh et al. () learning regimes, as impacted by batch size, learning rate, and regularization, is explored in relation to catastrophic forgetting, or the seminal work of Kaplan et al. () where scaling laws are introduced as a means to optimize a cost-performance tradeoff in large language models.

The NTK has emerged as a key concept for understanding these dynamics, particularly in the context of wide neural networks where it facilitates a tractable analysis of gradient descent dynamics.
Often, this work studies alignment effects in the NTK and what role these play in the learning taking place (Atanasov et al., ), or in identifying dynamics regimes (Lewkowycz et al., ).
Indeed, a landmark result utilizing theoretical scaling laws studies came in the form of $\mu$-transfer (Yang and Hu, ; Yang et al., ), that can be used to initialize width-scaled networks without losing accuracy.

Each of these studies highlights the importance of scaling laws and network dynamics in the field’s current state.
However, in each case, they do so by studying raw loss metrics as a function of changing network size and identifying steady changes in their performance, which are characterized by a scaling law.
In this work, we explore an alternative means of representing training dynamics, namely, the collective variables of the NTK.
In doing so, we isolate fundamental learning mechanisms such as information compression and structure formation and highlight their role and evolution during training.

### 1.2 Contribution Statement

We identify several relevant contributions in this work:

- •
We provide a detailed empirical analysis of NTK evolution across various neural network architectures and tasks, highlighting universal patterns and architecture-specific behaviors.
- •
We introduce collective variables derived from the NTK, such as entropy and trace, to quantify the diversity and effective learning rates of neural network updates.
- •
We demonstrate the applicability of our methods to modern machine learning models, including transformers, graph neural networks, and reinforcement learning agents, showcasing the utility of NTK-based analysis for understanding complex learning dynamics.
- •
We find that large models universally show an increase in entropy, in contrast to small models. This allows us to quantitatively define a deep learning regime.
- •
We relate existing scaling parameters of width and depth with our entropy and trace collective variables.

By extending the analysis of NTK to encompass a wide range of network architectures and tasks, we offer new tools and perspectives for scaling and comparing neural networks at an information-theoretic level.
This work enhances our theoretical understanding of neural network dynamics and provides practical insights for optimizing architecture design and training protocols.

## 2 Theory and Methods

### 2.1 Neural Tangent Kernel for Network Dynamics

The neural tangent kernel was first introduced to the machine-learning community by Jacot et al. () to investigate the infinite-width limits of neural networks.
In this work, we explore the evolution of neural networks and how this can be better understood by studying measures on this matrix. However, to understand why this is possible, consider a neural network, $f[x_{i},\bm{\theta}]$ acting on a single input data point, $x_{i}$ and parametrized by $\bm{\theta}$.
During training, each element of parameter, $\theta_{k}$, will be updated via some variant of stochastic gradient descent

$$ $\centering\theta_{k,t+1}=\theta_{k,t}-\eta\sum\limits_{i}\frac{\partial \mathcal{L}(f[x_{i},\bm{\theta}_{t}],y_{i})}{\partial\theta_{k}},\@add@centering$ (1) $$

where $\eta$ is the learning rate, $\mathcal{L}$ is the chosen loss function used in the minimization, and $y_{i}$ is the target value associated with the input data, $x_{i}$.
When performing gradient descent, the size of the step forward is determined by the learning rate, which is set under the assumption that the loss surface will not change drastically over the course of the update.
However, to perform a more rigorous analysis, we take this further and move into continuous time by defining

$$ $\centering\dot{\theta}_{k}=\lim_{\eta\rightarrow 0}\frac{\theta_{k,t+1}-\theta _{k,t}}{\eta}=-\sum\limits_{i}\frac{\partial\mathcal{L}(f[x_{i},\bm{\theta}_{t }],y_{i})}{\partial\theta_{k}},\@add@centering$ (2) $$

where the dot denotes a time derivative.
The problem with this formalism is that the network architecture itself is not expressed in the evolution, only the parameters of this architecture.
To understand how the neural network representations evolve, we compute the time derivative of the function itself, $\dot{f}[x_{i},\bm{\theta}]$ by the chain rule

$$ $\centering\dot{f}[x_{p},\bm{\theta}]=\sum\limits_{k}\frac{\partial f[x_{p},\bm {\theta}]}{\partial\theta_{k}}\dot{\theta}_{k}.\@add@centering$ (3) $$

Substituting Equation 2 into Equation 3 and expanding the loss derivative into a parameter and function component, we find

$$ $\centering\dot{f}[x_{p},\bm{\theta}]=-\sum\limits_{k}\frac{\partial f[x_{p}, \bm{\theta}]}{\partial\theta_{k}}\sum\limits_{i}\frac{\partial f[x_{i},\bm{ \theta}_{t}]}{\partial\theta_{k}}\cdot\frac{\partial\mathcal{L}(f[x_{i},\bm{ \theta}_{t}],y_{i})}{\partial f[x_{i},\bm{\theta}_{t}]}.\@add@centering$ (4) $$

Letting $\frac{\partial f[x_{p},\bm{\theta}_{t}]}{\partial\theta_{k}}\cdot\frac{
\partial f[x_{i},\bm{\theta}]}{\partial\theta_{k}}=\Theta_{pi}$, and $\frac{\partial\mathcal{L}(f[x_{i},\bm{\theta}_{t}],y_{i})}{\partial f[x_{i},
\bm{\theta}_{t}]}=\partial_{f_{i}}\mathcal{L}_{i}$, we find

$$ $\centering\dot{f}[x_{p},\bm{\theta}]=\sum\limits_{i}\Theta_{pi}\partial_{f_{i} }\mathcal{L}_{i},\@add@centering$ (5) $$

where the matrix $\Theta_{pi}$ is referred to as the neural tangent kernel or NTK.
In the evolution of the neural network representations, the loss derivative term plays the role of the minimization constraint, guiding the network parameters into an optimum.
The NTK, however, also plays a role in how the network evolves during training, specifically, how the architecture impacts this evolution.
Unfortunately, the NTK is often very large, and therefore, alternative tools are required to study it and its evolution.

### 2.2 Collective Variables

The NTK matrix is built from the inner products of the gradient vectors of a neural network evaluated at different data points.
It provides information on the degree of correlation in these gradient vectors, i.e., whether the representations of these points will evolve similarly or not.
To see the role this plays in the evolution, consider the eigendecomposition of the NTK

$$ $\centering\Theta_{pi}=\sum\limits_{n}D_{pn}\Lambda_{nn}\left(D^{-1}\right)_{ni },\@add@centering$ (6) $$

where $D_{pn}$ is the matrix of eigenvectors and $\Lambda_{ij}$ is the diagonal matrix of eigenvalues.
Substituting Equation 6 into Equation 5, yields (see also Krippendorf and Spannowsky ())

$$ $\centering\dot{f}[x_{p},\bm{\theta}]=\sum\limits_{i}\sum\limits_{n}D_{pn} \Lambda_{nn}\left(D^{-1}\right)_{ni}\partial_{f_{i}}\mathcal{L}~{}.\@add@centering$ (7) $$

Evaluating Equation 7 from right to left, we can see that during an update of the network, the loss derivative of the $i^{\text{th}}$ data-point is projected along the eigenvectors of the NTK and scaled by its eigenvalues, thus introducing effective learning rates and direction changes due solely to the architecture and its instantaneous prior over the data.
This description can be simplified into

$$ $\centering\dot{f}[x_{p},\bm{\theta}]=\sum\limits_{n}D_{pn}\tilde{p}_{n},\@add@centering$ (8) $$

where $\tilde{p}_{n}=\lambda_{n}D^{-1}_{ni}\partial_{f_{i}}\mathcal{L}$, referred to here as the $\lambda$-scaled projection factor, acts as a scaling factor on the update to the representation of data point $x_{p}$.
Regarding learning, this formalism argues that a neural network starts with a prior of the data, indicated by the spectrum of the NTK matrix, which correlates and scales modes based on its architecture.
During an update, the fundamental modes present in the neural tangent kernel are tuned based on the loss function.
This can occur either via scaling from the magnitude of the loss gradients, strengthening or weakening the degree of correlation between modes.
Alternatively, reorientation via the dot product of the gradient vector can force a change in the direction of the eigenvectors.
Thus, two data points that the network considered correlated before training, highlighted by high similarity in the NTK matrix, can either be separated via an orthogonal gradient vector or strengthened via an aligned gradient with a large eigenvalue.
The question remains: how can we discuss this mechanism taking place inside the neural network?
In their 2023 paper, Tovey et al. () introduced measures on the NTK to describe its current state and used this description to describe the role of data selection in neural network training.
By measuring the collective variables of their networks before training, they showed that it could provide insight into the quality of the prior data distribution and, thus, the expected performance of the training.
Specifically, they introduced the entropy of the NTK, computed from the normalized eigenvalues as

$$ $\centering S=-\sum\limits_{i}\lambda_{i}\log\lambda_{i}\@add@centering$ (9) $$

providing a measure of the diversity of the matrix, i.e., how many independent degrees of freedom exist in the update, and the trace of the NTK,

$$ $\centering\sum\limits_{i}\Theta_{ii},\@add@centering$ (10) $$

describing the magnitude of the scaling factor on the learning rate.
This work extends their analysis to the evolution of the neural network during training to understand what processes it undergoes.

### 2.3 Datasets and Architectures

Several datasets have been studied to demonstrate the use of the collective variables for model interpretation and to better highlight the dynamics’ universality.
Table 1 outlines each dataset, including their citations, the type of problem being learned, and the network architecture used in the training.

**Table 1: Data-sets along with the model architecture, task, and citation of the different experiments used throughout the study.**
| Name | Task | Architecture | Citation |
| --- | --- | --- | --- |
| UD Treebank Ancient Greek | NLP | Transformer | (Eckhoff et al., ; Bamman and Crane, ) |
| OGBG MolPCBA | Classification | GNN | (Hu et al., ) |
| Atari | RL | CNN | (Bellemare et al., ), |
| CIFAR10 | Classification | ResNet18 | (Krizhevsky and Hinton, ) |
| Fuel Efficiency | Regression | FFNN | (Quinlan, ) |
| MNIST | Generative | FFNN / CNN | (Lecun et al., ) |

Training for each network required unique optimizers, learning rates, and batch sizes and was performed over different data sets.
The training was distributed on several compute clusters utilizing NVIDIA 4090, 3090, and L4 GPUs.

### 2.4 ZnNL

Most network training and evaluation are performed using the ZnNL Python library throughout this work.
ZnNL utilizes the neural-tangents library (Bradbury et al., ) for NTK calculations and the Flax library (Heek et al., ) for neural network definition and training.
The novelty methods were trained using the Flax library using customized scripts outside the ZnNL framework.

## 3 Collective Variable Evolution

This work aims to understand the learning dynamics of neural networks and identify any universal properties of this evolution.
To this end, networks of different architectures are trained on different datasets and their properties, entropy, trace and loss, are measured at each epoch.
This has been performed on standard classification and regression problems in simple networks but also scaled to more modern architectures and data structures, including generative models, language problems with transformers, graph neural networks, and reinforcement learning.
First, we discuss the most comprehensive results of the MNIST data scan; afterwards, we introduce the additional architectures in more detail.

### 3.1 Architecture Scans

In the first study, the MNIST dataset is used to train dense and convolutional neural networks, along with the Fuel Efficieny dataset used to train a dense network, the architectures of which are outlined in Table 1.
Figure 1 displays the evolution of the collective variables for certain combinations of widths, depths, and activation functions for the convolution model.
Further results for dense MNIST and the Fuel Efficiency task are shown in Appendix figures 4 and 5 respectively, but are qualitatively similar.

Figure: Figure 1: Time evolution diagrams for the convolutional MNIST study. The top row shows the evolution difference due to a changing activation function. The second row is width scaling, and the third is depth scaling. The tuple over the plots highlights the architecture being studied. It is structured as (Width, Depth, Activation), where an $x$ indicates that this property is being changed during the study.
Refer to caption: x1.png

Each plot shows the variables computed on the test and the train data, differentiated by colors.
Furthermore, each row corresponds to the network’s varying architecture property while the other properties are held fixed, similar to performing measurements in a specific ensemble in statistical physics.
In addition to the evolution diagram, a more comprehensive scaling study was performed on the dense network architecture.
This was done by training dense neural networks of varying depths and widths and computing the loss, entropy, and trace evolution.
The results of this study at the start of training and after 100 epochs are displayed in Figure 2.

Figure: Figure 2: Architecture scaling sweep of a dense neural network trained on the MNIST dataset. The top row a) shows the entropy, trace and loss as a function of network width and depth at initialization. The bottom row, b), shows the same data after training the model.
Refer to caption: x2.png

To fully digest these results, we will discuss each variable individually, beginning with the loss.

#### Loss:

In the case of both test and training loss, we see drastic improvements with changes to the architecture.
Namely, using deeper networks, TanH over ReLU activation functions, and training wider networks results in lower test and train losses.
This trend is broken in the depth scaling as we see the test loss, and therefore, generalization improves with increasing depth despite a larger train loss, in line with previously reported scaling behaviour (Yang and Hu, ; Bahri et al., ).
These results are further strengthened in Figure 2 b) in the case of a dense network, as we see the deeper networks achieving significantly improved losses compared with their shallower counterparts.

#### Entropy:

Entropy evolution highlights two distinct learning dynamics that relate to different learning regimes arising directly from the information content in the neural network.
In general, entropy undergoes one of two evolutions: information compression and structure formation.
Information compression is signified by a reduction in the NTK’s entropy, which indicates that gradient vectors are beginning to align and the network is relating many data points to one another.
An increasing entropy indicates structure formation, suggesting that the network separates data to identify better representations.
What we observe in Figure 1 and in the dense network counterpart in Figure 4 is that depending on the size of the network, a different mechanism will become dominant.
In most cases, we see the entropy drop at the beginning of training, coinciding with the sharp drop in the loss early in the optimization procedure.
For smaller networks that are limited in dimension and, therefore, have limited flexibility to construct complex representations, the entropy remains low throughout the rest of the training.
These networks also typically attain a worse train and test loss than their larger counterparts, although they can still train.
After the compression phase of training, the larger networks undergo structure formation, indicated by increasing entropy.
This relationship, however, becomes complex when depths vs width are considered.
It is clear from the results that increasing the network dimension via width increases the entropy of the models.
However, adding layers appears to reduce this entropy.
Intuitively, this makes sense as adding layers does not increase the projected dimension but applies additional non-linearities to the representations.
However, we also see that deeper networks can achieve higher entropies than their shallower counterparts when scaled further in width and after training.
Thus, while adding layers to a network of one width can reduce its entropy, scaling it to a higher width increases its capacity for a higher entropy.
To explore this mechanism further, the higher resolution architecture scans of the dense neural network architecture are presented in Figure 2 a) at the start of training and Figure 2 b) at the end.
In these plots, increasing the network depth reduces entropy at initialization.
However, the deeper networks form a region of high entropy after training, corresponding with a better test loss.
This trend is not reproduced in the TanH case, indicating that these mechanisms are sensitive to the activation function.
Further work could explore how the scaling behavior for each activation can be used to perform targeted architecture selection.

#### Trace:

The final property to study is the trace of the networks.
Recall the trace corresponds to an effective learning rate, e.g., how far we will be pushed along a specific direction due to the network architecture.
In each study, trace increases with training time to a very large value, often in step with the loss evolution.
From a learning perspective, this indicates that the network continually increases its effective steps along specific directions as training goes on, similar in concept to the network’s confidence that the solution it is forming is correct.
From a physical perspective, this is very similar to the evolution of kinetic energy asymmetrically to the decrease in a potential.
Indeed, if one were to consider the kinetic energy of a network representation $f(x_{i})$, as $E^{k}=\frac{1}{2m_{i}}|\dot{f}(x_{i})|^{2}$, a simple expansion of variables would show

$$ $\centering E^{k}=\frac{1}{2m_{i}}\Theta_{ii}\sum_{p}\left(\frac{\partial \mathcal{L}(x_{p})}{\partial\theta_{j}}\right)^{2},\@add@centering$ (11) $$

where $m_{i}$ is some point mass and $\Theta_{ii}$ is indeed the trace of the NTK plotted here.
Representation of this value as kinetic energy is a convenient metaphor for a real effect within a neural network during training.
Namely, as one minimizes the loss, the network appears to become very sensitive to changes in these losses, as evidenced by this effective scaling parameter becoming large.
This indicates that if fictitious data were added to the data set late in training, driving the loss term up, the network would train on this data point with a significant degree of confidence as opposed to at the start of training, where it could possibly be ignored or saturated out by the other data points.
The conclusion here is that adding malicious data or even very new data to a data set at the end of training will result in a larger change in parameters than it would at the start of training, even if the loss computed on this value were the same at both times, simply due to the state of the neural network.
This result has significant ramifications for adversarial attacks on trained models and may help explain why small additions of noise to a model can lead to a drastic change in its outputs.
Further investigation of this property may lead to a reduction in model sensitivity and could also be applied to more effective continual learning strategies.
In the architecture sweeps shown in Figure 2, we see a similar trend in the trace at initialization and, after training, only shifted.
Specifically, increasing depth and width results in a larger trace value further exasperated by the training process.
Thus, supporting the work of Yang and Schoenholz (), deeper and larger networks are more sensitive than their shallower counterparts, not just at initialization but also after training.

## 4 Novelty Models

While exploring the evolution of collective variables on test problems provides deep insight into scaling relations and dynamics, it does not directly relate to the current state of machine learning.
To extend our analysis and argue that these dynamics are, in fact, a universal phenomenon, we apply this analysis to several so-called novelty models taken from various fields of machine learning.
The results of these evolution studies are outlined in Figure 3 and discussed below.
It should be noted throughout this discussion that due to the size of these architectures and data sets, the collective variables have not been computed at each epoch but rather at steps of between 10 and 1000 epochs.
Further, the NTK matrix was often subsampled during the calculations using 200 samples of 20 data points to construct a measurement.

Figure: Figure 3: Collective variable evolution for the novelty architectures. In each case, the raw value, along with a running average, is shown. (a) A resnet18 trained on the CIFAR10 data set. (b) A transformer trained on part-of-speech tagging of ancient Greek. (c) A CNN-based reinforcement learner trained on the arcade game Pong. Note the use of reward instead of loss on the y-axis. (d) A graph neural network trained on molecular property prediction. (e) A simple variational auto-encoder trained to reproduce the MNIST dataset.
Refer to caption: x3.png

Before discussing the individual models and the unique features emerging in the dynamics, their measurements’ global properties appear consistent with the test cases presented above.
In all cases, the entropy of the models increases throughout training, indicating that we are always within a structure formation regime.
This is not unusual, as each network is far larger than those required to induce constrained, compressive learning.
A notable difference is entropy’s complete lack of a compression phase.
This is likely because the dataset size is such that after even a single batch, so much back-propagation has been performed that the network moves directly to a structure formation phase.
Such a process indicates that the parameters of the network are free to perform sophisticated structure formation, thus learning fundamental features of the data.
Given the dominance of the entropy increase in the latter portion of the training, it is an open question whether the initial decrease is a required process or simply an artifact of non-ideal initialization.
One explanation would be that those models that perform only compression are not performing what the community would call deep learning. Only those models that can perform this structure formation should be considered to be performing deep learning.
Further, the trace in the later time consistently increases, in line with previous measurements.
Below, we examine four models more closely and interpret some emergent features.

### 4.1 NLP

In the NLP problem, a transformer was trained to perform part-of-speech tagging for ancient Greek text from Eckhoff et al. (); Bamman and Crane ().
We see a story similar to the test cases discussed in the entropy and trace evolution.
The trace increases steadily as the loss decreases.
The difference in magnitudes is interesting compared with the previous examples.
While the accuracy of the transformer reached 99 %, the loss, due to the specific one chosen, does not get too small, resulting in a more reasonable trace value and, thus, less "confidence" during the updates.

### 4.2 Reinforcement Learning

The next novelty model study looked at how the entropy and the trace of an agent trained via deep actor-critic reinforcement learning evolved.
To compute the NTK, a dataset of environments was constructed by letting the agent play the game hundreds of times at three stages of learning: untrained, occasionally successful, and end of training.
In this way, the dataset on which the NTK is computed covers the full range of possible configurations the agent could see.
The entropy, trace, and reward curves in Figure 3 are all plotted using a log-x scale to highlight exactly where certain transitions in this evolution occur.
The reason for this is the sharp increase in the entropy and trace after approximately 1000 training episodes.
What is interesting about this transition is that the entropy and trace, particularly of the actor, increase in the episode before the reward of the agent follows suit.
Upon finer examination, we see that this jump occurs in the episode directly before the agent begins to be able to perform the task successfully and achieve positive rewards.
This indicates that the entropy, and to a lesser extent, the trace, indicated that the agent could perform its task better given its current set of parameters.
This is a very natural interpretation of structure learning in a neural network, as seen through the eyes of an agent playing a game.
At some stage during training, the neural network begins to identify structure in the data and can use this structure to make better moves in the game.
The training then continues very quickly, partly due to the onset of structure formation but perhaps aided by the new scaling from the trace, which increases more significantly in the actor model.
In the case of RL, the concept of confidence also becomes quite interpretable as the agent begins to trust more that the structure it is beginning to learn will be successful. Therefore, it can make larger steps in its updates, arising from the larger trace.

### 4.3 Graph Prediction and Generative Modelling

In graph prediction and generative models, similar dynamics are realized during training.
Interestingly, in both cases, due to the small loss range, the trace evolution can be observed with a greater resolution, resulting in a large dip at the beginning of training.
This dip occurs over a small range compared with the fluctuations and values observed in other models, but it is by no means a simple fluctuation.
However, under the kinetic energy framework introduced above, where these dips arise could be accounted for by a missing term in the energy equation, perhaps related to repulsion between data points.

## 5 Conclusion

We have motivated the use of entropy and trace of the empirical neural tangent kernel as suitable measures of a neural network’s state.
These variables were then applied to understand architecture scaling and evolution on test problems in image classification and regression using dense and convolutional networks of various widths and depths.
We identify two dominant regimes in entropy during learning: compression, where entropy decreases and latent space representations begin to collapse, and a structure formation regime, where the entropy increases again later in training, and the network better differentiates the data.
We further identified that the compression regimes sometimes dominated the learning process in smaller architectures, albeit with degraded performance compared to the structure-forming networks.
This quantitative handle to distinguish between the small and deep networks provides us with a classification of what it means to be in the deep learning regime.
Further, we highlighted the trend of increasing trace during training, particularly in late time evolution.
As shown, this trace can be associated with an effective learning rate during model updates, thus indicating that adding data with a large loss to the network late in training could result in a very sharp change in parameters despite the model being otherwise well-trained.
To demonstrate the further use of these variables, we study the collective variable evolution on a set of so-called novelty models, which are more closely aligned with the current state of the art.
In each case, we see a similar trend in the entropy and trace, with a notable, short-time difference in the trace of the autoencoder and graph neural network studies.
Overall, the dynamics observed appear universal among architectures and shed light on how neural networks collect data in their latent space.
Further, we indicate avenues for improved stability by exploring the reduction of trace late in training.

## 6 Ethics Statement

The authors declare no conflicts of interest or other ethical concerns.

## 7 Acknowledgements

C.H and S.T acknowledge financial support from the German Funding Agency (Deutsche Forschungsgemeinschaft DFG) under Germany’s Excellence Strategy EXC 2075-390740016.
C.H and S.T acknowledge financial support from the German Funding Agency (Deutsche Forschungsgemeinschaft DFG) under the Priority Program SPP 2363, “Utilization and Development of Machine Learning for Molecular Applications - Molecular Machine Learning” Project No. 497249646.
C.H and C.L acknowledge funding by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Project Number 327154368-SFB 1313.
The authors acknowledge support by the state of Baden-Württemberg through bwHPC and the German Research Foundation (DFG) through grant INST 35/1597-1 FUGG.
The authors acknowledge support from the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) Compute Cluster grant no. 492175459.

## 8 Appendix

Figure: Figure 4: Time evolution diagrams for the dense MNIST study. The top row shows the evolution difference due to a changing activation function. The second row is width scaling, and the third is depth scaling.
Refer to caption: x4.png

Figure: Figure 5: Time evolution diagrams for the dense MPG regression dataset study. The top row shows the evolution difference due to a changing activation function. The second row is width scaling, and the third is depth scaling.
Refer to caption: x5.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: I would like to see very similar behaviour of entropy and trace across a number of tasks to be convinced about this two stage phenomena.

And where it doesn't occur, I expect some effort in trying to figure out why but appropriately titrating aspects of the task / dataset / network between the showing up of two stages vs one stage and figuring out exactly why and when this occurs. Similarly, a better analysis of why the trace shows a dip versus not during learning.

Line 253-254: "increasing entropy indicates structure formation" -- this is counter to the usual physics entropy where entropy decrease with structure formation. I would call the second phase "structure separation" or "structure refinement".
soundness: 2
strengths: The idea of analysing the entropy and scale of the NTK during training seems novel.

The authors have done a fair number of experiments across relevant architecture-task combinations.
summary: The authors use entropy and trace of the neural tangent kernel to understand the learning dynamics of neural network esp in regard to architecture scaling. They claim to see initial entropy decrease termed information compression and later entropy increase termed structure formation. They try to identify a deep learning regime. They also show how entropy and trace chance for various architecture-task combinations.
weaknesses: In section 3.1, drawing conclusions on scaling of entropy and trace of the NTK, just from architecture scans on MNIST seems quite premature. When even changing from ReLU to tanh changes results (line 353), it is unclear if results may change based on hyperparams requiring tuning for each architecture, i.e. each architecture needs its own learning rate, initialization, etc. Furthermore, quirks of MNIST might lead to some effects. These conclusions must hold across datasets and architectures to be justified -- the authors show this for dense and convolutional networks and MNIST and Fuel Efficiency tasks and say that results are qualitatively similar. The figures between MNIST and Fuel Efficiency dataset do not seem qualitatively similar to me. Main text says fuel efficiency dataset is Fig. 5 in Appendix, but in Fig 5, it says MPG regression dataset (which is not called this in the table either, but since it is the only task with regressions, this seems to be the Figure - please be consistent in nomenclature) In any case, this doesn't look qualitatively similar to the MNIST figures.

I'm not convinced that the trace gives any extra insight into why malicious data will affect the dataset more in the training. One could just as well use the argument used here in the usual weight update equation. If losses are low, and a new data point has a high loss in particular a high loss whose gradient with respect to current weights is high, then it'll have a large effect later in training.

The lack of a compression phase in models of section 4 is interesting, but is not studied, despite prominently mentioning this as information compression in the MNIST case. The authors just say this may be because of large batch /dataset size. This could have been easily studied by varying the batch size or restricting the dataset in some way.

The trace also varies very differently between different domains and no attempt is made to explain this except some conjectures on missing term in the energy equation.

Minor:
Line 185, use \cite instead of \citep to make authors of a work the subject of a sentence.

---

confidence: 2
contribution: 2
initial_rating: 8
presentation: 4
questions: Suggestions for improvement:
* Can the term "deep learning regime" be explicitly defined, and citations provided for wherever this definition matches or deviates from prior literature.
* Please explicitly provide information on network size/architecture/nonlinearities and initialization for every case considered. 
* For all tasks, how has the data been preprocessed/scaled/normalized? 
* L396: How was "200 samples of 20 data points" chosen? Justify this choice
* L406: Clarify "non-ideal initialization"; what would be ideal?
* L418: Unclear "loss, due to the specific one chosen…". What is being said here?
* L423: Missing details, what algorithm and hyperparameters were used for Deep RL?
* For the RL task, could some classic-control RL tasks be used instead of (or in addition to) a task like Atari CNN
* The idea of kinetic-energy introduced around L367 is quite interesting but seems underdeveloped. It should be left in, but more empirical details about this over training evolution could be provided. 
* The authors should consider connecting their findings to previous work on curriculum learning, specifically algorithms where loss on training examples are used for building automatic training schedules.
soundness: 3
strengths: Strengths
* Useful contribution: the problem studied and the point of view presented is practically useful to a large segment of the ML community. The basic claims were well studied/supported, with only minor concerns. 
* Paper is clear and well written, and as such has high pedagogical value.
summary: The authors explore neural network learning dynamics and scaling by using recently introduced entropy and trace measures on the empirical neural tangent kernel (NTK) spectrum. 
These metrics reveal insights into network representations and their evolution across diverse architectures and datasets, including a toy MNIST problem and more contemporary "novelty" architectures. 
Two main mechanisms are seen to emerge (1) information compression, characterized by entropy reduction and (2) structure formation, characterized by entropy increase especially in larger networks (noted by the authors to indicate a "deep learning regime"). While the novelty of the work is somewhat limited (because it's an empirical extension of work by Tovey et al (2023) who introduced the entropy and trace measures), I found the manuscript to be rigorous and very well written.
weaknesses: Weaknesses
* Limited novelty (as mentioned earlier), but by no means a deal breaker
* Some missing details (see below)

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Figure 1 is very hard to grasp. Please use similar colors for a given value of the quantity of interest, and separate "test" from "train" with e.g. dotted vs solid, or anything else that would make it easier to separate the values.

Should there be a sum over k in the definition of big-theta_pi in line 133? Where does the k go?

l. 179 is hard to parse.

What are the dark splotches near the sides of the panels in Figure 2?

It's not clear how Eq. 11 was derived. The original equations relate f_dot(xp) to dloss / df, but Eq. 11 show dloss/dtheta_j instead?  What is the j variable here?
soundness: 3
strengths: Understanding the dynamics of neural network training is an important problem.

The measures are reasonably well described.

The experiments appear thorough.
summary: The authors introduce two quantities to describe the training of neural networks, based on the neural tangent kernel (NTK). These measures are the entropy of eigenvalues in the NTK, and its trace.

They then assert that entropy decreases and increases correspond to "information compression" and "structure formation", respectively. The trace measures aggregated learning rate multipliers over the separate modes of the NTK.

They train multiple neural networks of different types and observe the changes in loss as well as in both entropy and trace of the NTK, trying to identify patterns.
weaknesses: It is not quite clear what message we're supposed to get from this work. The authors introduce these measures, and propose some kind of intuitive description of what they measure, which is not necessarily wrong, but what exactly it means for training is not obvious.

The authors also try to discern patterns in the differences in the evolution of these measures with network architecture, but there doesn't seem to be anything consistent beyond "trace increases" and "entropy always increases at some point unless the model is very small". For example,  the authors suggest that smaller networks have only decreasing entropy due to inability to create structure, but then have to admit that this is only for depth, not for width (line 261-262). Even more puzzling are the very large changes caused by changing the activation function, which seem to be at least as large as those caused by architecture changes.

Basically there does not seem to be any clear message from the data.

The parallel with kinetic energy at line 363 is unclear.

The authors make one prediction, if a rather weak one: that introducing noisy data should produce more changes in the parameters at the end of training than at the start of training (line 376). This prediction doesn't seem to be tested.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: -	In all the examples shown, it seems the models studied are succeeding (at least somewhat) at solving their tasks. So, we see relatively universal trends, and this is, so far, a reasonable result. However, I wonder what does it look like if the models are too small, or badly initialized, or have the wrong architectures for their tasks? What would these curves look for regimes in which the training dynamics are failing to solve the task? 
   - Could you, for example, take some of the models used in figure 3 and handicap them in some way to force them to be bad? What would the curves look like then?
    - Furthermore, for the example shown in figure 3, all these models are large enough such that their entropy is always increasing, as argued by the authors should be the case for over-parametrized, deep models. Could you then make the same study for these models as done in Figure 2 by making them smaller or shallower and recreate the transient decreasing entropy phase? 
    - By doing such an analysis, can you relate your results to the cost/performance trade offs for models at different scales?
soundness: 3
strengths: Originality: The paper extends known methods, namely computing the trace and entropy of NTK spectra, to a broader scope than previous applications.

Quality: The paper is well motivated, and the background material is introduced with good detail with relevant literatures cited.

Clarity: It is also written clearly, in good scientific English and is relatively comfortable to follow. The key results of the paper are also introduced well and discussed in a way that was easy to appreciate why the authors argue their significance. The figures in the paper are also well formatted.

Significance: The interpretation of the results has the potential to condense the analysis of training dynamics into universal patterns of collective variables, an insight that could help diagnose and analyze a broad variety of models.
summary: The authors present a study on collective variables of empirical Neural Tangent Kernels (NTK) of various scales, architectures, and tasks in order to describe the learning dynamics of these models. Two major results are described: 1) the evolutions of the entropy of the NTK spectrum during training can undergo two different modes, a decrease in entropy or increase. A decrease in entropy is described as an 'information compression' phase, where gradients from different inputs are aligning. This process seems to occur mostly for smaller networks in both width and depth and is often transitory. Conversely, an increase in entropy is described as ‘structure formation’, where different inputs result in more diverse gradient vectors, indicating more selective representations. An increase in the NTK spectrum’s entropy is almost always observed for large (or over-parametrized) networks and coincides with later stages of training when the models are converging to higher performance. The authors argue that the latter phase of increasing entropy can be argued as when a ‘deep learning’ phase is actually occurring. 2) The trace of the NTK increases as the loss decreases during training, where the trace can be interpreted as an ‘effective step size’ or learning rate. This increase in effective step size, even as the loss decreases, implies a heightened sensitivity of the model to loss, or, as an increase in confidence in the direction of the gradient vectors. The authors argue that this indicate that fictitious data has a stronger negative effect if shown at later stages of training compared to earlier, even if the loss value is the same. Finally, the authors show that these results share some universality across a variety of models and argue that this behaviour might indeed be descriptive of universality in learning dynamics and, in particular, deep learning.
weaknesses: -	(Clarity) For the figures, the colour coding/line-styles aren’t well grouped to my eyes. It might benefit to have the smaller and larger models have similar shades and then have different line styles for train vs. Test. I found myself having to constantly check the legends for each panel which was a bit tedious to interpret each figure.
-	(Significance and Quality) The interpretation of the changes in the entropy of the NTK spectrum as “information compression” vs “structure formation” is, as it’s written, somewhat convincing, but I’m not actually sure if this interpretation is simply a good analogy or if that is what is exactly happening. Perhaps this criticism is due to my own ignorance on the topic, but I think making these statements more convincing would empower this paper quite a bit as these are the central results of the paper. I’m not entirely certain what would make this point more convincing unfortunately, but I wonder if there could be custom datasets that one knows a priori that the data points are somewhat redundant and that the entropy should only decrease versus datasets that require more sophisticated representations and so the entropy should increase, as an example. Doing so might also help relate these results to the results from Tovey et al. (2023) and would help elucidate the contribution. Furthermore, expanding on how these methods are explicitly different from the Tovey et al. (2023) would also clarify the significance and contributions of this paper.

Meta-review: This submission investigates properties of the empirical neural tangent kernel (NTK), and demonstrates that the properties change predictably over the course of learning in two qualitative phases.

Empirical studies of the NTK may provide insight into the learning dynamics of models that follow NTK dynamics. However, I concur with Reviewer 3ejh that the descriptions afforded by the NTK trace and entropy are correlational in that they do not provide more than quantities that correlate with loss, even if they do so predictably, and that the phases identified are also observational rather than causal properties of the dynamics. To make the case as claimed in the manuscript that this empirical analysis "enhances our theoretical understanding of neural network dynamics and provides practical insights for optimizing architecture design and training protocols," the authors should answer the question of what to do with these descriptions—for example, to answer outstanding questions in neural network theory or, more practically, provide a way to attenuate a certain downstream property—which is not yet done in the manuscript.

Additionally, the authors nowhere discuss deviations of neural network training from the NTK regime; see [Geiger et al. (2019)](https://arxiv.org/abs/1906.08034) for example, and the discussion on μ-transfer confounds these regimes. Since feature learning is attested to underlie deep learning (see [this blog](https://kempnerinstitute.harvard.edu/research/deeper-learning/infinite-limits-of-neural-networks/) for a pedagogical overview), this may limit the explanatory power of any quantities derived from the NTK, including those derived here (even if reevaluated empirically throughout training).

Decision: Reject