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

• URL: http://arxiv.org/abs/2602.00576v1
• Date: Sat, 31 Jan 2026 07:40:36 GMT
• Title: Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs
• Authors: Tushaar Gangavarapu, Jiping Li, Christopher Vattheuer, Zhangyang Wang, Baharan Mirzasoleiman,
• Abstract summary: We show, for the first time, that a lower simplicity bias induces a better generalization.<n>Motivated by this insight, we demonstrate that the training data distribution by upsampling or augmenting examples learned later in training similarly reduces SB and leads to improved generalization.<n>Our strategy improves the performance of multiple language models including Phi2-2.7B, Llama3.2-1B, Gemma3-1B-PT, Qwen3-0.6B-Base-achieving relative accuracy gains up to 18% when fine-tuned with AdamW and Muon.
• Score: 60.68927774057402
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining [55.262510814326035]
  Existing reweighting strategies primarily focus on group-level data importance.<n>We introduce novel algorithms for dynamic, instance-level data reweighting.<n>Our framework allows us to devise reweighting strategies deprioritizing redundant or uninformative data.
  arXiv  Detail & Related papers  (2025-02-10T17:57:15Z)
• What Do Learning Dynamics Reveal About Generalization in LLM Reasoning? [83.83230167222852]
  We find that a model's generalization behavior can be effectively characterized by a training metric we call pre-memorization train accuracy. By connecting a model's learning behavior to its generalization, pre-memorization train accuracy can guide targeted improvements to training strategies.
  arXiv  Detail & Related papers  (2024-11-12T09:52:40Z)
• Efficient Continual Pre-training by Mitigating the Stability Gap [68.49269649759005]
  We study the behavior of Large Language Models (LLMs) during continual pre-training. We propose three effective strategies to enhance LLM performance within a fixed compute budget. Our strategies improve the average medical task performance of the OpenLlama-3B model from 36.2% to 40.7% with only 40% of the original training budget.
  arXiv  Detail & Related papers  (2024-06-21T02:28:37Z)
• Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization [12.472871440252105]
  We show that sharpness-aware minimization (SAM) learns different features more uniformly, particularly in early epochs. We propose a method that (i) clusters examples based on the network output early in training, (ii) identifies a cluster of examples with similar network output, and (iii) upsamples the rest of examples only once to alleviate the simplicity bias.
  arXiv  Detail & Related papers  (2024-04-27T03:30:50Z)
• Model Extrapolation Expedites Alignment [135.12769233630362]
  We propose a method called ExPO to expedite alignment training with human preferences.<n>We demonstrate that ExPO boosts a DPO model trained with only 20% steps to outperform the fully-trained one.<n>We show that ExPO notably improves existing open-source LLMs on the leading AlpacaEval 2.0 and MT-Bench benchmarks.
  arXiv  Detail & Related papers  (2024-04-25T17:39:50Z)
• Take the Bull by the Horns: Hard Sample-Reweighted Continual Training Improves LLM Generalization [165.98557106089777]
  A key challenge is to enhance the capabilities of large language models (LLMs) amid a looming shortage of high-quality training data. Our study starts from an empirical strategy for the light continual training of LLMs using their original pre-training data sets. We then formalize this strategy into a principled framework of Instance-Reweighted Distributionally Robust Optimization.
  arXiv  Detail & Related papers  (2024-02-22T04:10:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.