联邦轻量级微调

# 联邦轻量级微调
来源: https://arxiv.org/html/2607.18343
11institutetext:思科系统公司, 美国
11email:\{rachanta,wilreed\}@cisco\.com###### 摘要

联邦微调受到通信瓶颈的制约：FedAvg 和伪梯度方案传输的载荷随模型规模缩放，而梯度压缩仅将其缩小一个常数因子。我们另辟蹊径。映射网络通过一个固定的仿射投影，从一个小的可训练隐变量生成网络权重；由于映射是共享且仿射的，对隐变量取平均*等价于*对生成权重取平均。我们通过两项改动将其转变为实用的低带宽联邦通道：一个*低秩、种子可重生成*的投影因子分解（将生成器内存从约 80GB 降至约 10MB），以及一个*增量*公式 θ=θ_pre+UV^⊤z，它在共享的中心预训练基座周围学习一个加性修正——即联邦微调，这正是该方法在大规模下有效的原因。一个冻结的正交分类头进一步将头部移出传输载荷，同时提升准确率。在 CIFAR-100 上使用 ResNet-18+GroupNorm，我们的方法（FLITE，联邦低秩迭代训练引擎）每轮每客户端通信 1,280 个浮点数（约 5KB）——降低了 8,718 倍——并达到 74.67%，与全权重 FedAvg 相差约 0.5 个百分点。平均恒等式保持到浮点精度（6×10^{-8}）；该方法在带宽-准确率帕累托前沿上比 PowerSGD 和 top-k 低一到两个数量级；在强非独立同分布偏差下，它匹配或超过全权重 FedAvg。int4 隐变量每轮仅需 648 字节，准确率不变，而 int4 全权重 FedAvg 则会崩溃至随机猜测。

## 1 引言

联邦训练和分布式训练允许众多工作者在不汇集数据的情况下改进共享模型 [21, 9]，但它们付出了高昂的通信代价。主导范式直接交换模型状态：FedAvg [21] 传输模型本身，而 DiLoCo [5] 等伪梯度方案传输一个全尺寸的权重增量作为外循环更新。这两者的每轮带宽均为 Θ(|W|)，即模型参数量，因此对于现代网络来说，即使是一个小型 ResNet 的消息也达到数十兆字节，主导了联邦开销。

这些方法隐含一个结构假设：沿*频率*轴降低带宽——即减少通信频率——而每条消息的大小被视为固定。即使专门的梯度压缩方法（低秩投影 [29]、top-k 稀疏化 [1]、量化 [2]）也只能将固定消息缩小一个常数因子，仍然传输一个大小与 |W| 成正比的对象，而降低同步频率则用摊销成本换取客户端漂移，这在数据异质性最强时最为有害。本文采用正交的*大小*轴：我们保持频繁通信，但使每条消息成为一个微小、与模型无关的隐变量，同时保留在权重空间中的精确 FedAvg 平均语义。

#### 映射网络：另一条途径。

一个近期且非常不同的想法——*映射网络*——训练一个小的*隐变量*向量 z，通过一个固定（冻结）的投影生成一个更大网络的所有权重。隐变量是模型状态的紧凑载体。关键在于，当生成器是仿射且共享时，*对隐变量取平均等价于对生成权重取平均*——因此联邦客户端只需通信小的隐变量，而服务器聚合保持普通 FedAvg。这个想法很有吸引力，但按原有方案存在两个障碍：(i) 投影矩阵本身巨大（对于 ResNet-18 来说约为几十吉字节），(ii) 隐变量无法从头训练驱动大型网络（会崩溃至接近随机猜测的准确率）。

#### 我们的方法。

FLITE（联邦低秩迭代训练引擎）通过解决这两个障碍并将方法重新定位到其优势区域，将映射网络隐变量转变为实用的低带宽通信通道。首先，我们用*低秩、种子可重生成*的因子分解替换密集投影，将生成器内存从约 80GB 降至约 2GB（如果从种子重生成则约 10MB），重生成计算量也相应降低。其次，我们不要求隐变量编码整个网络，而是用它编码*围绕共享、中心预训练冻结基座的加性增量*——即联邦*微调*，而非从头训练。这保留了精确的平均恒等式，绕过了从头训练的失败，并利用了这样一个事实：良好训练模型所需的修正经验上是低维的。最后，我们冻结一个*正交初始化*的分类头，将其移出通信载荷，并且如消融实验所示，这提升了准确率。

#### 结果。

在 CIFAR-100 上使用 ResNet-18+GroupNorm，我们的隐变量每轮每客户端通信 1,280 个浮点数（约 5KB）——降低了 8,718 倍——并达到 74.67%，与全权重 FedAvg 相差约 0.5 个百分点，且统计上与集中式基线相当。维度描绘了一个可控的带宽-准确率前沿；该方法在帕累托前沿上比 PowerSGD 和 top-k 低一到两个数量级；在强非独立同分布偏差下，它匹配或超过全权重 FedAvg，并且在客户端数量上保持稳定。冻结的正交头在零通信下增加了 +0.54 个百分点。由于维度低，隐变量对量化高度鲁棒：int4 隐变量每轮仅需 648 字节，准确率不变，而 int4 全权重 FedAvg 则会崩溃至随机猜测。

#### 贡献。

- • 一个*增量低秩映射* θ=θ_pre+UV^⊤z，它保留*精确* FedAvg 平均恒等式（验证至 6×10^{-8}），同时将生成器内存减少约 40 倍（通过种子重生成可进一步降至约 10MB）。
- • 一个*低带宽联邦微调协议*，每轮传输约 5KB 的隐变量，准确率相当，在带宽-准确率帕累托前沿上优于 PowerSGD 和 top-k，并在异质性和小客户端数量下具有更强的鲁棒性。
- • *冻结正交分类器*作为免费、可组合的设计选择，将头部移出载荷并提升准确率；隐变量通道的 *int4 量化鲁棒性*，每轮仅 648 字节，而同等比特数下的全权重 FedAvg 崩溃。
- • *负面结果*：表明映射网络作为独立训练方法在 ResNet-18 尺度下无论维度、秩或热启动都会崩溃，从而确定联邦微调是该想法能奏效的范畴（§5.7）。

本文其余部分介绍方法（§3）、实验设置（§4）、结果（§5）以及关于范围与局限性的结论（§6）。

## 2 相关工作

#### 联邦平均及其变体。

FedAvg [21] 建立了主导范式：客户端本地训练，服务器对其模型求平均。后续工作改进了对异质性和客户端漂移的鲁棒性——FedProx [16] 添加了近端项，SCAFFOLD [10] 使用控制变量，Matched Averaging [30] 在平均前对齐神经元，自适应服务器优化 [25] 稳定聚合。所有这些每轮通信一个大小为 Θ(|W|) 的对象；它们改变的是*如何*聚合权重，而不是消息的*大小*。我们的方法是互补的：它精确保留 FedAvg 平均规则，但将消息缩减为一个隐变量。我们进一步指出一个设置上的区别：FedAvg、FedProx 和 DiLoCo 通常作为*从头*联邦训练运行，服务器初始化 θ_0，联邦本身产生训练好的模型。我们的方案在*联邦微调*范式中运行，其中每个客户端——无论是我们的还是 FedAvg 基线——都从相同的共享预训练检查点开始。§5 中所有的带宽和准确率比较都使用这个匹配的设置，因此对比的是两种不同带宽下的微调协议，而不是微调方法与从头方法之间的对比。

#### 低频和伪梯度方法。

Local SGD [27,12] 和 DiLoCo [5] 通过减少同步频率来降低通信，在合并之间执行多次本地步骤。这降低了摊销成本，但每次合并传输全尺寸更新，并面临客户端漂移风险，尤其在异质性下。我们的隐变量足够小，可以实现频繁同步，我们证明这在固定字节预算下转化为鲁棒性优势（§5.5）。

#### 通信压缩。

大量文献对传输的梯度/更新进行压缩：低秩投影（PowerSGD [29]）、幅值稀疏化（top-k [1]、深度梯度压缩 [20]）和量化（QSGD [2]、signSGD [3]），通常配合误差反馈以保持收敛 [11]。这些方法将载荷缩小一个常数因子，但压缩后的对象仍然派生自 |W| 并随其缩放。我们直接与 PowerSGD 和 top-k 比较，并展示在匹配准确率下我们的隐变量便宜一到两个数量级（§5.3）；此外，这两种方法是可以组合的（隐变量本身可以被量化）。

#### 子空间训练、低秩适应和映射网络。

在低维随机子空间中训练已知足以满足许多目标，这一观察由内在维度研究 [14] 形式化，并被 LoRA [8] 用于参数高效微调。映射网络 [26] 更进一步，通过一个冻结的、核心为仿射投影的生成器从低维隐变量生成网络的完整权重。我们采用映射网络生成器的*仿射特例*，并在三个方面有所不同：(i) 我们将投影分解为低秩、种子可重生成的形式（原始方案物化密集投影），消除了其巨大的内存和计算消耗；(ii) 我们在共享预训练基座上以增量模式运行，这是该方法在大规模下奏效的原因（从头映射失败，§5.7）；(iii) 我们将其用作联邦*通信*通道，利用仿射映射允许的精确平均恒等式。与 LoRA 不同，我们的因子是冻结且共享的，使得对隐变量求平均等于对权重求平均——这是联邦聚合所需的性质。

#### 以优化为中心的分散式联邦学习和函数空间方法。

并行工作通过优化、拓扑和个性化改善分散式联邦学习，同时保持在参数空间中 [32,13,28,18,4]；它们的消息仍然随 |W| 缩放，而我们的隐变量通道是互补的。另一条不同的线在*函数空间*中协作，在共享探测集上交换预测（FedMD [15]、FedDF [19]）。这些方法将消息与 |W| 解耦，并容纳模型异质性，但优化的是不同的（蒸馏）目标，不能直接与精确参数空间平均比较。我们的通道是权重平均，在低维隐变量中精确执行。

#### 冻结分类器和归一化。

固定分类头已知大致无害 [6]，并且在使用正则多面体/ETF 分类器 [24]（受神经坍缩启发 [23]）等良好条件化几何结构时可能有益；FedBABU [22] 在联邦期间冻结头部。我们使用冻结正交头作为可组合组件，将分类器移出载荷并提升准确率（§5.6）。BatchNorm 统计量在异质客户端间无法自然平均，这促使在联邦设置 [7] 中使用 GroupNorm [31] 或本地 BatchNorm [17]；我们采用 GroupNorm 并接受绝对准确率低于 BatchNorm 最优水平。

## 3 方法

我们首先回顾映射网络生成器及其对联邦平均有吸引力的性质（§3.1）。然后识别出使其无法直接在现代网络规模下使用的两个障碍（§3.2），并介绍消除这些障碍的增量低秩参数化（§3.3）。最后，我们描述可选的种子重生成投影（§3.4）、冻结正交分类器（§3.5）以及完整的联邦协议（§3.6）。

### 3.1 映射网络与平均恒等式

映射网络 [26] 通过一个固定（冻结）的生成器，从一个小的可训练*隐变量*向量 z 生成目标层的权重。原始提案通过调制和非线性增强了此映射；我们研究其仿射特例，这是允许精确平均恒等式的部分。对于一个具有 P 个参数的层，且 z ∈ ℝ^d，d ≪ P：

θ = b + W_m z,    (1)

其中 b ∈ ℝ^P 是固定基向量，W_m ∈ ℝ^{P×d} 是固定投影；只有 z 被训练。在原始公式中，W_m 是一个*密集*随机矩阵，完整物化。由于映射 z ↦ θ 是仿射的，并且 b、W_m 在所有客户端间*共享*，对隐变量取平均等价于对生成权重取平均：

1/K ∑_{k=1}^K (b + W_m z_k) = b + W_m (1/K ∑_{k=1}^K z_k).