发表时间: 2024-04 · arXiv:2404.10830 (AWS AI Labs, ICML 2024)
文章标题:减少截断可改善语言建模
作者/机构:Hantian Ding, Zijian Wang, Giovanni Paolini, Varun Kumar, Anoop Deoras, Dan Roth, Stefano Soatto
本文探讨了大型语言模型(LLM)预训练中一个普遍存在但被忽视的问题:输入数据的处理方式。当前,为了提高训练效率和避免使用填充(padding)令牌,普遍采用“拼接后切分”(concatenate-then-split)的方法,即将多个文档拼接在一起,然后切分成固定长度的序列。
核心问题与研究目标:
该拼接方法虽然高效,但严重损害了数据的完整性。它不可避免地将许多完整的文档切分成不完整的片段,导致过度的截断。这种截断会阻碍模型学习构建基于完整上下文、逻辑连贯且事实一致的内容,并可能使模型更容易产生幻觉。研究目标是提出一种新的数据分组方法,既能保持拼接法的高训练效率,又能最大限度地保护文档的完整性,减少不必要的截断。
创新点与主要贡献:
为解决上述问题,本文提出了Best-fit Packing,一种可扩展且高效的数据分组方法,通过基于长度的组合优化来将文档打包到训练序列中。
截断破坏了上下文的完整性。一篇写得好的完整文档天然是连贯且自包含的。文档中的事实陈述通常通过引用、蕴含或更复杂的推理逻辑上依赖于其前述上下文。我们将用于建立这种依赖关系的关键上下文跨度称为“基础上下文”(grounding context)。在进行下一词元预测学习时,如果基础上下文缺失,模型将被迫预测实际上无法从观察到的部分上下文中推导出的词元。因此,在推理时,模型有更高的几率忽略(即使提供了)基础上下文,并生成与给定上下文矛盾或无法验证的内容,这被称为“封闭域幻觉”(closed-domain hallucination)。我们在图2中说明了这一点。
代码截断导致未定义变量和病态模式学习。图2(a)展示了一个Python示例。尽管原始代码是正确的,但将变量定义和相应的使用分割到两个不同的训练序列中会引入语法错误。由于自注意力机制不跨越序列边界,DecoderModel和config在后一个训练序列中实际上是未定义的。以这种碎片化的方式格式化数据会使模型学习到病态模式,可能导致下游任务中的幻觉。例如,在程序综合任务中,模型可能会直接使用config而没有其定义。更糟糕的是,模型可能会忽略提供的上下文并捏造一个不相关的名称:如果我们打算实例化一个EncoderModel,并在上下文中指定import EncoderModel,模型仍可能生成model=DecoderModel(...),这是因为它学习到了model和DecoderModel之间的虚假关联。
图2. 文档截断导致幻觉或知识损失的示例。(a) 变量定义(蓝色)被截断,随后的使用调用导致未定义的名称(红色)。(b) 关键上下文信息被截断(蓝色),使得摘要不忠实(红色)。(c) 由于截断,模型不知道ICML 2024在哪里举行。
自然语言截断损害摘要的忠实性。图2(b)展示了自然语言中同样的问题,即截断损害了忠实性。摘要中的短语“Monday morning”在同一训练序列的上下文中找不到任何依据,因此变成了一种捏造。这种不完整的样本会降低模型的上下文感知能力(即关注上下文的能力),并导致不忠实的生成或无意义的推理。
截断阻碍知识获取。除了加剧幻觉,截断还会阻碍训练过程中的知识获取,因为知识的文本表示通常以完整的句子或段落形式出现,这很容易受到碎片化的影响。例如,在图2(c)中,模型将无法学习到ICML的举办地点,因为会议名称和其举办地位于不同的训练序列中。
引入简化随机过程进行直观分析。作为一个额外的直观来源,我们描述了一个简化的随机过程$(X_n)_{n \in \mathbb{N}}$,通过这个过程我们可以分析性地证明,即使在训练数据无限的情况下,一个在截断序列上训练的模型其序列建模准确性也严格差于一个在完整序列上训练的模型。虽然为Transformer模型建立关于截断如何影响学习的严谨理论很困难,但我们希望在分析探索上做出首次尝试,以更好地推动我们的提议。
过程定义。与语言建模类似,我们可以将$X_n$看作是二元词汇表{0, 1}中的词元。我们的过程是递归定义的,从一个伯努利变量$X_0$开始,它以0.5的概率取值为0,否则为1。对于$n \ge 1$,变量$X_n$以概率$p$取$X_0$的值,以概率$1-p$取$1-X_0$的值,其中$p \in (0.5, 1)$是固定的。与此过程相关的图模型将是一棵以$X_0$为根,$X_1, X_2, \dots$为叶的树。
模型对比设置。我们现在比较一个在序列$X_{0:L} := (X_0, X_1, \dots, X_{L-1})$上训练的“模型A”,和一个在序列$(X_0)$和$X_{1:L} := (X_1, X_2, \dots, X_{L-1})$上训练的“模型B”。因此,模型B的训练受到截断的影响。我们假设有足够的数据量让模型完美拟合训练序列。
模型A的损失。对于$m \ge 1$,模型A在词元$X_m$上实现的期望分类损失由条件熵给出:
其中$q = 1-p$。
模型B的损失。另一方面,如果我们将一个观测序列$(x_0, \dots, x_{m-1})$输入模型B,它对下一个词元的预测等于:
由于过程的简单性,这个分布可以被解析计算(见附录A),从而使我们能够确定模型B的期望损失。
分析结果与讨论。图3显示了模型B相对于模型A的损失相对增加量,作为词元位置$m$的函数。我们看到,即使在数据充足的假设下,模型B总是更差。随着$m$趋于无穷大,相对损失增加量总是收敛到0,但收敛速度取决于$p$,即取决于可见词元对被截断词元的依赖强度。这例证了一个在现实世界语言建模中也存在的效果:如果一个关键信息被截断但很快又被重复(高$p$),那么第一次提及的缺失不会产生持久影响;另一方面,如果只有关于被截断概念的模糊相关信息可用(低$p$),那么截断的影响会持续更长时间。
图3. 模型B(在截断序列上训练)相对于模型A(在完整序列上训练)的期望损失相对增量,作为词元位置m ≥ 1的函数,针对不同的p值。
Best-fit Packing方法概述。我们提出一种新方法,在高效分组训练数据的同时消除不必要的截断,如图1所示。给定模型的最大序列长度$L$,我们首先将每个文档分割成长度最多为$L$个词元的块。注意,这是由上下文长度限制的最小必要截断。然后,为了构建每个训练序列,我们选择若干个文档块,尽可能地填满$L$个词元的空间,而不破坏任何一个块。我们称这种选择策略为打包算法,将在§3.1中讨论。
方法面临的挑战。这里存在两个挑战:首先,由于并不总是能完全填满$L$个词元的序列,需要使用填充词元,这导致与拼接方法相比训练序列数量增加。这种增加需要每个epoch进行更多的训练步骤。因此,打包后的序列必须足够紧凑,以最小化填充词元的使用,从而保持训练效率。其次,该算法必须具有足够的可扩展性和速度,以便能够处理数十亿文档的数据集。
问题形式化定义。我们将Best-fit Packing形式化为一个组合优化问题。给定一组文档块$C = \{c_1, \dots, c_N\}$,其中$l(c)$是块$c$的词元长度且$l(c_i) \le L$,将这些块打包到训练序列中等价于确定$C$的一个分区$S = \{s_1, \dots, s_M\}$,满足$\sum_{c \in s_i} l(c) \le L$。一个训练序列由拼接$s_i$中的所有块构成。我们的目标是找到一个尺寸最小的分区$S$,这在实践中意味着生成最少数量的训练序列。
装箱问题与近似算法。上述优化问题被称为装箱问题(bin packing problem)【2, Bernhard, K. and Vygen, J. Combinatorial optimization: Theory and algorithms. 2008】,即必须将$N$个不同大小的物品装入有限数量的箱子中,每个箱子有固定的容量,目标是最小化使用的箱子数量。计算上,该问题是NP-难的。存在几种近似算法,其中First-Fit-Decreasing(FFD)和Best-Fit-Decreasing(BFD)【16, Eilon, S. and Christofides, N. The loading problem. 1971】是平衡效率和准确性方面最受欢迎的。我们在算法1中简要描述了这些启发式方法。
时间复杂度分析与优化。通常,FFD和BFD都需要$O(N \log N)$的排序时间和$O(N \log N)$的打包时间。查找$j^*$的步骤通常用一个大小为$O(N)$的平衡二叉树来实现,该树跟踪所有现有的箱子。然而,在我们的案例中,注意到序列长度总是在$[1, L]$范围内的整数,其中$L \ll N$。这通过计数排序将排序成本降低到$O(N)$,更重要的是,允许对打包部分进行进一步优化。由于在BFD中我们不区分具有相同剩余容量的箱子,因此只需跟踪剩余容量值而不是实际的箱子,这有效地将树的大小减少到$O(L)$,从而将打包时间减少到$O(N \log L)$。然而,同样的方法不适用于FFD,因为箱子的顺序很重要。
使用段树实现快速搜索。在实践中,我们使用一个段树在BFD中实现了上述快速搜索,定义如下:
* 树有$L$个叶节点。如果存在至少一个剩余容量为$i$的箱子,第$i$个叶子的值为$i$,否则为零。初始时,所有叶节点都设置为零,除了最后一个设置为$L$。
* 每个内部节点的值是其子节点的最大值。
为了找到最合适的箱子,我们从根节点查询树。在每个内部节点,如果左子节点的值不小于物品重量,我们就向左走,否则向右走。最终我们会到达一个叶节点,其值就是最合适的容量。我们使用一个容量到箱子的映射来检索最合适的箱子。最后,我们更新树以恢复上面列出的两个属性。更详细的说明请参考附录B。
运行时间对比。表1展示了在不同数据规模下,上下文长度为2048时,优化的Best-Fit Decreasing(OBFD)算法与标准First-Fit Decreasing(FFD)的运行时间比较,数据通过对约10亿文档的RefinedWeb数据集进行上/下采样得到。结果表明,我们的优化BFD在10亿规模下节省了60%的运行时间,并且相对加速比(FFD/OBFD)随数据大小呈对数增长。通过这种高效实现,我们的方法能够扩展到更大的数据集,因为渐进时间复杂度仅与数据大小线性相关。
表1. 我们在100万到20亿个文档上,以2048的上下文长度,对FFD和我们优化的BFD(OBFD)的运行时间(秒)进行了基准测试。两种算法都用Python实现,并在单个CPU线程上运行。基线FFD使用与优化BFD类似的段树,但有N个叶子对应所有箱子。我们优化的BFD通过线性可扩展性显著提高了打包效率。
紧凑性分析。打包的另一个重要方面是生成的训练序列有多紧凑。理论上,FFD和BFD都被保证渐近地使用的箱子数量不超过最优数量的11/9【23, Johnson, D. S., et al. Worst-case performance bounds for simple one-dimensional packing algorithms. 1974】。然而,在实践中,大多数文档相对于上下文长度来说都较短,如图4中的虚线所示。大量小尺寸物品的存在使得打包变得容易得多。因此,我们观察到打包产生的训练序列几乎和拼接方法得到的一样紧凑。如表2所示,无论是在2k还是8k的上下文长度,以及跨NL和PL数据集,Best-fit Packing产生的额外训练序列都不超过0.01%。从另一个角度看,表2也表明使用Best-fit Packing时,填充量微不足道。结果证明,以相同计算量处理的非填充词元数量来衡量,Best-fit Packing实现了与拼接大致相同的训练效率。
表2. Best-fit Packing的紧凑性。我们展示了由拼接方法生成的训练序列数量,以及由Best-fit Packing相对于拼接方法生成的序列数量(或百分比)的增量。两种方法之间的差异可以忽略不计。因此,Best-fit Packing保持了与拼接相同的训练效率。
截断减少效果分析。我们研究了Best-fit Packing在多大程度上缓解了截断问题。我们在NL和PL数据集中统计了每个文档被截断的次数,并按文档长度进行汇总,如图4所示。请注意,大多数文档包含的词元少于$L=2048$;因此,由拼接引起的截断主要发生在此范围内。通过完全消除文档长度低于$L$的截断,Best-fit Packing有效地保护了绝大多数文档的完整性,仅在绝对必要时才对更长的文档进行截断。
图4. 在2k或8k最大序列长度下,每个文档长度的文档计数和截断计数。使用Best-fit Packing后,截断数量显著减少。上图:自然语言。下图:编程语言。
模型架构:
数据集:
硬件配置:
软件与训练配置:
本文通过零样本和少样本提示对模型在多种下游任务上进行评估。结果表明,Best-fit Packing在多个任务中均提升了性能,尤其是在阅读理解(+4.7%)、自然语言推理(+9.3%)、上下文遵循(+16.8%)和程序综合(+15.0%)方面。此外,该方法有效减少了封闭域幻觉。在所有评估任务中,该方法要么优于基线,要么与之持平,从未出现统计上显著的性能下降。
阅读理解(Reading Comprehension)
自然语言推理(Natural Language Inference)
上下文遵循(Context Following)
摘要(Summarization)
常识与闭卷问答(Commonsense and Closed-book QA)
程序综合(Program Synthesis)
预训练数据。预训练数据对语言模型的质量至关重要。已有多个高质量的公开预训练数据集,如C4 【60, Raffel et al., 2020b】, Pile 【18, Gao et al., 2021】, RefinedWeb 【57, Penedo et al., 2023】, RedPajama 【11, Computer, 2023】, 以及The Stack 【28, Kocetkov et al., 2022; 39, Lozhkov et al., 2024】。在此基础上,多篇论文提出了各种过滤策略以提高数据质量。我们的工作可以广泛应用于这些预训练数据集之上。
语言模型训练中的数据分组。近期的Transformer语言模型采用了不同的策略来将训练数据分组为批处理序列,以解决文档长度可变的问题。对于仅编码器模型,RoBERTa 【36, Liu et al., 2019】首次研究了数据格式化,表明在同一训练序列中拼接来自多个文档的句子对性能影响很小。Krell等人【30, Krell et al., 2021】提出了一种基于近似的组合打包方法来加速BERT训练,但并未提高下游性能。值得一提的是,对于编码器模型,文档截断问题不那么令人担忧,原因有二:首先,它们通常在相对较短的文本跨度(128-512个词元)上训练,只尊重句子边界,这种情况下文档级截断不可避免。其次,它们不用于开放式生成,因此幻觉不是问题。
解码器模型中的数据分组策略。仅解码器语言模型主要采用“拼接后切分”策略以最大化训练效率【4, Brown et al., 2020; 8, Chowdhery et al., 2022; 59, Rae et al., 2021; 81, Zhang et al., 2022; 75, Touvron et al., 2023b; 69, Scao et al., 2022】。最近,Shi等人【71, Shi et al., 2024】提出将语义相关的文档拼接在同一训练序列中,这在下游任务中取得了显著改进。然而,该方法作为拼接的一种变体,仍然存在过度截断的问题,并且与我们的方法是正交的(可能互补)。
与LLM训练框架的集成。Best-fit Packing在数据层面操作,可以作为离线过程处理。因此,它不需要对训练实现进行任何更改,并且可以轻松集成到像Megatron-LM【72, Shoeybi et al., 2019】或DeepSpeed【62, Rasley et al., 2020】这样的通用分布式训练框架中。
语言生成中的幻觉。随着大规模生成语言模型的快速发展,幻觉问题吸引了越来越多的关注,因为它会影响性能并用捏造的事实误导用户【20, Ji et al., 2022】。已提出多种方法来解决此问题,包括检索增强生成、提示工程、上下文感知解码和监督微调。然而,在预训练阶段缓解幻觉在很大程度上被忽视了,我们是首批探索这一方向的研究之一。
当前语言模型训练中普遍采用的“拼接后切分”数据分组方法不可避免地导致文档碎片化。本文证明了这种截断效应会损害模型遵循上下文的能力,更糟糕的是,会使模型更容易产生幻觉。
基于这些发现,本文提出了Best-fit Packing,一种新的数据分组方法,它能最大限度地保留单个文档的完整性。该算法可扩展至数十亿文档的数据集,并保持与拼接方法同等水平的紧凑性。
通过在不同规模、跨文本和代码领域的模型上进行实验,本文实证了减少截断的有效性。具体而言,通过消除不必要的截断,Best-fit Packing在广泛的任务中表现出色,且未在其他任务上妥协性能。此外,它还有效地减少了语言生成中的封闭域幻觉。
尽管本文的实验主要集中在预训练阶段,但Best-fit Packing同样广泛适用于微调阶段。这项工作为开发更有效、更可靠的语言模型做出了贡献。
模型A和B准确性计算描述。我们现在描述计算§2.1中描述的玩具过程中模型A和B准确性所需的计算。
对于一个固定的$m \ge 1$,令$x = (x_0, \dots, x_{m-1})$为一个观测序列。下一个词元$X_m$服从伯努利分布,以概率$p$取值$x_0$,以概率$q=1-p$取值$1-x_0$。这正是模型A预测的分布。因此,模型A在第$m$个词元上的期望分类损失由伯努利分布的熵给出,并且实际上不依赖于$m$或$x_0$的值(因为交换$p$和$q=1-p$熵保持不变)。
模型B的推理过程。然而,模型B“认为”观测值$x$是根据过程$(X_n)_{n \in \mathbb{N}}$抽取的序列$(\tilde{x}_0, x_0, \dots, x_{m-1})$的一部分。我们可以将模型B的推断看作两个步骤:首先,预测隐藏词元$\tilde{x}_0$的概率分布;然后,预测序列中下一个词元的概率分布,从模型B的角度看,这是$X_{m+1}$。公式表示为:
隐藏词元的分布可以使用贝叶斯法则计算:
其中$k$是序列$x$中0的数量。我们可以代入(1)式,得到模型B预测的显式表达式:
模型B的期望损失计算。模型B的期望分类损失(对于给定的观测序列$x$)是真实分布(一个以概率$p$赋给$x_0$,概率$q$赋给$1-x_0$的伯努利分布)与模型B预测的分布(由(2)式给出)之间的交叉熵:
其中第一项对应结果0,第二项对应结果1。注意(3)式的右侧仅依赖于$x_0$和$k$,所以我们可以称之为$loss_B(x_0, k)$。最后,模型B的期望损失是$loss_B(x_0, k)$在所有可能观测$x$上的平均值:
其中$h$是子序列$(x_1, \dots, x_{m-1})$中0的数量。
算法数据结构。我们用图来说明所提出的优化的Best-Fit-Decreasing算法。宏观上,我们维护三个数据结构:(i) 一个箱子到物品的表(bin-to-items table),用于跟踪每个箱子当前分配的物品;(ii) 一个空间到箱子的表(space-to-bins table),允许我们根据给定的剩余空间检索一个箱子;(iii) 一个段树(segment tree),使我们能够以$O(\log L)$的时间找到最合适的容量。在这个例子中,我们假设最大序列长度(或最大箱子容量)为8。
初始化。开始时,箱子到物品的表和空间到箱子的表都是空的。段树中所有节点都为零,除了最右侧路径上的节点值为8,因为最初我们只有最大容量的空箱子。
图5. 初始化
运行示例。下面我们演示在算法的中间状态下如何打包一个新物品。假设我们已经打包了4个重量分别为8、6、6、4的物品,并达到以下状态。现在我们准备打包一个重量为3的新物品。
图6. 打包四个重量为8、6、6、4的物品后的状态。箱子到物品的表显示第0个物品(重8)已放入箱0,第1个物品(重6)放入箱1,依此类推。空间到箱子的表显示箱1和箱2在打包一个重6的物品后各剩2个空间,箱3剩4个空间。在段树中,由于我们有剩余空间为2、4或8的箱子,从左到右的第2、4、8个叶子分别被赋值为2、4、8。所有其他叶子为零。内部节点被递归更新为其子节点的最大值。
PROTECTED_IMAGE_23____PROTECTED_IMAGE_25
图7. 为了打包一个重为3的新物品,第一步是使用段树找到最合适的容量。我们从根节点查询树。在每个内部节点,如果左子节点的值不小于物品重量,就向左走,否则向右走。我们最终到达一个叶节点,其值就是最合适的容量(本例中为4)。
PROTECTED_IMAGE_26____PROTECTED_IMAGE_28
图8. 我们检索到一个剩余空间为4的箱子,即箱3。这个新的第5个物品(重3)被放入箱3。我们相应地更新箱子到物品的表和空间到箱子的表。最后,我们从下到上递归更新段树,以恢复§3中陈述的两个属性:(i)如果存在至少一个剩余容量为i的箱子,则第i个叶子的值为i,否则为零;(ii) 每个内部节点的值是其子节点的最大值。这就完成了一次打包迭代。
PROTECTED_IMAGE_29____PROTECTED_IMAGE_30
模型与注意力机制。我们使用与LLaMA 7B/13B【74, Touvron et al., 2023a】相同的模型架构。当一个序列包含多个文档时,我们遵循标准实践,在使用拼接方法时允许注意力跨越文档边界【4, Brown et al., 2020; 8, Chowdhery et al., 2022; 75, Touvron et al., 2023b】,但在使用Best-fit Packing时则屏蔽掉跨文档注意力,这样Best-fit Packing就等同于将每个文档视为一个独立的训练样本,尽管方式更为紧凑。得益于旋转位置嵌入(RoPE)【73, Su et al., 2021】的相对性,我们不调整模型输入的位置ID。我们在§C.3中对跨文档注意力进行了额外的消融研究。
优化器与训练设置。我们使用AdamW优化器【38, Loshchilov & Hutter, 2019】训练所有模型。我们使用3e-4的学习率和余弦学习率调度器,并在前3000步进行预热。全局批大小为2M词元。我们使用FlashAttention2【12, Dao, 2023】来加速训练。我们发现FlashAttention2中对文档级块状注意力的实现与标准的全序列注意力相比,没有产生可感知的开销。所有模型都在一个由256个A100 GPU组成的集群上进行训练。
数据流水线。Best-fit Packing的数据流水线实现如下。在离线阶段,我们将原始数据分词成最大序列长度的文档块,并运行BFD算法获取每个训练序列的块索引。在训练(在线阶段)期间,对于每个序列,我们通过块索引获取相应的文档块,并动态构建训练序列。这节省了在磁盘上拼接单个块的昂贵成本。
数据集。对于The Stack数据集,我们使用了7种流行编程语言的子集:Python、Java、C#、JavaScript、TypeScript、C和C++。
XSUM上的指标选择。我们没有使用SummaC【31, Laban et al., 2022】和QAFactEval【17, Fabbri et al., 2022】来评估XSUM的忠实性,原因有二。首先,SummaC假设忠实摘要中的每个句子都可以从文章的另一个句子中推断出来。这通常不适用于要求用一句话总结文章的XSUM。如果这唯一的一句话只是文章中一个句子的引申,那么摘要的覆盖范围将非常有限。其次,根据文献报道【17, Fabbri et al., 2022】,这两种方法在基于XSUM构建的摘要不一致性检测基准(即XSF【44, Maynez et al., 2020】)上的准确性低于在CNN/DailyMail的其他基准上的表现。
CNN/DailyMail上的指标选择。在CNN/DailyMail上,我们没有报告FAVA二元分数,因为它没有考虑摘要的长度。该指标偏向于较短的生成结果,因为它们产生幻觉的机会较小。如表6所示,用拼接和打包训练的模型生成的句子数量有显著差异。相反,在XSUM上,所有模型都能遵循指令生成单句摘要。
总结。总的来说,现有的评估摘要忠实性的方法有其局限性,这很大程度上是由于人类语言的复杂性。相比之下,依赖于程序分析的代码幻觉检测往往更可靠。
实验设置。Best-fit Packing将每个段落视为独立的,因此在训练中屏蔽了跨文档注意力(附录C.1)。为了研究这一选择的影响并凸显我们打包方法的重要性,我们预训练了一个2k NL模型,该模型使用拼接方法并启用了跨文档注意力掩码。在表10中,我们通过验证集上的困惑度(PPL)和下游任务的性能来比较这三个模型。注意,验证集中的每个序列只包含一个文档块。
实验结果。我们发现,尽管困惑度的增益可以归因于注意力掩码,但Best-fit Packing对于在下游任务中实现最佳性能至关重要。特别地,在传统拼接方法之上使用注意力掩码甚至会降低模型在摘要任务中的性能。
表10. 跨文档注意力的消融结果。我们额外训练了一个2k长度的模型,使用拼接和跨文档注意力掩码。我们报告了验证集上的困惑度(PPL),以及在阅读理解(RDC)、自然语言推理(NLI)、上下文遵循(CTX)、摘要(SUM,以ROUGE-2计)和常识(CMS)任务中的平均性能。