[论文解读] Distributionally Robust Federated Averaging

Yuyang Deng, Mohammad Mahdi Kamani+1 more|arXiv (Cornell University)|Feb 25, 2021

Stochastic Gradient Optimization Techniques被引用 20

一句话总结

本文提出分布鲁棒联邦平均（DRFA），一种通信高效的联邦学习算法，通过最小化异构客户端数据分布下的最坏情况损失，实现分布鲁棒性。DRFA 通过周期性全局平均与自适应采样，并结合一种新颖的快照机制来近似混合参数的历史信息，从而在凸与非凸设置下实现可证明的快速收敛，同时显著减少通信轮次。

ABSTRACT

In this paper, we study communication efficient distributed algorithms for distributionally robust federated learning via periodic averaging with adaptive sampling. In contrast to standard empirical risk minimization, due to the minimax structure of the underlying optimization problem, a key difficulty arises from the fact that the global parameter that controls the mixture of local losses can only be updated infrequently on the global stage. To compensate for this, we propose a Distributionally Robust Federated Averaging (DRFA) algorithm that employs a novel snapshotting scheme to approximate the accumulation of history gradients of the mixing parameter. We analyze the convergence rate of DRFA in both convex-linear and nonconvex-linear settings. We also generalize the proposed idea to objectives with regularization on the mixture parameter and propose a proximal variant, dubbed as DRFA-Prox, with provable convergence rates. We also analyze an alternative optimization method for regularized cases in strongly-convex-strongly-concave and non-convex (under PL condition)-strongly-concave settings. To the best of our knowledge, this paper is the first to solve distributionally robust federated learning with reduced communication, and to analyze the efficiency of local descent methods on distributed minimax problems. We give corroborating experimental evidence for our theoretical results in federated learning settings.

研究动机与目标

• 为解决联邦学习中数据异构性带来的挑战，即标准 FedAvg 模型在非独立同分布（non-iid）客户端数据上泛化性能差的问题。
• 设计一种通信高效的算法，在无需频繁更新混合权重的情况下，保持分布鲁棒性。
• 在最小化通信开销的同时，通过周期性平均与自适应采样，实现在多样化客户端数据分布下的鲁棒泛化能力。
• 理论上分析局部下降法在分布式极小极大优化问题中的收敛速率，填补先前研究的空白。
• 通过近端变体 DRFA-Prox 将框架扩展至正则化目标，提供可证明的收敛保证。

提出的方法

• 提出 DRFA，一种联邦算法，仅在同步轮次更新全局混合参数 λ，且周期性地同步本地模型。
• 引入快照机制，以近似混合参数 λ 的历史梯度累积，补偿混合参数更新频率较低的缺陷。
• 基于当前 λ 值对客户端进行自适应采样，以提升平均过程中的收敛性与鲁棒性。
• 在凸-线性和非凸-线性设置下分析收敛性，建立 O(1/T) 的次优性界。
• 针对正则化目标，提出 DRFA-Prox，采用近端更新机制，并在强凸-强凹与非凸（PL 条件）-强凹设置下完成收敛性分析。
• 采用新颖的技术工具，控制延迟 λ 更新带来的误差，包括 λ 估计误差的几何衰减与梯度方差的控制。

实验结果

研究问题

• RQ1 我们能否设计一种通信高效的联邦学习算法，在数据异构性下仍保持分布鲁棒性？
• RQ2 在混合参数更新频率较低的条件下，如何为分布式极小极大优化中的局部下降法实现可证明的收敛性？
• RQ3 在部分客户端参与和周期性平均的条件下，DRFA 在凸与非凸设置下的收敛速率如何？
• RQ4 正则化如何影响分布鲁棒联邦学习的收敛性？我们能否设计一种具有理论保证的近端变体？
• RQ5 我们能否从理论上证明，在通信高效设置下，使用快照机制近似混合参数历史梯度是合理的？

主要发现

• DRFA 实现了 O(1/T) 的收敛速率，次优性间隙显式依赖于通信频率 τ 和数据异构性。
• 算法的收敛性受初始 λ 估计误差、条件数 κ 和梯度方差项的约束，其衰减速率取决于 τ。
• 在非凸设置下，当满足 Polyak-Łojasiewicz（PL）条件时，DRFA 仍实现 O(1/T) 收敛，表明其在凸性之外也具备鲁棒性。
• 近端变体 DRFA-Prox 在强凸-强凹与非凸（PL）-强凹设置下均实现 O(1/T) 收敛，将框架扩展至正则化目标。
• 实验结果验证了理论发现，表明 DRFA 在异构联邦数据上相比 FedAvg 展现出更优的鲁棒性与泛化能力。
• 分析表明，即使 λ 更新频率较低，通信效率仍能保持，这得益于快照机制对历史梯度趋势的有效捕捉。