🤖 machine learning

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

本文提出了一种用于联邦学习的子空间优化方法（SSF），该方法通过在低维子空间中执行优化并采用回填式更新以保留残差信息，从而缓解由数据异构性引起的漂移，进而相较于现有方法在显著降低通信和内存开销的同时实现高精度。

原作者： Shuchen Zhu, Zhengyang Huang, Yuqi Xu, Peijin Li

发布于 2026-04-29

📖 1 分钟阅读☕ 轻松阅读

原作者： Shuchen Zhu, Zhengyang Huang, Yuqi Xu, Peijin Li

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一个庞大的小组项目，数百名学生（客户端）试图共同解决一个巨大的拼图，但由于隐私规则，他们无法分享实际的拼图碎片。相反，他们只向老师（服务器）发送关于他们认为拼图应为何种样子的笔记。

这就是联邦学习。通常，他们使用一种称为"FedAvg"的方法，即每个人都发送自己的最佳猜测，老师将其取平均。但存在一个问题：因为每个学生拥有不同的拼图碎片（异构数据），他们的猜测会逐渐偏离。他们开始完全解决不同的拼图，导致最终结果混乱不堪。

为了解决这个问题，聪明的研究人员发明了一种名为SCAFFOLD的方法。这就像老师给每个学生发一张“修正笔记”，让他们保持在同一条轨道上。然而，这些修正笔记非常巨大——就像每次更新都要发送一本 100 页的手册。如果学生们使用的是小型旧手机（资源受限设备），他们无法携带这些沉重的“手册”，互联网连接也会因此拥堵。

新方法的登场：SSF（子空间-SCAFFOLD）。

以下是 SSF 的工作原理，通过一个简单的类比来解释：

“素描本”与“完整蓝图”

想象学生们正在绘制一张巨大、详细的城市地图（大型 AI 模型）。

旧方法（SCAFFOLD）： 每次学生做出更改时，他们都会向老师发送一份完整的、高分辨率的、100 页的城市蓝图。老师检查后，发回一份巨大的修正笔记，学生据此更新他们的绘图。这很准确，但对于他们的背包和互联网来说太重了。
“子空间”方法（FedSub）： 为了节省空间，学生只发送一张仅 5 页的城市主要道路的小素描。这既快又轻。但是，如果老师试图基于这张小素描发送修正笔记，学生会感到困惑，因为素描没有显示公园或建筑物的细节。如果素描每周都在改变形状，旧的修正笔记就会变得无用，学生会迷失方向。
SSF 方法： 这是一个巧妙的折中方案。
1. 素描： 学生只将 5 页素描（低维子空间）发送给老师。这节省了海量的数据和电池。
2. 隐藏记忆： 这里的魔法在于：尽管他们只发送素描，但学生会在脑海中（或后台的硬盘上）保留完整的 100 页蓝图。
3. “回填”技巧： 当老师基于素描发送修正时，学生将该修正应用于素描，并使用一种特殊的“回填”技术来更新隐藏的完整蓝图。
4. 结果： 学生保持在正确的轨道上（就像沉重的 SCAFFOLD 方法一样），但只需携带轻便的素描本进行通信。

为什么这很重要？

该论文声称，SSF 解决了现代 AI 中的“三重威胁”问题：

计算： 因为它在小型素描上进行数学运算，而不是在巨大的地图上，所以速度更快。
内存： 由于繁重的工作是在后台完成的，而不是在活跃内存中，因此它占用的设备空间更少。
通信： 它发送的是微小的消息，而不是巨大的文件。

“稳定性”测试

研究人员通过两种场景测试了这一点：

数学玩具问题： 他们模拟了拥有非常不同数据的学生。他们发现，当素描变得太大或变化太频繁时，虽然“仅素描”方法（FedSub）最终会陷入困惑并崩溃（发散），但SSF 保持稳定并持续改进，效果几乎与沉重、缓慢的方法一样好。
真实图像识别（CIFAR-100）： 他们在识别图像的真实任务中进行了尝试。SSF 是表现第二好的方法，击败了标准方法（FedAvg）和“仅素描”方法，尽管它略逊于沉重、缓慢的方法（Full-SCAFFOLD）。

结论

该论文认为，SSF 兼具两者之长。它允许学生在小型设备上高效协作，而不会丢失那些防止他们偏离轨道的“修正笔记”。它证明你不必在“快速/轻量”和“准确/稳定”之间做出选择；你可以通过将“沉重”的信息隐藏在后台，只发送“轻量”版本，从而同时拥有两者。

该论文未声称的内容：

它不声称这适用于医疗诊断或临床用途。
它不声称这将解决未来所有的 AI 问题。
它严格专注于使联邦学习更快、更轻量同时保持准确性的数学和计算机科学。

以下是 Zhu 等人论文《非同质数据下高效联邦学习的子空间优化》的详细技术总结。

1. 问题陈述

本文解决了大规模模型时代（参数从数亿到数十亿）联邦学习（FL）面临的“三重挑战”：

计算：本地训练步骤计算量巨大。
内存：在资源受限的边缘设备上存储完整模型状态和辅助变量是不可行的。
通信：在有限带宽下传输全维更新效率低下。

这些挑战因**数据非同质性（Non-IID）**而加剧。在非同质设置中，本地目标函数与全局目标函数发生偏离，导致“客户端漂移”。

现有解决方案及其局限性：
- 异质性校正（例如 SCAFFOLD）：使用控制变量来校正漂移，但需要存储和通信全维度的辅助状态，从而抵消了大规模模型的效率增益。
- 压缩（例如量化）：减少了传输的比特数，但通常仍需要全维度的误差缓冲区，无法减少内存占用。
- 子空间/低秩方法：通过降低维度来提高效率，但通常缺乏校正客户端漂移的原则性机制。简单的组合往往在子空间变化时丢弃历史校正信息，导致不稳定性。

差距：目前尚无算法能够在低维子空间内原生执行异质性校正，同时保留像 SCAFFOLD 这样的全维度方法的鲁棒性。

2. 方法论：子空间 SCAFFOLD（SSF）

作者提出了SSF，这是一种将异质性校正直接集成到低维优化子空间中的算法。

核心机制

子空间投影：
- 在每一轮 $t$ ，生成一个共享的随机正交投影矩阵 $P_t \in \mathbb{R}^{r \times d}$ （其中 $r \ll d$ ）。
- 优化（本地更新和聚合）严格在 $r$ 维子空间内进行。
- 梯度被投影： $g_{proj} = P_t g$ 。
投影控制变量（异质性校正）：
- SSF 模仿 SCAFFOLD 机制，但在投影量上运行。
- 客户端维护本地控制变量 $c_i$ ，服务器维护全局控制变量 $c$ 。
- 本地更新方向利用投影后的本地与全局控制变量之差进行校正： $P_t g - c_{i, proj} + c_{proj}$ 。
残差保持回填（关键创新）：
- 与丢弃活跃子空间之外信息的标准子空间方法不同，SSF 维护全维度控制状态（ $c_i$ 和 $c$ ）。
- 更新规则：当子空间发生变化时，控制变量按以下方式更新：
  $c^{t+1} = (I - P_t^\top P_t)c^t + P_t^\top P_t (\text{新梯度信息})$
- 这意味着控制变量的**正交补（残差）**保持不变，仅刷新活跃子空间分量。
- 回填：子空间更新后，通过将新的投影坐标与保留的残差相结合，重构全空间模型： $x^{t+1} = P_t^\top x^{t+1}_{proj} + x^t_{res}$ 。

算法流程

分解：服务器将模型和控制变量分解为投影部分和残差部分。
本地步骤：客户端利用广播的子空间模型和服务器的残差重构完整模型。它们在完整空间中计算梯度，进行投影，并在子空间内执行异质性校正更新。
聚合：服务器在子空间内聚合投影后的端点。
回填与刷新：服务器将模型提升回全空间并更新控制变量，同时保留残差分量。

3. 主要贡献

算法设计（SSF）：
- 首个在低维子空间内执行原生异质性校正的联邦学习算法。
- 引入了一种回填式更新，保留控制变量的残差分量，确保当子空间旋转或变化时，历史漂移校正信息不会丢失。
理论分析：
- 证明在标准平滑性和有界方差假设下，SSF 实现了非渐近收敛速率 $\tilde{O}(1/T + 1/\sqrt{NKT})$ 。
- 该速率与全维度 SCAFFOLD 的线性加速相匹配，表明子空间约束不会降低关于客户端漂移的渐近收敛行为。
- 该分析成立的前提是不假设随机投影器与随机梯度之间相互独立。
实证验证：
- 在受控的矩阵回归玩具问题和深度学习基准（CIFAR-100 搭配 ResNet-110）上进行了演示。
- 表明 SSF 显著降低了内存、通信和计算成本，同时保持了与全维度 SCAFFOLD 相当的精度。

4. 实验结果

玩具问题（矩阵回归）

鲁棒性：在不同异质性水平（低、中、高）下，SSF 始终优于FedSub（一种没有残差保留的基准子空间方法）和FedAvg。
稳定性：由于正交信息的丢失，FedSub 在较高的子空间维度（ $r=50$ ）下表现出数值不稳定和发散（NaN 错误）。SSF 在所有维度下均保持稳定。
精度：在子空间比率 $r/d = 0.2$ 时，SSF 达到的相对误差非常接近全维度 SCAFFOLD，显著优于全维度 FedAvg。

深度学习（CIFAR-100 / ResNet-110）

排名：全维度 SCAFFOLD > SSF > FedAvg > FedSub。
性能：SSF 最终测试准确率达到45.43%，大幅优于 FedAvg（34.35%）和 FedSub（23.17%）。
效率：SSF 成功降低了通信和内存开销，同时保留了 SCAFFOLD 对抗数据异质性的鲁棒性。

5. 意义与影响

弥合差距：SSF 解决了鲁棒性（处理非同质数据）与效率（处理大规模模型）之间的根本矛盾。它证明不需要全维度辅助状态即可有效校正客户端漂移。
可扩展性：通过将通信和内存占用从 $O(d)$ 降低到 $O(r)$ （其中 $r \ll d$ ），同时保持理论收敛保证，SSF 使得在资源受限的边缘设备上部署大规模模型联邦学习成为可能。
理论洞察：“残差保持”机制为子空间优化提供了一种新范式，表明在动态子空间环境中，为控制变量维护全维度内存（即使仅部分访问）对于稳定性至关重要。

总之，SSF 为部署大规模、非同质联邦学习提供了一种实用且理论完备的解决方案，实现了“鱼与熊掌兼得”：子空间方法的效率与控制变量校正的鲁棒性。