Latent State Models of Training Dynamics
Authors: Michael Y. Hu
Angelica Chen
Naomi Saphra
Kyunghyun Cho
Venue: Accepted by TMLR ()

Abstract: The impact of randomness on model training is poorly understood. How do differences in data order and initialization actually manifest in the model, such that some training runs outperform others or converge faster? Furthermore, how can we interpret the resulting training dynamics and the phase transitions that characterize different trajectories? To understand the effect of randomness on the dynamics and outcomes of neural network training, we train models multiple times with different random seeds and compute a variety of metrics throughout training, such as the $L_2$ norm, mean, and variance of the neural network's weights. We then fit a hidden Markov model (HMM) over the resulting sequences of metrics. The HMM represents training as a stochastic process of transitions between latent states, providing an intuitive overview of significant changes during training. Using our method, we produce a low-dimensional, discrete representation of training dynamics on grokking tasks, image classification, and masked language modeling. We use the HMM representation to study phase transitions and identify latent "detour" states that slow down convergence.

Published in Transactions on Machine Learning Research (12/2023)
Delays, Detours, and Forks in the Road: Latent State Models of Training Dynamics
Michael Y. Hu
michael.hu@myu.ed
New York University
Angelica Chen New York University
ac5968
Naomi Saphra New York University
ns40088myn.edu
Kyunghyun Cho
kyunghhyu.cch wwyw.
New York University
Prescient Design, Genendech
CIFAR LMB
Reviewed on OpenReview: https: // openreview. net/ forum? id= NE2xXWoOLF
Abstract
The impact of randomness on model training is poorly understood. How do differences in data order and initialization actually manifest in the model, such that some training runs outperform others or converge faster? Furthermore, how can we interpret the result- ing training dynamics and the phase transitions that characterize different trajectories? To understand the effect of randomness on the dynamics and outcomes of neural network train- ing, train models multiple times with different random seeds and compute variety of metrics throughout training, such as the L2 norm, mean, and variance of the neural net- work's weights. We then fit hidden Markov model (HMM; Baum & Petrie 1966) over the resulting sequences of metrics. The HMM represents training as a stochastic process of transitions between latent states, providing an intuitive overview of significant changes during training. Using our method, we produce a low-dimensional, discrete representation of training dynamics on grokking tasks, image classification, and masked language modeling. We use the HMM representation to study phase transitions and identify latent "detour" states that slow down convergence.
1
Introduction
We possess strong intuition for how varions tuned byperparameters, such as learning rate or weight decay,
affect model training dynamics and outcomes (Galanti et al. 2023 Lyu et al. 2022 For example, a larger
lcarning rate may lead faster convergence at the cost of sub-optimal solutions (Hazan 2019 Smith et al.,
2021
Wu al. 2019). However, we lack similar intuitions for the impact of randomness. Like other
hyperparameters, random seeds also have significant impact training (Madhyastha & Jain 2019 Sellam
ct al. 2022 but we have limited understanding of how randomness in training actually manifests in the model.
In this work, we study the impact of random seeds through a low-dimensional representation of training dynamics, which we use to visualize and cluster training trajectorics with different parameter initializations and data orders. Specifically, we analyze training trajectories using a hidden Markov model (HMM) fitted set of generic metrics collected throughout training, such the means and variances of the

---

Published in Transactions on Machine Learning Research (12/2023)
neural network's weights and biases. From the HMM, we derive visual summary of how learning occurs for task across different random seeds.
This work is a first step towards a principled and automated framework for understanding variation in model training. By learning a low-dimensional representation of training trajectories, we analyze training at higher level abstraction than directly studying model weights. We use the HMM to infer Markov chain over latent states in training and relate the resulting paths through the Markov chain training outcomes.
Our contributions:
We propose to use the HMM as a principled, automated, and efficient method for analyzing variabil- ity model training. We fit the HMM a set of off-tho-shelf metrics and allow the model to infer latent state transitions from the metries. We then extract from the HMM "training map," which visualizes how training evolves and describes the important. metrics for each latent state (Section 2)
To show the wide applicability our method, we train HMMs on training trajectories derived from grokking tasks, language modeling, and image classification across a variety of model architectures and sizes. In these settings, we use the training map to characterize how different random seeds lead to different training trajectories. Furthermore, we analyze phase transitions in grokking by matching them to their corresponding latent states in the training map, and thus the important metrics associated with each phase transition (Section 3.1
2.
We discover detour states, which are learned latent states associated with slower convergence. We identify detour states using lincar regression over the training map and propose our regression method as a general way to assign semantics onto latent states (Sections 2.3 3.4).
To connect detour states to optimization, discover that classification by destabilizing the optimization process and, conversely, remove detour states in grokking by stabilizing the optimization process. By making few changes that are known to stabilize neural network training, such adding normalization layers, we find that the gap bet memorization and generalization in grokking dramatically reduced. Our results, along with prior work from Liu et al. (2023), show that grokking can be avoided by changing the architecture or optimization of deep networks (Section 3.3
induce detour states in image
Our code is available at hhttps://githun http:///igcabb....oiiiaaaaaa//// com http://ghhti..ooommmoooaoooooooo training
2
Methods
A(w)
B(wD
[(ww))
L(ws)
L(w))
0.05
0.07
Loss
0.95
0.93
1.0
Figure 1: From training runs we collect metrics, which are functions of the neural networks' weights. We then train hidden Markov model to predict the sequences metrics generated from the training runs, The hidden Markov model learns a discrcte latent state over the sequence, which we use to cluster and analyze the training trajectory.
In this work, we cluster training trajectories from different random seeds and then analyze these clusters better understand their learning dynamics and how they compare to each other. To cluster trajectorics, we assign each model checkpoint to discrete latent state using HMM. We choose the HMM because it is
2

---

Published in Transactions on Machine Learning Research (12/2023)
simple time series model with discrete latent space, and we specifically pick discrete latent space because previous works (Nanda et al. 2023 Olsson et al. 2022) have shown that learning can exhibit a few discrete, qualitatively distinct states.
Let WI.7 E RXT be the sequence of neural network weights observed during training. Each We is a model checkpoint. In this work, we use the Gaussian HMM label each checkpoint WI. with its own latent state, 51:7. Fitting the HMM directly over the weights is computationally infeasible, because the sample complexity of an HMM with O(D²) parameters would be prohibitively high. Our solution to this problem is compute small number of metrics
from WI-7; where d < D and fi RR - R, 1 i sis d.
2.1 Training an HMM over Metrics
In this work, we focus on capturing how the computation of the neural network changes during training by modeling the evolution of the neural network weights. To succinetly represent the weights, we compute various metrics such as the average layer-wise L1 and L2 norm, the mean and variances of the weights and biases in the network, and the means and variances of each weight matrix's singular values. A full list of the 14 metrics we use, along with formulae and rationales, is in Appendix B
To fit the HMM, we concatenate these metries into an observation sequence 21.T. We then apply z-score normalization (also known as standardization), adjusting each feature have a mean of zero and a standard deviation of one, as HMMSs are sensitive to the scale of features. We thus obtaín the normalized sequence ar. To bound the impact of training trajectory length, compute z-scores using the estimated mean and variance of (up to) the first 1000 collected checkpoints.
fi(w)))
[f((w)) ((i((w:)))ll(((((wm)))
2t
2,
fa(we)
[[fa(we) -(fddm::r)))/(((ddd
In total, we collect N observation sequences (=1 H from N different random seeds, normalize the distribu- tion of cach metric across training for given sced, and train the HMM over the sequences (ž1:) H using the Baun- Welch algorithm (Baum et al. 1970). The main hyperparameter in the HMM is the number of hidden states, which is typically tuned using the log- lieliiil,,ood, Akaike information criterion (AIC), and/or Bayesian information criterion (BIC). (Akaiko) 1998 Schwarz 1978) Here, we hold out 20% of the trajectorics as validation sequences and choose the number hidden states that minimizes the BIC. We use BIC because BIC imposes stronger preference for simpler, and thus more interpretable, models. Model selection curves are in Appendix H
2.2 Extracting the Training Map
Next, we use the HMM to describe the important features of each hidden state and how the hidden states relate to cach other. We convert the HMM into a "training map," " which represents hidden states as vertices and hidden state transitions as edges in state diagram (see Figure 2p.
First, extract the state diagram's structure from the HMM. The learned HMM has two sets of param- eters: the transition matrix p(s)(s:-1) between hidden states, and the emission distribution p(ž = k) N(u,, Ex), where Pa and Σ are the mean and covariance of the Gaussian conditioned on the hidden state k. respectively. The transition matrix is a Markov chain that defines the state diagram's structure: the hidden states and the possible transitions, or edges, between hidden states a priori. We set edge weights in the Markov chain zero if the edge does not appear in any of the N hidden state trajectorics inferred by the HMM.
3

---

Published in Transactions on Machine Learning Research (12/2023)
We annotate the hidden states $1.7 by ranking the features & posterior's 8 partial derivative with respect to &
according to the absolute value of the log
logp k(():))
0&
The absolute value of this partial derivative vector. Intuitively, if the ith index in this vector is large, then changes in the feature [ strongly influence the prediction that 8, k. We compute the closed form of this derivative in Appendix
To characterize edges in the state diagram where the hidden state changes j + k, we use this derivative, along with the learned means Hjj and μıg- A hidden state change from j to k means the new observation has moved closer to pig. Thus, we can summarize the movement features from j to & using the difference vector Ms Pj. However, not all these changes are necessarily important for the belief that 8, k. To account for this, we can rank these changes by our measure of influence, computed from partial derivatives of the postcrior.
In the results to follow, when examining stale transition j + k al timestep 4, we report the 3 most influential features for hidden state k. To aggregate across runs, we average the absolute value vectors.
In summary, we can obtain a training map from an HMM by extracting:
The state diagram structure from pruned transition matrix.
Edge labels from 1) differences between learned means and 2) partial derivatives of the posterior.
2.3 Assigning Semantics to Latent States
From the HMM's transition matrix, we obtain a training map, or the Markov chain between learned latent states of training. We then label the transitions in the training map using the HMM's learned means and partial derivatives of the posterior. But what do we learn from the path training run takes through the map? In particular, what impact does visiting a particular state have on training outcomes?
In order to relate HMM states to training outcomes, we select metric and predict it from the path a training run takes through the Markov chain. To do so, we must featurize the sequence latent states, and in this work we use unigram fcaturization, or a "bag of states" model. Formally, let 81, . 8r be the latent states visited during training run. The empirical distribution over states can be calculated as:
16, E, T
1s; =
P(==k)=
where represents a particular state and T is the total number of checkpoints in the trajectory. This distribution can be written d-dimensional vector, which is equivalent to unigram featurization.
In this work, we investigate how particular states impact convergence time, which we measure as the first timestep that evaluation accuracy crosses a threshold. We set the threshold to be a value slightly smaller than the maximum evaluation accuracy (see Scction 3.4h. We use lincar regression to predict convergence time from P. Here, we are not forecasting when a model will converge from earlier timesteps; rather, we are simply using linear regression to learn a function between latent states and convergence time.
After training the regression model, we examine the regression coefficients to see which states are correlated with slower faster convergence times. If the regression coefficient for a state is positive when predicting convergence time, then training run spending additional time in that state implies longer convergence time. Additionally, if that same state is not visited by all trajectories, then we can consider it detour, because the trajectories that visit the optional state are also delaying their convergence time.
1 'Computing feature importances using partial derivatives of the posterior was suggested by Nguyen Hung Quang and Khoa Doan. Previous versions of this paper used different. computation method.
4

---

In modular arithmetic, some training runs converge thousands of epochs earlier than others. Examining the modular addition training map, we find several paths of different lengths: some training runs take the shortest path through the map to convergence, while others do not. We fcature three such paths Figure Published in Transactions on Machine Learning Research (12/2023)
Definition.
A learned latent state is a detour state if:
Some training runs converge without visiting the state. This indicates that the state is "optional."
Its linear regression coefficient positive when predicting convergence time. This indicates that a training run spending more time in the state will have a longer convergence time.
Our method for assigning semantics to latent states can be extended to other metrics. For example, one might regression to predict mcasure of gender bias, which vary widely across training runs (Sellam) et al. 2022), from the empirical distribution over latent states. The training map then becomes map of how gender bias manifests across training runs, We recommend computing the p-value of the linear regression and only interpreting the coofficients when they are statistically significant.
3 Results
To show the applicability of our HMM-based method across a variety of training settings and model ar- chitoctures, we perform experiments across five tasks: modular addition, sparse parities, masked language modeling, MNIST, and CIFAR-100. For hyperparameter details, see Appendix [D] In this work, we ignore embedding matrices and layer norms when computing metrics, as we are primarily interested in how the function represented by the neural network changes.
Modular arithmetic and sparse parilies are tasks where models consistently exhibil grokking (Power el al. 2022|, a phenomenon where the training and validation losses seem to be decoupled, and the validation loss drops sharply after a period of little to no improvement. The model first memorizes the training data and then generalizes to the validation set. We call these sharp changes "phase transitions," " which are periods in training which contain an inflection in the loss (i.e., the concavity of the loss changos) that then sustained (no return to chance performance).
We study modular arithmetic and sparse parities to how phase transitions are represented by the HMM's discrete latent space. We complement these tasks with masked language modeling (Appendix P. and image classification. In Sections 3.1 and 3.2. we use the training map to examine the characteristics of slow and fast-converging training runs in the grokking settings and image classification. In Section 3.3 we show that variation in convergence times between runs can be modulated by changing training hyperparameters or model architecture. Finally, in Scction 3.4 we formalize the observations of Soction 3.3 using lincar regression connect detour states with convergence lime.
3.1 Algorithmic Data: Modular Arithmetic and Sparse Parities
Modular Arithmetic: Figure 2
In modular addition, we train a one-layer autoregressive transformer to predict 2 (z + +) mod 113 from inputs z and y. We collect trajectories using 40 random seeds and train and validate the HMM on random 80-20 validation split, split that use for all settings. This is replication of the experiments in Nanda et al. (203)).
All runs initialize in state 1 and achieve low loss in state 3, but there are several paths from to 3. The longest path (1 +1>3 5 2 + 3) coincides with the longest time to convergence of the three featured runs, and the shortest path (1 + 3) with the shortest.
Using the HMM, we further dissect this variability by relating the edges exiting state to how fast or slow generalizing runs differ with respect to model internals. The results of this examination are in the table of Figure [ Here, we take the top features of states 2, 5, and 3 via the learned covariance matrices, and quantify the feature movements of the top features by subtracting the learned means (recall 2) between these states and state 1. We find that the fast-generalizing path (1 + 3) is characterized by a "just-right"
5

---

Published in Transactions on Machine Learning Research (12/2023)
10
Train
Validation
1.957
0-905
0.096
0.906
0.01
0.005
0.005
Loss
(1,001
0.001
4
0.907
0.001
0.004
0.001
2000
4000
6000
8000
10000
Epoch
12.5
10.0
0. 1.987
0.965
0.096
0.996
0
Loss
0.01
0.006
0.005
5.0
0.002
0.001
0.007
2.5
0.001
0.004
0.001
0.0
0
2000
4000
G000
8000
10000
Epoch
12.5
10.0
0.987
0.995
0.096
0.996
C
Loss
7.5
01.01
CL005
0.005
0.002
0.001
5.0
0.997
0.001
0.004
2.5
0.001
2000
6000
8000
10000
Epoch
Edge
Top 3 important feature changes (z-score)
Transition frequency
Mean convergence epoch
1
((w) (1.99, L2 11.68, L1 [1.83
2 40
1810 + 71.5
L2 10.59, L1 40.88, p(A) f1.29
34 / 40
2950 + 104
1
L2 12.08, L1 22.25, a (w) 12.16
4 / 40
5450 + 313
Figure 2: One-layer transformer trained on modular addition. Edges exiting the initialization state all have different mcan convergence epochs. Feature changes are orderod by importance from most to least. For example, " "L 40.59" means that state 2 has learned La norm that is 0.59 standard deviations lower than state 1, and the Ta norm is the most important feature for state See Appendix B for glossary of metrics and Section 2.2 for how we identify important features.
drop in the La norm (11.68, see table). The slower-generaliiznng runs + + 3) and (1 + 5 2 3) are characterized by either smaller (10.59) larger (12.08) drops in Le norm.
We can also connect our training map results to phase transitions found in modular addition by prior work Nanda el al. (2023); Power et al. 2022):: State 1 encapsulates the memorization phase transition: the training loss drop to near-zero in state 1, while validation loss increases. Thus, according to the training map, the epoch which the generalization phase transition happens is affected by how fast the L2 norm drops immediately after the memorization phase transition. "just-right" drop the Lg norm is correlated with the quickest onset of generalization.

---

Published in Transactions on Machine Learning Research (12/2023)
Sparse Parities: Figure 8 in Appendix F
Sparse parities is similar rule-based task modular
addition, where multilayer perceptron must learn to apply AND operation to 3 bits within a 40-length bit vector; the crux of the task is learning which 3 of the 40 bits are relevant. We again collect 40 training runs.
Similar to modular arithmetic, path variability through the training map also appears at the beginning of training in sparse parities. Slow-generan hlowggeneaalizing runs take the path + - 5), while fast-generalizing runs take the more direct path (2 - 5). The L2 norm remains important here, with the edge (2 + 0) characterized by increase in the L2 norm and the edge (2 + 5) characterized by decrease. Once again, the speed at which the gencralization phase transition occurs is associated with specific change in the L2 norm immediately after the memorization phase transition.
3.2 Image classification: CIFAR-100 and MNIST
CIFAR-100: Figure 3 Training neural networks on algorithmic data is a nascent task. As counterpoint to the grokking scttings, consider image classification, well-studied task in computer vision and machine learning. We collect 40 runs of ResNet18 (He et al. 2016 trained on CIFAR-100 (Krizhevsky 2009 and find that the learning dynamics are smooth and insensitive to random seed. Unlike our results from the prior section, the training map for CIFAR-101 is linear graph, and the state transitions all tend to feature increasing dispersion in the weights. We show the top 3 features for cach state transition in the table of Figure The L, and L2 norms are increasing monotonically across all state transitions.
Train Validation
689
0.968
0.982
0.986
0.986
0
0.032
0.018
0.014
0.014
1000
2000
3000 Step
4000
5000
6000
Edge
Top 3 important feature changes (z-score)
4-11
L, 10.62, 14(A) 10.56, (A) 10.70
+ 3
(a(A) 10.75, Le 10.76, L, 10.76
+ 2
L 10.80, s4(A) 10.82, (A) 10.77
+ 0
L2 10.72, ((A) 10.75, (w) 10.81
Figure 3: ResNet18 trained on CIFAR-100. All 40 training runs we collected from CIFAR-100 follow the same path, although individual runs can spend slightly different amounts of time in each state, As shown by the table, the training dynamics of CIFAR-100 are similar between states.
The dynamics of MNIST are similar to that of CIFAR-100. We collect 40 training runs of two-layer MLP learning image classification on MNIST, with hyperparameters based on Simard al.] (2003). The training runs of MNIST again follow single trajectory through the training map. We examine several state transitions throughout training and find that the transitions are also characterized by monotonically increasing changes between features.
MNIST: Figure
Appendix G
3.3 Destabilizing Image Classification, Stabilizing Grokking
From the previous two scctions, we observe that the training dynamics neural networks learning algorithmic data (modular addition and sparse parities) are highly sensitive random seed, while the dynamics of networks trained image classification are relatively unaffected by random seed. We will now show that
7

---

Published in Transactions on Machine Learning Research (12/2023)
this difference in random seed sensitivity is due to hyperparameter and model architecture decisions within the training setups that we chose to replicate. Variability in training dynamics is not. a necessarily feature of the task, and it is not a feature of the tasks we examine in this paper. Grokking is also affected by model architecture and optimization hyperparameters, and small changes to training can both close the gap between memorization and generalization grokking and make training robust. changes in random seed. Furthermore, removing improvements to the image classification training process can induce variability training where it previously did not exist.
First, examine the training dynamics of ResNets without batch normalization (loffe & Szegedy 2015
and residual connections. Residual connoctions help ResNets avoid vanishing gradients (He et al. (016)) and smooth the loss landscape (Li et al. 2018). Batch norm has similarly been shown to add smoothness to the loss landscape (Santurkar et al., 2018 and also contributes to automatic learning rate tuning (Arora et al. 2019). We remove batch norm and residual connections from ResNet18 and train the ablated networks from scratch on CIFAR-100 over 40 random sceds.
Train Validation
Loss
0.002
0.005
0.988
0.995
2
0.035
1
2000
4000
6000
8000
10000
12000
Step
0.98%
4
Logs
3
BU002
0,005
0.98N
0.985
0
2000
4000
6000
8000
10000
12000
Step
Edge
Top important feature changes (z-score)
Transition frequency
Mean convergence step
3
maA) 10.17, L1 40.18, +. $ 2.12 μu(2) 1 0.67, L2 1 0.74, L1 t 0.63
29 / 40
3530 460
3 +
12 / 40
7260 3010
Figure 4: Without residual connections and batch normalization, ResNet training becomes unstable, causing convergence times to differ significantly. Slow-generalizing runs take the state transition (3 > 1), while fast- generalizing runs take the state transition (3 + 2). (Runs can take the path (3 - + > 2), so transition frequencies do not sum to 40). The variability induced by removing residual connections and batch norm occurs at the beginning of training.
In this experiment, we show that changing the training dynamics of task also changes the training map. Without batch norm and residual connections, ResNet18's training dynamics become significantly more sensitive to randomness. See Figure 4 Depending on the random seed, the model may stagnate for many updates before generalizing. This increase random variation is visible the learned training map, which now forks when exiting state 3, the initialization state. There now exists slow- Now-generalizing path (3 1) and fast-generaliving path (3 2), characterized by feature movements in opposite directions.
8

---

Published in Transactions on Machine Learning Research (12/2023)
If removing batch normalization destabilizes ResNet training in CIFAR-100, then adding layer normalization (which was removed by Nanda et al. (2023)) should stabilize training in modular addition. Thus, we add layer normalization back in and train over 40 random seeds. We also decrease the batch size, which leads SGD to flatter minima (Keskar et al.
2017) These modifications to training help the transformer converge around times laster on modular addition data. Furthermore, sensitivity random seed disappears the training map for modular addition in Figure becomes linear graph.
5
Train
Validation
4
Loss
3
0.799
0.9
Q
0
2
0.201
0.1
2
1
0
0
50
100
150 Epoch
200
250
300
Edge
Top 3 important feature changes (z-score)
+ 2
L1 40.93, Amar 10.93, (A) 41.52
+
(A)) 12.00, o(w) 11.56, μ(A) 41.11
Figure 5: With layer normalization and a lower learning rate, the one-layer transformer quickly learns the modular arithmetic task, with a convergence time stable across random seed. This stability is captured by the linear training map. Critically, the map still reflects the grokking phase transitions: memorization, which occurs in state 0, and generalization, which occurs in state 2.
From this soction, we draw two conclusions. First, that model training choices can amplify or minimize the grokking effect. Second, that using different hyperparamcters or architectures can result in different training maps for the same lask. In training setups sensitive to random seed, the HMM associates differences in training dynamics with different latent states.
3.4 Predicting Convergence Time
In Section 3.1 we identified latent states visited by slow-generalizing runs that skipped by fast- generalizing runs. We use our framework for assigning semantics to latent states from Section 2.3 to identify these skipped latent states as detour states, states that slow down convergence. The first step in our framework is to use paths through the training map as features in linear regression to predict convergence time. We define convergence time as the itcration where validation accuracy is greater than some threshold, and we take this threshold be 0.9 modular addition and sparse parities, 0.6 for the stable version of CIFAR-100, 0.4 for destabilized CIFAR-100, and 0.97 for MNIST. We set these values be slightly less than the maximum evaluation accuracy for each task, respectively. To visualize the variance in convergence times, see Appendix !
In Table 1 we find that lincar regression predicts convergence time from given training run's empirical distribution over latent states very accurately, as long as the training map contains forked paths. If the training map is instead linear, training follows similar paths through the HMM across different random seeds. We formalize this intuition of trajectory dissimilarity as the expected Wasserstein distance W(.+) [Kantorovich] 1939 Vaserstein 1969 between any two empirical distributions p and g, sampled uniformly over the N random seeds.
2 Trajectory dissimilarity == LE[W N(N - 1)
W W qj)
(1)
9

---

Published in Transactions on Machine Learning Research (12/2023)
Datasct
R²
p-value
Dissimilarity
Forking
Modular addition
0.977
<0.001
0.496
Modular addition, stabilized
0,514
<0.001
0.038
CIFAR-100
0.094
0.469
0.028
CIFAR- 100, destabilized
0.905
<0.001
0.806
Sparse parities
0.961
<0.001
0.183
MNIST
0.049
0.611
0.063
Table 1: Predictability of convergence epoch using unigram model of states. Dissimilarity is provided per Equation and the training maps are marked as forking unless they are lincar.
With statistically significant (p < 0.001) regression models for modular addition, sparse parities, and desta- bilized CIFAR-100, we can use the learned regression coefficients to find detour states. In Table - we highlight these detour states, defined any state with positive regression coefficient that is only visited by strict subset of training trajectories. In our tasks with linear graphs, there are no detour states, because every training run visits every lalent state. Our regression analysis largely confirms observations drawn from looking at the training maps and trajectories in sections prior: states 2 and 5 are detour states in modular addition, state 0 is a detour state in sparse parities, and state is a detour state in destabilized CIFAR-100.
State
Coefficient
State
Coefficient
0
-0.15
0.77
State
Coefficient
1
0.98
0.41
0.66
2
1.19
2
0.98
1
1.20
3
-0.20
3
-0.23
2
0.28
4 0.18
4
0.58
3
1.91
5
0.95
5
1.13
4
1.12
(a) Modular addition
(b) Sparse parities
(c) CLFAR-100, destabilized
Table 2: Learned linear regression coefficients. If the value is positive, then the time spent in the state is correlated with increased convergence time, and vice versa. Detour states are bolded.
Detour states signal that the outcome of training is unstable: they appear in training setups that are sensitive to randomness, and they disappear in setups that are robust to randomness. By adding layer norm and decreasing batch size, we remove detour states in modular addition, and the training map becomes a linear graph. Conversely, removing batch norm and residual connections destabilizes the training of ResNets, thereby inducing forks in the training map that lead to detour states.
4 Related Work
Prior works have examined the effect of random seed on training outcome (Sellam et al. 2022 Picard 2023
Fellicious et al. 2020). To our knowledge, this is the first work to 1) analyze random seed using a probabilistic model and 2) show how random seed manifests as specific changes in metrics during training. Weiss et al. (2018 2019) model the computation of neural networks as deterministic finite automata (DFA), which bears some similarity to the annotated Markov chain we extract from training runs. Williams (1992) use an extended Kalman filter (EKF) to train recurrent neural network and note the similarity between EKF and the real-time recurrent learning algorithm (Marschall et al. 2020). In contrast to the existing literature, we use state machines to understand the training process rather than the inference process. Measuring the state of a neural net work using various metrics was also done in Frankle et al. (2020).
Analyzing time scries data using probabilistic framework has been successfully applied to many other tasks in machine learning (Kim et al. 2017 Hughey & Krogh 1996. Bartollcci et al. (014). In similar spirit to work, Batty et al. (2019) use autoregressive HMM (ARHMM) to segment behavioral videos
10

---

Published in Transactions on Machine Learning Research (12/2023)
ResNet, CIFAR-100
Transformer, Modular Addition
0UIGS
U1982
0.966
0.966
3.0
0.799
QUV
10
Stable
0,014
0.201
Batch norm Residual connections
+
Layer norm
Decrease
batch size
0.052
0.987
0.965
0:905
0 0.906
0.006
0.01
0.005
0
0.005
Unstable
0
BABS
DLIIE
0.002
0.001
0.006
0.05
0.001
0.044
0.003
1.0
0.001
Figure 6: Training maps express variability in training dynamics as more densely connected graph. For stable training setups, the HMM learns linear graph as the training map. Training dynamics can be stabilized or destabilized by changing hyperparameters (batch size) or architecture (normalization layers, residual connections).
into semantically similar chunks. The ARHMM can capture both discrete and continuous latent dynamics, making it an interesting model to try for future work.
Our work substantively inspired by the progress measures literature, which aims to find metrics that can predict discontinuous improvement or convergence in neural networks. Barakk et al. (2022) first hypothesized the existence of hidden progress measures. Olsson al. (2022) found progress measure for induction heads in Transformer- based language models, and Nanda et al. (2023) found progress mcasure for grokking in the modular arithmetic lask.
The L2 norm is also known to be both important to and predictive of grokking, thereby motivating the use of weight decay to accelerate convergence in grokking settings (Nanda et al. 2023 Power et al. 2022 Thilak el al. 2022). Liu et. al. (2023) highlight the importance of the L2 norm by correcting for grokking via projected gradient descent within fixed-size L2 ball; conversely, they also induce grokking on new datasets by choosing disadvantagoous L2 norm. Our results show that grokking has other available remodies, beyond ones that directly manipulate the L2 norm. Merrill al. (2023) and Nanda et al. (2023) show that grokking in sparse parities and modular arithmetic (respectively) can be explained by the emergence of a sparse subnetwork within the larger network.
Finally, this work relates broadly to the empirical study of training dynamics. Much of the literature treats learning as a process where increases in training data lead to predictable increases in test performance
(Kaplan et al. 2020 Razeghi et al. 2022 and in model complexity (Choshen et al. 2022 Mangalam & Prabhu 2019 Nakkiran et al.) 2019 However, this treatment of training ignores how heterogeneous the factors of training can be. Different capabilities are learned at different rates (Srivastava et al. 2022 different layers converge at different rates (Raghu et al.) 2017). and different latent dimensions emerge at different rales (Jarvis el al. 2023 Saxe et. al. 2019). While early stages in training can be modeled nearly exactly through simple methods (Hu et al. 2020: Jacot et al. 2018), early stages are notably distinct from later stages, and simple models can often belic common training phenomena (Fort ct al.] (022)). Consoquently, methods like ours which treat training as a heterogencous process are crucial in understanding realistic training trajectories.
11

---

Published in Transactions on Machine Learning Research (12/2023)
5 Discussion
The training maps derived from HMMs are interpretable descriptions of training dynamics that summarize similarities and differences between training runs. Our results show that there exists low-dimensional, discrele representation of training dynamics. Via the HMM, this representation is generally predictive of the next set of metrics in the training trajectory, given the previous metrics. Furthermore, in some cases this low-dimensional, discrete representation can even be used to predict the itcration in which models converge.
5.1 Grokking and the Optimization Landscape
We conjecture that grokking is the consequence of a sharp optimization landscape. Consider the edits we performed to siguificantly decrease the grokking effect: adding layer normalization and decreasing batch size. Normalization layers and decreasing batch size have boen documented in the litcrature as increasing smoothness in the loss landscape (Santurkar et. al.] 2018 Arora et al. 2019 Keskar al.) 2017]. Image classification is well-studied task with many tricks for improving the efficiency of training; perhaps learning algorithmic data will become just efficient in the future, such that grokking is no longer concern.
5.2 Progress Measures and Phase Transitions
By modeling convergence time in grokking settings, we analyze phase transitions. We find that the gen- eralization phase transition can be sped up by avoiding detour states. These detour states are generally characterized by specific requirements in metrics such the L2 norm. For example, in the modular arith- metic setting, avoiding detour states without changing the training setup requires "just-right" decrcase in the Le norm not too little, and not too much. This observation aligns with the hypothesis from Liu et al. 2023), where the authors posit that grokking occurs because the weight norm is slow to reach a shell of particular L2 norm in weight space, previously callod the "Goldilocks zone" (Fort & Scherlis 2018
Our automated approach can be a complement to the progress measures literature, which in previous works has found measures predictive of phase transitions by hand (Barak et al. 2022 Nanda et al. 2023). In this work, instead of carefully choosing a single metric, compute variety of metrics and use unsupervised learning to find structure amongst them. We then use the learned latent representation to analyze phase transitions.
5.3 The Impact of Random Seed
We recommend that rescarchers studying training dynamics experiment with large number of training sceds. When claims are based on a small number of runs, anomalous training phenomena might be missed, simply due to sampling. These anomalous phenomena can be the most elucidating, as in our grokking experiments, where a small number of runs converge faster than the rest. The role of random variation has been highlighted for the performance and generalization of trained models (MeCoy et al., 2020 Sellam al., 2022 Juneja et al. 2023), but there are fewer studies on variation in training dynamics. We reconiendd studying training across many runs, and possibly relying on state diagrams like ours to distinguish typical and anomalous training phenomena.
5.4 Limitations and Future Work
Our work assumess that training dynamics can be represented by a linear, discrete, and Markovian model. Despite the successes of our approach, higher-powered model might capture even more information about training dynamics. Relaxing the assumptions of the HMM is likely fruitful area for future work. Addi- tionally, in this work we perform dimensionality reduction via hand-picked metrics. We use these metrics as interpretable features for our training maps, but a fully unsupervised approach without explicit metrics also deserves exploration. For very large models, training an HMM across many random seeds may be infeasible. A possible follow-up work could look whether models of training dynamics can generalize zero-shot across architectures and architecture sizes (Yang et al. 2021). If this were the case, then one could reuse dynamics models to interpret training.
12

---

Published in Transactions on Machine Learning Research (12/2023)
Finally, findings suggestive for future work on hyperparameter search. We demonstrate that training instability to random seed is highly dependent hyperparameters, and 2) instability manifests early in training. Thus, it may be more efficient to measure early variation across a few seeds to quickly evaluate a hyperparameter sctting, rather than waiting to mcasure the final evaluation accuracy on the trained model.
6 Conclusion
We make two main contributions. First, propose directly modeling training dynamics for interpretability and training dynamics research. We show that even with a simple model like the HMM, we can learn representations training dynamics that are predictive of key metrics like convergence time. Sccond, we discover detour states of learning, and show that detour states are related to both how quickly models converge and how sensitive the overall training process is to random seed. Detour stales can be removed by finding more efficient training hyperparameters or model architectures.
avenue
Acknowledgements
We would like to thank Nguyen Hung Quang, Khoa Doan, and William Merrill for their insightful comments and suggestions. We particularly thank Quang, who generously reached out about a significant error in a previous version (see Section 2.2) and whose suggested fix has been incorporated into the paper. MYH is supported by an NSF Graduate Research Fellowship. This work was supported by Hyundai Motor Company (under the project Uncertainty in Neural Sequence Modcling), the Samsung Advanced Institute of Technology (under the project Next Generation Deep Learning: From Pattern Recognition to AT), and the National Science Foundation (under NSF Award 1922658).
References
Hirotogu Akaike. Information Theory and an Extension of the Mazimum Likelihood Principle, pp. 199-213. Springer New York, New York, NY, 1998. ISBN 978-1-4612-1699-00 doi: 10.1007/973:1-4612266699 15. URL (ttpps///001.00r//00000////1//
Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. In International Conference Learning Representations, 2019. URL httpss///pparrriiee.nnttoooooo
id=r
Boaz Barak, Benjamin T.. Edelman, Surbhi Goel, Sham M. Kakade, eran malach, and Cyril Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
URL https://openrevier
Francesco Bartolucci, Alessio Farcomeni, and Fulvia Pennoni. Latent markov models: a review of a general framework for the analysis of longitudinal data with covariates. TEST, 23:433-465, 2014.
Eleanor Batty, Matthew Whiteway, Shreya Saxena, Dan Biderman, Taigs Abe, Simon Musall, Winthrop F. Gillis, Jeffrey E. Markowitz, Anne K. Churchland, John P. Cunningham, Sandeep Robert Datta, Scott W. Linderman, and Liam Paninski. Behavenet: nonlinear embedding and bayesian neural decoding of behav- ioral vidcos. In Neural Information Processing Systems, 2019.
Leonard E. Baur and Ted Petrie. Statistical Inference for Probabilistic Functions of Finite State Markov Chains. The Annals of Mathematical Statistics, 37(6))1554 1563, 1966. doi: 10.1214/aoms) //11//////1177711477 URL httpp::/ http:::///oooor//0.... 1214/aams/1111/ /t/////mm////1911177/
Loonard E. Baum, Ted Potric, George W. Soules, and Norman Weiss. A maximization technique occurring in the statistical analysis of probabilistic functions of markov chains. Annals of Mathematical Stafistics,
41:164-171, 1970.
13

---

Published in Transactions on Machine Learning Research (12/2023)
Leshem Choshen, Guy Hacohen, Daphna Weinshall, and Omri Abend. The Grammar-Learning Trajectories of Neural Language Models. arXiv:2109.000 0900009609000666 [cs], March 2022. URL http:///axxiv..o//bbbasss
06096 arXiv: 2109.06096.
Jacob Devlin, Ming-Wor Wei Chang, Kenton Loe, and Kristina Toutanova. BERT: Pre-training of deep bidi- rectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapler of the Associalion for Compulational Linguisties: Human Language Technologies, Vol- Minneapolis, Minnesota, June 2019. Association for 1423
ume (Long and Short Papers), pp. Computational Linguistics. doi: 10. (0.18653/v/ 1119-N199112 URL hctps://aclaathooogyy
Christofer Fellicious, Thomas Weissgerber, and Michael Granitzer. Effects of random seeds on the accuracy of convolutional neural networks. In Giuseppe Nicosia, Varun Ojha, Emanuele La Malfa, Giorgio Jansen, Vincenzo Sciacca, Panos Pardalos, Giovanni Giuffrida, and Renato Umeton (eds.), Machine Learning, Optimization, and Data Science, pp. 93- Springer International Publishing, 2020. ISBN 978-3-030- 64580-9.
Stanislav Fort and Adam Scherlis. The goldilocks zone: Towards better understanding of neural network loss landscapes. CoRR, 0bs/1807.0255, 2018. URL
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geom- etry and the time evolution of the neural tangent kernel. In Hugo Larochelle, Marc' Aurelio Ran- zato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Newral Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020. URL https:://rroceedings htwpg/g......ris..c///////////20 hash/ 460076699666-4881177777788888883 Abstratttt html
Jonathan Frankle, David J. Schwab, and Ari S. Morcos. The early phase of neural network training. In International Conference on Learning Representations, 2020. URL https://openrevier: wtw.rrrrrree.ee/oouuuiiiiii
Hk1 11RNFS
Tomer Galanti, Zachary S. Siegel, Aparna Gupte, and Tomaso Poggio. Sgd and weight decay provably induce low-rank bias in neural networks, 2023.
Mina Ghashami, Edo Liberty, Jeff M. Phillips, and David P. Woodruff. Frequent directions: Simple and deterministic matrix sketching. STAM Journal on Computing, 45(6):1762-11792 2016. doi: 10.1137/ 15M1009718. URL httpt:ttt//ooo httt:p:://oi..or/100.1111 1ttp:/00..117//11000077777
Elad Hazan. Introduction to online convex optimization. CoRR, abs/ 19tp//999000007, 2019. URL http://arxiv..
org/ hot/abbs/1909. 05207
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016. doi: 10.1109/CVPFR CCVp//21PR..22..0
Wei Hu, L. Xiao, Ben Adlam, and Jeffrey Pennington. The Surprising Simplicity of the Early- Time Learning Dynamics of Neural Networks. ArXiv, 2020.
Richard Hughoy and Anders Krogh. Hidden markov models for sequence analysis: extension and analysis of the basic method. Computer applications in the biosciences: : CABIOS, 12 2:95 107, 1996.
Niall P. Hurley and Scott T. Rickard. Comparing measures of sparsity. IEEE Transactions Information Theory, 55:4723-47411 2008.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 448-456, Tille, France, 07 09 Jul 2015. PMIR. URL http://cccccccconooooppsppsss/ss
mwl.iinnnmmm....mmmmmmmm..m.m. mlr Vtress//37//offe55...
14

---

Published in Transactions on Machine Learning Research (12/2023)
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and Generaliza- tion in Neural Networks. arxi 1806. 07572 fes, math, stat|, June 2018. URL http://arxiv. org/ /bs// 1806.07572 arXiv: 1806.07572.
Devon Jarvis, Richard Klein, Beujamin Rosman, and Andrew M Saxe. On the specialization of neural modules. In The Eleventh International Conference on Learning Representations, 2023. URL https:
Jeevesh Juneja, Rachit Bansal, Kyunghyun Cho, João Sedoc, and Naomi Saphra. Linear connectivity reveals generalization strategies. In The Bleventh International Conference on Learning Regresenlalions, 2023.
URL https://openneeviez wwtp:/ooottdddddoooonoooo
Dan Jurafsky and James H. Martin. Speech and language processing introduction to natural lan- guage processing, computational linguistics, and speech recognition. In Prentice Hall series in artificial intelligence, 2023. URL htttp:://eett stanford. wtt.:ee.sttafor..ddd/aaaaaaaaaaa edtu-jjraafsyy/slp33/ pdf
Leonid V Kantorovich, Mathematical methods of organizing and planning production. Management science, 6(4), 1939.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling Laws for Neural Language Models. arXi::2001.0631 fes, slatj, January 2020. URL
08361
arXiv: 2001.08361.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On batch training for deep learning: Generalization gap and sharp minima. In International Con- ference on Learning Representations, 2017. URL htttps://openreviex http://ee///oooo///ddooooo/la
Bomin Kim, Kevin H. Lee, Lingzhou Xue, and Xiaoyue Niu. A review of dynamic network models with latent variables. Statistics surveys, 12:105 135, 2017.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University of Toronto, 2009. URL
Hao Li, Zheng Xu, Gavin Taylor, and Tom Goldstein. Visualizing the loss landscape of neural nets, 2018. URL https://openreriew..ne het.://eee....t.ooooroo/enaaaaaa
Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data, In The
Eleventh International Conference on Leurning Representations, 2023. URL httpp://openreviee net/
2orrm??d-2DiHo I
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the gencralization benefit of normalization layers: Sharpness reduction. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advonces in Neural Information Processing Systems, 2022. URL https://oppeereeiee net/forummt/
1d=xp5V0BxTxZ
Pranava Madhyastha and Rishabh Jain. On model stability as a function of random scod. In Proceedings of the 23rd Conference on Computational Natural Language Learning (CoNLL), pp. 929-939, Hong Kong, China, November 2019. Association for Computational Linguisties. doi: 10. 10.18653/v /K/1//007.. URL httpp:///cclaattaooggg..rog kttp://1000.////1111019110
Karttikeya Mangalam and Vinay Uday Prabhu. Do deep neural networks learn shallow learnable examples first? ICML 2019 Workshop Identifying and Understanding Deep Learning Phenomena, 2019. URL
https://openreviie http:///pgn...eee..eeeeo.oortraa meet/forumniddHHatttvaa
Owen Marschall, Kyunghyun Cho, and Cristina Savin. A unified framework of online learning algorithms for training recurrent neural networks. J. Mach. Learn. Res., 21(1), jan 2020. ISSN 1532-4435.
Romit Maulik and Gianmarco Mengaldo. Pyparsvd: A streaming, distributed and randomized singular- valne-decomposition library. In 2021 7th International Workshop on Data Analysis and Reduction for Big Scientifie Data (DRBSD-7) pp. 19 25, 2021. doi: 10.1109/RRBSD7554563.... 2021 00007..
15

---

Published in Transactions on Machine Learning Research (12/2023)
R. Thomas McCoy, Junghyun Min, and Tal Linzen. BERTs of feather do not generalize together: Large variability in generalization across models with similar test set performance. In Proceedings of the Third BlackborNLP Workshop on Analyzing and Interpreting Neural Networks for NLP, pp. 217 227, Online, November 2020. Association for Computational Linguistics. doi: 10.18653 (0.18653/v1/2022 www.bbbblbbboo--... 1.21. URL
httpps://aclantholoogg org/
blackboxnlp-1122 orr.2000..laackooppp-111111
William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking competition of sparse and dense subnetworks, 2023.
Prectum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang, Benjamin L. Edelman, Fred Zhang, and Boaz Barak. SGD on Neural Networks Learns Functions of Increasing Complexity. arXiv:1905. 11604 fes, 11604 arXiv: 1905.11604.
stat), May 2019. URL
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Sleinhardt. Progress measures for grokking via mechanistic interpretability. In The Fleventh International Conference on Learning Repre- sentations, 2023. URL http:://ppperrrreieeeeeettttoore
Catherine Olsson, Nelson Elhage, Neol Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Halfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah, In-context learning and induction heads, 2022.
David Picard. Torch.manual seed(3407) is all you need: On the influence of random seeds in deep learning architectures for computer vision, 2023.
Alethea Power, Yuri Burda, Harrison Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datascts. CoRR, Abs/2201.02177 2022. URL
Maithra Raghu, Justin Gilmer, Jason Yosinski, and Jascha Sohl-Dickstein. SVCCA: Singular Voctor Canon- ical Correlation Analysis for Deep Learning Dynamics and Interpretability. arXiv: (ll:1706.0500 05806 fes, slat], June 2017. URL http://arxivv org
arXiv: 1706.05806.
Yasaman Razeghi, Robert L. Logan IV, Matt Gardner, and Sameer Singh. Impact of Pretraining Term Frequencics on Few-Shot Numerical Reasoning. In Findings of the Association for Computational Lin- guistics: EMNLP 2022, pp. 840 854, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL
org/2022 wwtp:///...ininnnee..o.e..o pendding--mnlp....
Audrey Repetti, Mai Quyen Pham, Laurent Duval, Emilie Chouzenoux, and Jean-Christophe Pesquet. Fuclid in taxicab: Sparse blind deconvolution with smoothed regularization. IEEE signal processing letters, 22 (5):539 543, 2014.
Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normalization help optimization? In Proceedings of the 32nd International Conference on Neural Information Processing Systems, N8PS'18, pp. 2488 2498, Red Hook, NY, USA, 2018. Curran Associates Inc.
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116(23):11533)) 11537-1166, Junc 2019. ISSN 0027-8424, 1091-6490. doi: 10.1073/ (..1073/pnns.18202261 16. URL http//ttt////ww/ pnas org/ content/ 116/23/11537
Publisher: National Academy of Sciences Section: PNAS Plus.
Gideon Schwarz. Estimating the dimension of model. Anmals of Statistics, 6:461-400 4644 1978.
Thibault Sellam, Steve Yadlowsky, lan Tenney, Jason Wei, Naomi Saphra, Alexander D' A'our, Tal Linzen, Jasmijn Bastings, Tulia Raluca Ture, Jacob Eisenstein, Dipanjan Das, and Ellic Pavlick. The multiBRTTs: BERT reproductions for robustness analysis. In International Conference on Learning Representations, 2022. URL Utt.aateeee//ggaaereeeeeeeeeerer
16

---

Published in Transactions on Machine Learning Research (12/2023)
Patrice Y. Simard, Dave Steinkraus, and John Platt. Best practices for convolutional neural networks applied to visual document analysis. In Seventh International Conference on Document Analysis and Recognition, 2003. Proceedings., pp. 958 963, 2003. doi: 1AA1109/ICDAR.200...2.
Samuel L Smith, Benoit Dherin, David Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. In International Conference on Learning Representations, 2021. URL
https://openreviers httpp///pppppeeeeeee.eeeerrrrrrr
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R. Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonso, Agnieszka Kluska, Aitor Lewkowyez, Akshat Agarwal, Alethea Power, Alex Ray, Alex Warstadt, Alexander W. Kocurek, Ali Safaya, Ali Tazarv, Alice Xiang, Alicia Parrish, Allen Nie, Aman Hussain, Amanda Askell, Amanda Dsouza, Ambrose Slone, Ameet Rahane, Anantharaman S. Tyer, Anders Andreassen, Andrea Madotto, Andrea Santilli, Andreas Stuhlmüller, Andrew Dai, Andrew La, Andrew Lampinen, Andy Zou, Angela Jiang, Angelica Chen, Anh Vuong, Animesh Gupla, Anna Gottardi, Antonio Norelli, Anu Venkatesh, Arash Gholamidavoodi, Arfa Tabassum, Arul Menezes, Arun Kirubarajan, Asher Mullokandov, Ashish Sabharwal, Austin Herrick, Avia Efrat, Aykut Erdem, Ayla Karakas, B. Ryan Roberts, Bao Sheng Loo, Barret Zoph, Bartdomiej Bojanowski, Batuhan Özyurt, Behnam Hedayatnia, Behnam Neyshabur, Ben- jamin Inden, Benno Stein, Berk Ekmekei, Bill Yuchen Lin, Blake Howald, Cameron Diao, Cameron Dour, Catherine Stinson, Cedrick Argueta, César Ferri Ramírez, Chandan Singh, Charles Rathkopf, Chenlin Meng, Chitta Baral, Chiyu Wu, Chris Callison-Burch, Chris Waites, Christian Voigt, Christopher D. Man- ning, Christopher Polis, Cindy Ramirez, Clara E. Rivera, Clemencia Siro, Colin Raffel, Courtney Ashcralt, Cristina Garbacea, Damien Sileo, Dan Garrette, Dan Hendrycks, Dan Kilman, Dan Roth, Daniel Free- man, Danicl Khashabi, Daniel Levy, Daniel Moseguí González, Danielle Perszyk, Danny Hernandez, Danqi Chen, Daphne Ippolito, Dar Gilboa, David Dohan, David Drakard, David Jurgens, Debajyoti Datta, Deep Ganguli, Denis Emelin, Denis Kleyko, Deniz Yuret, Derek Chen, Derek Tam, Dieuwke Hupkes, Diganta Misra, Dilyar Buzan, Dimitri Coelho Mollo, Diyi Yang, Ho Lee, Ekaterina Shutova, Ekin Dogus Cubuk, Elad Segal, Eleanor Hagerman, Elizabeth Barnes, Elizabeth Donoway, Ellic Pavlick, Emanuele Rodola, Emma Lam, Eric Chu, Eric Tang, Erkut Erdem, Ernie Chang, Ethan A. Chi, Ethan Dyer, Ethan Jerzak, Ethan Kim, Eunice Engefu Manyasi, Evgenii Zheltonozhskii, Fanyue Xia, Fatemeh Siar, Fernando Martínez-Pluni Plumed, Francesca Happé, Francois Chollet, Frieda Rong, Gaurav Mishra, Genta In- dra Winata, Gerard de Melo, Germán Kruszewski, Giambattista Parascandolo, Giorgio Mariani, Gloria Wang, Gonzalo Jaimovitch-López, Gregor Betz, Guy Gur-Ari, Hana Galijasevic, Hannah Kim, Hamnah Rashkin, Hannanch Hajishirzi, Harsh Mchta, Hayden Bogar, Henry Shevlin, Hinrich Schütze, Hiromu Yakura, Hongming Zhang, Hugh Mee Wong, Ian Ng, Isaac Noble, Jaap Jumelet, Jack Geissinger, Jack- son Kernion, Jacob Hilton, Jaehoon Lee, Jaime Fernández Fisac, James B. Simon, James Koppel, James Zheng, James Zou, Jan Kocoń, Jana Thompson, Jared Kaplan, Jarema Radom, Jascha Sohl-Dickstein, Ja- son Phang, Jason Wei, Jason Yosinski, Jekaterina Novikova, Jelle Bosscher, Jemnifer Marsh, Jeremy Kim, Jeroen Taal, Jesse Engel, Jesujoba Alabi, Jiacheng Xu, Jiaming Song, Jillian Tang, Joan Waweru, John Burden, John Miller, John U. Balis, Jonathan Berant, Jörg Frohberg, Jos Rozen, Jose Keenandez-Orallo, Joseph Boudeman, Joseph Jones, Joshua B. Tenenbaum, Joshua S. Rule, Joyce Chua, Kamil Kanclerz, Karen Livescu, Karl Krauth, Karthik Gopalakrishnan, Katerina Ignatyeva, Katja Markert, Kaustubh D. Dhole, Kevin Gimpel, Kevin Omondi, Kory Mathewson, Kristen Chiafullo, Ksenia Shkaruta, Kumar Shridhar, Kyle McDonell, Kyle Richardson, Laria Reynolds, Leo Gao, Li Zhang, Liam Dugan, Lianhui Qin, Lidia Contreras-Ochandoo, Louis Mooni--Pillippe Morency, Luca Moschella, Lucas Lam, Lucy Noble, Ludwig Schmidt, Luheng He, Luis Oliveros Colón, Luke Metz, Lütfii Kerem Senel, Maarten Bosma, Maarten Sap, Maartje ter Hoeve, Maheen Farooqi, Manaal Faruqui, Mantas Mazeika, Marco Baturan, Marco Marelli, Marco Maru, Maria Jose Ramírez Quintana, Marie Tolkichn, Mario Giulianelli, Martha Lewis, Martin Potthast, Matthew L. Leavitt, Matthias Hagen, Mátyás Schubert, Medina Orduna Baitemirova, Melody Arnaud, Mclvin McElrath, Michael A. Yee, Michacl Cohen, Michael Gu, Michael Ivanitskiy, Michael Starritt, Michael Strube, Michal Swędrowski, Michele Bevilacqua, Michihiro Yasunaga, Mihir Kale, Mike Cain, Mimee Xu, Mirac Suzgun, Mo Tiwari, Mohit Bansal, Moin Aminnaseri, Mor Geva, Mozhdeh Gheini, Mukund Varma T, Nanyun Peng, Nathan Chi, Naycon Lec, Neta Gur-Ari Krakover, Nicholas Cameron, Nicholas Roberts, Nick Doiron, Nikita Nangia, Niklas Deckers, Niklas Muennighoff, Nitish Shirish Keskar, Niveditha S. Tyer, Noah Constant, Noah Fiedel, Nuan Wen, Oliver Zhang, Omar Agha, Omar Elbaghdadi,
17

---

Published in Transactions on Machine Learning Research (12/2023)
Omer Levy, Owain Evans, Pablo Antonio Moreno Casares, Parth Doshi, Pascale Fung, Paul Pu Liang, Paul Vicol, Pegah Alipoormolabahii, Peiyuan Liao, Percy Liang, Peter Chang, Peter Eckersley, Phu Mon Htut, Pinyu Hwang, Piotr Milkowski, Piyush Patil, Pouya Pezeshkpour, Priti Oli, Qiaozhu Mei, Qing Lyu, Qinlang Chen, Rabin Banjade, Rachel Etta Rudolph, Raefer Gabriel, Rahel Habacker, Ramón Risco Del- gado, Raphaäl Millière, Rhythm Garg, Richard Barnes, Rif A. Saurous, Riku Arakawa, Robbe Raymaekers, Robert Frank, Rohan Sikand, Roman Novak, Roman Sitelew, Ronan LeBras, Rosanne Liu, Rowan Jacobs, Rui Zhang, Ruslan Salakhutdinov, Ryan Chi, Ryan Lee, Ryan Stovall, Ryan Teehan, Rylan Yang, Sahib Singh, Saif M. Mohammad, Sajant Anand, Sam Dillavou, Sam Shleifer, Sam Wiseman, Samuel Gruetter, Samuel R. Bowman, Samuel S. Schoenholz, Sanghyun Han, Sanjeev Kwatra, Sarah A. Rous, Sarik Ghaz- arian, Sayan Ghosh, Sean Casey, Sebastian Bischoff, Sebastian Gehrmann, Sebastian Schuster, Sepideh Sadeghi, Shadi Hamdan, Sharon Zhou, Shashank Srivastava, Sherry Shi, Shikhar Singh, Shima Asaadi, Shixiang Shane Gu, Shubh Pachchigar, Shubham Toshniwal, Shyam Upadhyay, Shyamolima, Debnath, Siamak Shakeri, Simon Thormeyer, Simone Melzi, Siva Reddy, Sncha Priscilla Makini, Soo-Hwan Lec, Spencer Torene, Sriharsha Hatwar, Stanislas Dehaene, Stefan Divic, Stefano Ermon, Stella Biderman, Stephanie Lin, Stephen Prasad, Steven T. Piantadosi, Stuart M. Shieber, Summer Misherghi, Svetlana Kiritchenko, Swaroop Mishra, Tal Linzen, Tal Schuster, Tao Li, Tao Yu, Tariq Ali, Tatsu Hashimoto, Te-Lin Wu, Théo Desbordes, Theodore Rothschild, Thomas Phan, Tianle Wang, Tiberius Nkinyili, Timo Schick, Timofei Kornev, Timothy Telleen-Lawton, Titus Tunduny, Tobias Gerstenberg, Trenton Chang, Trishala Neeraj, Tushar Khot, Tyler Shultz, Uri Shaham, Vedant Misra, Vera Demberg, Victoria Nyamai, Vikas Raunak, Vinay Ramasesh, Vinay Uday Prabhu, Vishakh Padmakumar, Vivek Srikumar, William Fedus, William Saunders, William Zhang, Wout Vossen, Xiang Ren, Xiaoyu Tong, Xinran Zhao, Xinyi Wu, Xudong Shen, Yadollah Yaghoobzadeh, Yair Lakretz, Yangqiu Song, Yasaman Bahri, Yejin Choi, Yichi Yang, Yiding Hao, Yifu Chen, Yonatan Belinkov, Yu Hou, Yufang Hou, Yuntao Bai, Zachary Seid, Zhuoye Zhao, Zijian Wang, Zijie J. Wang, Zirui Wang, and Ziyi Wu. Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models, June 2022. URL http: http:://aaxi..oorr/bb//222.... 04615. Number: arXiv:2206.04p0 arXiv:2206.04615 [cs, stat).
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon, 2022.
Leonid Nisonovich Vaserstein. Markov processes over denumerable products of spaces, describing large systems of automata. Problemy Peredachi Informatsii, 5(3):64 72, 1969.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Extracting automata from recurrent neural networks using queries and counterexamples. In Jennifer Dy and Andreas Kranse (ods.), Proceedings of the 35th Inter- national Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 5247-5256. PMLR, 10-15 Jul 2018. URL
Gail Weiss, Yoav Goldberg, and Eran Yahav. Learning deterministic weighted automata with queries and counterexamples. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL GBT5677666601Vee5trrdddaeeebe--2
neurips wttppp...c.ppp....pppp filess/paper/219//iiii
pdf.
R.J. Williams. Training recurrent networks using the extended kalman filter. In (Proceedings 1992] IJCNN International Joint Conference on Neural Networks, volume 4, pp. 241-246 vol.4, 1992. doi: 10.1109/ 1ICNN.199..22733.
Y. Wu, L. Liu, J. Bae, K. Chow, A. lyengar, C. Pu, W. Wei, L. Yu, and Q Zhang. Demystifying learning rate policies for high accuracy training of deep neural networks. In 2019 IEEE International Confer- ence on Big Data (Big Data), pp. 1971-1980, Los Alamitos, CA, USA, dec 2019. IEEE Computer Soci- cty. doi: 10. BiggData47090.2019..00 BigData47090. 2019. 9006104
httpss////// weweccopptersocceet....o.. org.yy....g..... 1109/
Greg Yang. Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter
18

---

Published in Transactions on Machine Learning Research (12/2023)
transfer. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (ods.), Advonces in Neural Information. Processing Systems, 2021. URL. htttps://operreieet http:///tn////eeeee.eeeeee//eonn .et./forrurricccBokkkkU0022
Wenjian Yu, Yu Gu, and Jian Li. Single-pass pea of large high high-dimensional data. In Proceedings of the 26th International Joint Conference Artificial Intelligence, IJCAT17, pp. 3350 3350-3356. AAAT Press, 2017. ISBN 9780999241103.
19

---

Published in Transactions on Machine Learning Research (12/2023)
A Derivation
We use the log posterior because it has a simplified form for Gaussians. The log posterior is:
p(s, k((z1e) p(žcl0; = k)p(s; = k; 21:e-1)
p(ε11)
logp(s k((z+)) = logp(E/10> = k) + togp(0t = k, m1--1) log p( (11a)
We take the derivative of these three terms separately:
i log p( k)
Ex'(px (Mx - 2t)
8z.
i log p(s = k, E1:1-1) =0
0E,
log p( (1:1)
1
8p(31:)
0%
p((1a)
tz
1
[Enea p( = -h)p(s: = h, 81.e-1)
p(ā1a)
1
Op(2r|5c = h)
p(s, h, 212 1)
p((1a)
t2
hEn
1
2 Plgg(Ea h)
p(s, h, 2 1))(21)
h))
p(
03.
hEB
p(st
Ō1 )p( Ee Stt 11) Sxm (14 1 - 21)
p(
hER
Σ
p(st 11(311)
(Jun - 21)
kER
So, for timestep and hidden state h;
2 log p(s: k(
Di 1 (px - 22) Ep(ss
p(s: = h1 (un 2t)
0z.
hen
p(s: = h h/)1:) can be efficiently computed using the forward algorithm. chapter for reference.
See Jurafsky & Martin (2023)
20

---

Published in Transactions on Machine Learning Research (12/2023)
B
Metrics
The chart in this section lists the 14 statistics we computed for each model checkpoint. We use these statistics to capture either 1) how the neural network weights weights are dispersed in space or the 2) properties of the function computed by a layer. For example, the L2 norm measures dispersion because it describes how far away the weights are from the origin. The spectral norm helps capture the function computed by a neural network because it describes the maximum amount that a vector might change as it passes through layer.
Of course, 1) and 2) are related, and thus the statistics we compute are also related; the L2 matrix norm upper bounds spectral norm. Our philosophy (and rocommendation) is to choose variety of metrics when modeling training dynamics to allow for interactions between metrics.
For metrics that become infeasible to compute during training at large model sizes, we recommend using streaming algorithms, matrix sketching algorithms, or other approximations such as random projections to make computation more efficient. For example, singular values can be computed using streaming algorithmss
MMaulik & Mengaldo 2021 Yu et al. 2017) or on a matrix sketch of reduced size (Ghashami al. 2016).
Name 1) Li
Description
The Li-orr, averaged over matrices. K ||l|| = k 2i=1 | where K is the number of weight matrices in the neural network. We average over matrices so that models with different depths are comparable.
1) L2 Ia
The L2-norm, averaged over matrices. = = v w
1) L2
Measures the sparsity of the weights (Repetti et al. 2014]. Zg
EE LT
1
which is the metric 4 averaged over the K weight
matrices. Lower is more sparse. For example, a one-hot vector is fully sparse and has code sparsity of 1. See [Hurley & Rickard (2008) for a discussion on measures of sparsity.
1) n(w)
Sample mean of weight. + Ei' W/, where N is the number of parameters in the network.
1) median(w)
Median of the weights, treated set concatenated together.
1) (w)
Sample variance of weights without Bessel's correction.
N (w:-w)²
1) a(b)
Sample mean of the biases. We treat the biases separately be-
cause they have a distinct interpretation from the weights.
1) median(b)
Median of the biases, treated as a set concatenated together. Sample variance of biases without Bessel's correction.
1) (bb)
2) trace
The average trace over K weight matrices. R Σ tr(W)),
where W₂ is the kth weight matrix.
2) Amax
The average spectral norm. ΣK W|
2) trace
Average trace over spectral norm. K E 1Wal
((()
Amux
2) P(A)
Average singular value over all matrices.
2) (A)
Sample variance of singular values over all matrices.
Table 3: A glossary of metrics. The "Name" column contains how the metrics appear in the text. The label
1) means the statistic intends to capture how the neural network weights weights are dispersed in space.
The label 2) means the statistic intends capture properties of the function computed by layer.
21

---

Published in Transactions on Machine Learning Research (12/2023)
C Baselines
We compare the performance of the full HMM, trained on all 14 statistics discussed in Appendix with two baselines:
1. K-means clustering, which learns a discrete latent space similar to the HMM but does not capture temporal structure.
HMM-1, which the HMM trained to model only the L2 norm. We chose this baseline because we found L2 norm be of the most important metrics throughout all settings (see Sections 33. and 3.2). and the L2 norm has also been noted metric predictive of model qualities in prior works
3.1
Liu et al. 2023 Nanda et al. 2023 Thilak et al. 2022
For each setting, we perform model selection and choose the optimal number of components according BIC. Below, list the number of components in the best model, along with its BIC. We find that K-means and HMM-1 tend to use consistently more components compared to the base HMM. We consider this undesirable because more components dilutes the interpretation of each individual component. In particular, K-means tends to use the maximum number of clusters we allowed to cluster the given sequence.
(NB: BICs are not comparable across models; we provide them for comparison in case the reader trains a model of the same class.)
Dataset
k-means
HMM-1
HMM
Components
BIC
Components
BIC
Components
BIC
Modular
16
103700
-1077
-5724
Modular, stabilized
10
3864
166.0
759.9
CIFAR-100
16
19360
10
-1851
5
-124400
CIFAR-100, destabilized
16
23210
15
-1432
5
-59080
Sparse parities
16
23660
13
-2965
-49530
MNIST
16
3244
-1064
-101200
Table 4: Dataset Scores
D Training Hyperparameters
For the MultiBERTs training.
(Sellam al. 2022
use the open-source training chockpoints without any additional
Hyperparameter
Value
Learning Rate
le-1
Batch Size
32
Training data size (randomly generated)
1000
Test data (randomly generated)
100
Architecture
Multilayer perceptron
Number of hidden layers
Model Hidden Size
128
Weight Decay
0.01
Seed Optimizer
through 40
SGD
Table Sparse parities, replicating Barak al. 2022
22

---

Published in Transactions on Machine Learning Research (12/2023)
Hyperparameter
Value
Learning Rate
1e-3
Batch Size
2048
Training data size
3831 (30% of all possible samples)
Architecture
Transformer, no layer normalization
Transformer Number Layers
1
Transformer Number of Heads Model Hidden Size
4
128
Model Head Size
32
Weight Decay
1.0
Seed
0 through 40 AdamW
Optimizer
Table 6: Modular addition, replicating Nanda et al. (2023). To stabilize training (Figure 5 we reduced the batch size from 2048 to 256 and added back layer normalization.
Hyperparameter
Value
Learning Rate
1e-3
Batch Size
256
Training data size
50000 (splits downloaded from PyTorch)
Architecture
ResNet18
Weight Decay
1.0
Seed Optimizer Data preprocessing
through 40
AdamW
Random crop, random horizontal flip, and normalization
Table 7: CIFAR-100. To destabilize training (Figure
we removed batch normalization and residual
connections.
Hyperparameter
Value
Learning Rate
le-3
Batch Size
256
Training data size
60000 (splits downloaded from PyTorch)
Architecture
MT.P
Number of hidden layers
Hidden size
800
Weight Decay Seed
1.0
through 40
Optimizer
AdamW
Data preprocessing
Flatten vector
Table 8: MNIST. MLP hyperparameters based on Simard et al. (2003
23

---

Published in Transactions on Machine Learning Research (12/2023)
E
Language Modeling: MultiBERTs
3.00
1e-5
0.75
Pretraining loss
10
2.75
Average weight
Average weight
8
0.1
0.25
2.50
0.75
2.25
0.25
2.00
4
0.1
0.25
1.75
0.75
0.0
0.5
1.0
1.5
2.0
Step
1e6
3.00
1e-5
0.75
2.75
Average weight
0.1
0.25
2.50
LoSS
0.75
0,8
1.0
2.25
Q
0.25
2.00
0.1
0.25
1.75
0.75
0.0
0.5
1.0 Step
1.5
2.0
1e6
Edge
Top 3 important feature changes, by z-score
Transition frequency
2 +
median w) 11.69, s(w) 11.70, Amar 1 1.14
2/3 5
2 3
Amar t 1.11, median (w) 41.33, ((() 11.30
5/5
Figure 7: MultBEETs. The average weight N ΣN w, initially decreases in of the five MultiBERT runs (2 -+ 0) and increases in the other three (2 - 3). However, all runs eventually converge to roughly the same average weight. The HMM represents this difference in runs the states 0 and 3. Critically, this difference is imperceptible from the pretraining loss.
To study variation in masked language model training, we use the five released training trajectories from the MuluBERTSs (Sellam et al. 2022), which are replications of the original BERT model (Devlin et al. 2019 trained under different random seeds. MultiBERT's differs from the other settings we consider because its training occurs over the course a single epoch, rather than over multiple epochs.
The most notable feature of the MultiBERTs training map is the fork state 2. MultiBERT's models all converge around 3.7 10-3, but the paths that the five models take to get there be clustered into two different trajectorics. For the path including (2 0), the average weight increases during states 2 and vero and then decreases during state 4, while the opposite is true for paths including (2 + 3). Understanding this difference between MultiBERTs models could be fruitful area for future work. Critically, this difference in model internals is imperceptible from the pretraining loss, which decreases at roughly the same rate for all five MultiBERTs runs. However, the MultiBERTs exhibit significant variation in transfer learning performance and gender bias Sellam et al. (2022). so these paths may indicate differences in behavior under specific distribution shifts and settings.
The average weights of the
24

---

Published in Transactions on Machine Learning Research (12/2023)
F
Algorithmic Data: Sparse Parities
1.5
Train
Validation
0.978
1.0
0.001
0,022
0.024
0.044
0,026
0.5
0.974
0.019
0.0
1.0
50
200
250
Epoch
1.25
Train
1.00
Validation
0.955
0.978
0.976
0
0.001
0.022
0.024
Loss
0.75
0.50
0.974
U.04
0.025
0.25
0.00
1.0
50
100
150
200
250
300
Epoch
Edge
Top important feature changes, by z-score
Transition frequency
L1 10.61, L2 10.11, L1 J0.32
39 40
- 5
L2 10.19, ((w) 4 0.74, (A( 1.20
1 1/ 40
Figure 8: Sparse paritics. Faster generalization in sparse paritics occurs with an carly decrease in the L2 norm. The norm ratio Li is metrie for dispersion, and il decreases as the vector becomes more sparse. For example, a one-hot vector completely sparse and is the minimum of Lg
L2
25

---

Published in Transactions on Machine Learning Research (12/2023)
G
Image Classification: MNIST
1.5
Train
0.942
0.978
0.981
Validation
0.058
0.022
3
4
1.0
Loss
0.019
0.5
0.025
0.021
5
2
0.0
1.0
0.975
0.979
200
400
600
800
1000
Step
Edge
Top 3 important feature changes, by z-score
3 + 4
L2 10.62, (w) 10.58, L1 10.61
0 + 2
(w) 10.70, L2 10.69, L1 10.70
5 3-1
L2 10.46, a (w) 10.50, L1 10.48
Figure 9: MNIST. All 40 training runs collected from MNIST follow the same path, although individual runs can spend slightly different amounts of time in each state. As shown by the training map and accom- panying annotations in the table, the training dynamics of MNIST are similar between states.
26

---

Published in Transactions on Machine Learning Research (12/2023)
H
Model Selection Curves
-100000
1000
3000
Log density
0
Log density
500
AIC/BIC
2000
AIC/BIC
-20000
50000
1000
-40000
-500
0
-1000
1000
8
2
8
Number of components
Number components
(a) Modular addition.
components.
(b) Modular addition, stabilized. 3 components.
40000
density
20000
50000
Log density
60000
AIC/BIC
-100000
AIC/BIC
0
50000
-20000
120000
-40000
50000
Number of components
Number of components
(c) CIFAR-100. components.
(d) CIFAR-100, destabilized. components.
30000
0
Log density
1000
-200
800
AIC/BIC
Log density
20000
-20000
AIC/BIC
-400
600
10000
40000
4 Number of components
2
Number of components
(e) MultiBRRTs. 5 components.
(f) Sparse parities. 6 components.
55000
80000
Log density
50000
AIC/BIC
900000
45000
100000
40000
2
8
Number of components
(x) MNIST. components.
Figure 10: Model selection curves for choosing the number of hidden states in the HMM, We choose the model with minimum BIC in all cases.
27

---

Published in Transactions on Machine Learning Research (12/2023)
Convergence Time Histograms
20
40
15
30
Count
Count
10
20
10
0
2000
4000
6000
8000
10000
0
2000
4000
6000
8000
10000
Convergence epoch
Convergence epoch
(a) Modular addition. Threshold 0.9.
(b) Modular addition, stabilized. Threshold 0.9.
40
20
30
Count
Count
20
10
10
0
2000
4000
6000
8000 10000 12000
0
2000
4000
6000
8000
10000 12000
Convergence step
Convergence step
(c) CIFAR-100. Threshold: 0.6.
(d) CIFAR-100, destabilized. Threshold: 0.4.
15
10
Count
Count
10
5
950
50
75
100
125
150
175
200
400
450
500
550
Convergence epoch
Convergence step
(e) Sparse parilies. Threshold: 0.9.
(I) MNIST. Threshokd: 0.97.
Figure 11: Visualization of convergence times. Convergence time here is defined as the first time model crosses some threshold of evaluation accuracy, and choose the threshold to be value slightly less than the performance that the final model achieves. For example, our fully trained models generally achieve perfect accuracy on modular addition, so we choose a threshold of 0.9.
28


Reviews:
[{"review_id": "Wo3uBiS7jn", "reviewer": "Reviewer_JCA7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I enjoyed the application of a probabilistic model to the training process. After all, researchers are subject to high-dimensional dynamic systems all the time during the attempts to make a new model work. Still, curricula are very short on the subject and knowledge about it is passed on as lore from the elder's mouth to the disciple's ear. Hence, applying the fields own tools to the fiels own problems seems like a good idea.\n\nThe paper contains lots of experimental evidence and offers a systematic statistical approach to most questions that pop up. Nevertheless, I am not convinced by the overall \"package\": it left me mostly confused. The first sentence of the conclusion sums it up nicely: \"We make several main contributions.\" I must say, the authors claim many different contributions in a rather unstructured and ad hoc way. A lack of focus emerges and that left me wondering what exactly I learned from the text. The feeling prevailed even after I spent considerable time on the manuscript.\n\nFurther, I am not convinced that many of the results are unknown. For example, we know that certain architectural tweaks (e.g. batch normalization) improve learning–the papers introducing these techniques generally do a good job at providing evidence. Something that I personally did not consider before was whether the sensitivity to the initialization implies something about generalization. But, to be honest, the evidence for this claim is either rather weak or not presented well/extensively in the manuscript.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors describe a series of experiments where Gaussian HMMs are applied to trajectories of neural network training. The trajectories itself are made of features of the weights of neural networks. Subsequently, certain structures emerge, which the authors trace back to different favourable or less favourable training runs and architectural choices. The authors then identify \"detour\" states, which they deem hurtful for convergence time.\nAdditionally, they employ various statistical experiments, such as predicting convergence time or the variety of training runs through a dissimilarity measure induced by the initialisation of the networks' weights.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The choice of HMM needs to be justified further. I think the model is not adequate, as the unobserved dynamics do not follow discrete states. What makes an HMM better than, say, a switching LDS, apart from the authors (and readers!) possibly being more familiar with it? I might even go further and say \"detour states don't exist, it's better explained by a flat region\". An HMM implies a detour state is a region from which it escapes randomly, while a flat region would mean it needs to work its way through it. Does this have any implications for my day to day work? I don't know, but I wish this paper would tackle such questions.\nYet, since all models are wrong and some are useful, the HMM might still be a reasonable choice. However, this would need support. A section on whether an HMM even is an adequate model–something like, \"let's check if this is a good model for the trajectories\". For example, I would like to see a comparison of samples from the model and the empirical data? \n\nTo take a view point from the other side, let us look at Figure 2, where we have different learning curves. Of course, an HMM captures the differences, but so does a simple graphical illustration. Do we even need an HMM for analysis? Don't we already see all that is necessary by noting the plateaus during training? Isn't the message \"some training runs plateau, other's dont–check your hyper parameters and run more experiments to find the best set!\" sufficient? Does this approach bring anything to the table? I am not convinced at this stage!\n\nThe contributions and findings need to be worked out better. Previewing more of the machinery in the earlier sections would certainly help. \nOverall, I have problems finding a proper thread. For example, the trajectory similarity measure and the regression are not mentioned in the preview at the end of the introduction. They take up quite some space though, and somewhat come as a surprise during the first reading. \n\nThe figures are very large sometimes and make lots of page turning necessary. Further, the formatting of those–especially the fonts–is not consistent. \n\nProposition 1 seems like the definition of a derivative to me. I don't think putting this fact into a formal framework helps. The same goes for the definition of the detour state.\n\nFurther, I think the evidence needs to be presented better. I must repeat that I felt very confused at times and felt that I am overloaded with unstructured information. I wish I could provide more constructive feedback at this point with concrete measures.", "broader_impact_concerns": "I don't think the submission requires a broader impact statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xFhLKqmlXj", "reviewer": "Reviewer_5mBE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n- The paper is nicely presented, with well-chosen figures that clearly convey the concepts.\n- Experiments are carefully constructed.\n\n# Weaknesses/Questions\n\nMy main issue with this paper is there are no baseline methods for any experiments.\n\nHere are some section specific comments:\n\nSection 2.2\n- What does it mean for an edge to be \"unused by the HMM for all training trajectories\"?\n- Should the $\\arg\\max$ be over $i$ instead of $\\tilde{z}[i]$? The states in Figure 1 are labelled with integers.\n- Intuition is lacking as to why the difference in means represents a reasonable quantity to use as a label for the edges.\n\nSection 2.3\n- What is meant by \"probabilistic reasoning\"?\n\nSection 3\n- I can understand the motivation to leave out embedding matrices when computing metrics, but not layer norms. Isn't the technique part of the neural network architecture choice?\n\nAppendix A\n\nLemma\n- I don't understand why \n$p(\\tilde{z}\\_{1:t}) = \\sum\\_{i=1}^t \\alpha\\_i(s_i) = \\sum\\_{i=1}^t p(s_i, \\tilde{z}\\_{1:i})$. If the summation was to marginalize out $s_i$, then it would make sense. But it's not clear why this is true when summing over $i$ .\n\nProposition 1\n- What happened to $E$ when going from the 2nd to the 3rd equality? I understand $E$ won't affect optimizations that use the log likelihood as an objective, but that isn't the same as saying the 2nd and 3rd line are equal.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a study on the effects of randomness in neural network training.\nTo accomplish this, they fit an HMM where observations are vectors of hand-chosen metrics computed from the checkpoints of a neural network during training. The number of hidden states is a hyper-parameter.\n\nThe main contribution is the identification of so called \"detour states\", which, in their words are \"latent states associated with slower convergence\".", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical for acceptance.\n- The authors should provide baselines for all experiments, or justify why a baseline can not be provided.\n- Please make the derivations more obvious (verbose).", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "jgbOx7E2E7", "reviewer": "Reviewer_ZHnM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n1. Whether the training process of the neural network is robust to randomness or not is an important question to study but so far not too much analysis has been done. This paper proposes an interesting method using HMM to analyze the training dynamics of neural networks. The authors suggested running each model with more training seeds to control the variation in the training processes.\n\n2. This paper provides some new insight into the grokking phenomenon. The experiments showed that adding layer normalization, residual connections, and decreasing batch size may help us smooth the loss landscape and avoid the grokking phenomenon. \n\nWeakness:\n1. There is no explanation of the choices of different metrics in Appendix B. Why can these metrics interpret the features of the neural networks during training? For very large models, some of these metrics cannot be efficiently recorded during training. \n\n2. The method proposed in this paper using HMM may be computationally inefficient for large models. This method requires repeatedly training neural networks many times with different random seeds, which becomes expensive for analyzing large models in practice. Any comments for this on training HMM and extracting the training map?\n\nMinor and question:\n1. Before section 2.1, $d<<D$ should be $d\\ll D$; $f:\\mathbb{R}^D\\to \\mathbb{R}^d$ should be $f_i:\\mathbb{R}^D\\to \\mathbb{R}^d,$ $1\\le i\\le d$.\n2. Notation \\{$z_{1:T}$\\}$1^N$ on page 3 section 2.1 is confusing. I understand you consider $N$ random seeds and for each seed, you have one $z_{1:T}$.\n3. In conclusion, the authors claim that stability to randomness is linked to generalization. Can you explain this more from the experiments in sections 3.-3.33? For instance, in Figure 2, although we observe the grokking phenomenon, what would be the final generalization error for each case? Can we get similar and comparable test errors if we train the neural network sufficiently longer time? Additional explanations may be needed for these experiments.\n4. Here are some additional references that may be related to this work. [1,2] also studied the training dynamics of the neural networks using different metrics. [3] also studied the grokking phenomenon in sparse parity tasks. [4,5] are some other references for the effect of randomness in neural networks.\n\n[1] Deep learning versus kernel learning: An empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel.\n[2] Spectral evolution and invariance in linear-width neural networks.\n[3] A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks.\n[4] Effects of Random Seeds on the Accuracy of Convolutional Neural Networks.\n[5] torch.manual seed(3407) is all you need: On the influence of random seeds in deep learning architectures for computer", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studied the effect of randomness on the training dynamics of neural networks. The author used a sequence of metrics on the training trajectories and established a hidden Markov model (HMM) to analyze the variations from random seeds during different training processes. Specifically, the authors presented that modular arithmetic and sparse parity experiments have grokking phenomena and are affected by randomness, while the image classification experiments with batch normalization and residual connections are insensitive to random seeds. The authors proposed detour states in HMM and showed that detour states appear when the model is sensitive to randomness and otherwise disappears.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper confronts an important and difficult problem in deep learning practice with a new analysis approach using a state space model (HMM) representation of the learning dynamics. Any attempt to progress the understanding of learning dynamics should be very much welcomed so even though this paper is merely a step on the way, it still holds enough insights to warrant publication despite a bit of misgivings from the referee.

The paper is accepted as is. But please remember to remove any highlighting in the final version.
Decision: Accept (poster)