高效训练大语言模型的新方法
在自然语言处理会议EMNLP 2024上,研究人员提出了一种训练框架,可将使用混合数据分布训练大语言模型(LLMs)或其他基于神经网络模型的计算成本降低高达91%。该方法同时还能提高最终模型的质量。
分布编辑模型
传统训练方法(如指令微调)通过称为网格搜索的方法选择训练数据分布的最佳混合比例,这种方法需要大量时间和资源。为应对这些限制,研究人员提出先在对应不同任务的数据分布上对预训练模型进行微调,然后从微调模型的参数值中减去原始模型的参数值。参数值的差异称为分布向量,通过将分布向量的加权和添加到原始模型的参数中来生成复合模型。
由此产生的模型称为分布编辑模型(DEM),以突出利用权重向量算术进行模型编辑的特点。权重基于每个微调模型的困惑度,即其参数值可从原始模型预测的概率。
这种方法依赖两个关键观察:
- 在每个数据集上分别训练模型可以更好地建模每个数据集的基础特性
- 困惑度可以在验证数据上通过单次前向传播计算,比网格搜索高效得多
方法步骤
- 个体分布训练:通过标准训练程序在个体数据分布上训练原始模型
- 分布向量计算:通过从微调模型的参数中减去预训练模型的参数来计算分布向量
- 合并系数优化:基于验证集上的困惑度找到组合数据分布向量的最佳系数
- 分布向量合并:通过可自定义权重的线性组合创建统一模型
- 灵活性和可扩展性:DEM支持在引入新数据集时进行增量更新,无需完全重新训练
评估与未来工作
评估研究表明,DEM在指令微调阶段将训练成本降低高达91%,同时在传统数据混合策略基础上实现高达16.1%的质量改进。该方法在MMLU、BBH和HELM等流行基准测试中表现出色,并在MathQA、SNI和CoT等数据集上展示了跨领域有效性。
DEM在不同模型规模(3B、7B和13B)上都显示出性能改进,证明了该方法的可扩展性。未来研究可能会探索该框架在其他训练场景中的有效性,以及将其扩展到其他模型架构,如编码器-解码器框架或专家混合模型。