想象一下,你拥有一座巨大且极其复杂的图书馆(一个大语言模型),你想教它一项新技能,比如回答关于特定主题的问题。通常情况下,要教导这座图书馆,你需要派出一支编辑团队,逐页阅读每一本书并记录下如何修改文本。这被称为“反向传播”(backpropagation),它需要消耗大量的内存和能量。如果图书馆规模过大,你的团队就会因为没有足够的空间记录笔记而导致过程崩溃。
问题所在:“盲目”的方法
一些研究人员尝试了一种不同的方法,叫做**零阶(Zeroth-Order, ZO)**优化。他们不再通过阅读每一页来寻找精确的错误,而是直接进行猜测。他们对书中的内容进行微小的、随机的改动,观察故事是否变得更好;然后进行另一种随机改动,再观察一次。这就像是在尝试通过随机摇晃一把巨大的钥匙扣上的钥匙,直到找到正确的那一把。
这种“盲目”猜测的问题在于,面对数以十亿计的参数,这些猜测会变得极其嘈杂且低效。你可能会去摇晃一把根本无关紧要的钥匙,从而浪费时间和精力。
解决方案:ZO-Act(“智能地图”方法)
论文介绍了一种更聪明的猜测方式——ZO-Act。与其随机摇晃任何一把钥匙,ZO-Act 首先观察图书馆对一些样本问题(称为“激活值”)的自然反应。
你可以这样理解:
- 单次成图(The One-Shot Map): 在开始教学之前,你会先问图书馆几个练习题。你注意到,当人们询问关于“历史”的问题时,图书馆内部的齿轮大多会以一种特定的、可预测的模式旋转。于是,你绘制了一张仅针对这些旋转齿轮的地图。
- 固定子空间(The Fixed Subspace): 你决定从现在起,只调整你地图上的那些齿轮。你冻结了图书馆的其他部分,以免它们被弄乱。
- 轻量级教练(The Lightweight Coach): 你不再试图直接移动那些沉重的、被冻结的齿轮,而是在你的地图上安装一个小型、轻便的杠杆(一个“系数矩阵”)。你只需要摇动这个小小的杠杆,就能带动沉重的齿轮向正确的方向移动。
为什么它是游戏规则的改变者
- 更少的噪声: 因为你只是在摇动一个小杠杆,而不是整个机器,所以你的猜测要准确得多。这就像是通过调节音量旋钮来调好收音机,而不是试图重建整个天线。
- 更快的学习: 由于“杠杆”很小,你可以使用标准的、强大的工具(如 Adam 优化器)来高效地摇动它。
- 适用于微型设备: 因为主图书馆保持冻结状态,你只需触碰这个小小的杠杆,所以你不需要超级计算机。即使是在内存非常有限的设备(量化模型)上,你也可以完成这项工作,这就像是用一本小笔记本而不是一部百科全书来教导图书馆。
论文的研究发现
研究人员在大型模型(如 Llama-3 和 OPT)上进行了测试,结果发现:
- 它比随机猜测效果更好: 在语言理解、回答问题和解决逻辑谜题等任务中,ZO-Act 的表现始终优于其他“盲目”方法。
- 它可以在小型计算机上运行: 即使当图书馆被压缩以适应小型设备(INT4 量化)时,ZO-Act 仍然表现出色。
- 地图是稳定的: 他们检查了最初绘制的“地图”是否依然有用。结果显示,这张地图依然有效!最重要的齿轮在整个训练过程中始终保持着相同的旋转方式,这证明了最初的“单次成图”是一个明智的选择。
代价
论文承认,虽然 ZO-Act 在节省内存和避免需要反向传播方面表现出色,但如果拥有无限的内存,它仍然无法达到传统“全能编辑”(一阶方法)那样的速度或完美程度。这是一种权衡:你得到的是一个极其高效、低内存的解决方案,它“足够好”,并且通常优于其他低内存技巧,但它本质上仍然依赖于猜测,而非直接看到确切答案。
总结
ZO-Act 就像是给一位蒙着眼睛的学生提供了一张智能地图,指明了巨大建筑中最关键的房间。他们不再是盲目摸索,而是只通过一个小型、易于操作的工具,来调整这些特定房间里的家具。这节省了能量,减少了错误,甚至让他们能在狭小的公寓里学习新技能。
技术摘要:ZO-Act
问题陈述
微调大语言模型(LLMs)通常依赖于一阶优化(反向传播),但随着模型规模扩展到数十亿参数,由于对激活值、优化器状态和梯度计算的存储需求,这会变得内存成本极高。零阶(ZO)优化提供了一种仅通过前向评估来估计更新方向的前向方案,从而避免了反向传播。然而,现有的 ZO 方法面临着重大挑战:
- 高方差: 对全参数空间进行扰动(如 MeZO)会导致梯度估计的高方差,从而导致收敛缓慢且不稳定。
- 次优子空间: 虽然一些方法将扰动限制在低维子空间内(例如 LOZO、SubZero),但它们通常依赖于随机构建的子空间或可能与实际梯度流不一致的随机扰动。
- 优化器兼容性: 许多 ZO 方法难以有效地利用标准的基于动量的优化器(如 Adam),因为其扰动结构无法暴露显式的、轻量级的可训练变量。
- 量化约束: 微调量化后的 LLM(例如 INT4)非常困难,因为这些方法需要实例化全权重扰动。
方法论:ZO-Act
作者提出了 ZO-Act,一种单次(one-shot)、由激活信息驱动的 ZO 微调方法。其核心洞察在于:线性层中权重梯度的主要成分位于输入激活矩阵的右奇异向量所张成的子空间内。
核心机制
- 单次初始化: 对于每个线性层,ZO-Act 通过在一个校准批次(calibration batch)上执行单次前向传播来计算输入激活矩阵 X。随后,通过薄奇异值分解(SVD)计算 X 的前 r 个右奇异向量 (Vr)。
- 固定子空间参数化: 权重更新被限制在由这些冻结的激活向量定义的低秩子空间内。有效的权重更新参数化为:
ΔW=VrB
其中 Vr∈Rm×r 是冻结的激活基底,B∈Rr×n 是轻量级的、可训练的系数矩阵。原始预训练权重 W 和基底 Vr 在整个训练过程中保持冻结。
- 在系数空间进行优化: ZO-Act 不再扰动全权重矩阵 W,而是在低维系数空间中采样随机扰动 Z (Z∼N(0,I))。有效的扰动权重变为 Weff=W+Vr(B+μZ)。
- 仅前向估计: B 的梯度通过前向差分损失评估来估计:
g^B=μL(B+μZ)−L(B)Z
该估计梯度随后用于使用标准的一阶优化器(如 Adam)来更新 B。
理论分析
论文将 ZO-Act 分析为针对受限系数空间 β(维度为 k)而非全权重空间(维度为 d)的零阶优化。
- 方差缩减: ZO 估计器的方差随扰动维度线性缩放。通过将维度从 d(全权重)降低到 k=∑rnℓ(系数,其中 r≪m),ZO-Act 显著降低了与方差相关的收敛项和有限差分误差。
- 偏差-方差权衡: 将更新限制在激活信息驱动的子空间内会引入子空间近似偏差。然而,作者认为 LLM 的激活值和梯度表现出低秩结构,这意味着主要的更新方向可以被输入激活的前奇异向量很好地捕捉。因此,这种偏差是受控的,并通过数据驱动的子空间特性得到了缓解。
核心贡献
- 激活信息驱动的子空间: 与使用随机或静态低秩矩阵的先前子空间方法不同,ZO-Act 直接从输入激活中导出扰动子空间,确保该子空间与层的占主导地位的激活方向相一致。
- 显式系数优化: 通过冻结基底并仅优化系数矩阵 B,ZO-Act 将 ZO 微调转化为一个在低维空间上的显式优化问题。这使得直接应用基于动量的优化器(如 Adam)成为可能,并避免了全权重扰动的实例化。
- 量化兼容性: 该方法天然支持量化后的 LLM(如 INT4),因为低比特权重保持冻结,且仅更新低比特的系数矩阵。
- 理论保证: 论文提供了收敛性分析,表明与全权重 ZO 相比,ZO-Act 在降低梯度估计方差和有限差分误差的同时,付出的代价是可控的子空间偏差。
实验结果
作者在 Llama-3-8B、OPT-13B 和 INT4 量化的 Llama-3-8B 上评估了 ZO-Act,涵盖了语言理解、问答和常识推理任务。
- 性能: ZO-Act 一致优于强力的 ZO 基线方法(包括 MeZO、LOZO、SubZero、AGZO 和 ZO-Muon)。
- 在 Llama-3-8B 上,它在 SST-2、RTE、BoolQ 和 WiC 任务上取得了最佳的 full-precision ZO 结果。值得注意的是,它将 RTE 的性能从 81.2 (ZO-Muon) 提升至 87.0,将 CB 从 69.6 提升至 89.3。
- 在 OPT-13B 上,它同样在 SST-2、RTE、CB 和 BoolQ 上取得了顶尖结果。
- 在 INT4 量化模型上,ZO-Act 仍能与全精度 ZO 基线保持竞争力,证明了其在内存受限微调中的有效性。
- 效率: 在固定的前向查询预算下,ZO-Act 的内存使用量在对比方法中最低(Llama-3-8B 约为 16.1 GB,对于 INT4 则降至 5.8 GB),且实现了最快的运行时间(42.0 分钟,而其他方法为 44.5 分钟以上)。
- 秩敏感性: 实验表明,对于全精度模型,秩 r=1 通常就足够了,能在最小化扰动维度的同时获得最佳精度。对于量化模型,较高的秩(r=32)则有助于补偿容量的降低。
- 稳定性: 诊断研究证实,激活信息驱动的子空间在微调过程中保持稳定,并能捕捉到很大一部分全梯度能量,验证了“单次(one-shot)”初始化的设计。
意义与主张
论文声称,ZO-Act 代表了使零阶微调成为大规模及量化 LLM 的实用且高性能替代方案的重要进步。其意义在于:
- 弥合差距: 它通过使用数据驱动的子空间减少估计方差,缩小了 ZO 与一阶微调之间的性能差距。
- 赋能优化器: 它通过暴露显式的、低维的可训练变量,解决了将 Adam 等标准优化器应用于 ZO 方法的难题。
- 可扩展性: 它提供了一个内存高效的方案用于微调量化模型,这对于在资源受限的硬件上部署至关重要。
作者承认了局限性,指出由于 ZO 估计固有的噪声和子空间偏差,ZO-Act 仍未能完全消除与一阶方法的性能差距。此外,由于每次更新需要多次前向传递,ZO-Act 的实际运行时间(wall-clock time)仍比基于反向传播的方法慢。然而,该方法被定位为在反向传播因内存限制而无法使用或不可用的场景下的优选方案。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。