Training Dynamics
Implicit regularization and emergence in deep learning — loss landscapes of deep linear networks, lazy vs rich regimes, grokking, and dynamical mean field theory and training-time phase transitions.
By Guillaume Corlouer (Moirai)
What you’ll learn
- Understand the concepts of implicit regularization
- Know of key results about loss landscape of Deep Linear Networks (DLNs)
- Gain some intuition about non trivial loss landscape (DLNs) and their symmetries
- Know about the edge of stability phenomena
- Derive gradient flow equations in DLNs
- Explain NTK
- Understand the importance of symmetries and conserved quantities
- Know that under assumptions we can interpret gradient flow in DLNs as minimizing a free energy
- Derive the simplicity bias in simple DLNs example (saddle to saddle)
- Explain the effect of initialization on training regime in DLNs (Lazy vs Rich)
- Know that implicit biases of stochasticity depends on the training regime
- Discuss why we care about implicit regularization for AI safety
- Know that there is a puzzle: emergent capabilities
- Know about toy models and examples of emergence: silent alignment in DLNs, grokking as delayed generalization, induction heads, emergent misalignment
- Explain the concept of phase transitions
- Explain the idea of mean field theory
- Understand the idea of dynamical mean field theory at a high level as a theory of emergence during training
- Explain grokking as a lazy to rich transition during training using DMFT in a toy model
- Discuss the importance of emergence for AI safety: early detection of phase transition (early warning signals), deriving order parameters that help us interpret the model (and simplifying the dynamics)
Overview
This module discusses some key results of learning dynamics in deep neural networks, focusing on implicit regularization and emergence. The implicit regularization section examines how the training process, specifically stochastic gradient descent, can have an implicit bias towards simple solutions. The topics include loss landscape geometry, the edge of stability, simplicity bias, the neural tangent kernel (NTK), and the influence of initialization on training regimes (lazy vs. rich) in deep linear networks as an important toy model of deep learning. The emergence section investigates the puzzle of unexpected "emergent capabilities" (e.g., grokking, induction heads, silent alignment). It uses theoretical concepts like phase transitions to explain these phenomena, such as viewing grokking as a transition from the lazy to the rich regime. One application of this phase transition perspective is to understand and enable the early detection of emergence for AI safety.
Prerequisites
-
Basics of deep learning (SGD, network, loss function)
-
Transformer architectures and LLMs
-
Statistics
-
Linear algebra (singular value decomposition)
-
Ordinary differential equations
-
May help but not strictly required: Group actions, Stochastic differential equations, Statistical physics and statistical field theory
Content
Fast track
-
Read Andrew Saxe paper.
-
Read the emergence slides.
-
If you have more time, read the key readings.
Main content
The main content is in the github repository. Go through the slides, exercises, and solutions.
Learn more
Go through the readings in the github repo.