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

TL;DR

This paper explores how modifying training data distribution can guide optimizers toward solutions with better generalization in large language models, supported by theoretical analysis and extensive experiments.

Contribution

It reveals that sharpness-aware minimization reduces simplicity bias and demonstrates data augmentation strategies that improve LLM generalization.

Findings

SAM induces lower simplicity bias, enhancing generalization.

Data augmentation reduces simplicity bias and improves accuracy.

Strategies yield up to 18% accuracy gains on multiple LLMs.

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.