Auto-FedAvg: Adaptive Federated Averaging

• Auto-FedAvg is a federated averaging variant that automatically learns aggregation weights and hyperparameters to tackle data heterogeneity.
• It employs techniques like softmax and Dirichlet parameterizations alongside meta-learning to adapt server parameters and client contributions.
• Empirical evaluations in medical imaging, CIFAR benchmarks, and heterogeneous deployments reveal enhanced convergence, accuracy, and robustness.

Auto-FedAvg refers to data-driven generalizations of the federated averaging (FedAvg) paradigm, where server aggregation weights and/or hyperparameters are dynamically adjusted throughout training, often via learnable procedures or direct gradient-based optimization. Initially formulated to address the challenges of non-i.i.d. client data distributions, especially in multi-institutional or cross-silo FL deployments, Auto-FedAvg variants have since expanded to encompass adaptive step-size, per-client weighting, and automatic hyperparameter selection strategies, with applications in medical imaging, heterogeneous benchmark tasks, and large-scale cross-device learning. Major instantiations include approaches with learnable aggregation weights (Xia et al., 2021), extrapolation-based adaptation of server step-sizes (Jhunjhunwala et al., 2023), doubly adaptive coordinate-wise step-sizes (Takakura et al., 16 May 2025), lightweight online estimation of weights under unknown participation (Wang et al., 2023), and algorithms leveraging federated automatic differentiation to meta-learn server and aggregation parameters (Rush et al., 2023).

1. Motivation and Limitations of Standard FedAvg

The classic FedAvg algorithm employs static client aggregation weights, typically proportional to local dataset sizes, when averaging model updates on the server. This design is principled under i.i.d. client data, but is often sub-optimal in real-world cross-silo and cross-device FL. Domain shifts (e.g., different medical imaging equipment, protocols, or population demographics) induce highly heterogeneous local objective landscapes. In such regimes, large or "easy" clients may dominate aggregation, biasing training and degrading both convergence and generalization. Furthermore, as participation rates vary across clients—often unknown a priori in realistic deployments—fixed weights can permanently skew the global model toward more active or larger participants.

Auto-FedAvg strategies seek to resolve these issues by adaptively learning or inferring the most effective aggregation weights or server-side optimization parameters, thus improving robustness to data and participation heterogeneity (Xia et al., 2021, Wang et al., 2023).

2. Mathematical Formulations and Algorithms

At the core of many Auto-FedAvg approaches is the relaxation of fixed-weighted aggregation to a parametrized, learnable objective. A canonical instance is the simultaneous minimization over shared model parameters $\theta$ and global aggregation weights $\alpha=(\alpha_1,\ldots,\alpha_K)$: $$\min_{\theta,\alpha} \mathcal{L}(\theta,\alpha) = \sum_{k=1}^K \alpha_k F_k(\theta) \quad \text{s.t.}\quad \sum_{k=1}^K \alpha_k = 1,\, \alpha_k \ge 0.$$ Here, $F_k$ denotes the expected local loss for client $k$, with $\theta$ representing the neural network parameters (e.g., a 3D U-Net in medical segmentation).

Auto-FedAvg algorithms typically alternate between:

1. Local model update ($\theta$): Each client receives the current global model, performs several SGD steps, and returns its locally trained model.
2. Aggregation parameter update ($\alpha$): The server (sometimes with client assistance) optimizes $\alpha$ under simplex constraints, often using indirect gradient signals, softmax or Dirichlet parameterizations, and validation loss feedback.
3. Global model aggregation: The server recombines local models with the latest $\alpha$ weights, updating $\theta$ accordingly (Xia et al., 2021).

A noteworthy extension is the "FedDuA" framework (Takakura et al., 16 May 2025), which recasts the FedAvg server update as mirror descent and selects both the server step-size $\eta_g^t$ and diagonal preconditioning matrix $G_t$ in a minimax-optimal manner, adapting to inter-client and coordinate-wise heterogeneity: $w_{t+1} = w_t + \eta_g^t G_t^{-1} v_t,$ where $v_t$ is the aggregated pseudo-gradient and $\eta_g^t$ is computed to minimize the worst-case Bregman divergence under an approximate projection condition.

Pseudocode for these methods is directly available in each primary reference (Xia et al., 2021, Takakura et al., 16 May 2025, Wang et al., 2023, Jhunjhunwala et al., 2023, Rush et al., 2023).

3. Learning and Optimization of Aggregation Weights

Several mechanisms for learning $\alpha$ have been proposed:

• Softmax parameterization: $\alpha_k = \exp(\beta_k) / \sum_i \exp(\beta_i)$, with $\beta$ updated via gradient signals.
• Dirichlet-based stochasticity: $\alpha \sim \mathrm{Dirichlet}(\beta)$, leveraging reparameterization gradients as in pathwise estimators (Xia et al., 2021).
• Validation loss feedback: Gradients for $\beta$ (or, more generally, aggregation-relevant parameters—server learning rate, momentum, per-client exponents, etc.) are computed with respect to the validation loss after each global aggregation, facilitating meta-learning of weighting schemes.
• Cutoff-interval estimators: In settings with unknown and heterogeneous client participation rates, each client maintains running averages of its inter-participation intervals (with a cut-off $K$ to balance bias–variance), resulting in an online estimator $\omega_t^n$ that approaches the optimal inverse-frequency weighting (Wang et al., 2023).

A summarizing table of aggregation weight learning techniques:

Method	Weight Parametrization	Adaptation Signal
Softmax/Dirichlet	$\alpha$, $\beta$ (global)	Validation loss gradient
Inverse participation	$\omega_t^n$ (per-client)	Local interval statistics
Server step-size	$\eta_g^t$	Pseudo-gradient alignments
FAD meta-learning	$\theta$, $\alpha$, $\eta$	Federated hypergradients

By disentangling local model parameters from global aggregation weights or server hyperparameters, Auto-FedAvg variants enable dynamic adaptation to both static and evolving heterogeneity in FL.

4. Theoretical Properties and Optimization Guarantees

Rigorous convergence analysis exists for multiple Auto-FedAvg variants:

• Learnable aggregation weights: Under convexity and L-smoothness, alternating gradient descent yields faster convergence and improved generalization, provided aggregation parameter updates are frequent and communication-efficient (Xia et al., 2021).
• Step-size adaptation (FedExP): Adaptive $\eta_g^t$ confers accelerated convergence in overparameterized convex regimes, with theoretical justifications via analogies to extrapolated POCS algorithms and direct minimization of the global model's distance to the optimum (Jhunjhunwala et al., 2023).
• Doubly adaptive mirror descent (FedDuA): Minimax-optimal step-sizes and preconditioners decrease the worst-case Bregman divergence to any putative optimum under approximate projection conditions. Coordinate-wise adaptation (FedDuAdagrad/FedDuAdam) enables dimension-free convergence under anisotropic update structures (Takakura et al., 16 May 2025).
• Adaptive weighting for participation heterogeneity (FedAU): When participation probabilities are unknown, online averaging of participation intervals ensures that the weights $\omega_t^n$ converge to the optimal $1/p_n$, and the global FedAvg iterate solves the true federated objective $\frac{1}{N}\sum F_n$ with classical $O(1/\sqrt{NIT})$ rate, plus an extra $O(\log^2 T/T)$ term due to weight estimation error. Empirically, bias from mis-weighted aggregation is dominated by the variance bound of the cutoff-geometric estimator (Wang et al., 2023).

All these approaches are compatible with standard FL assumptions (full or partial participation, bounded variance, and L-smooth local objectives).

5. Communication, Computation, and System Considerations

While augmenting FedAvg with learnable weights or server-side adaptivity introduces additional computation and communication overhead, these are modest compared to the gains in robustness and performance:

• Aggregation parameter updates: Procedurally, learning $\alpha$ requires all clients to evaluate and, occasionally, communicate $O(K)$ model parameters for the weighted sum, though this is amortized by performing $\alpha$-updates every $t_0 \gg 1$ rounds and by the small size of $\beta$ or $\alpha$ vectors (Xia et al., 2021).
• Extrapolated step-sizes/dually adaptive methods: Only aggregate statistics (e.g., $\|\Delta_k\|^2$) or running sums of pseudo-gradients are needed, requiring $O(1)$ extra communication per client (Jhunjhunwala et al., 2023, Takakura et al., 16 May 2025).
• Online participation estimators: Memory requirements are minimal ($O(1)$ per client for FedAU), with all computation local and no need for participation statistics to be relayed to the server (Wang et al., 2023).
• Federated automatic differentiation (FAD): When meta-learning hyperparameters or aggregation weights, an extra round of federated sum (for validation loss gradients) is required. Mixed-mode FAD keeps communication efficient by broadcasting only small Jacobians and using sum reductions compatible with privacy-preserving primitives (Rush et al., 2023).

Compatibility with secure aggregation protocols and differential privacy is preserved, as most additional signals are either scalar statistics or use sum-only primitives.

6. Empirical Results and Applications

Empirical studies across heterogeneous vision, language, and synthetic tasks demonstrate that Auto-FedAvg variants systematically outperform baseline methods:

• Medical imaging (COVID-19, pancreas segmentation): Auto-FedAvg improves global test Dice by 1–2% over FedAvg/FedProx on strongly non-i.i.d. multi-institutional datasets, with significance confirmed via Wilcoxon tests (Xia et al., 2021).
• CIFAR-10/100, FEMNIST, Shakespeare: Adaptive methods (Auto-FedAvg, FedDuA, FedExP) accelerate convergence and/or increase final accuracy by several percentage points over vanilla and momentum-augmented FedAvg (Takakura et al., 16 May 2025, Jhunjhunwala et al., 2023).
• Cross-participation heterogeneity (FedAU): On benchmarks with highly heterogeneous and unknown client participation, adaptive online reweighting (FedAU) consistently yields higher accuracy and unbiased solutions, outperforming fixed-weight and state-heavy variance-reduction baselines (Wang et al., 2023).
• Meta-learned hyperparameters (FAD): Federated hypergradient methods reduce the cost of hyperparameter sweeping by an order of magnitude and adaptively interpolate between uniform and example-proportional client weighting strategies (Rush et al., 2023).

Dynamic α-adaptation is especially impactful in early training, where small, fast-converging clients are initially upweighted, then gradually downweighted as larger clients catch up, as observed in both learned weights and sharpened Dirichlet posteriors (Xia et al., 2021).

7. Limitations, Extensions, and Future Directions

Auto-FedAvg methods, while robust and performant, exhibit several operational caveats:

• Assumptions of persistent, always-online clients restrict some algorithms (notably those needing all clients to participate in α-learning rounds) to cross-silo settings. Extensions to cross-device scenarios would require dropout-tolerant or asynchronous aggregation schemes (Xia et al., 2021).
• For extremely large client populations, even modest per-client communication overhead during aggregation parameter updates may become non-negligible.
• Current theoretical guarantees for step-size and weight adaptivity focus on convex and full-participation regimes; extensions to non-convex objectives and stochastic or partial participation are active research directions (Jhunjhunwala et al., 2023, Takakura et al., 16 May 2025).
• Opportunities remain to integrate richer aggregation mechanisms (cluster-wise, layer-wise, or personalized weights), differential privacy, and local drift-correction (e.g., SCAFFOLD) with automatic global adaptivity for further gains.

The combination of federated automatic differentiation with data-driven aggregation adaption suggests broad applicability wherever federated optimization must adapt to both data and system heterogeneity, with minimal infrastructure changes and strong privacy compatibility (Rush et al., 2023).