[ACM MM 2024] FedBCGD：打破大模型联邦学习的通信壁垒

TL;DR

在联邦学习（FL）中，由于客户端上行带宽极其珍贵，传输如 ResNet-18 甚至 Transformer 这样的大模型往往成为性能瓶颈。本文提出的 FedBCGD (Federated Block Coordinate Gradient Descent) 框架，让每个客户端只负责上传模型的一个“切片”（Block），却能通过巧妙的动量补偿和全量本地更新策略，在不损失精度的情况下将通信效率提升数倍至十倍以上。

背景定位

传统的联邦优化算法（如 FedAvg）要求客户端每次都上传整个模型。随着模型参数量跌破亿级，这种“全量上传”模式在移动设备或边缘计算场景下已难以为继。虽然常见的压缩方法（如量化、剪枝）能减少体积，但它们大多是针对梯度本身的。FedBCGD 则另辟蹊径，从算法优化路径入手，首次将**块坐标下降（BCD）**思想引入水平联邦学习。

核心痛点：为什么简单的“采样上传”行不通？

如果像传统 BCD 那样，本地只更新一部分参数，其他参数保持冻结（Freezing），会导致严重的“块漂移”（Parameter Block Drift）。

• 现象：不同客户端更新不同的 Block，合并后发现各个 Block 之间失去了协同性，模型无法收敛。
• 作者直觉：本地训练必须更新所有参数来捕获完整的特征交互，但上传时可以只传增量最为关键的部分。

方法论详解：FedBCGD 的三大支柱

1. 结构化划分与共享层

作者将参数 $x$ 分为 $N$ 个非重叠块和 1 个极小的共享块 $x_s$。

• 共享块（Shared Block）：通常是深度模型的最后几层（如分类器）。研究发现这些层虽然参数量仅占 0.01%，但对模型稳定性至关重要。
• 客户端分工：不同的客户端被分配到不同的 Block，并行优化。

图 1：FedBCGD 框架示意图，展示了客户端如何分块处理参数。

2. 本地更新不冷冻（No Freezing）

与传统分块下降不同，FedBCGD 客户端在本地迭代时使用 SGD 更新所有层，确保特征提取器与分类器之间不会脱节。

3. 服务器端动量补偿与变差缩减（FedBCGD+）

为了弥补缺失块的信息，服务器端引入了动量项 $v_{(j)}$。针对 Non-IID（非独立同分布）数据导致的异构性问题，高级版 FedBCGD+ 引入了两套对照变量（Control Variates）：

• 客户端漂移控制：修正本地梯度方向，使其靠近全局最优。
• 随机方差缩减：参考 SVRG 技术，压制随机梯度带来的噪声。

实验战绩：全线超越 SOTA

在多个视觉任务上，FedBCGD 展示了恐怖的加速比：

• 收敛速度：在 CIFAR-100 + LeNet-5 实验中，FedBCGD 仅需 77 天（模拟时长）即可达到 40% 精度，而 FedAvg 需要 558 天，提速 7.3 倍。
• 大模型支持：在 ViT-Base 模型上，面对 Tiny ImageNet 数据集，收敛速度比 FedAvg 快 11.5 倍。
• 泛化性能：出人意料的是，即使通信量更小，FedBCGD 的最终精度往往比全量更新的分布式 SGD 更高，这暗示该算法有助于模型跳出尖锐局部极小值，获得更好的泛化能力（Generalization）。

表 1：CIFAR-100 数据集上各算法性能对比，FedBCGD 系列在收敛天数（d）上优势明显。

深度洞察

FedBCGD 的成功不仅仅在于减少了传输量。从优化理论看，它将原本 $O(1/ϵ)$ 的通信复杂度，在 $N$ 个 Block 的设置下降低到了之前的 $1/N$。这意味着：模型分得越细，通信收益越高（在一定范围内）。

总结与展望

FedBCGD 证明了在联邦学习中，“少即是多”的哲学是可行的。它通过算法层面的重构，解决了上行带宽的顽疾。 局限性：Block 的划分目前仍依赖人工经验（如手工按层划分）。 未来研究方向：如何利用神经网络的结构信息进行“自动动态分块”，以及如何自适应地选择每个轮次中贡献最大的 Block 进行上传。

本文基于 ACM Multimedia 2024 论文重构。