通过调整数据分布，让优化器学得更好泛化，无需昂贵的SAM。
Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs
- 用数据重采样让优化器少学简单特征，提升泛化能力。
- 在数学推理任务上，模型准确率最高提升18%。
- 适合想低成本提升LLM泛化性能的研究者和工程师。
训练大型语言模型时，能否通过调整数据分布引导优化器获得更好的泛化能力？本文首次在上下文线性回归模型与多头线性自注意力结构中，理论分析了梯度下降（GD）与尖锐度感知最小化（SAM）的训练动态。尽管SAM具有更优泛化性，但其计算成本极高，难以用于中等规模模型训练。研究发现，SAM能降低简化偏差（SB），即优化器早期优先学习简单特征的倾向，这是其泛化优势的关键原因。受此启发，我们提出通过重采样或增强训练后期学习的样本，同样可降低简化偏差并提升泛化性能。大量实验表明，该方法在多个模型（Phi2-2.7B、Llama3.2-1B、Gemma3-1B-PT、Qwen3-0.6B-Base）上有效，采用AdamW和Muon优化器微调后，在数学推理任务上相对准确率提升最高达18%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。