Efficient Training of Language Models with Compact and Consistent Next Token Distributions

📄 arXiv: 2407.02819v1 📥 PDF

作者: Ashutosh Sathe, Sunita Sarawagi

分类: cs.CL, cs.LG

发布日期: 2024-07-03

备注: ACL 2024

💡 一句话要点

提出紧凑一致的下一Token分布，加速并提升语言模型训练效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 语言模型 预训练 n-gram 正则化 模型压缩

📋 核心要点

  1. 现有语言模型训练方法计算成本高昂，直接使用n-gram统计进行正则化会显著降低训练速度。
  2. 论文提出一种紧凑的下一token分布表示，在期望上与完整n-gram分布对齐，并降低小批量方差。
  3. 实验表明，该方法在模型质量和收敛速度上均优于现有方法，并能扩展到更大的数据集和模型。

📝 摘要（中文）

本文针对语言模型预训练中最大化下一个token似然这一目标，提出了一种通过预先聚合语料库的折叠n-gram分布来更快地训练更好模型的方法。虽然之前的研究已经提出了语料库级别的n-gram统计作为正则化项，但如果以朴素的方式构建和查询这些n-gram，其代价高昂，会显著降低训练速度，从而限制了它们在现代大型语言模型预训练中的应用。本文介绍了一种替代的、紧凑的下一token分布表示，该表示在期望上与完整的n-gram分布对齐，同时显著降低了小批量之间的方差，优于标准的下一token损失。实验结果表明，与现有方法相比，n-gram正则化模型和本文提出的近似方法在模型质量和收敛速度方面均有显著提高。此外，与直接的n-gram正则化方法相比，本文的近似方法有助于将增益扩展到更大的数据集和模型。

🔬 方法详解

问题定义：现有语言模型预训练依赖最大化下一个token的似然，计算量大。直接使用n-gram统计作为正则化项可以提升模型性能，但朴素的n-gram构建和查询方法计算成本很高，严重影响训练速度，难以应用于大规模语言模型。

核心思路：论文的核心思路是使用一种紧凑的、近似的下一token分布来替代完整的n-gram分布。这种近似分布在期望上与完整n-gram分布对齐，但显著降低了小批量之间的方差，从而加速训练过程并提高模型质量。通过预先聚合语料库的n-gram信息，减少了训练过程中的计算量。

技术框架：该方法主要包含两个阶段：1) 预处理阶段：对语料库进行n-gram分析，构建紧凑的下一token分布表示。2) 训练阶段：使用构建的紧凑分布作为目标，训练语言模型。在训练过程中，使用该紧凑分布计算损失，替代传统的下一token预测损失。

关键创新：关键创新在于提出了紧凑且一致的下一token分布表示，它是一种对完整n-gram分布的有效近似。与直接使用n-gram正则化相比，该方法显著降低了计算复杂度，并减少了小批量之间的方差，从而提高了训练效率和模型性能。这种近似方法使得n-gram信息能够应用于更大规模的数据集和模型。

关键设计：论文的关键设计包括：1) 如何构建紧凑的下一token分布表示，使其在期望上与完整n-gram分布对齐。2) 如何设计损失函数，利用该紧凑分布来指导语言模型的训练。3) 如何平衡紧凑分布的近似程度和计算复杂度，以获得最佳的训练效果。具体的参数设置和网络结构细节在论文中未详细描述，属于未知信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文实验结果表明，提出的紧凑下一token分布方法在模型质量和收敛速度方面均优于现有方法。与直接的n-gram正则化方法相比，该方法能够更好地扩展到更大的数据集和模型，在大型数据集上的性能提升更为显著。具体的性能数据和提升幅度在摘要中未提及，属于未知信息。

🎯 应用场景

该研究成果可应用于各种需要大规模语言模型预训练的场景，例如机器翻译、文本生成、对话系统等。通过提高训练效率和模型质量，可以降低训练成本，并提升下游任务的性能。该方法尤其适用于资源受限的环境，能够以更低的计算成本训练出更好的语言模型。

📄 摘要（原文）

Maximizing the likelihood of the next token is an established, statistically sound objective for pre-training language models. In this paper we show that we can train better models faster by pre-aggregating the corpus with a collapsed $n$-gram distribution. Previous studies have proposed corpus-level $n$-gram statistics as a regularizer; however, the construction and querying of such $n$-grams, if done naively, prove to be costly and significantly impede training speed, thereby limiting their application in modern large language model pre-training. We introduce an alternative compact representation of the next token distribution that, in expectation, aligns with the complete $n$-gram distribution while markedly reducing variance across mini-batches compared to the standard next-token loss. Empirically, we demonstrate that both the $n$-gram regularized model and our approximation yield substantial improvements in model quality and convergence rate compared to existing methods. Furthermore, our approximation facilitates scalability of gains to larger datasets and models compared to the straightforward $n$-gram regularization method.