FedAdamW: A Communication-Efficient Optimizer with Convergence and Generalization Guarantees for Federated Large Models

AdamW has become one of the most effective optimizers for training large-scale models. We have also observed its effectiveness in the context of federated learning (FL). However, directly applying AdamW in federated learning settings poses significant challenges: (1) due to data heterogeneity, AdamW often yields high variance in the second-moment estimate $\mathit{v}$; (2) the local overfitting of AdamW may cause client drift; and (3) Reinitializing moment estimates ($\mathit{v}$, $\mathit{m}$) at each round slows down convergence. To address these challenges, we propose the first \underline{Fed}erated \underline{AdamW} algorithm, called $\mathtt{\text{FedAdamW}}$, for training and fine-tuning various large models. $\mathtt{\text{FedAdamW}}$ aligns local updates with the global update using both a $\mathbf{\text{local}}\mathbf{\text{correction}}\mathbf{\text{mechanism}}$ and decoupled weight decay to mitigate local overfitting. $\mathtt{\text{FedAdamW}}$ efficiently aggregates the $\mathtt{\text{mean}}$ of the second-moment estimates to reduce their variance and reinitialize them. Theoretically, we prove that $\mathtt{\text{FedAdamW}}$ achieves a linear speedup convergence rate of $\mathcal{O}(\sqrt{(L\mathrm{\Delta }{\sigma }_l^2)\mathrm{/}(SKR{\epsilon }^2)}+(L\mathrm{\Delta })\mathrm{/}R)$ without $\mathbf{\text{heterogeneity}}\mathbf{\text{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. We also employ PAC-Bayesian generalization analysis to explain the effectiveness of decoupled weight decay in local training. Empirically, we validate the effectiveness of $\mathtt{\text{FedAdamW}}$ on language and vision Transformer models. Compared to several baselines, $\mathtt{\text{FedAdamW}}$ significantly reduces communication rounds and improves test accuracy. The code is available in https://github.com/junkangLiu0/FedAdamW.