RegMix-用回归任务解决大模型数据混合问题

发布于 2024-7-8 07:38
浏览
0收藏

写在前面

大型语言模型在预训练过程中,如何选取数据的混合比例(利用较少的Tokens来实现较小的Loss从而加速预训练过程)是一个复杂和关键的问题。手动确认数据集中各个组成的比例是不可扩展的,并且很可能不是最优选择。

今天给大家介绍一个用回归任务解决大模型数据混合问题的方法-RegMix。其核心思想是,利用不同的数据混合比例先训练多个小模型并获取其结果,在利用这些样本训练一个回归模型,再遍历所有比例利用回归模型找到最优的数据混合比例,最后用最优数据混合比例训练更大的语言模型。

Paper: https://arxiv.org/abs/2407.01492
Github: https://github.com/sail-sg/regmix

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

通过训练512个1M的小模型,拟合回归模型,找到top64的数据混合比例,训练1B模型,最优数据混合比例训练的模型的验证集loss也是最低。

方法

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

整体流程如上图所示,

  • 生成随机数据混合比例,按照比例采用混合数据并训练小模型;
  • 利用数据混合比例作为特征值,模型训练的目标值作为标签,拟合回归模型;
  • 在模拟更大数据混合比例空间,利用回归模型预测最佳目标值,以获取最佳混合比例;
  • 使用模拟出的最佳混合比例的数据训练更大的模型。

训练小模型时越多越好,但为了节约成本需要尽量减少小模型训练次数,那么在初始化数据混合比例时就需要时多样化的,并且每个数据领域需要都存在极端值,数据采用过程主要是基于Tokens(chunk-level)分布的狄利克雷分布来实现。

详见:mixture_config/synthesize_mixture.py

同时在拟合回归模型时,采用了线性回归LightGBM两种回归模型。

结果

数据集采用Pile dataset中不涉及版权的17个数据集,如下表所示,

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

512个1M小模型在1B Tokens训练得到的回归模型,与在25B Tokens数据下训练的1B模型,排序具有97.12%的高相关性,如下表所示,

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

同时训练次数要比训练的总Token数要重要,更影响回归模型的效果,并且采用LightGBM建模要比线性回归建模要好。

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

PS:跟作者@乾神交流过,512个样本训练回归模型会不会数据量太少,乾神说他们做过1024的实验,但并回归模型效果无明显提高,并且从成本考虑,那么512最佳。

不同的数据混合比例对下游任务结果影响较大,在Lambada数据集上最好和最差的效果相差14.6%,如下表所示,

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

同时发现了一个与传统理解不一致的结果,一般我们任务维基数据质量很高,是评估大型语言模型最具代表性的数据集。但实验结果发现,网络数据集上评估的效果,更能体现模型在下游任务上的好坏,如下图所示,可以发现Pile-CC数据集作为验证时损失值与下游任务的相关性更强。

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

并且RegMix可以发现各领域数据之间是如何相互作用的,数据领域之间复杂的相互作用利用人类固有经验很难直接区分。

RegMix-用回归任务解决大模型数据混合问题-AI.x社区

本文转载自 NLP工作站​,作者:刘聪NLP

收藏
回复
举报
回复
相关推荐