Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs

TL;DR

The paper addresses how training data distribution can guide optimizers toward better generalization in large language models by analyzing an in-context linear regression setup with multi-head linear self-attention. It shows that Sharpness-Aware Minimization (SAM) reduces Simplicity Bias (SB) compared to gradient descent (GD) and develops an ODE-based account of early feature learning, revealing SAM slows learning of simple features while accelerating harder ones. To emulate SAM's benefits without its computational cost, the authors propose a data-centric approach that upsamples examples learned later in training (hard examples) based on proxy-model loss trajectories, yielding improved generalization across multiple LLMs on mathematical reasoning tasks (up to 18% relative gains) and across optimizers like AdamW and Muon. The work provides both theoretical insights into optimizer dynamics and a practical, scalable data-augmentation strategy that leverages data distribution to steer learning toward more generalizable solutions.

Abstract

Can modifying the training data distribution guide optimizers toward solutions with improved generalization when training large language models (LLMs)? In this work, we theoretically analyze an in-context linear regression model with multi-head linear self-attention, and compare the training dynamics of two gradient based optimizers, namely gradient descent (GD) and sharpness-aware minimization (SAM), the latter exhibiting superior generalization properties but is prohibitively expensive for training even medium-sized LLMs. We show, for the first time, that SAM induces a lower simplicity bias (SB)-the tendency of an optimizer to preferentially learn simpler features earlier in training-and identify this reduction as a key factor underlying its improved generalization performance. Motivated by this insight, we demonstrate that altering the training data distribution by upsampling or augmenting examples learned later in training similarly reduces SB and leads to improved generalization. Our extensive experiments show that our strategy improves the performance of multiple LLMs-including Phi2-2.7B , Llama3.2-1B, Gemma3-1B-PT, and Qwen3-0.6B-Base-achieving relative accuracy gains up to 18% when fine-tuned with AdamW and Muon on mathematical reasoning tasks.

Figures (5)

• Figure 1: Comparison of training loss for (left) GD, (right) SAM, and (middle) GD + upsampling (our method). We perform 25 independent runs and report the median test loss across all 25 runs on top of each plot. Both SAM and GD+upsampling exhibit lower simplicity bias, i.e., more uniform feature learning, (demonstrated by abrupt drops in the loss) than GD, and obtain lower test loss.
• Figure 2: Comparison of preconditioned f-SAM and AdamW. f-SAM achieves lower evaluation loss compared to AdamW but incurs around 2x GPU memory usage and 6x training time, when fine-tuning Phi2-2.7B on the MathInstruct dataset. This makes it prohibitively expensive for training even medium-sized LLMs (we couldn't finish fine-tuning with f-SAM after 6 days on 4$\times$ NVIDIA A40 GPUs, hence reporting validation loss).
• Figure 3: Upsampling improves performance across models and math benchmarks. (Left) Each model is finetuned with AdamW on the original dataset and with targeted upsampling, and evaluated using zero-shot greedy decoding. (Right) Targeted upsampling is optimizer-agnostic, yielding consistent performance gains under Muon. For each model, we report the final-epoch checkpoint with the highest average accuracy across all benchmarks. Detailed dataset-level results are provided in \ref{['tab:adamw-math-upsampling']}, Appendix \ref{['app:aux_results']}. (Different scales are used across subfigures to accurately show performance gains while accounting for variability in individual model performance.)
• Figure 4: Data upsampling ablations. Each subfigure lists the target model in its title; detailed dataset- and model-level results are provided in Appendix \ref{['app:aux_results']}. (a) Performance is comparable when using the lightweight Pythia-70M as the proxy instead of the target model. (b) Identifying hard examples using loss trajectories from all $15$ proxy checkpoints performs better than using only the final checkpoint. (c) No clear trend emerges as a function of the number of proxy checkpoints used in loss trajectories for clustering. (d) Under compute-matched finetuning, upsampling hard examples yields the strongest generalization, while upsampling easy examples degrades performance. (e) Variational problem synthesis offers a promising strategy to further improve model generalization on hard examples.
• Figure 5: Toy example of Shannon entropy in quantifying uniformity in feature learning. We consider the extreme case of only 3 features. Following the $\texttt{ATTN}_S$ setting, the training can be separated by features. For each feature, there exists an early training phase (in green), and $t_1, t_2, t_3$ mark the times spent in this regime. More uniform learning is achieved when these values are closer to each other (bottom case), and after sum-normalization, this naturally corresponds to a distribution that resembles the uniform distribution (e.g. $(1/3, 1/3, 1/3)$ over three objects) and hence has higher entropy

Theorems & Definitions (18)

• Lemma 3.1: Equivalence to three-layer convolutional network (CNN); from $\S$4.1 of zhang2025training
• Lemma 4.1: Manifolds of saddle points
• Theorem 4.2: Early Training Dynamics of SAM
• Theorem 4.4: SAM has lower simplicity bias than GD
• Definition 2.1: Global vs. layer-wise SAM regimes
• Theorem 2.2: SAM accelerates training for merged KQ
• proof
• Theorem 2.3: Escape Time
• proof
• Lemma 2.4