解决联邦学习中二阶优化器因数据异构导致的训练发散问题。
Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data
- 提出预条件器对齐与修正框架,分离参数聚合与几何同步。
- 在非独立同分布数据上实现稳定训练,图像任务最高提升5.8%准确率。
- 适用于视觉与语言任务,尤其适合追求高精度的联邦学习场景。
二阶优化器可显著加速大规模训练,但其直接应用于联邦学习时,在非独立同分布(non-IID)数据上常不稳定甚至发散。我们发现关键原因是预条件器漂移:客户端的二阶训练导致各异的曲率定义几何结构(即预条件器坐标系),而服务器端在不兼容度量下进行模型平均,破坏了全局下降方向。为此,我们提出 exttt{FedPAC} 框架,通过显式解耦参数聚合与几何同步来应对这一几何不匹配:(i) 对齐(将本地预条件器聚合为全局参考,并用全局预条件器热启动客户端);(ii) 修正(使用全局预条件方向引导本地预条件更新,抑制长期漂移)。我们在部分参与条件下提供了带漂移耦合的非凸收敛保证,并实现线性加速。实验表明, exttt{FedPAC} 在视觉与语言任务中均显著提升稳定性与精度,于使用 ViTs 的 CIFAR-100 上取得高达 5.8% 的绝对准确率提升。代码已公开于 https://anonymous.4open.science/r/FedPAC-8B24。
原文摘要 · Abstract (English)
Second-order optimizers can significantly accelerate large-scale training, yet their naive federated variants are often unstable or even diverge on non-IID data. We show that a key culprit is \emph{preconditioner drift}: client-side second-order training induces heterogeneous \emph{curvature-defined geometries} (i.e., preconditioner coordinate systems), and server-side model averaging updates computed under incompatible metrics, corrupting the global descent direction. To address this geometric mismatch, we propose \texttt{FedPAC}, a \emph{preconditioner alignment and correction} framework for reliable federated second-order optimization. \texttt{FedPAC} explicitly decouples parameter aggregation from geometry synchronization by: (i) \textbf{Alignment} (i.e.,aggregating local preconditioners into a global reference and warm-starting clients via global preconditioner); and (ii) \textbf{Correction} (i.e., steering local preconditioned updates using a global preconditioned direction to suppress long-term drift). We provide drift-coupled non-convex convergence guarantees with linear speedup under partial participation. Empirically, \texttt{FedPAC} consistently improves stability and accuracy across vision and language tasks, achieving up to $5.8\%$ absolute accuracy gain on CIFAR-100 with ViTs. Code is available at https://anonymous.4open.science/r/FedPAC-8B24.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。