针对平均模型输出优化训练,提升大模型性能
Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models
- 将平均模型优化设计为最优控制问题,动态调整训练过程
- 新方法PACE在多种任务中显著降低平均模型误差,提升效果
- 适合需要高精度平均模型的微调与预训练场景
许多现代语言模型采用迭代平均(如指数移动平均)作为最终输出,而非直接使用最后一个迭代点。这引发一个核心问题:若最终返回的是平均模型,应如何调整训练过程以提升该平均模型的性能?本文将优化器设计为针对迭代平均估计器的最优控制问题,在连续时间随机二次模型下求解最小化返回平均误差的控制策略,并引入对干预强度的惩罚。基于此控制器的实用近似得到PACE——一种轻量级的AdamW封装,通过剪裁后的逐坐标控制力,将当前权重拉向其指数移动平均。理论证明,简化版PACE在凸优化中达到标准速率,且在二次设定下可严格降低迭代平均估计器的极限平方误差,某些情况下甚至提升幅度可任意大。实验表明,无论是在1-2B参数模型的监督微调,还是GPT-2在FineWeb上的预训练中,PACE均在广泛学习率、衰减策略等超参数设置下优于AdamW及经EMA评估的AdamW。
原文摘要 · Abstract (English)
Many modern Language Model (LM) pipelines return an averaged model, such as an exponential moving average of the training iterates, rather than the final iterate itself. This raises a fundamental question: given that we will return an iterate average, how should we change training to improve the performance of this average? We study this question by formulating optimizer design for the iterate-average estimator as an optimal-control problem. In a continuous-time stochastic quadratic model, we solve for the control strategy that minimizes the error of the returned average subject to a penalty on the size of the intervention. A practical approximation to this controller yields PACE, a lightweight wrapper around AdamW that pulls the live weights toward their exponential moving average with a clipped, per-coordinate control strength. We prove that a stylized version of PACE converges at the standard stochastic convex optimization rate, up to a factor depending on the averaging rule, while in the quadratic setting it can strictly improve the limiting squared error of the iterate-average estimator and can do so by an arbitrarily large factor on some instances. Empirically, our results suggest that PACE improves over AdamW and EMA-evaluated AdamW in supervised fine-tuning of 1-2B parameter LMs and in GPT-2 pretraining on FineWeb for a wide range of learning rates, decay schedules, and other hyperparameters.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。