[Paper] 轻量客户端，完整精度：混合零阶和一阶分割联邦学习

Source: arXiv - 2601.09076v1

概述

本文解决了分割联邦学习（SFL）中的一个长期瓶颈：边缘设备在反向传播过程中必须承担的巨大内存和计算负载。作者提出了 HERON‑SFL，一种混合训练方案，它用廉价的零阶（ZO）近似取代客户端侧基于梯度的（一级）更新，同时保持服务器侧的更新为一级。其结果是一个能够在轻量设备上训练现代深度网络的系统，显著降低内存和计算需求，同时不牺牲模型精度。

关键贡献

• 混合优化框架 – 在客户端使用零阶更新，在服务器端使用一阶更新，保持整体训练保真度。
• 辅助网络辅助的 ZO 更新 – 使用轻量级“助理”网络生成扰动前向传播，实现无梯度更新，避免激活缓存。
• 理论保证 – 在低有效秩假设下证明收敛速率与模型维度无关，规避 ZO 方法常见的维度灾难。
• 实证验证 – 在 ResNet 图像分类和语言模型微调上展示 HERON‑SFL 在保持基准准确率的同时，将客户端峰值内存降低最多 64 %，每步计算降低最多 33 %。
• 可扩展性蓝图 – 展示该方法如何将 SFL 推广到先前对边缘设备不可及的模型（例如更大的 CNN、基于 Transformer 的语言模型）。

方法论

1. Split Architecture – 模型被划分为 client‑side 前端和 server‑side 后端。客户端处理原始数据，将中间激活（即 “cut‑layer” 输出）发送给服务器，并接收服务器的响应用于计算损失。

2. Zeroth‑Order Client Update

   • 与其对客户端网络进行反向传播，客户端会采样一小组随机扰动 $\delta$，仅使用 前向 计算 $\mathbf{x} + \delta$ 的损失。
   • 通过有限差分估计器（例如两点差分或高斯平滑），客户端为其参数构建 近似梯度，而无需存储激活值。
   • 一个辅助网络（远小于主客户端模型）生成扰动，并提供一种低成本的方式来计算 ZO 步骤。

3. First‑Order Server Update

   • 服务器接收中间激活，计算自身（更大）后端的真实梯度，并执行标准的 SGD/Adam 步骤。
   • 服务器端的更新会传回给客户端，完成一次全局迭代。

4. Hybrid Loop – 训练循环交替进行：客户端在本地执行廉价的 ZO 步骤；服务器聚合并应用 FO 更新。该过程重复直至收敛。

低有效秩假设（损失相对于客户端参数的雅可比矩阵位于低维子空间）使作者能够界定 ZO 估计器的方差，从而得到一个 不随参数数量增长而恶化 的收敛速率。

结果与发现

• 准确率持平 – 在所有基准测试中，HERON‑SFL 与完整的一阶基线相差不超过 0.2 %。
• 内存占用 – 客户端峰值内存显著下降，因为反向传播不需要存储激活图。
• 计算节省 – 每个客户端步骤仅需前向传播（加上廉价的扰动生成），从而降低每次迭代的 FLOPs。
• 可扩展性测试 – 在 Raspberry Pi 级别的设备上训练 ResNet‑50（≈25 M 参数）变得可行，而原始 SFL 因 OOM 而崩溃。

消融实验表明，辅助网络的规模和 ZO 样本数量在额外开销与估计方差之间存在权衡。

实际意义

• Edge‑AI 部署 – 公司现在可以将更复杂的模型（例如视觉 Transformer、中等规模语言模型）推送到物联网设备、可穿戴设备或智能手机，而无需重新设计模型架构。
• 降低带宽成本 – 更少的反向消息意味着上行流量更低，这对蜂窝或卫星连接的设备至关重要。
• 能源效率 – 仅前向计算消耗更少的电能，延长了设备端学习场景（个性化、持续学习）的电池寿命。
• 简化的 SDK – 开发者可以将 HERON‑SFL 集成到现有的联邦学习框架（TensorFlow Federated、PySyft）中，只需进行最小的更改：将客户端优化器替换为 ZO 包装器。
• 监管与隐私收益 – 通过在设备端保留更多计算并限制梯度泄漏，该方法很好地符合隐私设计原则的监管要求（GDPR、HIPAA）。

限制与未来工作

• Zeroth‑Order Variance – 虽然低秩假设可以缓解，但零阶估计仍会引入额外的随机性，可能影响在高度非凸任务上的收敛速度。
• Auxiliary Network Overhead – 辅助网络会增加参数量和推理成本；为不同硬件找到最佳规模仍是一个未解决的工程问题。
• Server Load – 服务器仍需对后端执行完整的反向传播，这在大规模部署时可能成为瓶颈。
• Theoretical Scope – 收敛性证明假设损失函数平滑且扰动有界；将其扩展到非平滑目标（例如量化模型）是未来的工作。
• Broader Benchmarks – 实验主要聚焦于图像分类和语言模型微调；在强化学习、图神经网络或多模态模型上评估 HERON‑SFL 将提升其通用性。

作者

• Zhoubin Kou
• Zihan Chen
• Jing Yang
• Cong Shen

论文信息

• arXiv ID: 2601.09076v1
• 类别: cs.LG, cs.DC, cs.IT, cs.NI, eess.SP
• 发表时间: 2026年1月14日
• PDF: 下载 PDF