Fed-ADE: Adaptive Learning Rate for Federated Post-adaptation under Distribution Shift

Heewon Park; Mugon Joe; Miru Kim; Kyungjin Im; Minhae Kwon

cs.LGMar 1, 2026

Fed-ADE: Adaptive Learning Rate for Federated Post-adaptation under Distribution Shift

Authors: Heewon Park, Mugon Joe, Miru Kim, Kyungjin Im, Minhae Kwon

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

Figure 1: Overview of the Fed-ADE. Each client receives a pre-trained model, splits it into shared ( $\psi_{c}$ ) and personalized ( $\phi_{c}$ ) layers, and adapts to unlabeled, distribution-shifting data using an adaptive learning rate.

		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 $\pm$ 1.0	76.5 $\pm$ 1.3	83.8 $\pm$ 0.1	74.5 $\pm$ 0.2	87.1 $\pm$ 0.2	84.7 $\pm$ 0.8	87.5 $\pm$ 1.2	88.2 $\pm$ 1.1	86.1 $\pm$ 0.4	89.1 $\pm$ 0.1
	Sin.	77.9 $\pm$ 0.8	76.8 $\pm$ 1.1	83.2 $\pm$ 0.5	74.5 $\pm$ 0.2	87.5 $\pm$ 0.4	84.8 $\pm$ 0.7	87.3 $\pm$ 1.4	88.0 $\pm$ 1.1	87.6 $\pm$ 0.6	88.9 $\pm$ 0.1
	Squ.	77.2 $\pm$ 0.8	78.5 $\pm$ 1.4	83.3 $\pm$ 0.1	74.7 $\pm$ 0.3	86.4 $\pm$ 0.9	83.0 $\pm$ 0.9	87.4 $\pm$ 1.2	88.2 $\pm$ 0.7	86.4 $\pm$ 0.5	88.9 $\pm$ 0.1
	Ber.	78.2 $\pm$ 1.1	77.6 $\pm$ 1.1	82.7 $\pm$ 0.3	73.8 $\pm$ 0.8	86.5 $\pm$ 0.7	83.8 $\pm$ 0.9	86.5 $\pm$ 1.7	87.8 $\pm$ 1.2	86.3 $\pm$ 0.6	88.7 $\pm$ 0.1

Table 1: Performance comparison of different online adaptation methods (average accuracy (%) and average wall time (sec.))

Panel of Figure 2 — Figure 2: Impact of learning rate bounds on Fed-ADE under label shift scenarios on Tiny-ImageNet and CIFAR-10 datasets.

Distribution Shift		Fed-ADE	FedCCFA	Fed-POE	UDA	UNIDA	ATLAS	FTH
Gaussian	Lin.	69.2 $\pm$ 1.7	62.3 $\pm$ 1.2	65.7 $\pm$ 2.1	20.1 $\pm$ 2.0	20.2 $\pm$ 1.8	25.3 $\pm$ 2.2	22.3 $\pm$ 1.7
	Sin.	68.3 $\pm$ 1.6	63.1 $\pm$ 1.5	64.8 $\pm$ 1.7	27.8 $\pm$ 2.4	27.7 $\pm$ 2.2	33.5 $\pm$ 2.3	23.9 $\pm$ 2.1
	Squ.	66.4 $\pm$ 1.5	61.9 $\pm$ 2.6	61.4 $\pm$ 2.2	20.3 $\pm$ 2.1	19.8 $\pm$ 1.7	23.1 $\pm$ 2.0	23.0 $\pm$ 1.8
	Ber.	67.1 $\pm$ 1.3	63.0 $\pm$ 2.5	60.9 $\pm$ 2.5	18.2 $\pm$ 2.3	17.7 $\pm$ 1.9	22.5 $\pm$ 2.1	21.6 $\pm$ 1.9
Exponential Decay	Lin.	67.4 $\pm$ 1.6	54.1 $\pm$ 3.3	61.1 $\pm$ 1.8	23.4 $\pm$ 3.1	25.4 $\pm$ 2.8	26.1 $\pm$ 1.5	22.3 $\pm$ 1.9
	Sin.	67.6 $\pm$ 1.6	54.2 $\pm$ 2.5	61.4 $\pm$ 2.5	23.9 $\pm$ 2.9	24.5 $\pm$ 2.3	25.6 $\pm$ 2.1	21.5 $\pm$ 1.7

Table 2: Performance with Pre-trained Model Trained on Various Distributions with CIFAR-10.

Measure	Lin.	Sin.	Squ.	Ber.
Fed-ADE (Cosine similarity)	73.8 $\pm$ 0.6	73.6 $\pm$ 0.5	72.2 $\pm$ 1.6	72.9 $\pm$ 2.2
KL divergence	63.2 $\pm$ 3.8	62.4 $\pm$ 2.3	61.1 $\pm$ 3.4	61.4 $\pm$ 3.6
Wasserstein	70.8 $\pm$ 0.6	69.9 $\pm$ 1.1	68.8 $\pm$ 1.3	66.5 $\pm$ 2.1
Bayesian CPD	71.5 $\pm$ 0.5	70.8 $\pm$ 0.8	69.7 $\pm$ 2.4	69.0 $\pm$ 1.1

Table 3: Ablation study on the choice of distribution-shift similarity measure in Fed-ADE under CIFAR-10 label-shift scenarios.

$S_{\text{unc}}$	$S_{\text{rep}}$	Variant	Label Shift (%)	Covariate Shift (%)
✓	✓	Fed-ADE (full)	$73.8\pm 0.6$	$64.4\pm 0.2$
✗	✓	w/o $S_{\text{unc}}$	$71.3\pm 0.3$	$64.2\pm 0.1$
✓	✗	w/o $S_{\text{rep}}$	$73.1\pm 0.4$	$64.0\pm 0.3$
✗	✗	Fixed LR	$70.8\pm 2.1$	$63.9\pm 0.3$

Table 4: Ablation of Fed-ADE estimators on CIFAR-10

Notation	Description	Notation	Description
$\mathbf{x}_{G}^{0}$	Pre-training data at server	$\mathbf{y}_{G}^{0}$	Pre-training labels at server
$\mathbf{Q}_{{\mathcal{G}},\mathbf{x}}^{0}$	Feature distribution of $\mathbf{x}_{G}^{0}$	$\mathbf{Q}_{{\mathcal{G}},\mathbf{y}}^{0}$	Label distribution of $\mathbf{y}_{G}^{0}$
$\mathcal{I}$	Set of class indices	$i$	Index for class
$t$	Timestep	$T$	Total number of timesteps
$\mathcal{C}$	Set of FL clients	$c$	Client index
$\theta$	Global model parameters	$\theta_{c}$	Model parameters of client $c$

Table 5: Notation Table

Figure 3: Visualization of how the parameter $\omega(t)$ changes over time for various shift types. The Lin. presents a steady and continuous increase or decrease, reflecting a gradual and predictable change. Both the Squ. and Sin. display periodic patterns. The Squ. alternates sharply between two values at regular intervals, while the Sin. oscillates smoothly in a wave-like manner. The Ber. introduces randomness at each timestep, resulting in a stochastic and less predictable trajectory for $\omega(t)$ .

Panel of Figure 4 — Figure 4: Visualization of heterogeneity among clients according to Dirichlet distribution factor $\alpha$ in CIFAR-10 dataset. The $x$ -axis represents the ID of the client, while the $y$ -axis represents the class of the data. The color of the circles varies with each class type, and the size of the circles indicates the size of the data. This level of heterogeneity is not dependent on the data set but on $\alpha$ values.

Hyperparameter	Tiny ImageNet	CIFAR-10	CIFAR-100
Number of clients	100	100	100
Communication rounds ( $R$ )	10	10	10
Local epochs	4	4	4
Optimizer	SGD	SGD	SGD
Batch size	128	32	32
Participant rate	10%	10%	10%

Table 6: Hyperparameters for Tiny ImageNet, CIFAR-10, and CIFAR-100.

Hyperparameter	CIFAR-10-C	CIFAR-100-C	LAMA
Number of clients	100	100	5
Communication rounds ( $R$ )	10	10	5
Local epochs	4	4	1
Optimizer	SGD	SGD	AdamW
Batch size	32	32	4
Participant rate	10%	10%	100%

Table 7: Hyperparameters for CIFAR-10-C, CIFAR-100-C, and LAMA

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

Table 8: Hardware specifications used for all simulations.

Panel of Figure 5 — Figure 5: Impact of learning rate bounds on Fed-ADE under label shift scenarios.

		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 $\pm$ 2.3	70.7 $\pm$ 2.1	75.1 $\pm$ 2.1	63.1 $\pm$ 2.3	81.5 $\pm$ 1.2	81.7 $\pm$ 1.8	81.5 $\pm$ 1.1	81.2 $\pm$ 1.3	78.1 $\pm$ 1.5	84.5 $\pm$ 1.1
	Sin.	73.7 $\pm$ 2.4	71.1 $\pm$ 1.9	74.8 $\pm$ 2.5	61.3 $\pm$ 2.1	82.1 $\pm$ 1.2	80.9 $\pm$ 1.1	81.9 $\pm$ 1.7	81.0 $\pm$ 2.1	77.5 $\pm$ 2.1	84.7 $\pm$ 0.7
	Squ.	72.6 $\pm$ 2.4	70.9 $\pm$ 1.8	74.6 $\pm$ 2.2	61.7 $\pm$ 2.1	81.5 $\pm$ 1.3	78.8 $\pm$ 1.9	82.4 $\pm$ 1.2	80.9 $\pm$ 1.6	77.7 $\pm$ 1.3	83.9 $\pm$ 0.9
	Ber.	72.1 $\pm$ 2.3	71.0 $\pm$ 2.0	74.7 $\pm$ 2.1	61.1 $\pm$ 1.7	80.9 $\pm$ 1.6	79.1 $\pm$ 1.1	82.1 $\pm$ 0.9	80.8 $\pm$ 1.7	77.8 $\pm$ 1.5	84.2 $\pm$ 0.4

Table 9: Performance comparison of different online distribution shift adaptation methods across datasets and shift types with pretrained model from Gaussian distribution (average accuracy (%) and average wall time (sec.))

		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 $\pm$ 1.0	70.2 $\pm$ 2.1	77.3 $\pm$ 0.8	69.0 $\pm$ 0.6	79.9 $\pm$ 0.3	82.1 $\pm$ 0.4	81.8 $\pm$ 0.5	82.5 $\pm$ 0.3	81.4 $\pm$ 0.2	83.1 $\pm$ 0.4
	Sin.	77.9 $\pm$ 0.8	71.4 $\pm$ 1.6	77.1 $\pm$ 0.6	69.4 $\pm$ 0.5	80.4 $\pm$ 0.6	81.5 $\pm$ 0.4	82.1 $\pm$ 0.2	82.7 $\pm$ 0.5	81.3 $\pm$ 0.3	82.7 $\pm$ 0.5
	Squ.	77.2 $\pm$ 0.8	71.1 $\pm$ 1.8	76.9 $\pm$ 0.6	69.4 $\pm$ 0.5	80.9 $\pm$ 0.5	81.2 $\pm$ 0.3	82.3 $\pm$ 0.4	82.1 $\pm$ 0.4	80.1 $\pm$ 0.3	82.6 $\pm$ 0.3
	Ber.	78.2 $\pm$ 1.1	70.7 $\pm$ 1.6	76.6 $\pm$ 0.7	69.7 $\pm$ 0.4	80.4 $\pm$ 0.5	80.4 $\pm$ 0.3	81.1 $\pm$ 0.4	82.1 $\pm$ 0.2	80.4 $\pm$ 0.3	81.8 $\pm$ 0.4

Table 10: Performance comparison of different online distribution shift adaptation methods across datasets and shift types with pretrained model from exponential decay distribution(average accuracy (%) and average wall time (sec.))

Adaptive Learning Rate

See all 17 papers

Beyond Non-IID: Learner--Client Distribution Mismatch in Federated Learning
Aug 27, 2026Yiming Xie, Lili Su, Ningfang MiClient Selection in Federated Learning Non-IID Federated Learning
Department of Electrical and Computer Engineering Northeastern University, Boston, MA, USA
FedeRage: Provably Convergent Agnostic Federated Learning under General Client Drift
Sep 17, 2026Herlock Rahimi, Dionysis KalogeriasFederated Learning Federated Learning Aggregation
FedA2L: Adaptive layer-wise learning rate adjustment in decentralized federated learning
Aug 10, 2026Van Truong Vo, Khoa Nguyen, Taehong KimNon-IID Federated Learning Federated Learning
School of Information and Communication Engineering, Chungbuk National University, Cheongju, South Korea

Aug 27, 2026J/K move · Enter open · S save

Beyond Non-IID: Learner--Client Distribution Mismatch in Federated Learning

Yiming Xie, Lili Su, Ningfang Mi

Department of Electrical and Computer Engineering Northeastern University, Boston, MA, USA

Client Selection in Federated Learning Non-IID Federated Learning Data Valuation

Open paper AskPDF arXiv

Federated learning systems are increasingly deployed to facilitate collaborative model training across a heterogeneous client population. Existing practice mostly implicitly assumes that the aggregated client data distribution is representative of the learner's target distribution or that learning from all available clients is uniformly beneficial for the learner distribution. However, such an assumption often does not hold in reality. Traditional client selection strategies in FL literature largely overlook such misalignment, while most existing work on multi-source transfer learning either requires direct access to local data or uses one-shot model/feature aggregation. In this paper, we take the initiative to understand and mitigate the impacts of such learner-client population misalignment. In particular, we consider the practical setting where the learner keeps a small proxy dataset. We observe that client contributions vary significantly across training rounds, and traditional technology is insufficient to identify beneficial sources under multi-source transfer diversity. Then, we propose a dynamic, influence-aware client selection framework that estimates each client's potential utility to the learner's optimization objective using proxy influence signals on a learner-specific proxy set. Via using leave-one-out evaluations, we prioritize the most informative sources of knowledge while controlling the negative impacts of statistical noise and data heterogeneity. Experiments on CIFAR-10 under heterogeneous data partitions demonstrate that our approach consistently outperforms static and dynamic baselines, achieving faster convergence and higher accuracy.

Abstract

Adaptive Learning Rate

Explore similar work

Beyond Non-IID: Learner--Client Distribution Mismatch in Federated Learning

FedeRage: Provably Convergent Agnostic Federated Learning under General Client Drift

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