Learning to Predict Learning Dynamics
In the book Cybernetics, Norbert Wiener characterized the sciences in a continuum based on their modeling decisions. At one extreme, Newtonian mechanics makes predictions based on abstractions over the current system, such as mass and velocity. Towards the other extreme lies a field like meteorology, where it’s not clear, to this day, exactly what information is relevant about the state of a current system. Predictions are statistical, with error and noise.
Wiener labels these two modes of science as operating in Newtonian and Gibbsian time, respectively. In Newtonian time, system dynamics are predictable in both the forward and reverse directions. Given the positions, masses, and velocities of billiards, I can predict ball positions in both the future and past. Conversely, in Gibbsian time the predictive model is statistical and imperfect in both directions. We can’t forecast the future with certainty, nor exactly reconstruct the past. Each present observation forces us to revise estimates of past states and future predictions.
Our study of learning dynamics for language models can proceed in a Newtonian way, with simplified mathematical models—for examples, see work from @blake__bordelon or the ICML keynote from @ShamKakade6. But in this post, I’ll discuss an alternative direction I’ve been pursuing with @kchonyc that treats the prediction of future states in language model training as a machine learning problem. Instead of building a simplified model, we choose a powerful function approximator, train it, and test the model on the basis of generalization.
The purpose for doing this is that we can often build good predictors without having clean models of phenomena. Language models themselves are a great example: linguists built elegant models of language with automata, the Minimalist program, et cetera, and language models don’t use any of that; they’re big function approximators. We can apply this same philosophy to learning dynamics: build a model of learning dynamics that scales with more data.
A key decision behind learning learning dynamics is deciding how to represent the state of your neural network. During language model training, we have a perfect, bit-level description of every weight matrix, activation, and optimizer state. We also have losses on the training set, validation set, and various evaluations—gigabytes and gigabytes of information. What is relevant from all this? Here, I’ll separate this information into two parts, 1) language model outputs and 2) the parameters of the training procedure. For the first part, I’ll discuss losses, and for the second, neural network weights.
Losses
Scaling laws take as input model size, dataset size, or compute and predict the average validation loss, typically using a power law. This is a surprisingly small representation of training that remains both predictive and smooth over several orders of magnitude. Billions of parameters modeling trillions of token-level predictions get summarized by a few scalar inputs.
However, average loss becomes an inadequate representation when predicting downstream model behavior, which can have a variety of different scaling trends. When you average losses over many documents into one number, you lose a lot of relevant signal that could be used to predict downstream model behavior for different tasks.
In Neural Neural Scaling Laws, we avoid this bottleneck by directly giving token-level losses to a neural network. We process token-level losses using a one-dimensional CNN, which can learn which patterns in the loss distribution are useful, instead of taking a strict average. We find that this learned neural model (NeuNeu) makes better predictions than logistic regression, the current standard model for predicting downstream task performance from training loss.
Weights
The challenge with analyzing weights of a neural network is that there’s too many of them, so some compression of the weights is necessary. Our hypothesis in Latent State Models of Training Dynamics is that, despite the high dimensionality of neural network weights, their evolution is governed by a low-dimensional process. We were particularly interested in modeling discrete, phase-transition-like phenomena like grokking, where a model first memorizes the training distribution, then suddenly generalizes.
Our solution was to model the evolution of neural network weights using a state space model—in this case, a hidden Markov model. We first computed generic statistics over neural network checkpoints, such as norms, means and variances, then trained an HMM to learn a discrete latent state over these sequences of statistics.
This turns a very high-dimensional sequence of weights into a small state machine that describes the trajectory of training. Training runs can follow the same road to convergence, fork into different roads, or enter optional “detour states” that slow down generalization. These weight-derived latent states contain sufficient information to predict differences in generalization behavior for different training runs. Generalization appears with respect to the validation set, but it is preceded and predicted by changes within model weights.
Foundation models of learning dynamics
Collect a big dataset (of learning dynamics perhaps), find a workable input representation, train a model, and it usually generalizes. I think of my own work and others such as G.pt and Configuration-to-Performance Law as precursors to foundation models of learning dynamics. As more training trajectories get open-sourced, we can train models that predict how neural network training will progress.
Such models could be used as world models for training. We could perform hyperparameter search and data ablations by rolling forward a learning dynamics model without necessarily executing the training run. Another application is scientific regret minimization. When a black-box model reliably predicts a phenomenon, we can improve our mathematical theories against it by minimizing regret to the best black-box model in hindsight.
Physics operates in both Newtonian and Gibbsian regimes. I think the same will be true for our study of learning dynamics.