✨ 要点🔬 技术摘要
想象一下,你有一个巨大的、极其聪明的机器人大脑(大型语言模型),你想教它一个新技巧。通常,为了教这个大脑,你需要一台拥有海量内存和大量时间的超级计算机。但如果想用成千上万个小型、性能较弱的设备(如智能手机或智能家居设备)来教这个大脑,且永远不把它们的隐私数据发送到中央计算机呢?这被称为联邦学习(Federated Learning) 。
问题在于,这些小型设备太弱了,无法处理“反向传播”(标准的 AI 教学方法),因为这种方法需要存储大量的临时数据。如果它们尝试这样做,就会耗尽内存并导致崩溃。
旧方案:零阶优化(Zeroth-Order Optimization, ZO)
为了解决内存问题,研究人员开发了一种叫做**零阶优化(ZO)**的方法。
类比: 想象你正试图在一个雾气弥漫的山谷中寻找最低点(即教导机器人的最佳方式),但你手里没有地图。传统的教学方法(反向传播)就像有一个 GPS,能准确告诉你哪边是下坡;而 ZO 方法则像是用棍子在地上试探 。你向随机方向戳向地面,看看它是上升了还是下降了。
代价: 为了获得关于哪边是下坡的准确信息,你必须在许多、许多个方向上进行试探。虽然这节省了内存(你不需要 GPS 地图),但由于你试探得如此频繁,过程会变得非常缓慢且计算成本极高。
新方案:FedSPZO
该论文的作者提出了一种名为 FedSPZO (联邦拆分扰动零阶优化)的新方法。他们发现了一种巧妙的方法,可以在不损失准确性的情况下,让这种“试探”过程变得快得多。
以下是他们是如何实现的,我们使用**“两阶段厨房”**来进行类比:
拆分厨房: 想象机器人大脑是一个有两个区域的厨房:
区域 A(准备站): 一个巨大的区域,用于切碎和混合食材。这是模型的“第一块(first block)”。
区域 B(烤箱): 一个较小的区域,用于最后的烹饪。这是“第二块(second block)”。
旧方法(低效): 为了弄清楚如何改进整个厨房,旧方法会随机改变整个 厨房(包括准备站和烤箱),然后品尝食物。接着它会再次改变,然后再改。因为厨房很大,每次做微小的改变后都要品尝一次食物,这需要花费很长时间。
FedSPZO 的做法(高效):
第 1 步: 他们只对准备站 (区域 A)进行轻微的改变。他们保持烤箱(区域 B)完全不变。
第 2 步: 他们将准备站的输出 送入烤箱,并在烤箱中进行许多、许多次 运行,每次仅对烤箱进行微小的随机改变。
神奇之处: 因为准备站没有改变,所以每当测试烤箱时,厨房都不需要重新“切碎”食材。他们可以直接重复使用已经切好的食材(中间激活值),并只专注于测试烤箱。
结果: 与之前的相比,他们能以更少的总“品尝”步骤,更准确地掌握如何改进烤箱(以及间接改进准备站)。
他们发现了什么?
研究人员在几个著名的 AI 模型(如 RoBERTa、OPT 和 LLaMA)上进行了测试,发现:
速度: 他们的这种新方法在计算量方面比其他类似的“试探”方法快了高达 3 倍 。
内存: 它仍然使用极少的内存,非常适合像手机这样的微型设备,就像原始的“试探”方法一样。
通信: 他们向中央服务器发送的数据量非常小(仅为数字,而不是整个大脑),这对于网络连接较慢的情况非常有利。
准确性: 机器人的学习效果几乎与那些沉重的标准方法一样好,仅有极小的性能下降。
核心结论
可以将 FedSPZO 看作是一种聪明的方法,利用一群小型、弱性能的设备来教导一个巨大的机器人。他们并没有要求每个设备都进行大规模、耗费内存的计算,而是将任务拆分为两个部分。他们完成一次繁重的体力活,然后重复利用这项工作来快速测试许多微小的变化。这使得他们在通常无法胜任的设备上训练强大的 AI 成为可能,从而节省了时间、电池和数据。
技术摘要:面向资源受限设备的语言模型高效零阶联邦微调
1. 问题陈述
联邦学习 (FL) 为跨分布式数据源微调大语言模型 (LLMs) 提供了一个保护隐私的框架。然而,在资源受限的边缘设备上部署这些模型面临着重大障碍:
内存限制: 一阶方法(例如标准反向传播)和参数高效技术(如 LoRA)需要存储用于梯度计算的中间激活值。这导致内存占用(例如 RoBERTa-large 约为 4 GB)超出了大多数边缘设备的能力范围。
通信开销: 传输全量模型参数甚至低秩自适应矩阵 (LoRA) 会产生巨大的上传成本,这在低带宽无线环境中尤为严重。
零阶 (ZO) 方法的局限性: 虽然零阶优化 (ZO) 方法(如 MeZO)通过前向传播和伪随机扰动消除了对反向传播的需求,但它们面临收敛速度慢和计算需求高的难题。为了获得稳定的梯度估计,现有的 ZO-FL 方法在每个训练步骤中需要大量的扰动,从而导致过多的仅前向计算 (FLOPs)。
这里存在一个关键的权衡:目前的 ZO 方法虽然内存效率高,但计算成本高;而一阶方法虽然计算效率高,但内存密集。
2. 方法论:FedSPZO
作者提出了 Federated Split-Perturbation Zeroth-order Optimization (FedSPZO) ,这是一个旨在降低 ZO-FL 计算负担,同时保留其内存和通信优势的新型框架。
核心机制:分裂扰动策略
FedSPZO 利用神经网络的顺序组合结构,将模型分为两个连续的块:f 1 f_1 f 1 (第一个较大的块)和 f 2 f_2 f 2 (第二个较小的块),使得 y = f 2 ( θ 2 ; f 1 ( θ 1 ; B ) ) y = f_2(\theta_2; f_1(\theta_1; B)) y = f 2 ( θ 2 ; f 1 ( θ 1 ; B )) 。
非对称扰动预算: FedSPZO 并非对整个网络应用相同数量的扰动,而是将更高的扰动预算分配给第二个较小的块 (f 2 f_2 f 2 )。
块 1 (θ 1 \theta_1 θ 1 ): 使用 P 1 P_1 P 1 个方向进行扰动。
块 2 (θ 2 \theta_2 θ 2 ): 使用 P 2 P_2 P 2 个方向进行扰动,其中 P 2 = 2 × P 1 × P s P_2 = 2 \times P_1 \times P_s P 2 = 2 × P 1 × P s 。
激活重用: 关键的效率提升来自于中间激活值的重用。
客户端在单个方向 (+ z 1 +z_1 + z 1 ) 上扰动 θ 1 \theta_1 θ 1 ,并通过 f 1 f_1 f 1 进行一次前向传播,生成中间输出 + y l +y_l + y l 。
该单一的 + y l +y_l + y l 随后被用作 P s P_s P s 个不同的 θ 2 \theta_2 θ 2 扰动(包括正向和负向方向)的输入。
该过程在 θ 1 \theta_1 θ 1 的负方向 (− z 1 -z_1 − z 1 ) 上重复进行,生成 − y l -y_l − y l ,并将其重用于另一组 P s P_s P s 个 θ 2 \theta_2 θ 2 的扰动。
梯度估计:
θ 2 \theta_2 θ 2 的梯度通过 P 2 P_2 P 2 个扰动的损失差异来估计,这些扰动共享来自 f 1 f_1 f 1 的相同输入。
θ 1 \theta_1 θ 1 的梯度通过平均跨多个 f 2 f_2 f 2 扰动的损失差异来估计,这有效地减少了由 f 2 f_2 f 2 扰动的随机性引入的噪声。
通信协议
FedSPZO 保持了通信高效的设计:
客户端侧: 客户端进行 K K K 步本地训练。它们不上传模型参数。相反,它们仅上传标量梯度值 (G 1 , G 2 G_1, G_2 G 1 , G 2 ) 和用于再生扰动向量的随机种子 (S 1 , S 2 S_1, S_2 S 1 , S 2 )。
服务端侧: 服务器通过从接收到的种子中再生扰动向量并应用标量更新来重建精确的客户端模型。这消除了客户端传输模型权重的需求,使上传开销比基于参数的 FL 降低了数个数量级。
3. 核心贡献
FedSPZO 框架: 一种新型 ZO-FL 方法,它将模型划分为两个块,并将较大的扰动预算分配给更深层的较小块。这利用了中间激活值的重用,以更少的前向评估次数来更新整个网络。
计算效率: 该方法与最先进的 ZO-FL 基准(特别是 Fang 等人,2022 年,以及 Li 等人,2025 年)相比,实现了高达 3 倍的计算量减少 ,同时保留了 ZO 的内存和通信优势。
全面评估: 在 RoBERTa-large、OPT-1.3B 和 LLaMA-3.2B 模型以及多个数据集(SST-2, RTE, WiC, MultiRC, SQuAD, BoolQ)上的广泛实验表明,FedSPZO 显著降低了总 GFLOPs,同时保持了与一阶反向传播方法相当的准确度(仅有轻微下降)。
消融分析: 作者进行了关于模型拆分、扰动预算以及独立块间梯度计算的研究,验证了分裂扰动策略优于独立块估计或均匀扰动方案。
4. 实验结果
评估将 FedSPZO 与一阶方法(FedAvg, FedAvg+LoRA)和其他 ZO 方法(FedZO, DecomFL)进行了对比。
内存效率: FedSPZO 保持了类似于推理的内存占用。对于上下文长度为 256 的 RoBERTa-large,它需要约 1.75 GB,而 FedAvg 需要约 4.9 GB,LoRA 需要约 3.9 GB。这使得 FedSPZO 对于 LoRA 和标准 FL 无法胜任的边缘设备而言是可行的。
通信开销: 与 LoRA 相比,FedSPZO 将上传成本降低了超过三个数量级 。LoRA 上传的是低秩矩阵(数百 MB),而 FedSPZO 仅上传标量梯度和种子(千字节到兆字节)。
计算成本:
与 FedAvg(LoRA) 相比:由于 ZO 方法固有的收敛速度较慢,FedSPZO 的总计算量更高(约 17 倍)。
与 其他 ZO 方法 相比:FedSPZO 的效率显著提高。在各种模型和数据集上,它比 DecomFL 和 FedZO 的计算量减少了 2.1 倍至 3 倍 。
准确度: FedSPZO 达到了与 FedAvg(LoRA) 相当的准确度,仅有轻微下降(例如在 SST-2 上 <1%,在 RTE 上约为 4%)。它的表现略逊于全量 FedAvg(反向传播),后者作为性能上限,但在考虑资源约束的情况下,这种差距被认为是微不足道的。
训练时间: 在模拟边缘环境(Jetson Orin Nano)中,FedSPFO(RoBERTa/SST-2 为 2.6 小时)比 LoRA(0.45 小时)慢,但比 DecomFL(6 小时)和 FedZO(166 小时)快得多。
5. 重要性与主张
本文将 FedSPZO 定位为针对资源受限环境的务实权衡 。
对于边缘设备: 当内存和通信带宽是主要瓶颈,且 LoRA 和标准反向传播无法运行的情况下,它被视为一种首选方案。
平衡之道: 作者承认 FedSPZO 并没有消除 ZO 的计算成本,而是显著减轻了这一负担。它为在存储激活值进行反向传播是不可能的、且传输全量模型更新过于昂贵的设备上进行 LLM 微调提供了一条可行路径。
未来展望: 作者指出,虽然 FedSPZO 提高了现有 ZO 方法的计算效率,但在原始速度上仍落后于一阶方法。未来的工作旨在探索在低精度仅推理加速器上的 ZO,以进一步缩小这一差距。
文章总结道,FedSPZO 成功地将高内存/高通信成本的一阶 FL 与高计算成本的传统 ZO 解耦,使得在以往被认为不可行的场景下进行 LLM 微调成为可能。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。