论文解读 | EMNLP2024 一种用于大语言模型版本更新的学习率路径切换训练范式

点击蓝字

关注我们

AI TIME欢迎每一位AI爱好者的加入！

点击 阅读原文 观看作者讲解回放！

作者简介

王志豪，厦门大学博士生

刘诗雨，厦门大学硕士生

内容简介

新数据的不断涌现使版本更新成为大型语言模型（LLMs）不可或缺的需求。LLMs的版本更新训练范式包括从头预训练（PTFS）和继续预训练（CPT）。初步实验表明，PTFS在预训练性能上表现更好，而CPT的训练成本较低。此外，随着版本更新的进行，两种范式的性能和训练成本差距逐渐扩大。为探究这一现象的根本原因，作者分析了学习率对CPT的两个阶段的影响：准备初始化参数（checkpoint）和基于该checkpoint的继续预训练。研究表明，在第一阶段中使用较大学习率以及在第二阶段中使用具有完整率衰减过程的学习率对于LLMs的版本更新至关重要。因此，作者提出了一种基于学习率路径切换的训练范式。该范式包括一条主路径和多条分支路径。LLMs在主路径上以最大学习率进行预训练，而分支路径则基于LLMs在主路径上的checkpoints使用新增数据进行版本更新。广泛的实验表明该范式的有效性和泛化性。特别是在训练4个版本的LLMs时，该范式在保持与PTFS相当的预训练性能同时，将总训练成本减至58%。

论文地址：https://arxiv.org/abs/2410.04103

研究动机

这篇论文的研究动机在于，随着新的预训练数据不断涌现，大型语言模型（LLMs）面临着版本更新的需求，以确保模型能力的持续提升。现有的版本更新范式包括从头预训练（Pre-Training From Scratch，简称PTFS）和继续预训练（Continue Pre-Training，简称CPT）。图中列出了这两种范式在使用cosine学习率调度策略时的学习曲线。可以明显看到，两种范式之间的主要区别在于更新过程中的学习率变化。这启发了作者从学习率角度出发，研究新的预训练范式。

先导实验

为了比较这两种范式在性能和成本上的差异，作者进行了一个先导实验。实验选择了训练LLMs时最常见的3种学习率调度策略，测试在更新4个版本LLMs情况下的性能和成本差距。需要注意的是，这里的性能用困惑度（Perplexity，PPL）来表示，数值越低代表LLMs性能越好。

从实验结果可以观察到，尽管CPT的版本更新成本远低于PTFS，但PTFS的性能优于CPT，而且这种性能差距随着版本数的增加而增大。

性能差距增大的原因

为了研究这种性能差距产生的原因，作者将CPT拆分成两个阶段。第一个阶段是为CPT准备初始参数（checkpoint），第二个阶段是基于初始checkpoint进行继续训练。紧接着，作者基于这两个阶段，分别设计两组实验来探索学习率对CPT两个阶段的影响。

在第一组实验中，作者采用具有不同衰减速度的学习率作为第一阶段的学习率曲线，并固定了第二阶段的学习率曲线。结果表明，当第一阶段的学习率固定为最大值时，初始checkpoint的模型性能最低，但最终性能却是最好的。

在第二组实验中，作者固定了第一阶段的学习率曲线，采用具有不同衰减速度的学习率作为第二阶段的学习率曲线。结果显示，当第二阶段学习率快速衰减到最小值时，对应的LLMs性能最佳。

基于上述两组实验我们可以得出如下结论：1.第一阶段的大学习率和第二阶段完整的学习率衰减过程对CPT的性能尤为重要。2. CPT无法兼顾不同版本LLMs的性能。完整的学习率衰减过程能确保当前版本的LLMs的最优性能，但后续版本的LLMs则需要以大学习率训练提供的初始化checkpoint，这是CPT无法同时满足的。

训练范式

为了解决CPT两阶段对不同学习率要求的冲突，作者提出了一种基于学习率路径切换的训练范式。该范式包括一条主路径和多条分支路径。在主路径上，LLMs以最大学习率从头开始预训练，为后续版本更新提供初始化checkpoint。当我们想获得新版的LLMs时，可以直接基于主路径的当前checkpoint继续预训练。在这个过程中，学习率会经历一个完整且快速的衰减过程，从而以较低的成本来保证新版LLMs的性能。同时，在主路径上LLMs仍然使用新增数据对当前checkpoint以最大学习率进行预训练，以便于后续的版本更新。

不同于PTFS和CPT，该范式还包含关键超参数α用于控制分支路径在训练步长中所占的比例。根据版本更新的总时间复杂度计算，该范式与CPT一样，确保了线性的复杂度。