A central challenge in modern machine learning is to understand \emph{why} and \emph{how}
gradient-based training of neural networks leads to models that generalize well and learn
meaningful internal representations.
This thesis approaches this challenge by studying the training dynamics of first-order
methods on a sequence of models of increasing complexity, from linear predictors through
linear networks to attention-only transformers.
In each setting, the goal is to obtain a closed-form, end-to-end characterization of the
training dynamics and to identify the mechanisms by which gradient-based learning gives
rise to structured feature representations and what components of the architecture and algorithm are responsible for such features.
In Chapter 2, we study dynamics of momentum based methods where the features are fixed and consider the least-squares regression problem in the non-strongly-convex
setting and propose an algorithm that adapts Nesterov's accelerated gradient descent to the
stochastic setting.
The algorithm provably achieves the optimal statistical rate while simultaneously accelerating
the forgetting of the initial conditions, resolving a tension between acceleration and variance
that had prevented prior methods from achieving both goals. Through explicit characterization of the learning dynamics, we provably show that the stochastic algorithms can be accelerated for non-strongly-convex functions.
In Chapter 3, we study two-layer linear networks trained by gradient flow and stochastic
gradient flow on the square loss and give a variational characterization of the weight matrices
learned at convergence.
A key result is a separation between the two algorithms: gradient flow preserves the rank of the
parameter matrix, while label-noise stochastic gradient flow drives the rank to zero regardless
of initialization, acting as an implicit regularizer that induces low-rank, sparse feature
representations.
In Chapter 4, we study a two-layer transformer trained to perform in-context prediction
on sequences drawn from n-gram language models.
We establish a landscape result showing that sub-n-gram estimators are near-stationary points
of the cross-entropy loss, explaining the prolonged plateaus observed during training.
By further restricting to sequences drawn from random permutation maps, we derive closed-form
gradient-flow trajectories and show how training progresses stage-by-stage through a sequence
of saddles, with each stage corresponding to the emergence of one component of the induction
head circuit.
In Chapter 5, we analyze a value-softmax model capturing the core structure of
self-attention and show that gradient flow on this parameterization inherently drives the outputs
towards low-entropy, sparse solutions.
This polarization effect provides a mechanism for empirically observed phenomena such as
attention sinks and massive activations, and establishes a direct link between the softmax
architecture and the sparse inter
