想象一下,你正试图用乐高积木搭建一座宏大而精巧的城堡。通常情况下,唯一的办法是让一个庞大的建筑团队同时在整个结构上协作。如果一个人掉落了一块积木,整个团队都必须停下;而且如果团队规模变得太大,房间会变得拥挤,项目也会变得极其昂贵且难以修复。这就是科学家们通常构建“大语言模型”(LLM)的方式——那些能写故事、回答问题并与我们聊天的超级智能计算机大脑。它们是作为一个巨大的、不可分割的整体进行训练的。但如果,你可以把城堡建成一个个独立的、较小的房间,让不同的团队独立地完成每个房间的建造,然后在最后将它们全部拼装在一起呢?这正是这篇论文所提出的核心问题:我们能否将智能人工智能的训练分解成更小、更易于管理的碎片,并在稍后重新组装它们,同时又不丢失那份“魔力”?
这项研究背后的研究人员正在开展一项名为“混合训练”(Mixture of Training,简称 MoT)的项目,他们决定测试这个想法。他们不仅仅是在凭空猜测;他们构建了一个拥有 13 亿参数的人工智能模型(一种“Gemma 式”模型),并尝试对其进行“切片式”训练。他们并没有一次性训练整个模型,而是将模型分成了两个大块。为了确保这些块稍后能够完美契合,他们使用了一个聪明的技巧:一个“脚手架”。想象一下,这是一套辅助训练轮或一个刚性的框架,在各个部分被建造时将其固定在原位。他们冻结了这个框架(称为“对齐器”,aligner)使其保持不变,然后在这个框架内分别训练每一块人工智能。通过这种方式,即使每一块是在分别训练的,它们也能学会使用相同的“数据语言”。
结果既有“成功了!”也有“视情况而定”的情况。团队发现,他们确实可以将这些独立训练的块拼接在一起,并得到一个可以运行的人工智能。然而,就在他们将这些块拼在一起的那一刻,人工智能变得有些困惑,犯错也变多了(困惑度得分从标准的 15.0 上升到了 19.3)。但有趣的部分在于,经过一段非常短的“微调”过程(即让他们让整个模型进行一段时间的自我对话),人工智能就变得和传统的“巨型整体块”方法一样出色了。事实上,他们发现了一种训练方法,其消耗的计算能力与标准方法相同,且能达到完全相同的质量。
但这里有一个限制。论文指出,这种方法并不是一个能让所有人立即省钱的“魔杖”。“脚手架”本身在最初建造时需要消耗大量的能量。如果你只建造一个人工智能,脚手架会让整个过程比传统方法更昂贵。然而,如果你重复使用同一个脚手架来建造许多个不同的人工智能(比如使用同一套辅助训练轮来建造三座或更多城堡),那么每座城堡的成本就会下降,这种方法也就成为了赢家。研究人员还发现,如果你试图将人工智能拆分成过多的微小碎片,它们就会变得难以拼凑,质量也会下降。因此,虽然这目前还不能完全取代那种“巨型一体化”的训练方法,但它证明了将人工智能训练分解为更小、可重用且可并行的任务是可行的。它为人工智能的实验开辟了一种新途径,使得修复错误、重启部分工作以及在未来训练更智能的模型变得更加容易,而不再需要一个规模如城市般巨大的超级计算机。
技术摘要:混合训练 (Mixture of Training, MoT)
问题陈述
大型语言模型 (LLM) 的预训练目前被组织为单一、单体式的端到端优化过程。这种架构造成了扩展瓶颈:所有层必须同时进行训练,这使得整个过程在面对故障时显得非常脆弱(单个错误会导致整个运行停止),并且在迭代研究方面缺乏灵活性(改进工作往往需要重启或继续一个庞大的耦合系统)。本文探讨了是否可以将预训练分解为更小的、可独立训练的任务,并能在不牺牲质量的前提下,随后将这些任务重新组合成一个更大型的连贯模型。
方法论:混合训练 (MoT)
作者引入了 混合训练 (MoT),这是一种脚手架式模块化预训练程序,旨在将目标 Transformer 划分为连续的层块。该方法分为三个阶段:
- 阶段 0(准备): 选择或准备一个预训练的对齐器 (aligner) (A)。目标 Transformer (F) 与对齐器都被划分为 K 个连续的层块。对齐器在形状上与目标模型兼容,共享全局宽度、注意力头维度、前馈宽度、Token 嵌入以及输出头。
- 阶段 1(脚手架训练): 将目标模型分解为 K 个子模型 (f1,…,fK)。对于每个子模型 fi,通过将 fi 插入冻结的对齐器结构中来构建脚手架网络 Si:
Si=aK∘⋯∘ai+1∘fi∘ai−1∘⋯∘a1
在此设置中,只有目标块 fi 是可训练的;所有的对齐器切片均为冻结状态。每个脚手架使用标准的下一 Token 预测损失进行独立优化。至关重要的是,在此阶段目标块之间不进行梯度交换,但它们都在由共享对齐器提供的稳定表示上下文中进行学习。
- 阶段 2(重组与适配): 丢弃对齐器。将训练好的目标块重组成完整模型 F^=fK∘⋯∘f1。这个“冷组合 (cold-composed)”模型可以直接进行评估,或者可以经过一段短时间的端到端适配过程,以解决可能存在的表示失配问题。
核心贡献
- 脚手架式模块化训练: 一种通过冻结的对齐器保留共享表示接口,从而独立训练目标深度切片的程序,实现了无需事后缝合层即可进行重组。
- 机制证明: 通过一项小规模研究,证明了独立训练的深度切片可以成功重组为一个可用的 1.3B 参数语言模型。
- 计算与调度核算: 对计算量 (FLOPs)、Token 暴露量以及理想化的层等效关键路径进行了详细分析。这包括对对齐器成本的显式处理,区分了“直接”训练成本与“全额计费 (fully charged)”成本(即将完整的对齐器成本分配给单次运行)以及多次运行下的摊销成本。
- 消融实验: 通过实证证据展示了对齐器的必要性、不相交数据流的影响,以及与增加拆分数量 (K) 相关的质量-效率权衡。
实验结果
作者在 C4 数据集上,针对一个 12 层、1.3B 参数的 Gemma 风格模型评估了 MoT。
- 质量对等性: 一个“质量对等”调度方案(通过延长子模型训练并包含 30k 步的适配过程)达到了与单体基准模型相同的困惑度 (15.0)。
- 计算权衡:
- “质量对等”调度方案处理了更多的总 Token 数 (47.1B vs. 33.6B),且具有更短的理想化关键路径(具有 1.7 倍的加速潜力),但当对齐器成本被全额计费到单次运行时,其所需的总 FLOPs 更多 (285.0 EFLOPs vs. 268.4 EFLOPs)。
- 然而,如果对齐器在 R 个独立的运行中被复用,则每次运行的有效成本会下降。论文指出,当 R≥3 时,有效成本会低于单体基准模型。
- 一个低计算量方案(MoT + 15k 适配)在 189.4 EFLOPs(全额计费)下达到了 15.9 的困惑度,这低于基准预算,尽管没有提供在相同预算下的等计算量单体基准进行直接比较。
- 消融实验:
- 对齐器的必要性: 在没有对齐器的情况下,冷组合质量崩溃(PPL 38.9),证实了对齐器能够防止接口失配。
- 数据流: 在存在对齐器的情况下,子模型使用不相交的数据流比使用共享数据(PPL 20.3)能获得更好的冷组合质量(PPL 19.3)。
- 拆分数量: 将拆分次数从 K=2 增加到 K=4 会减少计算量,但会降低冷组合质量(PPL 24.8),凸显了并行性与接口复杂度之间的权衡。
意义与主张
本文将 MoT 定位为一种用于研究脚手架式子运行能否作为可复用训练单元的小规模框架,而非作为单体预训练的通用替代方案。
作者提出了三个主要结论:
- 如果是在共享的对齐器脚手架内进行训练,独立训练的子模型可以重组为连贯的语言模型;如果没有对齐器,冷组合将会失败。
- 独立训练的块之间的失配在很大程度上是可恢复的;一段短时间的适配过程可以弥补大部分差距,而更长的调度方案可以达到与单体训练相当的水平。
- MoT 改变了预训练的工程形态,创造了更小、可独立调度、可重启且可消融的任务。这特别适用于小规模训练研究,使其可能更便宜、更具可重复性且更容易迭代。
作者明确指出,目前的证据展示的是模块化可组合性和摊销复用调度方案,而非建立统一的计算或等硬件速度优势。本研究在规模上是有意限制的(仅限于一种模型族、一个数据集、仅使用困惑度指标),旨在作为实现可复用的脚手架式子运行这一具体设计空间的机制证明。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。