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

arXiv: https://arxiv.org/abs/2602.00576 · HTML: https://mlsanity.com/paper/2602.00576/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 47 of its 56 references that are in the map.
- First posted on arXiv 31 Jan 2026.

## Builds on (most cited first, 30 of 47)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Gemma 3 Technical Report](https://mlsanity.com/paper/2503.19786/) (2025)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Reconciling modern machine learning practice and the bias-variance trade-off](https://mlsanity.com/paper/1812.11118/) (2018)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Deep Double Descent: Where Bigger Models and More Data Hurt](https://mlsanity.com/paper/1912.02292/) (2019)
- [MAmmoTH: Building Math Generalist Models through Hybrid Instruction Tuning](https://mlsanity.com/paper/2309.05653/) (2023)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [SGD on Neural Networks Learns Functions of Increasing Complexity](https://mlsanity.com/paper/1905.11604/) (2019)
- [A mathematical theory of semantic development in deep neural networks](https://mlsanity.com/paper/1810.10531/) (2018)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced](https://mlsanity.com/paper/1806.00900/) (2018)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Towards Understanding Sharpness-Aware Minimization](https://mlsanity.com/paper/2206.06232/) (2022)
- [NumGLUE: A Suite of Fundamental yet Challenging Mathematical Reasoning Tasks](https://mlsanity.com/paper/2204.05660/) (2022)
- [Neural Networks as Kernel Learners: The Silent Alignment Effect](https://mlsanity.com/paper/2111.00034/) (2021)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [How Does Sharpness-Aware Minimization Minimize Sharpness?](https://mlsanity.com/paper/2211.05729/) (2022)
- [The Surprising Simplicity of the Early-Time Learning Dynamics of Neural Networks](https://mlsanity.com/paper/2006.14599/) (2020)
- [Asymptotic theory of in-context learning by linear attention](https://mlsanity.com/paper/2405.11751/) (2024)

## Abstract (arXiv)

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.