早周期权重平均结合高学习率用于大语言模型预训练
机器学习
2023-12-13 v2 人工智能
计算与语言
摘要
训练大语言模型（LLMs）成本高昂；因此，任何加速模型收敛的策略都有助益。在本文中，我们研究了一个简单思路的能力：在训练轨迹上对已保存检查点进行平均，以在训练早期显著改善收敛性与泛化能力。我们在此表明，采用高学习率训练的模型因检查点平均而获得更高的收益。此外，当检查点在训练步数上以较大间隔采样时，这些收益会被放大。我们的训练方案优于常规训练以及流行的检查点平均基线，如指数移动平均（EMA）和随机权重平均（SWA）。我们通过预训练 LLMs 来评估我们的训练方案，其中由于极大的批大小，高学习率天然受到青睐。具体而言，我们在包含 9B tokens 的 OpenWebText 数据集上预训练了不同规模的 nanoGPT-2 模型：小（125M）、中（335M）和大（770M）。此外，我们给出了公开可用的 Pythia LLMs 的结果，规模从 1B 到 12B 不等，这些模型在包含 207B tokens 的 PILE-deduped 数据集上训练。
引用
@article{arxiv.2306.03241,
title = {Early Weight Averaging meets High Learning Rates for LLM Pre-training},
author = {Sunny Sanyal and Atula Neerkaje and Jean Kaddour and Abhishek Kumar and Sujay Sanghavi},
journal= {arXiv preprint arXiv:2306.03241},
year = {2023}
}
备注
17 pages, 13 figures, presented at NeurIPs 2023 WANT workshop