FedA2L: Adaptive layer-wise learning rate adjustment in decentralized federated learning

Decentralized federated learning (DFL) suffers from convergence inefficiency under data heterogeneity due to the use of a uniform learning rate (LR) that ignores layer-specific optimization needs. The paper states: Foundational layers are responsible for maintaining network consensus, while specialized layers adapt to local data characteristics, leading to conflicting gradients and degraded performance under non-IID conditions.

The authors identify a fundamental tension: "Low LRs help maintain consensus in foundational layers and stabilize the network, but they limit the ability of specialized layers to adapt to local data. In contrast, high LRs accelerate adaptation in specialized layers but destabilize foundational layers, causing parameter oscillations and degraded convergence."

The paper introduces FedA2L, a serverless method for dynamically tuning layer-wise LRs based on model divergence signals. The core innovation formalizes model states at distinct transitions of each DFL round: the base state (before local training), the trained state (after local gradient updates), and the aggregated state (after network consensus).

FedA2L derives two layer-wise metrics from these state transitions:

1. Weight divergence (σ): measuring local update intensity across nodes — defined as the normalized L2-norm of parameter change during local training: σ = ‖θ T - θ B‖2 / (‖θ B‖2 + ε)

2. Aggregation stability (ζ): measuring network consensus constraints — defined as the proportion of parameters that remain stable after aggregation: ζ = (1/θ) Σ I[θ A - θ T < τ]

The method uses Z-score normalization over a sliding temporal window to stabilize metrics, then fuses them into a fusion score: λ = β·e ω + (1-β)·e δ, where β ∈ [0,1] balances the two signals.

The layer-wise LR is computed as: η = η0 · (1 + tanh(log(λ))) · γ r, where γ r = (1+ξ·r)(-0.5) is a global decay factor.

Key properties: "Since λ > 0, log(λ) ∈ R and tanh(log(λ)) ∈ (-1, 1), which implies 1 + tanh(log(λ)) ∈ (0, 2). Together with γ r ≤ 1, the effective layer-wise LRs satisfy 0 < η < 2η0γ r, so FedA2L cannot arbitrarily amplify the step size."

1. FedA2L introduces a state-based formalization that characterizes model dynamics across three states of each DFL round

2. Adaptation of two layer-wise metrics (weight divergence σ and aggregation stability ζ) enables layer-wise dynamic LR adjustment

3. Layer-wise adaptive LRs based on local model divergence metrics can be seamlessly integrated into existing DFL algorithms without modifying the core aggregation logic or introducing additional communication overhead

4. "Extensive empirical validation demonstrates that FedA2L significantly accelerates convergence across six DFL algorithms and DFedHPO, spanning multiple neural network architectures and diverse datasets, achieving up to 4.94× faster convergence than vanilla DFL and up to 59% fewer communication rounds than the best-performing scheduler-based baselines"

The paper establishes convergence guarantees for both non-convex and strongly convex settings. For non-convex convergence, Theorem 1 shows that FedA2L achieves the standard O(1/√R) rate: matching the standard O(1/√R) non-convex rate of decentralized SGD. FedA2L does not worsen the convergence order; the bounded multipliers only modify the constants.

For strongly convex settings, the paper shows FedA2L preserves the linear convergence of decentralized SGD with contraction factor ρ0 ∈ (0,1) determined jointly by μ and the spectral gap φ.

Across 36 scheduler-strategy settings, FedA2L achieves the fastest convergence in 32 cases and the second-fastest in 3. Key results include:

• On TinyImageNet with ResNet-34 under FedAWA: FedA2L reaches 72 rounds versus 356 rounds on ResNet-34 TinyImageNet (4.94× speedup)

• Under FedProx on CNN CIFAR-10: FedA2L still achieves 2.24× speedup on CNN CIFAR-10 (127 vs. 285 rounds)

• On ResNet-18 CIFAR-100: FedA2L converges in 177 rounds compared to 491 rounds for the vanilla baseline under FedAvg (2.77× speedup)

FedA2L achieves the highest test accuracy across diverse settings, maintaining both model quality and training stability under non-IID conditions. For ResNet-18 on CIFAR-100, the best test accuracy drops from 49.50% at β = 0.6 to 49.23% at β = 0.0 and 49.10% at β = 1.0, demonstrating that σ and ζ provide complementary and non-redundant information.

• Severe heterogeneity (α=0.01): FedA2L achieves the fastest convergence (181 rounds for FedAvg, 187 for FedAWA) and maintains a substantial 1.66× speedup over the vanilla baseline

• Scalability (20-30 nodes): FedA2L sustains its performance advantage over all baselines, requiring 278 and 273 rounds under FedAvg and FedAWA, respectively, for 20 nodes, and 339 and 334 rounds for 30 nodes

• Sparse topologies (Ring): "Most baseline methods fail to converge entirely. Despite these severe structural constraints, FedA2L successfully converges, achieving a remarkable 1.55× speedup (318 rounds for FedAvg, 319 for FedAWA) compared to the best-performing converging baseline (OCLR)"

On ExchangeRate: FedA2L reaches target MSE in 11 rounds versus 16 rounds for the best scheduler (OCLR) under FedAvg, a 1.45× speedup that translates to 31.25% reduction in communication rounds and 26.37% reduction in total time.

The default configuration is: R warm = 10 (CNN) or 40 (ResNet), ρ = 5 (CNN) or 10 (ResNet), τ = 1e-3, β = 0.6, ξ = 0.3 (CNN) or 0.1 (ResNet).

Key findings: "Among the five parameters, ρ and τ show limited sensitivity across all tested architectures, while β remains stable within [0.4, 0.8], confirming that the dual-signal fusion is robust to reasonable variation in weighting. The primary configuration parameters are R warm and ξ, both of which interact directly with model depth."

FedA2L achieves per-round and per-node computation having complexity O(θ · D i · E) with adaptive computation overhead of O(θ + Lρ). Critically, FedA2L introduces zero additional communication overhead, requiring no extra messages, proxy datasets, or centralized coordination.

The paper notes limitations: some hyperparameters (e.g., R warm and ξ) depend on model depth and may require additional tuning across different architectures and the performance gains are less pronounced when the base DFL algorithm already incorporates strong drift correction or adaptive regularization (e.g., FedNTD and FedYogi).

Based on FedA2L, I can implement the following improvements to AI systems:

1. Adaptive Layer-Wise Learning Rate Controller for Distributed Training

• Add a module that tracks three model states (base, trained, aggregated) during each training round

• Compute weight divergence (σ) and aggregation stability (ζ) per layer

• Dynamically adjust per-layer learning rates using the fusion formula: η = η0 · (1 + tanh(log(λ))) · γ r

• Apply this to existing federated/decentralized training pipelines without changing aggregation logic

Improved capability: The system automatically increases learning rates for specialized layers that need faster adaptation while keeping foundational layers stable, resulting in 2–5× faster convergence under non-IID data distributions without manual LR tuning.

2. Self-Tuning Scheduler for Heterogeneous Data Environments

• Implement Z-score normalization over sliding windows to stabilize metric signals

• Use β=0.6 fusion weight to balance divergence and stability signals

• Apply global decay factor γ r = (1+ξ·r)(-0.5) for convergence guarantees

3. Zero-Overhead Drift Correction Module

• Add computation of σ and ζ with complexity O(θ + Lρ) per round

• Integrate into existing DFL algorithms (FedAvg, FedProx, FedAWA, etc.) without extra communication

• Apply layer-wise LR adjustments only, leaving aggregation logic untouched

4. Architecture-Aware Hyperparameter Auto-Configuration

• Implement adaptive warm-up period (R warm = 10 for CNN, 40 for ResNet)

• Adjust decay rate ξ based on model depth (0.3 for CNN, 0.1 for ResNet)

• Keep β stable within [0.4, 0.8] for robustness

5. Complementary Signal Fusion for Model Quality

• Use both σ (local update intensity) and ζ (consensus stability) as complementary signals

• Apply tanh(log(λ)) transformation to bound effective LR within (0, 2η0γ r)

Decentralized intelligence systems with heterogeneous devices and limited coordination increasingly rely on decentralized federated learning (DFL). However, DFL suffers from convergence inefficiency under data heterogeneity due to the use of a uniform learning rate (LR) that ignores layer-specific optimization needs. Foundational layers are responsible for maintaining network consensus, while specialized layers adapt to local data characteristics, leading to conflicting gradients and degraded performance under non-IID conditions. To address this fundamental tension, this work introduces FedA2L, a method that dynamically adjusts layer-wise LRs based on model divergence signals. By leveraging local update intensity and network consensus constraints, FedA2L seamlessly integrates into existing DFL protocols without additional communication or coordination. Extensive evaluations across DFL algorithms, various model architectures, and datasets demonstrate that FedA2L achieves up to 4.94 times faster convergence than vanilla DFL and reduces communication rounds by up to 59% compared to scheduler-based baselines. Furthermore, FedA2L exhibits resilience to severe data heterogeneity, larger network sizes, and sparse topologies, reducing communication overhead and establishing it as a versatile optimization tool for resource-constrained or large-scale distributed learning in edge and IoT deployments. The code is released at https://github.com/nclabteam/FedA2L.