FedVSSAM: Mitigating Flatness Incompatibility in Sharpness-Aware Federated Learning¶
会议: NeurIPS 2026
arXiv: 2605.09144
领域: 优化/理论
关键词: 联邦学习, 锐度感知最小化, 平坦度不兼容, 方差抑制, 非凸收敛
一句话总结¶
针对非独立同分布（Non-IID）数据与批次扰动友好对抗现象诱发的“平坦度不兼容”失效模式，FedVSSAM 构建参数服务器端维护的方差抑制调整方向，并在本地平坦度搜索、本地模型更新与全局聚合三阶段统一锚定该方向，无需上传额外辅助变量即实现更优的全局泛化与非凸线性加速收敛。
研究背景与动机¶
联邦学习（FL）通过协调各边缘设备在保护本地私有数据的前提下协同训练全局模型。然而，跨设备的数据异质性（Non-IID）会导致严重的设备漂移，使得本地更新轨迹与全局目标优化方向产生分歧。为了提升全局模型在异质甚至未知数据分布下的泛化能力，研究者将锐度感知最小化（SAM）引入联邦优化（如 FedSAM），旨在通过对抗扰动寻找损失曲面上的平坦极小点，避免模型落入过拟合的尖锐极小值。
但在联邦学习中直接应用 SAM 存在本质结构性缺陷：设备在本地搜索到的平坦区域往往与全局聚合目标所需的全局平坦区域严重冲突。这种局部与全局平坦度失配的根本原因来自两方面：一是跨设备数据异质性导致本地梯度与全局梯度存在系统性偏差；二是批次随机扰动中的“友好对抗（Friendly Adversary）”现象，即针对某一批次数据构建的最大化损失扰动，在另一批次梯度方向不一致的数据上反而变成了降低损失的“良性”扰动。这两类误差并非独立存在，本地扰动方向的不准确直接导致后续下降梯度评估点的错位，并在多轮本地迭代和部分设备抽样聚合中逐级放大，最终导致聚合后的全局模型脱离全局平坦盆地。
现有改进方案往往割裂处理优化漂移与锐度对齐，或者依赖额外上传控制变量与对偶状态，增加了设备上行通信负担。核心 idea：显式建模并解耦数据异质性与友好对抗带来的正交误差，通过在服务器端维护一个全局指数滑动平均调整方向，并在本地扰动搜索、本地扰动后更新以及全局模型聚合三个阶段全链路统一锚定该方差抑制方向，在不增加上行通信载荷的前提下彻底消除平坦度不兼容。
方法详解¶
整体框架¶
FedVSSAM 采用经典的“参数服务器-多设备”迭代架构，每一轮通信由参数服务器（PS）向均匀随机抽样（UAR）选出的设备子集广播当前全局模型 \(\theta^t\) 与方差抑制调整方向 \(h^t\)。各设备在本地执行 \(K\) 步局部迭代：首先结合历史全局方向与当前随机梯度构建方差抑制的混合扰动方向，寻找平坦邻域扰动点；接着在该扰动点评估梯度，并再次与全局方向混合生成方差抑制更新方向以执行参数下降；本地训练结束后仅上传最终模型参数，PS 聚合差分并以指数移动平均（EMA）更新全局调整方向 \(h^{t+1}\) 与全局模型 \(\theta^{t+1}\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
PS_Init["参数服务器 PS<br/>下发全局模型 θ^t 与调整方向 h^t"] --> Dev_Search["本地平坦度搜索<br/>混合方向 h_i^{t,k} 构造扰动参数 θ~_i^{t,k}"]
Dev_Search --> Dev_Update["本地扰动更新<br/>在扰动点求导并构造混合步 u_i^{t,k} 更新本地模型"]
Dev_Update --> Dev_Upload["设备上传最终本地模型 θ_i^t<br/>（零额外上行辅助变量）"]
Dev_Upload --> PS_Global["全局方差抑制更新<br/>聚合差分 g^{t+1} 并通过 EMA 演进 h^{t+1} 与 θ^{t+1}"]
PS_Global -->|进入下一轮通信迭代| PS_Init关键设计¶
1. 平坦度不兼容性解耦与本地平坦度搜索：消除友好对抗与设备漂移对扰动方向的扰动 在传统 FedSAM 中，本地扰动方向直接由单批次随机梯度归一化决定（\(\epsilon_i \propto g_i / \|g_i\|\)）。本研究在理论上定义了平坦度不兼容度指标 \(\Delta_{\mathrm{FI}}(\theta, \rho)\)，并证明该不兼容度可严格分解为跨设备数据异质性偏差 \(\delta_i(\theta) = \nabla F_i(\theta) - \nabla F(\theta)\) 与设备内部批次梯度的友好对抗方差 \(\zeta_{i,\phi_i}(\theta) = g_i - \nabla F_i(\theta)\)，二者在期望下相互正交。若直接采用本地批次梯度，二者会共同劣化扰动方向。为此，FedVSSAM 设计了凸组合的本地扰动构造规则： $\(h_i^{t,k} = (1 - \gamma_l) h^t + \gamma_l g_i^{t,k}, \quad \tilde{\theta}_i^{t,k} = \theta_i^{t,k} + \rho \frac{h_i^{t,k}}{\|h_i^{t,k}\|}\)$ 其中 \(\gamma_l \in (0, 1]\) 为本地混合系数，\(h^t\) 为服务器下发的全局稳定方向。较小的 \(\gamma_l\) 借由全局先验大幅削弱单批次梯度方差与异质性漂移，确保扰动方向探测的是全局一致的平坦盆地，而非孤立过拟合的本地尖锐波谷；同时保留的 \(\gamma_l g_i^{t,k}\) 分量保留了对本地曲率的必要敏感度。
2. 本地扰动更新的方差抑制：切断扰动误差向局部下降步的传播链路 在确定扰动位置 \(\tilde{\theta}_i^{t,k}\) 后，若直接沿该点评估的随机梯度 \(\tilde{g}_i^{t,k} = \nabla F_i(\tilde{\theta}_i^{t,k}; \phi_i^{t,k})\) 进行梯度下降，由于评估点本身存在潜在偏差，下降方向会再次放大异质性漂移。FedVSSAM 采用同构的混合更新方向： $\(u_i^{t,k} = (1 - \gamma_l) h^t + \gamma_l \tilde{g}_i^{t,k}, \quad \theta_i^{t,k+1} = \theta_i^{t,k} - \eta_l u_i^{t,k}\)$ 由于扰动构造与参数更新共享同一控制系数 \(\gamma_l\) 与全局锚点 \(h^t\)，更新方向条件方差被严格压缩至 \(\gamma_l^2 \mathrm{Var}[\tilde{g}_i^{t,k} \mid \mathcal{F}^{t,k}]\)。这种一致性设计避免了扰动生成与梯度下降“方向打架”，确保局部优化轨迹紧随全局优化流。
3. 全局方差抑制聚合：免额外上行载荷的高精度梯度追踪 完成本地 \(K\) 步更新后，设备仅将最终模型参数 \(\theta_i^t = \theta_i^{t,K}\) 上传至服务器。服务器利用模型差分反推有效伪梯度 \(g^{t+1} = \frac{1}{\eta_l K} (\theta^t - \frac{1}{S} \sum_{i \in \mathcal{S}^t} \theta_i^t) = \frac{1}{S K} \sum_{i \in \mathcal{S}^t} \sum_{k=0}^{K-1} u_i^{t,k}\)。随后利用全局混合系数 \(\gamma_g\) 执行指数移动平均，更新全局调整方向与模型参数： $\(h^{t+1} = (1 - \gamma_g) h^t + \gamma_g g^{t+1}, \quad \theta^{t+1} = \theta^t - \eta_g h^{t+1}\)$ 相较于 SCAFFOLD、FedGAMMA 等必须上传控制变量（Control Variates）导致上行带宽翻倍的方法，FedVSSAM 上行通信开销与最基础的 FedAvg 完全相同。理论证明，在非凸平滑假设下，\(h^t\) 与真实全局梯度 \(\nabla F(\theta^t)\) 的均方误差（MSE）被严格限制在 \(\mathcal{O}(\frac{L\Delta}{\beta T} + \beta \Sigma)\) 级别，能以极高保真度实时追踪全局目标。
损失函数 / 训练策略¶
优化目标基于全局经验损失 \(F(\theta) = \frac{1}{N} \sum_{i \in \mathcal{N}} F_i(\theta)\) 的邻域平坦度极小化。核心训练超参数包括：本地步长 \(\eta_l \le \frac{1}{\sqrt{2} \gamma_l K L}\)、全局步长 \(\eta_g \le \min\{\frac{1}{24L}, \frac{\gamma_l \gamma_g}{6L}\}\)，混合积 \(\beta = \gamma_l \gamma_g\)。针对中度异质性分布设置 \((\gamma_l, \gamma_g) = (0.4, 0.6)\)，针对高度异质性分布设置 \((\gamma_l, \gamma_g) = (0.1, 0.6)\)，扰动半径基准设置为 \(\rho = 0.05\)。
实验关键数据¶
主实验¶
在 CIFAR-10（ResNet-18）和 CIFAR-100（WRN-28-4）上设置 100/200 客户端环境，评测在狄利克雷分布（\(\alpha \in \{0.5, 0.1\}\)）与病态非独立同分布（Pathological，\(c \in \{6, 3\}\) 或 \(\{20, 10\}\)）下的 Top-1 准确率（%）及达到目标精度（CIFAR-10: 75%, CIFAR-100: 40%）所需的通信轮数（Rd.）。
| 数据集 | 异质性划分 | 指标 | FedVSSAM (本文) | 次优基线 | 性能增益 | FedAvg 基线 |
|---|---|---|---|---|---|---|
| CIFAR-10 (100设备, 10%采样) | Dirichlet \(\alpha=0.5\) | 最终 Acc. (%) / 达标轮数 | 83.56 ± 0.19 / 111 | 82.36 ± 0.13 (FedGAMMA) / 201 | +1.20% / 提速 44.8% | 77.05 ± 0.14 / 315 |
| CIFAR-10 (100设备, 10%采样) | Dirichlet \(\alpha=0.1\) | 最终 Acc. (%) / 达标轮数 | 79.57 ± 0.14 / 296 | 78.94 ± 0.14 (MoFedSAM) / 385 | +0.63% / 提速 23.1% | 72.57 ± 0.19 / 未达标 (–) |
| CIFAR-10 (100设备, 10%采样) | Pathological \(c=6\) | 最终 Acc. (%) / 达标轮数 | 83.57 ± 0.22 / 107 | 82.94 ± 0.16 (MoFedSAM) / 178 | +0.63% / 提速 39.9% | 77.55 ± 0.24 / 299 |
| CIFAR-10 (100设备, 10%采样) | Pathological \(c=3\) | 最终 Acc. (%) / 达标轮数 | 81.11 ± 0.23 / 196 | 80.76 ± 0.23 (FedSMOO) / 289 | +0.35% / 提速 32.2% | 74.63 ± 0.23 / 未达标 (–) |
| CIFAR-100 (100设备, 10%采样) | Dirichlet \(\alpha=0.5\) | 最终 Acc. (%) / 达标轮数 | 58.78 ± 0.30 / 237 | 56.87 ± 0.23 (FedSMOO) / 279 | +1.91% / 提速 15.1% | 46.52 ± 0.23 / 313 |
| CIFAR-100 (100设备, 10%采样) | Dirichlet \(\alpha=0.1\) | 最终 Acc. (%) / 达标轮数 | 51.22 ± 0.20 / 328 | 50.59 ± 0.28 (SCAFFOLD) / 450 | +0.63% / 提速 27.1% | 36.76 ± 0.23 / 未达标 (–) |
| CIFAR-100 (100设备, 10%采样) | Pathological \(c=20\) | 最终 Acc. (%) / 达标轮数 | 56.87 ± 0.19 / 258 | 55.78 ± 0.33 (SCAFFOLD) / 368 | +1.09% / 提速 29.9% | 42.59 ± 0.15 / 486 |
| CIFAR-100 (100设备, 10%采样) | Pathological \(c=10\) | 最终 Acc. (%) / 达标轮数 | 50.27 ± 0.18 / 337 | 48.60 ± 0.18 (FedSMOO) / 463 | +1.67% / 提速 27.2% | 36.67 ± 0.21 / 未达标 (–) |
| DBpedia-14 (100设备, 10%采样) | Dirichlet \(\alpha=0.1\) | 最终 Acc. (%) / 达标轮数 (90%) | 96.90 ± 0.07 / 71 | 96.44 ± 0.08 (FedACG) / 39 | +0.46% (收敛精度最高) | 92.56 ± 0.08 / 205 |
| DBpedia-14 (100设备, 10%采样) | Pathological \(c=3\) | 最终 Acc. (%) / 达标轮数 (90%) | 96.74 ± 0.11 / 46 | 96.21 ± 0.11 (FedSMOO) / 80 | +0.53% / 匹配最快收敛 | 91.09 ± 0.11 / 215 |
消融实验¶
在 CIFAR-10 数据集上评估各阶段方差抑制机制的有效性。以 \(\gamma_l = \gamma_g = 1\)（等价于无方差抑制的退化基线 FedVSSAM-base）为基准，分别单独激活本地平坦度搜索（LF, 式 14）、本地更新（LU, 式 15）、全局更新（GU, 式 16）的方差抑制。
| 配置 | 核心改动说明 | 收敛精度与加速表现 |
|---|---|---|
| FedVSSAM (Full) | 同时激活 LF + LU + GU 三阶段方差抑制 | 达到全局最高准确率（\(\alpha=0.5\) 下达 83.56%，\(\alpha=0.1\) 下达 79.57%），收敛轨迹最平稳 |
| w/o LF & LU (仅保留 GU) | \(\gamma_l = 1\), 仅在服务器端聚合时进行动量平滑更新 | 性能优于 base，但无法阻止本地多步迭代中漂移扰动的发散累积 |
| w/o LF (仅保留 LU + GU) | 本地扰动采用纯批次梯度，仅在下降更新中融合全局方向 | 扰动位置与下降方向不一致，损失了平坦区域探测的全局兼容性 |
| w/o GU (仅保留 LF + LU) | \(\gamma_g = 1\), 服务器直接算术平均更新模型 | 聚合伪梯度的方差未被平滑，导致下发至各设备的 \(h^{t+1}\) 存在随机震荡 |
| FedVSSAM-base | \(\gamma_l = 1, \gamma_g = 1\)（完全退化为 Vanilla FedSAM 架构） | 遭受严重平坦度不兼容，强异质性下精度掉点明显（低于完整模型 4%~6%） |
关键发现¶
- 全链路一致性的必要性：消融实验清晰证明，孤立修正扰动或仅在服务端做动量聚合均无法完全消除漂移。只有 LF、LU、GU 三者联动，才能切断“误差扰动 \(\to\) 错位求导 \(\to\) 偏倚更新 \(\to\) 漂移聚合”的恶性循环。
- 异质性越强，增益越突出：在 CIFAR-10 极端异质性（\(\alpha=0.1\) 与 \(c=3\)）场景下，传统 FedAvg 和 FedProx 无法在 800 轮内达到 75% 准确率，而 FedVSSAM 稳定在 200~300 轮内达标，最终精度领先 FedAvg 7 个百分点以上。
- 通信效率显著占优：在达到相同目标精度时，FedVSSAM 相比 FedAvg 节省了 30%~65% 的通信轮数，且因无需上传双倍维度的控制变量，实际物理上行带宽消耗仅为 SCAFFOLD 和 FedGAMMA 的一半。
亮点与洞察¶
- 平坦度不兼容问题的理论显式化：首次在联邦 SAM 框架下识别出“平坦度不兼容”失效机理，并通过引理将该失效在数学上精确解耦为跨设备数据异质性与设备内友好对抗两大正交成分。
- 单方向全局锚定机制：通过一个服务器端低成本维护的 EMA 向量 \(h^t\) 贯穿三阶段，在不增加客户端上行通信载荷的约束下，优雅地达成了全局收敛边界与扰动一致性。
- 即插即用的轻量级架构：该设计思路完全兼容现有的参数服务器架构，可直接推广至各类基于锐度感知或正则化平滑的分布式学习系统（如去中心化协作、差分隐私联邦学习）。
局限与展望¶
- 假设局限于标准光滑性与全梯队抽样：当前的非凸收敛保证依赖于 \(L\)-Lipschitz 平滑性、有界梯度方差以及独立均匀的客户端抽样，在系统异构（慢节点 Stragglers、非均匀掉线）下的理论分析有待扩展。
- 超参数耦合调节需求：混合系数 \(\gamma_l\) 和 \(\gamma_g\) 需要根据数据异质性先验进行粗调（中度设为 0.4/0.6，极度设为 0.1/0.6），未来可探索根据每轮聚合方差动态自适应调整衰减率的算法。
- 面向超大规模基础模型验证：目前实验主要基于 ResNet 与 TextCNN 等经典视觉与文本架构，尚未在百亿级参数大语言模型（LLM）的联邦指令微调场景下进行全尺度验证。
相关工作与启发¶
- vs FedSAM [6]: FedSAM 仅在本地做单批次梯度的简单 SAM 扰动，在 Non-IID 下各客户端极小化不同的局部平坦区，导致聚合后模型依然尖锐；FedVSSAM 引入全局方差抑制方向锚定，从源头压制了不兼容性。
- vs SCAFFOLD [22] & FedGAMMA [9]: 后两者利用控制变量纠正漂移，虽理论收敛性好但必须设备向服务器传输双倍上行向量（模型参数+控制变量）；FedVSSAM 客户端只传更新后模型，通过服务端差分反推完成 EMA 平滑，极大节省上行带宽。
- vs MoFedSAM [45]: MoFedSAM 仅在客户端引入动量，未能解决本地单批次扰动遭遇友好对抗时的扰动失效问题；FedVSSAM 在扰动构造阶段即注入全局方向，使得对抗邻域搜索更加鲁棒。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ 首次定义并解耦了联邦 SAM 的平坦度不兼容问题，三阶段统一方差抑制设计精炼有效。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 CVPR/NeurIPS 经典标准图像基准（CIFAR-10/100）及文本数据集，对比了 12 种代表性强基线。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导完备自洽，从问题定义、引理分析到算法提出逻辑闭环清晰。
- 价值: ⭐⭐⭐⭐⭐ 兼顾高性能泛化与极低通信开销，对边缘计算与分布式联邦优化具有重要的工程指导意义。