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

AI summary

Overview

• Research area: Federated learning (FL) and stochastic optimization for large-scale models, specifically adaptive optimizers for Transformer architectures.
• Technical level: Advanced. The paper combines algorithm design with non-convex convergence proofs, PAC-Bayesian generalization analysis, and Hessian structure analysis.
• Scope: Proposes and validates FedAdamW, a federated variant of the AdamW optimizer that reduces communication cost, mitigates client drift, and carries theoretical convergence and generalization guarantees for training and fine-tuning large vision and language models.

What This Paper Is About

Federated learning lets many clients train a shared model without sharing raw data, but the standard optimizer (SGD/FedAvg) converges slowly on modern Transformer models such as ViT, GPT, and BERT. AdamW works far better in centralized training, yet applying it naively in federated settings causes three problems: unstable second-moment estimates under non-i.i.d. data, local overfitting that pushes clients away from the global model (client drift), and repeated resetting of optimizer state that wastes the training history. This paper builds a federated AdamW that fixes all three while keeping communication cheap, and backs it with convergence and generalization theory.

Key Contributions

1. A principled FedAdamW algorithm. The authors integrate a global update estimate into each client's local update to align local and global directions, use decoupled weight decay to curb local overfitting, and carry the second-moment estimate across rounds instead of resetting it.
2. Communication-efficient block-wise aggregation of the second moment. Motivated by the near block-diagonal structure of Transformer Hessians, clients transmit only the mean of each block of the second-moment estimate rather than the full vector, cutting communication from billions of scalars to a small number of block means.
3. First federated adaptive optimizer proved without a gradient heterogeneity assumption. FedAdamW achieves a linear-speedup rate of O(√(LΔσ_l²/(SKRε²)) + LΔ/R), which is strictly better than the comparable rates for Local AdamW and FedLADA because it avoids the heterogeneity term σ_g².
4. PAC-Bayesian generalization analysis. A bound of order O(1/√n) is derived that formally explains why decoupled weight decay (λ > 0) improves generalization relative to Local Adam (λ = 0), matching the empirical λ sweep.

Main Findings

• AdamW clearly beats SGD locally in FL. Local AdamW converges substantially faster than Local SGD for ViT, GPT-2, and BERT, motivating its federated adaptation.
• Second-moment variance is the first bottleneck. Non-i.i.d. data inflates the variance of v across clients because squared gradients amplify heterogeneity; averaging the mean of v reduces this variance and stabilizes training.
• Client drift is real and measurable. Local AdamW drifts further from the global optimum than Local SGD, and the added global-update term Δ_G^r corrects it.
• Reusing v across rounds helps; reusing m does not. Initializing v from the previous round's aggregated mean accelerates convergence, while m is still reset to zero because it adapts quickly to recent gradients.
• Block-wise mean aggregation is the sweet spot. Aggregating only the mean of each v block matches or beats full v aggregation (38.15% vs 38.12% accuracy on ViT-Tiny) at one-third of the communication cost (5.7M vs 17.1M scalars).
• Best-in-class empirical results. FedAdamW reaches 66.12%/63.01% on CIFAR-100 with ResNet-18 (Dir-0.6/Dir-0.1), 42.56%/39.86% with ViT-Tiny, 85.85% and 85.23% fine-tuning Swin Transformer on CIFAR-100 and Tiny ImageNet, and 81.79% average accuracy across seven GLUE tasks with RoBERTa-Base + LoRA — ahead of FedAvg, SCAFFOLD, FedCM, FedAdam, FedLADA, Local Adam, and Local AdamW.
• Ablations confirm every component. Removing v aggregation, global alignment, or decoupled weight decay each degrades accuracy (37.51%, 37.42%, 38.25% vs 39.86% for the full method). The alignment coefficient α = 0.5 and weight decay λ = 0.01 are optimal, matching theory.
• Theory aligns with practice. The predicted trade-off in λ (a log term that grows and an inverse term that shrinks) is visible in the empirical weight-decay sweep, where Local Adam collapses to 4–29% accuracy at large λ while FedAdamW stays above 38%.

Methodology in Plain English

The authors start by diagnosing why AdamW misbehaves in federated settings. Because each client sees a different data distribution, its squared gradients — which feed the second-moment estimate v — vary wildly, making the adaptive scaling unstable. Meanwhile AdamW's aggressive updates push each client hard toward its own local optimum, so clients drift apart, and resetting m and v to zero every round throws away useful statistics.

Their fix has three parts. First, in addition to the local AdamW update, each client adds a term proportional to the previous round's global update direction, a "correction" that keeps local steps pointed toward the shared model. Second, they keep decoupled weight decay and tune it more aggressively, which the PAC-Bayesian bound justifies as trading off regularization against curvature. Third, instead of communicating the full second-moment vector (which would double communication), they exploit the observation that Transformer Hessian blocks are nearly block-diagonal: parameter groups such as query/key heads, MLP layers, and value heads each have fairly uniform curvature, so a single scalar mean per block captures the information. Clients send these B block means plus their model deltas; the server averages both and broadcasts the updated model, aggregated block means, and the new global direction.

For theory, the authors prove convergence under standard smoothness and bounded-gradient assumptions — notably without the usual bounded-heterogeneity assumption that competing analyses require, because the global alignment term absorbs client drift. They then apply PAC-Bayesian bounds to show how the generalization gap shrinks with data size and how λ governs a bias-variance-like trade-off.

Why This Matters

Research impact. This is the first federated AdamW with both convergence and generalization guarantees, and the first federated adaptive optimizer whose rate is independent of gradient heterogeneity. It opens a path for adapting other modern centralized optimizers (the authors name LAMB and Lion) to federated settings, and makes the case that optimizer design, not just aggregation strategy, is a first-class lever in FL research.

Real-world applications:

• Privacy-preserving fine-tuning of LLMs across institutions — hospitals, banks, or law firms can jointly fine-tune a shared RoBERTa/GPT-style model with LoRA without exchanging proprietary text.
• On-device keyboard and speech models — mobile clients can adapt large models locally with fewer communication rounds, saving bandwidth and battery.
• Medical imaging across hospitals — federated ViT or Swin training on non-i.i.d. patient populations, where heterogeneity is severe and data cannot leave the site.
• Cross-organization vision systems — e.g., manufacturing defect detection or autonomous-driving perception, where each site's data distribution differs sharply.

Industry relevance. Federated training of large models is bottlenecked by communication rounds and unreliable convergence under heterogeneous data. FedAdamW's block-wise aggregation directly reduces the bandwidth cost of shipping optimizer state, and its stability under Dir-0.1 (high heterogeneity) means fewer rounds to reach target accuracy — both are practical cost drivers in deployed FL systems.

Future Directions

• Extending the framework to other optimizers. The paper explicitly suggests LAMB and Lion as candidates for a federated treatment, but whether their moment structures admit the same block-wise aggregation is untested.
• Adaptive or learned block partitioning. The current block definitions are architecture-specific by hand (query/key, value, MLP, embeddings). Automating the partitioning, or learning it from per-layer Hessian statistics, could generalize the method to new architectures.
• Privacy guarantees beyond data locality. The paper communicates aggregated second-moment means; whether these leak information, and how they interact with differential privacy or secure aggregation, is not analyzed.
• Theory under weaker assumptions. The convergence proof assumes bounded gradients and β₁ = 0, λ = 0 in the stated theorem. Closing the gap between the analyzed setting and the practical configuration (β₁ = 0.9, λ = 0.01) remains open.
• Scaling to larger models and real cross-silo deployments. Evaluation covers ViT-Tiny, Swin, and RoBERTa-Base; behavior on truly large foundation models with heterogeneous clients is unknown.

Target Audience

Federated learning and distributed optimization researchers, especially those working on adaptive optimizers and heterogeneity-robust algorithms. Practitioners training or fine-tuning Transformers in cross-device or cross-silo settings will benefit from the practical communication savings and the concrete hyperparameter guidance (α = 0.5, λ = 0.01, v carried over, m reset). Readers need comfort with non-convex optimization and PAC-Bayesian theory to fully appreciate the theoretical sections, though the algorithmic ideas and empirical results are accessible to anyone familiar with Adam/AdamW and the FedAvg loop.

Authors’ abstract

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 $\boldsymbol{v}$; (2) the local overfitting of AdamW may cause client drift; and (3) Reinitializing moment estimates ($\boldsymbol{v}$, $\boldsymbol{m}$) at each round slows down convergence. To address these challenges, we propose the first \underline{Fed}erated \underline{AdamW} algorithm, called \texttt{FedAdamW}, for training and fine-tuning various large models. \texttt{FedAdamW} aligns local updates with the global update using both a \textbf{local correction mechanism} and decoupled weight decay to mitigate local overfitting. \texttt{FedAdamW} efficiently aggregates the \texttt{mean} of the second-moment estimates to reduce their variance and reinitialize them. Theoretically, we prove that \texttt{FedAdamW} achieves a linear speedup convergence rate of $\mathcal{O}(\sqrt{(L Δσ_l^2)/(S K R ε^2)}+(L Δ)/R)$ without \textbf{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. We also employ PAC-Bayesian generalization analysis to explain the effectiveness of decoupled weight decay in local training. Empirically, we validate the effectiveness of \texttt{FedAdamW} on language and vision Transformer models. Compared to several baselines, \texttt{FedAdamW} significantly reduces communication rounds and improves test accuracy. The code is available in https://github.com/junkangLiu0/FedAdamW.