Fed-ADE: Adaptive Learning Rate for Federated Post-adaptation under Distribution Shift
Organizations: Department of Electrical and Computer Engineering, Sungkyunkwan University, Republic of Korea · Soongsil University, Seoul, Republic of Korea
Abstract
Federated learning (FL) in post-deployment settings must adapt to non-stationary data streams across heterogeneous clients without access to ground-truth labels. A major challenge is learning rate selection under client-specific, time-varying distribution shifts, where fixed learning rates often lead to underfitting or divergence. We propose Fed-ADE (Federated Adaptation with Distribution Shift Estimation), an unsupervised federated adaptation framework that leverages lightweight estimators of distribution dynamics. Specifically, Fed-ADE employs uncertainty dynamics estimation to capture changes in predictive uncertainty and representation dynamics estimation to detect covariate-level feature drift, combining them into a per-client, per-timestep adaptive learning rate. We provide theoretical analyses showing that our dynamics estimation approximates the underlying distribution shift and yields dynamic regret and convergence guarantees. Experiments on image and text benchmarks under diverse distribution shifts (label and covariate) demonstrate consistent improvements over strong baselines. These results highlight that distribution shift-aware adaptation enables effective and robust federated post-adaptation under real-world non-stationarity.
Figures & tables
| Localized Learning | Federated Learning | ||||||||||
| Dataset | Shift | FTH | ATLAS | UNIDA | UDA | Fed-POE | FedCCFA | FixLR(Low) | FixLR(Mid) | FixLR(High) | Fed-ADE |
| (i) Label Shift Scenarios | |||||||||||
| Tiny ImageNet | Lin. | 78.2 1.0 | 76.5 1.3 | 83.8 0.1 | 74.5 0.2 | 87.1 0.2 | 84.7 0.8 | 87.5 1.2 | 88.2 1.1 | 86.1 0.4 | 89.1 0.1 |
| Sin. | 77.9 0.8 | 76.8 1.1 | 83.2 0.5 | 74.5 0.2 | 87.5 0.4 | 84.8 0.7 | 87.3 1.4 | 88.0 1.1 | 87.6 0.6 | 88.9 0.1 | |
| Squ. | 77.2 0.8 | 78.5 1.4 | 83.3 0.1 | 74.7 0.3 | 86.4 0.9 | 83.0 0.9 | 87.4 1.2 | 88.2 0.7 | 86.4 0.5 | 88.9 0.1 | |
| Ber. | 78.2 1.1 | 77.6 1.1 | 82.7 0.3 | 73.8 0.8 | 86.5 0.7 | 83.8 0.9 | 86.5 1.7 | 87.8 1.2 | 86.3 0.6 | 88.7 0.1 | |
| Distribution Shift | Fed-ADE | FedCCFA | Fed-POE | UDA | UNIDA | ATLAS | FTH | |
|---|---|---|---|---|---|---|---|---|
| Gaussian | Lin. | 69.2 1.7 | 62.3 1.2 | 65.7 2.1 | 20.1 2.0 | 20.2 1.8 | 25.3 2.2 | 22.3 1.7 |
| Sin. | 68.3 1.6 | 63.1 1.5 | 64.8 1.7 | 27.8 2.4 | 27.7 2.2 | 33.5 2.3 | 23.9 2.1 | |
| Squ. | 66.4 1.5 | 61.9 2.6 | 61.4 2.2 | 20.3 2.1 | 19.8 1.7 | 23.1 2.0 | 23.0 1.8 | |
| Ber. | 67.1 1.3 | 63.0 2.5 | 60.9 2.5 | 18.2 2.3 | 17.7 1.9 | 22.5 2.1 | 21.6 1.9 | |
| Exponential Decay | Lin. | 67.4 1.6 | 54.1 3.3 | 61.1 1.8 | 23.4 3.1 | 25.4 2.8 | 26.1 1.5 | 22.3 1.9 |
| Sin. | 67.6 1.6 | 54.2 2.5 | 61.4 2.5 | 23.9 2.9 | 24.5 2.3 | 25.6 2.1 | 21.5 1.7 | |
| Measure | Lin. | Sin. | Squ. | Ber. |
|---|---|---|---|---|
| Fed-ADE (Cosine similarity) | 73.8 0.6 | 73.6 0.5 | 72.2 1.6 | 72.9 2.2 |
| KL divergence | 63.2 3.8 | 62.4 2.3 | 61.1 3.4 | 61.4 3.6 |
| Wasserstein | 70.8 0.6 | 69.9 1.1 | 68.8 1.3 | 66.5 2.1 |
| Bayesian CPD | 71.5 0.5 | 70.8 0.8 | 69.7 2.4 | 69.0 1.1 |
| Variant | Label Shift (%) | Covariate Shift (%) | ||
|---|---|---|---|---|
| ✓ | ✓ | Fed-ADE (full) | ||
| ✗ | ✓ | w/o | ||
| ✓ | ✗ | w/o | ||
| ✗ | ✗ | Fixed LR |
| Notation | Description | Notation | Description |
|---|---|---|---|
| Pre-training data at server | Pre-training labels at server | ||
| Feature distribution of | Label distribution of | ||
| Set of class indices | Index for class | ||
| Timestep | Total number of timesteps | ||
| Set of FL clients | Client index | ||
| Global model parameters | Model parameters of client |
| Hyperparameter | Tiny ImageNet | CIFAR-10 | CIFAR-100 |
|---|---|---|---|
| Number of clients | 100 | 100 | 100 |
| Communication rounds ( ) | 10 | 10 | 10 |
| Local epochs | 4 | 4 | 4 |
| Optimizer | SGD | SGD | SGD |
| Batch size | 128 | 32 | 32 |
| Participant rate | 10% | 10% | 10% |
| Hyperparameter | CIFAR-10-C | CIFAR-100-C | LAMA |
|---|---|---|---|
| Number of clients | 100 | 100 | 5 |
| Communication rounds ( ) | 10 | 10 | 5 |
| Local epochs | 4 | 4 | 1 |
| Optimizer | SGD | SGD | AdamW |
| Batch size | 32 | 32 | 4 |
| Participant rate | 10% | 10% | 100% |
| Component | Specification |
|---|---|
| CPU | AMD Ryzen 9 7950X, 16-Core, 32 Threads |
| GPU | NVIDIA GeForce RTX 3090X2 |
| RAM | 256 GB DDR5 |
| Storage | 1.8 TB NVMe SSD |
| Localized Learning | Federated Learning | ||||||||||
| Dataset | Shift | FTH | ATLAS | UNIDA | UDA | Fed-POE | FedCCFA | FixLR(Low) | FixLR(Mid) | FixLR(High) | Fed-ADE |
| (i) Label Shift Scenarios | |||||||||||
| Tiny ImageNet | Lin. | 74.1 2.3 | 70.7 2.1 | 75.1 2.1 | 63.1 2.3 | 81.5 1.2 | 81.7 1.8 | 81.5 1.1 | 81.2 1.3 | 78.1 1.5 | 84.5 1.1 |
| Sin. | 73.7 2.4 | 71.1 1.9 | 74.8 2.5 | 61.3 2.1 | 82.1 1.2 | 80.9 1.1 | 81.9 1.7 | 81.0 2.1 | 77.5 2.1 | 84.7 0.7 | |
| Squ. | 72.6 2.4 | 70.9 1.8 | 74.6 2.2 | 61.7 2.1 | 81.5 1.3 | 78.8 1.9 | 82.4 1.2 | 80.9 1.6 | 77.7 1.3 | 83.9 0.9 | |
| Ber. | 72.1 2.3 | 71.0 2.0 | 74.7 2.1 | 61.1 1.7 | 80.9 1.6 | 79.1 1.1 | 82.1 0.9 | 80.8 1.7 | 77.8 1.5 | 84.2 0.4 | |
| Localized Learning | Federated Learning | ||||||||||
| Dataset | Shift | FTH | ATLAS | UNIDA | UDA | Fed-POE | FedCCFA | FixLR(Low) | FixLR(Mid) | FixLR(High) | Fed-ADE |
| (i) Label Shift Scenarios | |||||||||||
| Tiny ImageNet | Lin. | 78.2 1.0 | 70.2 2.1 | 77.3 0.8 | 69.0 0.6 | 79.9 0.3 | 82.1 0.4 | 81.8 0.5 | 82.5 0.3 | 81.4 0.2 | 83.1 0.4 |
| Sin. | 77.9 0.8 | 71.4 1.6 | 77.1 0.6 | 69.4 0.5 | 80.4 0.6 | 81.5 0.4 | 82.1 0.2 | 82.7 0.5 | 81.3 0.3 | 82.7 0.5 | |
| Squ. | 77.2 0.8 | 71.1 1.8 | 76.9 0.6 | 69.4 0.5 | 80.9 0.5 | 81.2 0.3 | 82.3 0.4 | 82.1 0.4 | 80.1 0.3 | 82.6 0.3 | |
| Ber. | 78.2 1.1 | 70.7 1.6 | 76.6 0.7 | 69.7 0.4 | 80.4 0.5 | 80.4 0.3 | 81.1 0.4 | 82.1 0.2 | 80.4 0.3 | 81.8 0.4 | |