通过矩阵正交化加速联邦学习,减少通信轮次并缓解客户端漂移。
FedMuon: Accelerating Federated Learning with Matrix Orthogonalization
- 在本地优化中引入矩阵正交化,利用权重矩阵的几何结构提升更新效率。
- 在非独立同分布场景下,通过动量聚合与全局对齐,将通信轮次减少30%以上。
- 适合需要低通信开销的分布式训练场景,尤其适用于异构数据环境。
联邦学习的核心瓶颈在于通信轮次过多。现有方法多采用逐元素本地优化器(如Adam/SGD),忽视了权重矩阵的几何结构,导致局部更新时病态方向被放大,条件数恶化,收敛变慢。为此,我们提出局部优化器Muon,通过矩阵正交化优化矩阵结构参数。实验表明,在独立同分布(IID)设置下,局部Muon显著加快收敛速度,减少通信轮次。但在非独立同分布(non-IID)场景下,基于本地分布的独立矩阵正交化引发严重客户端漂移。为解决此问题,我们提出新型联邦优化器FedMuon,包含两项关键技术:(1)动量聚合,客户端使用聚合动量进行本地初始化;(2)本地-全局对齐,使本地梯度与全局更新方向一致,显著抑制客户端漂移。理论上,证明了FedMuon在无需异质性假设下实现线性加速收敛,其中S为每轮参与客户端数,K为本地迭代次数,R为总通信轮次。实证上,我们在语言和视觉模型上验证其有效性,相比多个基线方法,显著降低通信轮次并提升测试准确率。
原文摘要 · Abstract (English)
The core bottleneck of Federated Learning (FL) lies in the communication rounds. That is, how to achieve more effective local updates is crucial for reducing communication rounds. Existing FL methods still primarily use element-wise local optimizers (Adam/SGD), neglecting the geometric structure of the weight matrices. This often leads to the amplification of pathological directions in the weights during local updates, leading deterioration in the condition number and slow convergence. Therefore, we introduce the Muon optimizer in local, which has matrix orthogonalization to optimize matrix-structured parameters. Experimental results show that, in IID setting, Local Muon significantly accelerates the convergence of FL and reduces communication rounds compared to Local SGD and Local AdamW. However, in non-IID setting, independent matrix orthogonalization based on the local distributions of each client induces strong client drift. Applying Muon in non-IID FL poses significant challenges: (1) client preconditioner leading to client drift; (2) moment reinitialization. To address these challenges, we propose a novel Federated Muon optimizer (FedMuon), which incorporates two key techniques: (1) momentum aggregation, where clients use the aggregated momentum for local initialization; (2) local-global alignment, where the local gradients are aligned with the global update direction to significantly reduce client drift. Theoretically, we prove that \texttt{FedMuon} achieves a linear speedup convergence rate without the heterogeneity assumption, where $S$ is the number of participating clients per round, $K$ is the number of local iterations, and $R$ is the total number of communication rounds. Empirically, we validate the effectiveness of FedMuon on language and vision models. Compared to several baselines, FedMuon significantly reduces communication rounds and improves test accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。