Title: FedMuon: Accelerating Federated Learning with Matrix Orthogonalization
Primary area: optimization
Keywords: Federated Learning

Abstract: 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 (named \texttt{Local Muon}), which has matrix orthogonalization to optimize matrix-structured parameters. 
Experimental results show that, in IID setting, \texttt{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  \underline{Fed}erated \underline{Muon} optimizer (\texttt{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 of $\mathcal{O}(\sqrt{(L \Delta \sigma_l^2)/(S K R)}+(L \Delta)/R)$ 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 \texttt{FedMuon} on language and vision models. Compared to several baselines, \texttt{FedMuon} significantly reduces communication rounds and improves test accuracy. The code is available in \url{https://anonymous.4open.science/r/FedMuon-02F5}.