Optimizer-Induced Low-Dimensional Drift and Transverse Dynamics in Transformer Training
This work provides insights into how optimizer choice affects training geometry in transformers, which is incremental for understanding neural network optimization.
The study analyzed training trajectories in small transformers, finding that optimizer updates create a dominant drift direction capturing most parameter movement early on, with AdamW producing multi-dimensional drift while SGD variants yield nearly colinear evolution.
We study the geometry of training trajectories in small transformer models and find that parameter updates organize into a dominant drift direction with transverse residual dynamics. Using uncentered, row-normalized trajectory PCA, we show that a single direction captures a large fraction of cumulative parameter movement early in training, while remaining components encode oscillatory behavior in auxiliary probe performance. Instantaneous gradients exhibit little alignment with this dominant direction, indicating that it arises from accumulated optimizer updates rather than per-batch gradient structure. Comparing AdamW with SGD variants at matched loss levels reveals substantial differences in trajectory geometry: AdamW develops multi-dimensional drift structure, whereas SGD-family optimizers produce nearly colinear parameter evolution and weaker probe dynamics. Reheating selectively perturbs transverse components with minimal effect on the dominant drift coordinate. These findings suggest that optimizer choice shapes the effective dimensionality and structure of learning trajectories beyond what is apparent from loss values alone.