想象一下,你拥有一座巨大且极其复杂的图书馆(即大语言模型,LLM),它懂得如何写作、推理和聊天。你想教这个图书馆一项新的、特定的技能,比如编写法律合同或解决数学问题。这个过程被称为“微调”(fine-tuning)。
通常,教导这个图书馆需要一种非常昂贵的方法,叫做“反向传播”(backpropagation)。这就像一位老师,在学生写完每一句话后,都必须走遍整个图书馆,检查每一本书,并精确计算出应该如何调整学生的脑回路。这种方法很准确,但它需要海量的内存和能量,往往使得在标准计算机上完成这项工作变得不可能。
为了解决这个问题,研究人员开发了一种“零阶”(Zeroth-Order)方法(例如 MeZO)。这种方法不再是走遍整个图书馆,而是像一位蒙着眼睛的探险家。探险家做一个微小的尝试(扰动),观察结果是变好了还是变差了,然后朝着相反的方向再做一个微小的尝试。通过比较这两个结果,他们可以在从未见过全貌的情况下,判断出前进的方向。这节省了大量的内存。
问题:“一刀切”的地图
现有的蒙眼探险家使用一种固定的、随机的策略。他们根据一个标准的、枯燥的规则(比如掷一个公平的骰子)来猜测方向。这种方法效率低下。有时,图书馆中存在特定的“房间”(参数块),在这些房间里你需要非常精准;而在其他房间里,你可以更大胆一些。固定的策略并不知道这一点;它把图书馆的每个部分都视作同等对待。
解决方案:ZO 微调器(聪明的探险家)
作者创建了 ZO 微调器(ZO Fine-tuner),这是一个“学会如何学习”的系统。他们没有使用固定的规则,而是训练了一个微小的、极其高效的“教练”(一个小型的神经网络)来教导这位蒙眼的探险家如何更好地进行猜测。
以下是它的运作方式,使用了几个类比:
- 教练(ZO 微调器): 想象一位教练在观察探险家。教练不需要看到整个图书馆。相反,教练会观察一些简单的统计数据:“这个房间现在有多乱?”以及“上一次的尝试是有利还是有害?”基于这些信息,教练会告诉探险家:“在这个特定的房间里,迈出大胆的一步。而在那个其他的房间里,迈出小心翼翼的一小步。”
- 块策略(The Block Strategy): 图书馆规模庞大(拥有数十亿个参数)。如果教练试图为每一本书单独提供指令,速度会太慢且极其耗费内存。然而,研究人员注意到,图书馆是按“块”(类似于书籍的分区)组织的。因此,教练学会了为每个块提供一条指令。这就像是一位教练告诉整个团队的队员如何协同移动,而不是对着每一个单独的队员大喊大叫。这让内存占用保持在极低的水平。
- “训练一次,广泛复用”的魔力: 通常情况下,你必须为每个新任务训练一个新的教练。但作者发现了一个惊人的事实:无论你是在教数学还是在教写故事,图书馆的“形状”变化并不大。因此,他们仅在单个数据集(COPA)上训练了一次教练,然后就带着同一个教练去教导图书馆处理完全不同的任务(如情感分析或阅读理解),甚至是不同版本的图书馆。
- 结果: 在一个任务上训练出的教练,在所有其他任务上也表现得惊人地好。这就像是在特定的赛道上训练了一名司机,然后让他能在完全不同的公路上完美驾驶,而无需任何新的练习。
实验结果
论文在四个不同的语言模型和七个不同的任务上进行了测试。
- 性能: 在大约 82% 的案例中,这个新的“聪明探险家”比旧有的“固定规则”探险家更快地达到了更好的结果。
- 效率: 它不需要额外的太多内存。这个“教练”本身非常小(对于一个 300 亿参数的模型,其大小不到 2MB),就像是在一本 60GB 的百科全书里添加了一页笔记。
- 稳定性: 新方法对“学习率”(即步长的大小)的敏感度也较低。即使你告诉它要迈大步,它也不会像旧方法那样容易崩溃。
总结
这篇论文介绍了一种无需超级计算机即可教导巨型 AI 模型新技能的方法。他们用一个微小的、聪明的教练取代了僵化、随机的猜测策略,这个教练学会了如何针对模型的不同部分给出正确的猜测方向。一旦这个教练被训练好,就可以高效地复用来教导模型完成许多其他任务,从而节省了时间和计算机内存。
技术摘要:学习用于大语言模型微调的零阶优化器
问题陈述
使用标准的阶一阶优化器(如 Adam)对大语言模型(LLM)进行微调会产生极高的内存开销,由于需要存储优化器状态和反向传播激活值,其内存需求通常高达推理内存的 12 倍。虽然零阶(ZO)优化器(如 MeZO)作为一种仅依赖前向传播的内存高效替代方案已经出现,但现有方法通常采用手工设计的静态采样策略(例如标准正态分布)。这些静态策略无法适应不同模型或任务的具体结构特性,往往需要大量的超参数调优,并且无法充分利用损失函数的局部几何结构。
方法论:ZO Fine-tuner
作者提出了 ZO Fine-tuner,这是一种基于学习的零阶优化器,旨在通过“学习如何学习”(L2L)框架自动学习高效的扰动策略。其核心创新在于从静态、均匀的扰动转向自适应的块状(block-wise)扰动。
1. 块状扰动架构
受近期研究中关于 Transformer Hessian 矩阵呈现近似块对角结构的启发,该方法避免了坐标级参数化(这在大规模应用时会导致内存开销过大),转而采用块状自适应。
- PertNN:对于 LLM 的每个参数块 i,一个轻量级的辅助神经网络(PertNN)会预测一个共享的扰动方差 σt(i)。
- 输入:PertNN 的输入包括当前参数的紧凑统计量(块的均值和方差)、前一步的扰动方差以及前一步的损失值。
- 输出:网络输出更新后的方差 σt(i),用于构建对角方差矩阵 Σt=diag(σt(1)I,…,σt(n)I)。
- 归一化:为了将学习到的方差与有效学习率解耦,该方法对 Σt 进行归一化,使其 Frobenius 范数保持恒定(∥Σt∥F=d)。这确保了扰动幅度保持稳定,同时允许优化器学习相对的块间方差。
2. 通过“学习如何学习”(L2L)进行训练
优化器通过一个元目标(meta-objective)进行训练,其中 LLM 本身的微调轨迹充当监督信号。
- 元目标:目标是最小化单次 ZO 步之后(即 θ1 时刻)的后验损失 L(θ1)。损失相对于 PertNN 参数 ω 的梯度是通过重参数化技巧(采样 z∼N(0,I) 并设置 u=Σz)计算的,使得扰动过程是可微的。
- 训练策略:PertNN 在单个数据集(如 COPA)上使用一阶优化器进行训练,以生成模型状态。为了防止对低损失区域的过拟几,作者引入了周期性重置机制,在特定间隔后将 LLM 参数重新初始化为微调前的状态。
- 泛化性:一旦针对特定基座模型完成训练,ZO Fine-tuner 即可在各种下游任务和衍生检查点之间复用,无需进一步训练。
核心贡献
- 将 L2L 扩展至 LLM:论文证明了学习到的单一优化器(基于基座模型训练)可以有效地泛化到不同的下游任务和衍生检查点,验证了“一次训练,广泛复用”的工作流。
- 块状方差学习:作者提出了一种内存高效的块状参数化方法,具有每个块共享方差的设计。该设计在理论上基于 Transformer Hessian 的块对角结构,使得在基础模型规模下进行 L2L 成为可能,且不会产生过高的内存成本。
- 实验性能:在四种 LLM(LLaMA-3.2-1B, LLaMA-3.1-8B, Qwen2.5-14B, OPT-30B)和七个数据集上的测试表明,在收敛损失方面,ZO Fine-tuner 在 82.1% 的任务-模型组合中优于最强的零阶基线(MeZO, MeZO-AdamU, HIZOO, LOZO),平均准确率提升了 2.5%。
实验结果
- 收敛性与准确率:ZO Fine-tuner 比基线方法更快达到更低的损失。它展示了对学习率选择的卓越鲁棒性,通常在较低的学习率(例如 10−8)下也能实现与基线相当甚至更好的性能,而基线在此情况下往往无法收敛或变得不稳定。
- 迁移能力:在 COPA 数据集上训练的优化器成功迁移到了 SST-2, CB, SQuAD, WSC, BoolQ 和 DROP 等任务,以及数学推理任务(MetaMathQA)的模型微调中。
- 衍生检查点:该方法能从基座模型(LLaMA-3.1-8B)泛化到其指令微调版本(LLaMA-3.1-8B-Instruct),在损失和准确率上均优于 MeZO。
- 效率:
- 内存:与模型规模相比,内存开销微乎其微(例如 OPT-30B 上的辅助网络小于 2MB)。总内存占用与 MeZO 相当,且显著低于一阶方法(例如 OPT-30B 下,MeZO 约为 62GB,而一阶方法为 312GB)。
- 时间:查询 PertNN 引入的时间开销极小(对于较小的模型小于 3.4%,随着模型增大该比例降低),因为主要的成本仍然是零阶估计所需的前向传播过程。
重要性与主张
论文声称 ZO Fine-tuner 解决了 LLM 微调中关键的内存效率瓶颈,同时克服了静态零阶策略的僵化问题。通过利用“一次训练,广泛复用”的范式,作者认为模型创建者可以将预训练的学习型微调器随基座模型一同发布,从而使下游用户能够以接近推理时的内存成本,进行高效、高性能的微调。这项工作确立了学习自适应扰动分布在基础模型规模下的可行性,并为可扩展、内存高效的 LLM 适配提供了一条切实可行的路径。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。