想象一下人工智能的世界,就像一座巨大且繁忙的图书馆,最聪明的书正由一群巨大的、隐形的抄写员撰写。这些抄写员就是“大语言模型”(LLM),他们每天都在变得更大、更聪明。但问题在于:为了编写这些书,他们需要一个如此庞大的图书馆,以至于建造它需要耗费巨资,维持灯火通明则需要消耗如山般的电力。这些模型内部的“权重”就像是抄写员的记忆库;记忆库越大,他们能学习的东西就越多,但这也意味着他们变得越来越沉重且昂贵。
有一段时间,科学家们认为让这些模型变小的唯一方法是先完整地写出它们,然后再尝试缩小它们,就像是在行李箱已经装满后尝试压缩它一样。也有人尝试在旅途中只使用行李箱的一小部分,同时保留一个“全尺寸”的备份。但如果从一开始就只用最核心、最轻便的物品来打包行李呢?这就是“低秩”(low-rank)训练的梦想:从第一天起就将模型构建得自然而轻量且高效。问题在于,当科学家们尝试从零开始构建这些轻量化模型时,它们总是会崩溃。它们内部的数学逻辑会变得疯狂,数值会爆炸,训练过程也会崩溃,就像汽车引擎转速过高导致气缸爆裂一样。大问题在于:我们能否从头开始构建这些轻量化模型而不让它们发生爆炸,而且它们是否依然能像那些沉重、昂贵的模型一样聪明?
这篇论文介绍了一种名为 Spectron 的新方法,它给出了肯定的回答。作者发现,这些轻量化模型之所以崩溃,是因为它们内部的数值正在不受控制地增长,就像一个没有塞子的气球在被吹大一样。他们发现,模型内部发生的变化的“规模”(称为谱范数/spectral norm)增长得太快、太大了。为了解决这个问题,他们发明了一个安全阀。想象一下你正在驾驶一艘速度极快且非常轻盈的小船。如果你把方向盘转得太猛,船就会失控旋转。Spectron 就像是一个自动驾驶仪,每当船试图转向过急时,它都会轻轻地将方向盘拨回到一个安全的角度。它通过不断检查船只运动的“规模”,并将其按比例缩小到足以保持稳定、又不至于让船停止移动的程度来实现这一点。
研究人员在多种不同规模的模型上测试了这一方法,范围从约 9400 万参数的小型模型到 4.54 亿参数的大型模型。他们发现,通过使用 Spectron,这些轻量化模型可以在没有“沉重”备份权重的情况下,从零开始进行训练。不仅如此,它们不仅保持了稳定性,而且学习效果也和那些庞大、沉重的模型一样好。事实上,当他们将一个使用 Spectron 训练的 4.54 亿参数模型与一个 7.8 亿参数的沉重模型进行比较时,他们发现这个轻量化模型可以用显著更少的资源达到同样的智能水平。这就像是找到了一种方法,可以制造出一辆速度与重型卡车一样快、但油耗却只有一半的跑车。
论文还探讨了如何构建这些模型才能获得最佳效果。他们通过模拟实验,试图找出模型规模与阅读数据量之间的完美平衡点。他们发现,对于这些轻量化模型,其“甜点位”(最佳平衡点)与重型模型略有不同:拥有一个稍小一点的模型并让它阅读稍多一点的数据会更好。这表明在未来,我们或许能够构建出极其聪明的人工智能系统,它们运行起来成本更低,也更容易适配普通计算机,而不是必须依赖超级计算机。作者基于他们的实验结果对这些结论充满信心,展示了该方法在不同规模和任务中都能可靠运行,为构建下一代高效人工智能提供了一套稳定的方案。
技术摘要:稳定原生低秩 LLM 预训练
1. 问题陈述
基础模型取得了显著成功,但其不断增加的参数量造成了巨大的计算和内存瓶颈。虽然低秩分解(将权重矩阵表示为 W=AB⊤)在微调阶段(如 LoRA)已被证明有效,但在从头开始进行预训练时应用该技术仍然难以实现。
此前尝试以原生低秩模式(不使用辅助全秩权重)进行训练的方法面临严重的稳定性问题。这项工作所识别的核心问题是权重更新中谱范数(最大奇异值)的失控增长。与通过 AdamW 等优化器保持谱范数有界的稠密训练不同,低秩分解允许存在无限的等效表示(W=(λA)(λ1B))。这种缩放不变性使得缩放因子 λ 会无限制增长,从而导致激活值爆炸和训练发散。现有方法(例如 Wei et al., 2024a; Huh et al., 2024)通过维持辅助全秩权重或使用混合架构来规避此问题,但这会产生显著的内存开销,或者无法实现真正的端到端分解训练。
2. 方法论:Spectron
作者提出了 Spectron,一种能够在没有辅助全秩组件的情况下,实现稳定、端到端原生低秩预训练的方法。Spectron 通过结合**谱重归一化(Spectral Renormalization)和梯度正交化(Gradient Orthogonalization)**两种机制来解决谱不稳定性。
核心机制
该方法旨在将复合权重更新 ΔW 的谱范数限制在预设半径 η 内(由学习率导出)。
- 梯度正交化: 该方法使用 Newton-Schulz 迭代对因子 A 和 B 的动量更新进行正交化。这确保了更新位于局部约束半径内,防止了任意方向性的增长。
- 自适应谱重归一化: 为了满足全局约束 ∥ΔW∥2≤η,该方法根据 A 和 B 当前的谱范数动态调整因子的学习率。
因子的更新规则为:
ΔA←∥A∥2+∥B∥2+1η⋅Ortho(MA)
ΔB←∥A∥2+∥B∥2+1η⋅Ortho(MB)
其中,∥A∥2 和 ∥B∥2 通过单步幂迭代(power iteration step)进行高效估计。
实现效率
- 正交化: 通过 5 次 Newton-Schulz 迭代实现。
- 谱估计: 通过 1 次幂迭代步骤实现。
- 开销: 总计算开销小于 1% 的 FLOPs,相比于需要约 25% 额外 FLOPs 的“自我引导(self-guided)”训练方法,这是一个显著的改进。
3. 主要贡献
- Spectron 算法: 一种自适应低秩因子更新方法,它将正交化与谱重归一化相结合。它限制了产生的权重更新谱范数,使得能够从随机初始化开始进行稳定的端到端训练,而无需辅助稠密组件。
- 经验验证: 证明了 Spectron 训练的分解模型在 FineWeb 数据集上的三个模型规模(94M, 297M, 454M 参数)上,性能达到了与或超过了稠密基准模型的水平。
- 计算最优缩放法则: 通过 IsoFLOP 分析,作者推导出了原生低秩 Transformer 的特定缩放法则:
- 最优模型大小:Nopt∝C0.479
- 最优训练 Token 数:Dopt∝C0.521
这些指数与 Chinchilla 法则(N∝C0.49, D∝C0.51)略有不同,表明计算最优的低秩模型更倾向于采用较小的架构并配合比例更大的数据集进行训练。
4. 结果
- 稳定性: Spectron 消除了在朴素 AdamW 和自我引导训练中观察到的损失峰值和发散现象。它在整个训练过程中保持了有界的谱范数。
- 性能:
- 在 454M 参数模型上,Spectron 实现了 12.11 的验证困惑度(perplexity),优于自我引导基准(13.70)和朴素 AdamW(14.57)。
- 在相同 FLOPs 下训练时,一个 454M 的分解模型达到了 780M 稠密模型的验证损失,有效地回收了约 42% 更多参数的模型性能。
- 下游任务(HellaSwag, PIQA, ARC-easy)显示出相对于基准模型和同等参数量稠密模型的持续提升。
- 推理效率: 由于达到相同困惑度所需的参数量减少,Spectron 模型提供了显著的推理成本降低(在现代计算预算下,估计比 Chinchilla 最优的稠密 Transformer 降低高达 50%)。
5. 意义与主张
论文声称,当应用适当的谱约束时,原生低秩训练是稠密预训练的一种可行且稳定的替代方案。
- 科学洞察: 这项工作揭示了在优化过程中,全秩表示并非严格必要;此前观察到的不稳定性是由于不受控制的谱增长,而非低秩结构的固有局限性。
- 实际影响: 通过消除对辅助全秩权重的需求,Spectron 使基础模型开发更加民主化,有望降低内存消耗和硬件要求。
- 缩放行为: 推导出的缩放法则表明,对于固定的计算预算,低秩模型的最佳策略是在更多数据上训练更小的模型,从而产生具有内在推理效率的模型。
承认的局限性:
- 实验是在高达 1.5B 参数的模型上进行的,规模小于万亿级 Token 的前沿领域。
- 由于使用了未融合的算子(unfused kernel operations),目前的实现与稠密训练相比存在实际运行时间(wall-clock time)上的差距,但作者指出这可以通过定制的融合算子(fused kernels)来解决。
- 该方法专注于 LLM,尽管其原理可能扩展到多模态架构。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。