FedDUAL: A Dual-Strategy with Adaptive Loss and Dynamic Aggregation for Mitigating Data Heterogeneity in Federated Learning
Authors: Pranab Sahoo
Ashutosh Tripathi
Sriparna Saha
Samrat Mondal
Venue: Accepted by TMLR ()

Abstract: Federated Learning (FL) marks a transformative approach to distributed model training by combining locally optimized models from various clients into a unified global model. While FL preserves data privacy by eliminating centralized storage, it encounters significant challenges such as performance degradation, slower convergence, and reduced robustness of the global model due to the heterogeneity in client data distributions. Among the various forms of data heterogeneity, label skew emerges as a particularly formidable and prevalent issue, especially in domains such as image classification. To address these challenges, we begin with comprehensive experiments to pinpoint the underlying issues in the FL training process, such as gradient instability and the emergence of sharp minima in the global model, both of which contribute to performance inconsistencies. Based on our findings, we introduce an innovative dual-strategy approach designed to effectively resolve these issues. First, we introduce an adaptive loss function for client-side training, meticulously crafted to preserve previously acquired knowledge while maintaining an optimal equilibrium between local optimization and global model coherence. Secondly, we develop a dynamic aggregation strategy for aggregating client models at the server. This approach adapts to each client's unique learning patterns, effectively addressing the challenges of diverse data across the network. Our comprehensive evaluation, conducted across three diverse real-world datasets, coupled with theoretical convergence guarantees, demonstrates the superior efficacy of our method compared to several established state-of-the-art approaches.

Published in Transactions on Machine Learning Research (11/2025)
FedDUAL: A Dual-Strategy with Adaptive Loss and Dy- namic Aggregation for Mitigating Data Heterogeneity in Federated Learning
Pranab Sahoo
pronab 2021cs25@ilp.
Department of Computer Science Indian Institute of Technology Patna, India
Ashutosh Tripathi
ashutoshtripathi199 @ com
Electrical and Electronics Engg. Department Rajiv Gandhi Institute of Petroleum Technology, India
Sriparna Saha
sriporna@hitp.ac..
Department of Computer Science Indian Institute of Technology Patna, India
Samrat Mondal
sammal@hilp..cm ac. ac.in
Department of Computer Science
Indian Institute of Technology Patna, India
Reviewed on OpenReview: httpss//pppprrrreerrn///////ottd
Abstract
Federated Learning (FL) marks transformative approach to distributed model training by combining locally optimized models from various clients into unified global model. While FL preserves data privacy by eliminating centralized storage, it encounters significant chal- lenges such as performance degradation, slower convergence, and reduced robustness of the global model due the heterogeneity in client data distributions. Among the various forms of data heterogeneity, label skew emerges & particularly formidable and prevalent issue, begin especially in domains such as image classification. To address these challenges, with comprchensive experiments to pinpoint the underlying issues in the FL training pro- cess, such as gradient instability and the emergence of sharp minima in the global model, both of which contribute to performance inconsistencies. Based on our findings, we in- troduce an innovative duall-srategy approach designed to effectively resolve these issues. First, we introduce an adaptive loss function for client-side training. meticulously crafted to preserve previously acquired knowledge while maintaining an optimal equilibrium between local optimization and global model coherence. Secondly, develop dynamic aggregation strategy for aggregating client models at the server. This approach adapts to each client's 's unique learning patterns, effectively addressing the challenges of diverse data across the network. Our comprehensive evaluation, conducted across three diverse real-world datasets, coupled with theoretical convergence guarantees, demonstrates the superior efficacy our method compared to several established stato-of-the-art approaches. The code can be found at http://ttuubcoo///aaaabtttt/// FedDUUA
1 Introduction
Federated learning (FL) has revolutionized collaborative model training by enabling multiple clients contribute to global model without compromising the privacy of their local data (MeMahan al. 2017).

---

Published in Transactions on Machine Learning Research (11/2025)
This decentralized strategy avoids the need for sending data to central server, thus maintaining data privacy. As the digital landscape evolves, with an increasing number of distributed data sources emerging from mobile devices, healthcare institutions, and Internet of Things (IoT) networks, FL has emerged as a pivotal solution for training sophisticated deep networks across goographically dispersed and hetcrogencous environments (Bonawitz et al., 2016), Sahoo al. (2024b), (Hu et al. 2024). However, a significant practical the form of skewness in labels and quantity of the data across various clients (Kairouz et al., 2021), (Li al. 2020). Diverse user behaviors lead to significant heterogeneity in the local data of different clients, leading non-independent and identically distributed (non-IIT)) data. This variability can introduce biases in model training, leading to unstable convergence and potentially degrading the model's performance or making it counterproductive (Li et al. 2022). ZZhao el. al.) 2018). While FedAvg (McMahan el. al.) 2017) is effective and widely used, it often falls short in accuracy and convergence with static aggregation methods. These methods combine model updates from different clients fixed manner, failing to adapt hetcrogencous data distributions and client drift, as discussed in (Karimireddy el al. 2020).
obstacle encountered during federated training data heterogencity
Previous studies have addressed the issue of client drift by implementing penalties for deviations between client and server models (Li ct al. 2020). [Li et al.] (021a), employing variance reduction techniques during client updates [Karimireddy et al. 2020). (Acar et al. 2021), or utilizing novel aggregation methods on the
server side CCen et al. 2023), Chowdhury & Halder 2024
1.1 Motivation
Prior studies by Yashwanth al. (2024). Hu al. (2024) have demonstrated that in non-IID scenarios, federated models tend to converge to 'sharp minima', resulting in significant perfor- mance degradation and compromised generalizability. In this study, we investigate the root causes of this phenomenon and propose a novel solution to mitigate its effects. Our study begins with a detailed analysis of loss landscapes for FedAvg- trained models across IID and non-IID data distributions. Fig- ure visually depicts the loss landscapes of two models the FMNIST dataset with systematic parameter perturbations. The model trained on IID data exhibits notably smoother and wider valley in its loss landscape, suggesting greater robust- noss and better generalization. In contrast, the model trainod on non-IID manilests sharper peaks and narrower valleys, indi- cating higher sensitivity to parameter variations and potential overfitting. These visualizations offer strong evidence that in the presence of non-IID data, the FedAvg algorithm achieves suboptimal generalization Motivated by this observation, we investigate the underlying mechanisms by analyzing gradient norms to identify which parts of the neural network are most af- fected by data heterogeneity. Our findings, presented in Fig. 2 reveal notable pattern: non-IID scenarios, the gradient norms of the final layers, including the classification layer, ex- hibit significant amplification compared to their IID counter- parts. Such amplification leads to model instability, impedes convergence, and ultimately compromises the gencralizability of the federated model. Our investigation suggests that ef- fective federated training in non-lID environments necessitates targeted adjustments during server-side aggregation, particu- larly for these highly affected layers, achieve performance comparable to IID settings.
20.0
17.5
1500
12.5
10.0
7.5
5.0
1,00
-1.55.1 18,00.00 (..7% 9..0 60,77
0.25
10.25
Portucbution
0.75
1.00
04
Loss
10
0,50
-1.8.13.0 18..5.00.22.0 28.00 23.50 75
-3.5
Aeriraation
-0.75
+3.00
[M
Figure 1: Visualization of the loss surface for the global model trained the FM- NIST dataset using the FedAvg algorithm: (a) depicts the loss landscape when trained on IID data, while (b) illustrates the land- scape for non-IID data distribution.
2

---

Published in Transactions on Machine Learning Research (11/2025)
This prompts critical question: Can static aggregation methods effectively address severe non-IID data distributions across clients while maintaining higher convergence, perfor- mance, and gencralizability in federated models? The answer decidedly negative. Static aggregation methods inherently struggle with the dynamic heterogeneity present in lederated networks, where adjusting parameters based on client distributions and performance each communication round crucial. Although incorporating predetermined parameters into the aggregation process may provide some partial mitigation, these methods fail to address the complex challenges posed by non-IID data distributions. A more dynamic and nuanced approach is necessary to effectively manage these multifaceted issues. To address this challenge, we apply dynamic aggregation to the model's final layers, where gradient norms fluctuate significantly in non- IID scenarios, while using traditional aggregation (FedAvg) for the lower layers. For dynamic aggregation, we leverage the concept of Wasserstein Barycenter (Agueh & Carlier 2011), derived from optimal transport theory, to integrate clicnt-spocific lcarning behaviors in these affected layers. By minimizing discrepancies from non-IID data, the Wasserstein Barycenter helps to aligu gradients from diverse clients, offering precise model updates. This approach ensures fair aggregation, adapts to data heterogeneity, reduces bias, and enhances robustness, ultimately leading to more stable model convergence and improved generalization.
IID
Non 10
2000
302
422
1500
Gradent Nomn
#02
500
Layar
Layer
(a)
(b)
Figure Comparison of gradient norms between models trained on IID and non-IID datasets using the FedAvg algorithm. (a) FMNIST dataset using LeNet model. (b) CIFAR10 dataset using VGG16 model.
In addition to the server-side dynamic aggregation, we introduce an adaptive loss function for local training on the client side. This function allows clients to effectively explore the minima on their local datasets while preventing overfitting, thereby enhancing local optimization. Simultaneously, preserves the global knowledge of the federated model, ensuring that the benefits from all participating clients are integrated. By incorporating regularization parameter, B, the local loss function dynamically balances the trade-offs between local and global objectives. The contributions of this paper are follows:
We introduce FedDUAL, an innovative dual-strategy approach designed to effectively develop a robust and generalized federated model in highly heterogeneous data environments.
We introduce an adaptive loss function for client-side training to balance the trado-offs between local and global objectives.
Instead of straightforward server-side averaging, propose dynamic aggregation technique that uses Wasserstein Barycenter to reduce the effects of non-IID data by integrating the learning behav- iors of participating clients.
We conducted extensive experiments on three real-world datasets, demonstrating significant perfor- mance improvements over state-of-the-art methods and offering theoretical convergence guarantees for both convex and non-convex scenarios.
3

---

Published in Transactions on Machine Learning Research (11/2025)
2
Related Work
The landscape of FL research has been significantly shaped by efforts to address data heterogeneity chal- lenges, yielding a diverse array of innovative solutions. These approaches can be divided into three primary calegories: (1) client drift miligation strategies, which refine local client objectives to foster better align- ment with the global model
[Li et al., 2021a), (Karimiredyy et al., 2020), (Acar et al. 2021p. Luu et al, 2023) (2) aggregation scheme optimization, aimed at cnhancing server-side fusion of model updates (Hsu et al. 2019 (Lin et al. 2020), (Wang et al. 2020b), (Wang el al.) 2020a and (3) personalized
2021), (Li ct al.)
FL, which tailors models to individual clients (Fallah et al. 2020), (Sattler et al. 2020), (Bui al. 2019).
Our research primarily focuses on two interconnected aspects of FL: mitigating client drift and optimizing server-side aggregation, and we will discuss the same in the literature review.
McMahan et al. (017)) introduced FL. as an extension of local Stochastic Gradient Descent (SGD) Stich 2019), enabling increased local gradient updates on client devices before server synchronization and signif- icantly reducing commmication costs in identically distributed data settings. However, the method faces considerable obstacles when dealing with non-ITD scenarios. Since then, various methods have emerged to address the challenge of data heterogencity in FL (Li et al., 2019), (Yang ct al.] 2021), (Lin et al., 2018), (Hsu let al. 2019). FedProx (Li et al., (020) incorporates a proximal regularization term the optimization fune- tion to reduce model drift and addresses client stragglers. However, this term can also lead to local updates being biased towards the previous global model, which may result in misalignment between local and global optima. Building on previous work, |Acar et al. (2021) introduced a dynamic regularization term to align local updates more closely with global model parameters, effectively reducing client drift caused by local model overfitting. Sun et al.) (2023) further advanced the field with a momentum- -astdd algorithm that ac- celerates convergence by combining global gradient descent with a locally adaptive optimizer. Similarly, several studies use variance reduction techniques, such as SCAFFOLD (Karimireddy al. 2020). However, this approach often results in higher communication costs due to the transmission of additional control vari- ates (Halgamuge et al. 2009). FedPVR (Li et a... 2023) addresses these limitations by reassessing FedAvg's 's performance on deep neural networks, uncovering substantial diversity in the final classification layers. By proposing a targeted variance reduction strategy focused solely on these final layers, FedPVR outperforms several benchmarks. MOON Li et al. 2021a) introduces an innovative model-contrastive framework leverag- ing a contrastive loss to align local client representations with the global model, effectively mitigating client drift, and enhancing convergence, particularly in challenging non-IlD environments, Luo et al. (2021 intro- duced CCVR (Classifier Calibration and Variance Reduction), which employs classifier regularization and calibration method to enhance federated learning performance. CCVR's approach involves fine-tuning the classifier using virtual representations sampled from an approximated Gaussian mixture model. (Shi et al. (2023) introduced novel differentially private federated learning (DPFL) algorithm that integrates the Sharpness-Awail Minimization (SAM) optimizer to enhance stability and robustncss against weight pertur- bations. By generating flatter loss landscapes and reducing the impact of differential privacy (DP) noise, it mitigates performance degradation and achieves state-of-the-art results, supported by theoretical analysis and rigorous privacy guarantces. Fani et al. (2024) proposed FED3R, leveraging Ridge Regression pre- trained features to tackle non-IID data challenges, effectively mitigating client drift, enhancing convergence, and optimizing efficiency in cross-device settings.
Another line of research targets optimizing server-side aggregation in FL. For instance, Hsu et al. (0019) investigated the impact of non-IID data on visual classification by creating datasets with diverse distributions and found that increased data heterogeneity negatively affected performance, leading them to propose server momentum as potential solution. FedNova (Wang et al., 2020b) addressed the problem of objective inconsistency due to client heterogeneity in federated optimization by introducing a normalized averaging technique, which resolves this inconsistency and ensures rapid error convergence. Addressing the limitations of traditional parameter averaging methods, Lin et al. (2020) introduced ensemble distillation for model fusion. This approach allows for the flexible aggregation of heterogeneous client models by training a central classifier on unlabeled data, using the outputs from the client models as guidance. FedMRL (Sahoo et al. 2024a) introduced a novel framework by using a loss function that promotes fairness among clients and employed a multi-agent reinforcement learning for personalized proximal terms, and a self-organizing map to dynamically adjust server-side weights during aggregation.
4

---

Published in Transactions on Machine Learning Research (11/2025)
3 Definitions
In this subsection, summarize the key mathematical definitions used in the paper to ensure clarity.
Bullback-Leibles (KL) Divergence. The KL divergence measures how one probability distribution diverges from a second reference distribution. For two distributions P and Q over the same probability space, it defined as:
P(() log P(i) Q(1
DRL(PI1 Σ P(i
(1)
Wasscrstein Distance. The Wasserstein distance (also known as the Earth Mover's Distance) measures the optimal cost of transporting mass to transform one probability distribution into another. For two distributions and v, the p- Wasserstein distance is defined as:
inf Jx d(s, ((c.))d((a,)))) 1/P (ertiee
(2)
where v) v) denotes the set of all couplings with marginals μ and v, and d(.,-) is ground metric. We leverage the Wasserstein distance and its barycenter to aggregate client models in manner that accounts for data heterogeneity.
| Wasserstein Barycenter. Given distributions {ax k== and weights (Ae) the Wasserstein barycenter j is defined as the distribution minimizing the weighted sum of Wasserstein distances:
K
p = arg min
de We v).
(3)
This barycenter allows us to aggregate the last-layer representations of client models more elfectively than simple averaging, improving robustness to non-ITD data.
4 Methods and Materials
We consider practical FL scenario with non-IID data distribution among K independent clients, each with local training data (x, w), where (x, v) denoting the data points. We initialize the global model weights 0g and share it to the participating clients. The clients download the weighs from the server and train it using their local dataset Dz(x, y). The updated model parameters 0f from each client k for pth communication round are uploaded to the server to aggregate into a global model eg. Our objective is to develop a robust global model by collaboratively training local models across clients, even under varying heterogeneous conditions. To formalize, we define the optimal global model as follows:
0- = arg min F(0), F(0) = K
Σ fk()), K
(4)
where (f) defined in Eq.
fx(0) Ei D. v),,
(5)
where 0 represents the global model parameters, fa( is the model's prediction, and é is the loss function.
4.0.1 Client Side Update.
At the beginning of each round t, the server randomly selects a subset S, K of clients participate in the federated training process and subsequently shares the current global model 0g to these participating clients.
5

---

Published in Transactions on Machine Learning Research (11/2025)
Each client updates its local model by initializing with the global model parameters (0% = f%)) and then updates its local model by minimizing the local objective function. For local training, we have developed an adaptive objective function that balances local loss with the divergence between local and global models. The extent of this divergence is quantified using the Kullback-Leibler (KL) divergence (Csiszár 1975). which effectively compares the probability distributions of the local model weights f²(w) with the global model weights g(w). The KL divergence mathematically defined in Eq. 7 To obtain the probability distributions of the local and global model weights, we first flatten the weights and then apply the softmax function. This process yields the desired probability distributions (p), specified in Eq.
exp(flatten weights) p Σ exp(flatten weights)
(6)
p (w()
(KL(P"10)) = (nt(w) log
(7)
((ww)
where of and qi are the probabilities associated with the hth component of the weight vectors. The local model must excel on local data while maintaining alignment with the global model to enhance overall generalization. This balance between minimizing local loss and aligning with the global model is defined as local adaptive fanction (0) in Eq. g
(1 +3 (p
where Sa (0) is cross-entropy loss for htth client and is a regularization parameter and should be adaptive to account for the performance discrepancy between the local and global models. When the local model substantially outperforms the global model, 3 should increase to enforee greater alignment. Conversely, if the models perform similarly, ,9 should decrease, allowing the local model to focus more on local optimization. The definition of 33 is given in Eq.
Ale Agiobsal
where a is the sigmoid function, Afccal represents the local model accuracy, and A is the global model accuracy for client k. We calculated the global model's accuracy A for client k by evaluating it on the training data of client k prior to performing local updates in the current round. Incorporating the adaptive parameter B in Eq. 8 the adaptive loss function for client k is represented in Eq. 10
Ck adaptive = (1 - (o - Aliobal)) * Clocal +o Agtohal) (p
(10)
After defining the adaptive loss function for each client, we optimize the local model parameters using stochastic gradient descent (SGD). The gradient update for the local model weights wy based on the adaptive
loss function given in Eq.
w = w/ - V/ Cidiapaiv (wf)
(11)
where y is the local learning rate. Expanding the gradient term V Cadaptivee (wf), we obtain Eq.
12
V Lk
adaptive (wi) = (1 (o A Agioba)) Ck (wf) + a (A - Aglobal) V DKL (p ([g).
(12)
The KL divergence term, a (AAfecal V DKT. (x* [[q) in Eq. 12 acts as a regularizer to keep the local model gradients aligned with the global model gradients, thereby preserving model coherence despite non- IID data. The adaptive cocfficient / (Eq. Dh is dynamically computed as a function of the performance gap (Atrral Ak A. When the global model outperforms the local one, 8 tends toward (via the sigmoid function), thus increasing the weight on the local loss term (1 - (8) and enabling the client to focus more

---

Published in Transactions on Machine Learning Research (11/2025)
its local data. Conversely, when the local model performs better than the global, 3 increases toward 1, sering the KT, term to preserve global knowledge rather than blindly aligning the models. The rationale is that superior local accuracy often reflects overfitting to non-llI) client data. A higher 3 in such cases regularizes the local model by constraining it to remain close to the global parameter manifold, improving overall generalization. Conversely, when the global model generalizes better, a smaller 3 allows the client to emphasize local learning. For example, client trained only on classes 0-4 may achieve high local accuracy but would poorly generalize to unseen classes 5 9. The higher and associated KL regularization preserve the global model's multi-class knowledge, thereby preventing catastrophic forgetting.
Note that we employed KL divergence on the softmax of the parameters to capture distributional alignment between local and global models rather than direct numerical differences. By transforming tened pa- rameters into probability distributions through softmax, we interpret model weights as expressing relative importance rather than absolute magnitude. KT divergence thus quantifies how the structural configuration of model parameters diverges between clients and the server. Compared to alternatives, logit-based distances caplure only output-level similarily and overlook the underlying parameter drift that degrades generalization. Layer- wise representation distances would require computing activations on reference dataset, introducing computational overhead and ambiguity regarding data selection in heterogencous settings. Similarly, the L2 parameter distance (as in FedProx) reflects magnitude differences but not distributional structure which is an essential factor for understanding model behavior. Empirically, we find KT, divergence ad advantageous because its asymmetry aligns with the local to global adaptation objective and its bounded nature ensures stable gradients during optimization.
4.0.2 Server Side Update.
After obtaining the weights from the participating clients at round t, the server calculates the Wasserstein Barycenter to effectively aggregate the weights of the last layers of the client models. Computing exact Wasserstein Barycenter can be computationally expensive, so we have approximated it using the Sinkhorn- Knopp Knight 2008) algorithm for efficient computation. We consider the local model weights as distribu- tions and assign equal importance to each client in the computation of the Wasserstein Barycenter (π). This barycenter represents the distribution that minimizes the sum of Wasserstein distances the individual client gradient distributions, as formally defined in Eq. 13
K A D)
A = arg min
(13)
where A are weights corresponding to the importance or reliability of the client k. The Wasscrstein distance W W(),,u;)) between two gradient distributions μa and μ4 of clients and is defined in Eq. 14
W(uk)15)) = ( inf I da. dy(w,
1/p
(14)
(TT xxxx
where T(Hs, (5) denotes the set of all couplings (or joint distributions) 7 on X' x X with marginals dh and Hý respectively, and d(z, y) the distance between points v and y in the metric space X. After that, we use Sinkhom-Knopp algorithm to calculate the Wasscrstein Barycenter.
This barycenter is computed iteratively, starting by caleulating a scaling factor 7 using Eq. 15 followed by Eq. 16
W(F.P.)
7 exp
(15)
E,hrp
Pnew
Σ, Am
(16)

---

Published in Transactions on Machine Learning Research (11/2025)
where is the current estimate of the barycenter, ps refers to the tth client's gradient distribution, is small positive constant, and the iterations continue until convergence. After few iterations, we get the Wasserstein barycenter that is used to update the global model weights. We update the the global model weights for the last layers by substracting them from the calculated Wasserstein barycenter for effectively aggregating the updates from the last layers. Our gradient analysis (Fig. 2 indicates that non-IID data disproportionately amplifies gradient norms the final layers, leading to heightened instability compared to IID conditions. Traditional averaging methods aggregate client updates in Euclidean space without accounting for the underlying distributional differences across clients. The Wasserstein Barycenter addresses this limitation by operating in the space of probability distributions, finding the optimal aggregation point that minimizes distributional discrepancies (Wasserstein distances) to all client updates. This geometry- aware approach provides more robust aggregation by explicitly accounting for heterogeneous client learning behaviors, reducing the bias and instability caused by non-lID data. The algorithm of proposed method FedDUAL given in the Algorithm I The proof of the convergence for both convex and non-convex settings for the proposed method can be found in Section A of the Appendix.
Algorithm 1 FedDUAL
1: Input: Number of clients K, Number of communication rounds T, and Global model G. 2: Output: Trained global model g".
Define a mask e E (0, 1)3, where e = for the last few layers and for the rest layers. Let Snaive = {j e; = 0} and Sdynamic = (j :jjee = 1}.
Initialize global model weights 0
for t = 1 to T do
7: Sample a subset of clients S, S (1,..... K}
8: Initialize lists: local model weights W + D. gradients A + S, do Initialize local model Mr with global weights 89.
9: for each client
10:
11:
Train M& on local datasct Dr using adaptive loss function defined in Eq. 10 W +- WU [0x}
12:
Store local model weights Og
13:
Compute gradients Vn for Mr
14:
A + A {
Store gradients V
15:
end for
16:
for j E (1...., do
17:
if e = then Extract last layers' gradients {V&Di} from A
Layer belongs to Saymaiic
18:
19:
Compute Wasserstein Barycenter of last layer j gradients V Update global model's last layer j weights 0* 00 V
20:
21:
else
> Layer belongs to Snaive
22:
Perform Federated Averaging for layer j: 00 +- [..] EvES, Oxli]
23:
24:
end if
25:
end for
26: end for
27: g + 05
Final trained global model
5 Experimental Results
5.1 Experimental Setup
To assess the elfectiveness of the proposed FedDUAL approach, we conducted extensive experiments using three widely recognized classification benchmarks: CIFAR10 (Krizheskyy et al., 2009). CI- FAR100 (Krizhevsky 2009), and FMNIST (Xiao et al. 2017). To simulate real-world non-IID data dis- tributions, we employed a client-wise partitioning strategy based on the Dirichlet distribution (Hsu et al., 2019). This distribution is governed by concentration parameter a, which controls the degree of data
8

---

Published in Transactions on Machine Learning Research (11/2025)
Table 1: Top-1 accuracy (%) on CIFARIO, CIEARI0O, and FMNIST datasets. The values bold represent the highest accuracy achieved. *** denotes algorithms that failed to achieve convergence.
CIFAR10
CIFAR100
FMNIST
FedAvg
46.68 + 0.25
26.88 + 0.18
81.70 = 0.20
FedProx
47.58 + 0.30
26.89 + 0.22
80.54 + 0,28
FedNova
48.44 + 0.35
FedBN
26.88 + 0.19
81.36 + 0.23
FedDyn
43.97 + 0.40
18.27 + 0.32
71.86 + 0.45
MOON
46.57 + 0.28
28.50 + 0.25
80.09 + 0.27
SCAFFOLD
FedPVR
42.26 + 0.42
23.78 + 0.31
80.32 + 0.33
Proposed
48.70 = 0.20
29.15 + 0.24
81.99 + 0.21
Munn
Peiloy
Propoded
Figure Learning curves comparing the proposed method with baselines across various datasets: CIFAR-10, (b) CIFAR-100, and (c) FMNIST.
(a)
heterogeneity among clients. Lower values result in more skewed data distributions, closely mimicking uneven data partitions. In all experiments, we set 0.01 to simulate severe data heterogeneity, closely approximating real-world conditions. Throughout the communication rounds, each client retains a fixed lo- cal data partition. To evaluate the global model's classification performance, we use a separate test dataset maintained at the server, which remains unseen during training. For our experiments, we used LeNet (LeCun) jet al., 1998) for FMNIST dataset and a pre-trained VGG16 (Simonyan & Zisserman 2015) for CIFAR-10 and CIFAR-100 dataset, following the methodology outlined in (Hu et al. 2024). We applied the proposed dynamic aggregation mechanism only to the last two layers of these models. Our setup involved 100 clients, with 10% randomly sampled per communication round, and a batch size of 32. Each client performed three
Targot accuracy: 0.40
Target accuracy:
Targe accuracy: 0.70
(a)
(b)
(c)
Figure 4: Number of FL rounds required to reach the target accuracy for the proposed method and other baselines different datasets: (a) CIFAR-10, (b) CIFAR-100, and (c) FMNIST.

---

Published in Transactions on Machine Learning Research (11/2025)
local epochs of model updates. We have computed each result three times with different seed values and reported the mean value with standard deviation. To determine the optimal client learning rate for each experiment, we conducted a grid search over 0.05, 0.01, 0.2, 0.3. For the baseline FedProx, we tested proximal values of 0.001, 0.1, 0.4, 0.7 to find the optimal setting, and for FodNova, we evaluated proximal SGD values from 0.001, 0.003, 0.05, 0.1, following the recommendations in Li et al. (2024). Across all experiments, we used the Adam optimizer for consisteney. We have run cach algorithm three times and reported the average outcome. The experimental setup utilized an NVIDIA Quadro RTX 4000 GPU boasting 40GB of memory. The implementation was crafted using Python leveraging the TensorFlow framework X utilizing Windows 11.
5.2 Comparison with the -late-of-the-e Methods
5.2.1
Baseline.
We evaluate the proposed FedDUAL method against eight notable state-of-the-art (SOTA) FL baselines,
including FodAvg (McMahan et al. 2017), FedProx (Li ct al. 2020), FodNova (Wang et al. 2020b), SCAF-
FOLD (Karimiredy et al. 2020), FedBN (Li et al. 2021bb, FedDyn (Acar el al., 2021), MOON [Li et al.
2021a) and FedPVR @Li et al. 2023
5.2.2
Comparison of Accuracy.
The results, detailed in Table D reveal that many recent FL methods often fall short compared to the standard FedAvg baseline. In contrast, our proposed method consistently achieves SOTA performance, surpassing FedAvg along with other baselines across all evaluated scenarios. Furthermore, our approach exhibits remarkable adaptability across diverse datasets. Unlike some algorithms that excel on specific datasets but falter on others, the proposed FedDUAL consistently putperforms baselines across wide range of data environments. This improvement suggests that method addresses fundamental challenges in FL, potentially offering more generalizable solution to the issues posed by data heterogeneity in federated settings. We also observed that FedNova, FedBN, and Scaffold did not perform effectively in our experimental setup.
5.2.3 Comparison of Convergence.
Figure [ compares the learning curves of our method with baselines, while Fig. [0 in the Appendix includes the corresponding curves with error bars. Across all datasets, our method consistently converges faster and achieves higher final accuracy. Although the number of communication rounds varies by dataset, performance generally saturates by the final round. Notably, our method not only allains a more robust. final model but also displays markedly faster convergence across all datasets examined. This effectiveness is further highlighted in Fig. [ where it consistently reaches target accuracy with far fower communication rounds compared to bascline approaches. However, the proposed method incurs higher computation cost due to Wasserstein barycenter based aggregation at the server side. We have provided a thorough computational complexity analysis and runtime analysis in the Section C and the Section D of the appendix respectively.
5.3 Validation of the Motivation
To substantiate our claim that the proposed method yields models in flatter loss landscapes compared to FedAvg, we conducted comparative analysis. Using VGG-16 models trained on the FMNIST dataset under non-lID conditions (a 0.01), visualized their respective loss landscapes following the approach outlined in Li et al. 2018)). Figure 0 in the Appendix depicts these landscapes, with each model centrally located within its respective lerrain. The visualization reveals that our proposed method situates the model in a notably flatter region compared to FedAvg. This finding supports our assertion that our approach guides federated training towards more stable and generalizable solutions, characterized by flatter loss landscapes. The performance improvement of the proposed models stems from two key innovations: a Wasserstein
httpp://www.ppthhmee httpp:// https://ww.ttensolllo
10

---

Published in Transactions on Machine Learning Research (11/2025)
Barycenter-based aggregation for final layer gradients, mitigating client drift in heterogeneous data envi- ronments, and an adaptive loss function balancing local optimization with global consistency during client training. This synergistic approach preserves global knowledge while promoting client-specific optimization, addressing fundamental FL challenges.
6 Ablation Study
In our ablation study, we performed all experiments on the FMNIST dataset with a 0.01. The study comprised four types of experiments: (1) performance analysis of the individual modules, (2) assessment of the impact of dynamic aggregation across different neural network layers, (3) hyperparameter analysis, and (4) evaluation of various levels of data heterogeneity.
6.0.1 Performance Analysis of Individual Modules
To assess the effectiveness of the proposed adaptive loss and dynamic aggregation techniques, we conducted three ablation experiments across FMNIST, CIFAR-10, and CIFAR-100. The results for FMNIST are shown in Table 2 In the first experiment, we employed only the adaptive loss alongside standard server- side aggregation. Notably, this configuration underperforms FedAvg across all datasets, indicating that adaptive loss alone cannot effectively address data heterogeneity- due to its limited capacity to improve generalization despite fostering local-global aligument. The second experiment implemented our dynamic aggregation technique at the server, while retaining the conventional cross-entropy loss function locally, Finally, the third experiment combined both proposed methods: the adaptive loss function and the dynamic aggregation technique. As evidenced by Table 2 the integration of both proposed approaches in the third experiment yielded the highest accuracy, highlighting the impact of our dual stralegy on model performance. The learning curves for these experiments using FMNIST dataset are illustrated in Fig. 7 of the Appendix.
6.0.2 Impact of Dynamic Aggregation on Different Network Layers
To substantiate our decision to apply dynamic ag- gregation technique selectively to last layers, we ex- amined its impact across various layers of the neural network. Our earlier findings highlighted that data heterogeneity primarily affects last. layers of the net- work. Figure 5 illustrates that random utilization of the dynamic aggregation to all layers diminishes performance. Conversely, targeted implementation on layers proximal to the classifier yielded optimal accuracy and convergence. These outcomes validate our hypothesis and demonstrate the method's effi- cacy in mitigating heterogeneity-indune susceptible layers, we directly address the core challenge of data heterogeneity in fodcrated training, resulting in enhanced model performance and faster convergence.
Table Ablation study of FodDUAL across different datasets.
Adaptive Losss
Dynamic Agg.
Dataset
Acc. (%)
80.70 = 0.22
x
FMNIST
80.91 + 0.25
/
/
81.99 1: 0.18
/
X
41.05 :
X
/
CIFARI0
46.50 + 0.15
/
48.70 = 0.20
25.05 = 0.11
X
CIPAR100
27.01 + 0.17
29.15 0.24
issues, By focusing our dynamic aggregation technique on the most
6.0.3 Hyperparameter Analysis
In the proposed architecture, there are two key hyperparameters consider: the scaling factor (7) and the number of iterations used compute the Wasserstein Barycenter. The proposed FedDUAL approach utilizes dynamic server-side aggregation by applying the Wasserstein Barycenter concept to combine the weights of the final layers from local models. This iterative process involves a small positive constant (e) to determine the scaling factor (7). To optimize performance, we conducted two sets of experiments on the FMNIST dataset with a = 0.01, each exploring a range of values for these crucial hyperparameters. The hyperparam- eter e influences the sensitivity of the barycenter calculation to variations in Wasserstein distance. A smaller makes the barycenter more responsive to differences in Wasserstein distance, while larger diminishes this sensitivity. This impacts how the barycenter integrates each distribution according its distance from the current estimate. During the iterative update of the barycenter, affects the scaling factor 7 applied to each
11

---

Published in Transactions on Machine Learning Research (11/2025)
distribution. An excessively small result in slow potentially non-existent convergence due to mini- mal scaling factor, whereas a too-large e may cause oversmoothing, reducing the barycenter's s effect iveness in accurately representing the distributions. For this setting we have fixed the number of iterations to compute Wasserstein Barycenter as 150. Figure in the Appendix shows test accuracy across different values, indicating that larger can degrade performance or hinder convergence. Figure in the Appendix presents the corresponding learning curves for these settings. The number of iterations the Wasscrstein Barycenter function another critical hyperparameter that affects both the accuracy and efficiency of the barycenter computation. Generally, more iterations enhance convergence and accuracy, ensuring that the barycenter more closely approximates the optimal value. However, increasing the number of iterations also prolongs computation time, necessitating a balance between accuracy and efficiency. Finding the optimal number of iterations involves a Irade-ofT: too few iterations may result in suboptimal oulcomes, while loo many can yield diminishing returns accuracy. To achieve the best performance, begin with a reasonable default value, monitor convergence by obscrving changes the barycenter, and adjust itcratively basod on empirical results and available computational resources. For this setting, we fixed the epsilon value as 0.0001, which yields the highest results in previous experiment. Figure [T illustrates the test accuracy for different values of itera- tions to calculate Wasserstein Barycenter, suggesting that larger iterations may adversely affect performance. Figure 12 presents the corresponding learning curves for these settings. From both experiments, we observe that the highest performance achieved with = 0.00001 and 150 iterations. Therefore, to optimize perfor- mance, is advisable to set e to a smaller value while keeping the number of iterations between 100 and 150.
6.0.4 Experiment
on Different Level of Data Heterogeneity
Figure 13 illustrates the accuracy of the proposed method and various baselines across different levels of data heterogeneity on the FMNIST dataset. In this context, heterogencity is quantified by a, with lower values indicating greater data heterogeneity. The results show that as a decreases, the test accu- racy for all models increases, because data hetero- geneity among clients is decreased. Remarkably, the proposed method consistently achieves the highest test accuracy and cxhibits the slowest performance decline compared to other algorithms, retonstrat- ing superior performance of the proposed method on varying degrees of non-IID data partitioning. The learning curve is presented in Fig. 14
0.6
Test Accuracy
0.5
0.2
Layers:
Layers: 2 to
Layers: to
Layers: 6 (Proposed)
100
200
300
400
500
FL rounds
Figure 5: Illustration of the Dynamic aggregation method applied across various layers the neural net- work.
7 Conclusion
This research presents a novel approach to address the challenges posed by data heterogeneity among clients in the federated approach. We systemati- cally analyze the factors contributing to federated model performance degradation under severe data
heterogeneity and propose architecture incorporating dual-strategy innovations. First, we implement adaptive loss function for client-side training. Second, create dynamic aggregation strategy for server side aggregation, tailored to client-specific learning bchaviors. The proposed FodDUAL effectively overcomes the challenges of heterogeneous data, outperforming eight SOTA baselines. It demonstrates faster conver- gence and consistently improved performance, making it an excellent solution for large-scale FL applications in real-world scenarios. Our approach's flexibility paves the way for rescarch into hybrid federated learn- ing models that adapt to changing client environments and data. Future studies will focus on integrating personalized learning paths to enhance model adaptability and efficiency across various datasets.
12

---

Published in Transactions on Machine Learning Research (11/2025)
References
Durmus Alp Emre Acar, Yuc Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N. Whatmough, and Venkatesh Saligrama. Federated learning based dynamic regularization, 2021.
Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, (3(2):904-999 924, 2011.
Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for federated learning on user-held data. arXie preprint arXiv:1611.0445 2016.
Duc Bui, Kshitiz Malik, Jack Goetz, Honglei Liu, Seungwhan Moon, Amuj Kumar, and Kang G Shin. Federated user representation learning. Xie preprint 1rx: 12535, 2019.
Dengsheng Chen, Jie Hu, Vince Junkai Tan, Xiaoming Wei, and Enhua Wu. Elastic aggregation for federated optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognation (CVPR), pp. 12187 12197, June 2023.
Sujit Chowdhury and Raju Halder. Fedsat: A statistical aggregation approach for class imbalaced clients in federated learning, 2024. URL https:/ http//aaxvvv.o//b////////// http:///////////3/////3/ 03862
Imre Csiszár. T-divergence geometry of probability distributions and minimization problems. The annals of probability, pp. 146 158, 1975.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. arXiv preprint 2tl.20002.0744.. 2020.
Eros Fani, Raffaello Camoriano, Barbara Caputo, and Marco Ciccone. Accelerating heterogeneous federated learning with closed-fone classifiers. Proceedings of the International Conference on Machine Learning, 2024.
Malka Nisha Halgamuge, Moshe Zukerman, Kotagiri Ramamohanarao, and Hai L Vu. An estimation of sensor energy consumption. Progress In Electromagnelicss Research B, (12):259 295, 2009.
Tzu-Ming Harry Hsu, Hang Qi, and Malthew Brown. Measuring the elfects of non-identical data distribution for federated visual classification. Xrviv preprint arXie: 0039.0 2019.
Ming Hu, Yue Cao, Anran Li, Zhiming Li, Chengwei Liu, Tianlin Li, Mingsong Chen, and Yang Liu. Fedmut: Generalized federated learning via stochastic mutation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 11528-12537, 2024.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in foderated learning. Foundations and trends@ in machine learning. 14(1-2):1-210, 2021.
Sai Pranceth Karimireddy, Satyen Kale, Mchryar Mohri, Sashank Roddi, Scbastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic rolled averaging for federated learning. In International conference on machine learning, pp. 5132-5143. PMLR, 2020.
Philip A Knight. The sinkhorn knopp algorithm: convergence and applications. SIAM Journal on Matrix Analysis and Applications, 30(1):261 275, 2008.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun, Lčon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied docu- ment. recognition. Proceedings of the IEFE, 86(11):221 2324, 1998.
13

---

Published in Transactions on Machine Learning Research (11/2025)
Bo Li, Mikkel N Schmidt, Tommy S Alstrom, and Sebastian U Stich. On the effectiveness of partial variance reduction in lederated learning with heterogeneous data. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 3964 3973, 2023.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
Qinbin Li, Bingsheng Hc, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF conference on compuler vision and pallern recognition, pp. 10713 10722, 2021a.
Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng He. Federated learning on non-iid data silos: An experimental study. In 2022 IEFE 38th international conference data engineering (ICDF), pp- 965- 978. IEEE, 2022.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Amoet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429 450, 2020.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. arXiv preprinl urXiv:1907.02100 2019.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. arXiv preprint (xX::2102.07623, 2021b.
Xingyu Li, Zhe Qu, Bo Tang, and Zhuo Lu. Fedlga: Toward system-heterrogenetty of federated learning via local gradient approximation. IEEE Transactions on Cybernetics, 54(1):401-414, Jamuary 2024. ISSN 2275.2275. doi:
URL http://dx.doi.oo 10.1109/ 1tt/:100.111T........1 TCYB.
Tao Lin, Sebastian Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini bati-h, use local sgd. Xiv preprint Xiv:::1808. 07217, 2018.
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351 2363, 2020.
Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, and Jiashi Feng. No fear of heterogeneity: Classifier calibration for federated learning with non-iid data. Advances in Neural Information Processing Systems, 34:5972-555 5984, 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication--ffieenn learning of deep networks from decentralized data. In Artificial intelligence and statistics, pp. 1273-1282. PMLR, 2017.
Pranab Sahoo, Ashutosh Tripathi, Sriparna Saha, and Samrat Mondal. Fedmrl: Data heterogeneity aware federated multi-agent deep reinforcement learning for medical imaging, 2024a. URL https:/ hutpps//arxxvv.ogg /arxiv.orgg/ abss2440.005% 05800
Pranab Sahoo, Ashutosh Tripathi, Sriparna Saha, Samrat Mondal, Jyoti Prakash Singh, and Bhisham Sharma. Adafedprox: A heterogeneity-aware federated deep reinforcement learning for medical image classification. IEEE Transactions on Consumer Electronics, 2024b.
Felix Sattler, Robort Müller, and Wojciech Samck. Clusterod fedorated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE transactions on neural nelworks and learning systems, 32(8):3710 07222, 2020.
Yifan Shi, Yingqi Liu, Kang Wei, Li Shen, Xueqian Wang, and Dacheng Tao. Make landscape flatter in differentially private federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pallern Recognition, pp. 24552 24562, 2023.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition, 2015. URL https:/ hwtp//wwwwwwww/////////////////e
14

---

Published in Transactions on Machine Learning Research (11/2025)
Sebastian U. Stich. Local sgd converges fast and communicates little, 2019. URL https: htttss//arxivv.oo http://wxwwi/rrr///// abs/ 1805.09767
Yan Sun, Li Shen, Hao Sun, Liang Ding, and Dacheng Tao. Efficient. lederated learning via local adaptive amended optimizer with linear speedup. arXiv preprint orXiv:2308.0052. 2023.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conference on Learning Representations, 2020a. URL
https://openneview..ru wwt.///wweeeee... ftt///orummiccee htwp://wwoiii.etttooommmmddddddd Bid
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and Vincent Poor. Tackling the objective inconsistency
problem in heterogeneous federated optimization. Advances in neural information processing systems, 33: 2020b.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image datasel for benchmarking machine learning algorithms. arXiv preprint orXiv.170..07741. 2017.
Haibo Yang, Minghong Fang, and Jia Liu. Achieving linear speedup with partial worker participation in non-iid lederated learning. arXi preprint X1i::2001.1100, 2021.
M. Yashwanth, Gaurav Kumar Nayak, Harsh Rangwani, Arya Singh, R. Venkatesh Babu, and Amirban Chakraborty. Minimizing layerwise activation norm improves generalization in federated learning. Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), pp. 2287-2296, January 2024.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXio preprint arxiv:1806. 00582, 2018
Appendix
A Convergence Proof
Before presenting the main convergence theorems, we establish assumptions and several key lemmas.
Assumptions
1. L-Smoothness:
Each local loss function (a) is L-smooth:
(v)x - (O)(l LIj0- - 0'I,, V0,0' E W,
(17)
2. Unbiased Stochastic Gradients: For any client ls and parameter 0, the stochastic gradient is unbiased:
[V = (()))).
(18)
3. Bounded Variance: The variance of stochastic gradients is bounded:
E EEDD, (1) (() s V0, k
(19)
where o² max of.
4. Bounded Gradients: There exists G > 0 such that:
[V S WOEW,
(20)
KL Divergence Properties: The KT, divergence term and dynamic weighting satisfy:
|
[DDa.((((0))(((01 Gkl, v0,
(21)
15

---

Published in Transactions on Machine Learning Research (11/2025)
[V LB,
V8,
(22)
5 (0) - fmax < 1, 40, k:
(23)
6. Wasserstein Barycenter Approximation: For the final layers aggregated using Wasserstein Barycen- ters, the approximation error ewB is bounded as follows:
V WB
Vecact||
S EWB
(24)
where Visuct represents the exact gradient aggregation.
A.1 Key Lemmas:
Lemma 1: Smoothness of Modified Loss function:
Under Assumptions and 5, the modified loss function fc(() is f-smooth with
= L + Lo(G + GKL) + Smax GKL + LgGr,
(25)
where G5 = max; supa (0)] and GKL. = max, sup, (e)
Proof. Tet 0, 0' W. The complete gradient the modified loss function
fkk(0) 30 [1 [11-B8((0))))000++444))))Xx00010
(26)
+ VB,(0) (.V
Similarly for 8':
fi((t') = 3Vx
+ V3x(0')
(27)
Calculating the gradient difference results in Eq. 28
-
(VBa(0) V.x00)) fa(0))
+ [(1 -/(0)) (1
(28)
+ [V (p)((((()))) - V3x(0') : (px (0))l((0))
+ (PL ())) - (xx (p))
We bound each term separately using the triangle inequality:
Term 1: (0) fu(0) - (0) fa))
Adding and subtracting (0) (x))
(V)((0) (A( VBa(0') fa(0')l
$ (1) (f) - fa((')) + (V)))(() V3L(0)) fa(8')l
< (0)] - +1(00)) VBa(0'11 (1.x(())))
(29)
< L L3|10 GG
[s(L+
16

---

Published in Transactions on Machine Learning Research (11/2025)
Term 2: [11 (x(O)) (1 -
Adding and subtracting (1 - (.(0))
[[1 3x (O)) (1 - (=()))V/\((0I
- [[ - (v(0))(Vxx(0) - Vf((0')) ((()(())-((0)))(((
s (1 (x(0)) l1
(30)
s LI10 - 0'11 Lg|00 G
(L + LsC))10
Term 3: (Vx(0)) DEL (p)(0))l((()) (Ok (xL))((0))))(
Following similar decomposition:
(V)) (O)(()(()))(()))) - (Bk) (p)
LOGKLLLll0 l1| | LgGKr -
8'll
(31)
L 1)100 0'|l
Term 4: -
< (1Ba(00((Vxxx.xxxx(0)))))))) (p)
+ - (0
-
(32)
I Smax GKL 8'1 LgGKull0 l'l
L (Bmax + 1) 0'll
Final bound: Combining all terms:
[v)( Vi G,(L+G;)+
(33)
+ LLGKI 1) + LGGKL(Pmas + 1)]] - 0']]
For conservative and simplified bound, we can write:
== L Lz(G + GKI. Gf) + L2GKI. + L.G
(34)
Or more compactly, assuming G2 S G and using conservative bounds:
L = L Lo GKL.) + BmaxGGLLL + LoGg
(35)
Therefore, (0) is L-smooth.
Lemma 2: Bounded Variance of Modified Gradients: Under Assumptions 2, 3, and 5, the variance of stochastic gradients for the modified loss is bounded:
E (1Vvg(00:)) V s0²
(36)
Proof. We begin by analyzing the variance of the stochastic gradient of the modified loss function Fz( Recall that the per-sample stochastic gradient is defined as:
(.(0)))) (1 3
- (0) (a(0) ( (px (px (()))
17

---

Published in Transactions on Machine Learning Research (11/2025)
whereas the full-batch gradient is:
- =. (())
Subtracting the two, we obtain:
- (1- 3x) (Ve(0:)) ffx(
Here, the KL divergence term cancels out since it is deterministic and does not depend on the stochastic sample g. Taking the squared norm and expectation over the stochasticity of &, we have:
- f Es -
||]
(1 Bx(())² B& [V((0:)) (f))
< Eg [IV((0:)) v (since (1 3x((0)² < 1)
S o 0²²
where we have used Assumption upper bound the variance of the stochastic gradients by o².
Lemma 3: Local Update Analysis: Let of be the local model client l after E local updates initialized from the global model 0,. Then, under Assumptions 5, the expected squared deviation from the global model after local training satisfies:
E 22 2
E 1 - 0, + (EV0000))
<
E
|100 t. - 0.11] + Eo?
(37)
where gh-s denotes the local model on client k after local steps, is the smoothness constant of fa (Lemma 1, and 52 is the bounded variance of the modified stochastic gradient (Lemma 2).
Proof. We begin by expressing the full local model update as a telescoping sum over E local steps:
0% - 1
where &e denotes the data sample used in the e-th local step.
Rewriting this update in terms of the true gradient at the initial point O,, we add and subtract (a).))):
B-1
E-1
0-
Via - - Σ
Rearranging terms gives:
E-1 - 0, + (
Σ
Taking the norm squared and expectation:
E 110 0, +
= Σ Vi((0!= vi((o))
Applying Jensen's 's incquality (or the incquality
Σ aell 2 s EE 1|a
F
2 E
E
(o. (6) v v
e=0
18

---

Published in Transactions on Machine Learning Research (11/2025)
Now decompose the difference inside each term:
E
Fu(00
vix())]
- 2
Vi -ix |1]
+ 2E [i((0 vin(o,)
From Lemma 2, the variance of the stochastic gradient is bounded:
vi(00" (0 ] S.a
E
From Lemma 1, the gradient of fi I-Lipschitz:
VViieo---vii Vi((E))) 1 00
18.00 0.1 0.1?
Combining the two bounds:
E Fia(0"()E)) -V
=200 +2ï2E
|100-00 8, 0,112]
Substituting back:
E-1 Σ
o0.e - 0, -o1"") 25 + 2
E 0 - 02 + +EVVi
S E
Grouping constants:
E
0. -011 0,1²]
222 + 211² 1
E
Finally, simplifying constants and using 1 factor for future algebraic convenience:
2 E
B 11 0.
aal"] + F
Theorem 1 (Convex Convergence)
Suppose Assumptions 6 hold and F(0) is convex. Let 0* arg minsew F(0), and set the learning rate
S ALE'
Then, FedDUAL guarantees the following convergence bound:
2 11
nLEG2
E[F(Or) F(0))
+nLE*G +ewBG
rT
K
where Ōr = - ET-1 Oz and 5² is as defined in Lemma 2.
B Proof
Since F()) is i-smooth and convex, write the fundamental smoothness inequality:
19

---

Published in Transactions on Machine Learning Research (11/2025)
F(0.41) & F(O.) (VF(0)), 0441 0i) + 2
_ 110 0ell
(38)
This is the standard smoothness incquality. For any L-smooth function f, we have f(0) & +(x)+
2) + 1
The FedDUAL update consists of two phases:
(For early layers): Standard aggregation
gearly @early
7
fa (0,
(39)
K
(For final layers): Wasserstein Barycenter aggregation
6+1
(BB(O,,-PD)E) dWB
(40)
where lów S Ewn is the Wasserstein approximation error.
The total update be written
R-1 (0 ikk,e) + Stotal.
11+1 02
02
(41)
where ||ootail| < EWB. By substituting Eq.
41
into the inner product term defined in Eq. 38 we derive
Eq. 42
(V((O)), 0c)
E-1
VF(O).
7
ΣΣ
( :kk,e) + 8rotal
K
(42)
K Æ-1
(V(())) (0, :k.,)) + (VF(B,), Opota))
= K
c=
For the error term, we apply the Cauchy Schwarrz inequality to obtain:
(VF())), (ro((a))) S (V. S GEwB-
(43)
Here, we use Assumption 4, which states that |V(())) s G. We add and subtract (a,) from the first term of Eq. 42 and obtain Eq.
44
(V.)))) Vi(00 Ek,e)
(44)
(Vi(0)), fa.))) + (VF(O.), (aa (n.e)) (..(0))
The first term of right hand side of Eq.
44 gives us:
K
(((M)) (VF().).
Σ [o(0o)) 11
(45)
K
20

---

Published in Transactions on Machine Learning Research (11/2025)
By definition, VF(O.) K E (O)), and assuming uniform data distribution, this simplifies to K E ((.)
For the second term, we decompose:
Via00, - (0.))
(46)
[V fi((0 (V)(())))))) - (..))
Taking expectation and applying the Cauchy Schwarz inequality, we obtain:
BI (((O)... fo. (x.e) ((.))l
[II ;k,.o) Vi((0
(47)
+ [ (f+(0)111
Using Young's inequality with parameter a > 0:
ab - a² ah?
(48)
2a
2
Applying Eq.
48
into first term in Eq. 47
obtain:
E V (0 Ex,e) - [(a(A"0011]
E
F(O))
SEIIV :En f0, (8."
2a
2
(49)
E F(0))11)
ag²
2a
2
Here we have used Lemma 2 which bounds the variance of stochastic gradients by Ø2.
Similarly, write for the second term in Eq. 47
fa)(ll))]
(1 2a
ai?
(50)
Blllot- 0.11%).
We use the L-smoothness of fs from Lemma 1. By selectively combining the terms from Eq. 43
Eq.
44
Eq. 45 and Eq. 50 and substituting them into Eq. 42 we obtain Eq. 51 The choice of a in Young's 's inequality is standard optimization choice that balances the two terms in the bound. When applying Young's inequality ab < 2a + ab? setting = gives equal weight to both the gradient norm term and the variance terms, which minimizes the overall bound and leads to the subsequent steps.
B((FF(0)) 0,)
S 5-1B]]VF(0)191 π K
K E-1
E
ag2
+
ai2
EE
11
ghe
0,139
+ GEwB-
(51)
2a
2
Substituting 1. obtain:
--BBIIIF(001111 DE K
82
1 F
E|
B[]VF(())]
Blllo -0e113
GswB-
(52)
6=0
21

---

Published in Transactions on Machine Learning Research (11/2025)
Simplifying Eq.
52
obtain Eq.
53
7L'E
F
- 0ol1] nE
<
-
+ GewB-
(53)
4K
4
e=0
Again we starting from Eq. 41
K
E
0c+1- - 0g
=
(ie((O,esSssee)) Ex,,e) + Spotal
K
e
where ||totaal| I EWB the Wasserstein Barycenter approximation error.
Taking Norm Squared, we get below:
K E-1
|10+11 0,11
fa (6) Sk,e) + drotal
K R= c=0
Now we apply the below inoquality:
||a + 6|12 s 2||aa|| 2|1b11
So, we get below:
1 K E-1
2
|10+11 0,11 2 r
Σ
V
fla Skse)
2 2110eot 11
K
Now we take Expectation of the above inequality:
K E-1
27
E [11 - 0,113] s
2E
1
+ 2e
K
Using Jensen's inequality and bounded gradients, obtain Eq. 54
K E-1
27
s 2
ΣΣ
v
(0 En,e)
2E w
K
(54)
< E² 2E
From Lemma 3, we can derive the following.
122202
BI| 0.1191
+
(55)
2
Summing over all local steps yields:
1/02 E-1
E-1
-0,119-
<
Σ
2
Σe è
2K
2K
c=0
(56)
n E
T22
2K
2K
22

---

Published in Transactions on Machine Learning Research (11/2025)
We use Eezo e² /3 s E² for practical purposes. Substituting Eq.
56
into Eq.
53
we get Eq.
57
TTZE
F2
F2F.
nE
[[ - [[F( - -EIIVF(
4K
2K
2K
(57)
1 (2) FF + 2) + GewB-
Simplifying the higher-order terms and keeping dominant terms:
E "E]IVE((N11 nLEE
+ nnEE tWwnG.
(58)
K
From the definition convex functions, we have:
(. (VF(0)), - 07) 2 F(O.) F(0")
(59)
Applying the Cauchy-Schwand incquality along with 2ab I a² - b², we obtain:
|F(00) 2(F(0.) =(0')) 0*)
(60)
10, 1-
Subtracting F(0*) from both sides of Eq. 58 we obtain:
-
- +C,
(61)
where C = nLEs² + GLE &EwBG.
K
For convex lunctions, we leverage the inequality:
E|| =()))11 2(E F(()))
(62)
This yields the recursive bound:
- < [[ -(()))
(63)
Unrolling the recurrence gives:
B[F(0r)] < ( T (F((0) F(0)))
(64)
Using Jensen's inequality for the averaged iterate Ör DT-1 O::
T T-1
E[F(0r))
<
Σ
F(
T
1=0
(65)
2/ 11
nLEG²
+nf.F
+wwBG
n
K
Putting 17 ALE
in above equation and simplifying yields:
23

---

Published in Transactions on Machine Learning Research (11/2025)
B[F(Or) - -F(0*) - 8LEi|00 1-
Ø²
EG²
CWBG
(66)
T
4K
4
- = ( + constant
(67)
B.1 Non- Convex Convergence Analysis
Theorem 2: Given that Assumptions are satisfied and F(0) is non-convex, setting the learning rate
I ALE ensures that FedDUAL achieves the following convergence guarantee:
T B[IVP()))))) 2(F(0o) 2(F(0)) F(O')) n/T
NLEG*
K + nLEE EWBG
(68)
where 0 any global minimum of F(0) and 6² as defined in Lemnma 2. From the L-smoothness of F(0) (established in Lemma 1), we have:
Proof:
L
F(6.11) S ((,)) + (VF(0.), 0,) 2 |10 0.11²
(69)
The FedDUAL update rule is defined as follows, as presented Eq.
K
E-1
0,11 0,
"
(0 + Ototat,
(70)
where ||onotat| EwB represents the Wasserstein Barycenter approximation error.
Substituting Eq. 70 into the inner product in Eq.
69
get Eq.
71
(VF(0.), - 0₂)
17
Σ
(V))))
(0. (h)))) + (V.())), 8cotal)
(71)
K
Applying the Cauchy-Sehie Schwarz inequality to the error term in Eq. 71 obtain the following.
(VV 6totat)) < [[V I GewB-
(72)
For each gradient term in Eq. 71 we add and subtract the term (f.().):
(VF(O))
(V)( vi.) - (vF(0)) (0. -
(73)
The first term of Eq. 73 gives the following:
1 K
(F(O)) = - 1VF
(74)
K
k=1
For the sccond term Eq.
73
we further decompose:
1. gke))
-
- + -..)
(75)
24

---

Published in Transactions on Machine Learning Research (11/2025)
Taking expectations over the second term in Eq. 73 and applying Young's incquality with parameter a, derive (omitting details identical to the convex setting):
E (()))
E
Vi.0 Vi((ONM)
Bll
2
al
Bll10, - -.
(76)
2
Using Lemma 2 for the variance bound and setting a in Eq.
76
obtain the following deviation bound:
1
(V(()), (
+ E BlI0 1.11
(77)
4
4
By combining all the terms and substituting them into Eq.
we get the following bound:
E
1.2
El 8.41 0,)l -
7
+
Bllle 0c11
K
+ G.wB.
(78)
Simplifying Eq. 78
we arrive at:
nL'E
nEo2
((V - 0:)1 s --
B1110 R. - 0,
+ G.wB.
(79)
From Lemma. and using the fact that Eede 2²
F S F.2 for practical purposes, obtain:
R Elllo 0.113) 2K
E22202
E22
(80)
K
2K
From Eq. 7
applying Jensen's inequality yields:
B[[ 0.113] < 2E
v (a(o
+e S 222 + 2EWB.
(81)
K
J=
E=0
By combining all the above terms and substituting them into the smoothness incquality in Eq.
69
we obtain:
BIF(0(+1) S(F(())) - "E[IVF(O))
+ GLE + 4 e2 22 + GEwB-
(82)
K
Rearranging and using the learning rate condition 17 ALE
in Eq. 82
we get:
" E[F((
L + LEivE + GEwB-
(83)
2
K
Summing the above equation over t = 0, 1...... and dividing by nT. 2
2(F(00) E
-B(F(Or))
nLEaa
24 +
2GewB
+nLEZ +
(84)
T 1=
nf
K
7
25

---

Published in Transactions on Machine Learning Research (11/2025)
Since F(Or) 2 F(O*), and combining the error terms, we have:
2(F(0)) - =((0)) T B(I
T-1
nLE02 K
+nLE + EWBG
+
(85)
nf
where the term Lewn has been absorbed into the dominant sweG term for simplicity. To minimize the right-hand side, we choose the learning rate that balances the first and second terms:
2 (F(0o) F(OO))
LT2T
Substituting this choice of into the inequality gives:
2 ((((0))
2 (F(()) F((-))
E [[VE(0))1)) s
Lo
LwwG..
2(F(00)-F(00)
L
L0²T
Simplifying both terms, we obtain:
2 2(F(0n) - F(O2))
22(((0)) F(O
- Σ B [IVF(())
+ ewwG
VT
VT
3 2(F(0)) F(0)))
Σ
ewBG.
VT
Therefore, the convergence rate is:
T Σ E [IV -0 (0(2) + eWBG
B.2 Discussion on Convergence Bounds
Our thcoretical analysis establishes convergence guarantoes for both convex and non-convex settings:
Convex Setting: Theorem 1 shows that FedDUAL achieves
T +
(F(())-F((')) o
constant,
which matches the optimal rate for first-order methods under smoothness and convexity assump- tions. This indicates that FedDUAL preserves convergence efficiency while incorporating dynamic weighting and Wasserstein aggregation.
Non-Convex Setting: Theorem 2 proves that.
1 (
E [[IVF(())
$
+ ewBG,
consistent with the standard rate for non-convex optimization foderated learning. The additional CwB term quantifies the effect of Wasserstein Barycenter approximation, ensuring that its impact remains bounded.
26

---

Published in Transactions on Machine Learning Research (11/2025)
Comparison with state of the art methods: Standard algorithms such as FedAvg McMahan et al, 2017] and FedProx [Li et al., 2020) achieve similar asymptotic rates (O(1/T) for convex and (1/vTT) for non-convex), but suffer from significant degradation due to client drift under non-llD data. Specifically, Fe- dAvg incurs heterogenetty-deeendent terms of O({2K2/T) in non-convex settings, while FedProx improves this lo O({²K/T) through proximal regularization. Variance-reduced methods like SCAFFOLDD (Karim- ireddy et al., 2020) achieve superior bounds with heterogencity terms of only O(ς2/T) (independent of local steps K) and variance scaling as O(/nkKT) but require maintaining control variates at each client. Our bounds show that FedDUAL matches these theoretical rates while introducing Wasserstein aggregation and adaptive loss weighting, providing robust to heterogeneity without the additional memory and commu- nication overhead of control variates where I'' denotes the number of global communication rounds, K is the number of local SGD steps performed at each client wetween communications, n is the total number of participating clients, is the client drift parameter measuring data heterogeneity, n is the total clients, and ry the learning rate.
C Computational Complexity Analysis
We present a formal complexity analysis comparing the proposed method with baseline approaches for a single communication round involving K clients, a model of dimension d, and n denoting the number of parameters in the final layer.
Client-Side Complexity
FedAvg:
The computational complexity per communication round is O(E B d), where E denotes the number of local epochs, B the batch size, and the model dimension. This cost primarily arises from the forward and backward passes performed during local training on each client.
FedProx: The per-round computational complexity OEE B + d), where the additional O(d) term accounts for computing the proximal regularization term 10g- - 0,
SCAFFOLD: The per-round complexity is O(E B d + 2d), with the extra O(d) arising from storing and updating control variates used to correct client drift.
FedDUAL (Proposed): The per-round complexity is O(E B .B.da + d), where the O(d) term corresponds to KL divergence computation comprising weight flattening, softmax operations, and divergence calculation, all linear d. Thus, FedDUAL maintains the same asymptotic complexity as FedProx.
Please note that Our adaptive loss adds negligible client-side overhead compared local training cost (as E. B d).
Server-Side Side Complexity
The aggregation complexity per communication round is O(K d), corresponding to the simple averaging of model parameters across clients. FedProx/ CodProx/FedNNovee Foddova: The aggregation step also has com- plexity OKK d), as weight normalization and scaling operations remain linear the model dimension. SCAFFOLD: The aggregation complexity is O(K . d), dominated by the averaging of client updates along with control variate terms.
FedAvg:
FedDUAL (Proposed): The aggregation complexity is O(K (1K-++1 I K n²), where O(K d) accounts for standard aggregation of lower-layer parameters and O(T K -n²) arises from computing the Wasserstein barycenter for the final layers. Here, I denotes the number of Sinkhorn iterations (typically 100-150), and n represents the number of parameters in the last layers (n < d); for instance, in VGG-16, d z 138M total parameters, where the last two fully connected layers contain n x 67.2M parameters (two layers of 4096 4096 each), constituting approximately 48.7% of the total parameters.
Observation: While a represents substantial portion of the network, the Wasserstein aggregation complexity O K n²) kn?) remains tractable when applied selectively. For our setup:
VGG-16 on CIEAR- 10: n = 67,240,000 (last two FC layers), I = 150, = 10 (per round)
Wasserstein aggregation: 150 10 x (67,240,000) 6.78 10 operations
27

---

Published in Transactions on Machine Learning Research (11/2025)
Standard aggregation: 10 x 138 x 10 1.38 10 operations
Overhead ratio: ~ 4.9 x 108x x for the affected layers
However, this computational overhead is incurred only during the aggregation phase the server, not during client-side training. Moreover, practical implementations use approximate Wasserstein distance computa- tions (e.g,, Sinkhorn iterations) which significantly reduce this theoretical complexity while maintaining the benefits of permutation-invariat aggregation.
subsubsection Amortized Analysis Considering total time-to-connvergence:
Trotal = R x (Telient + Teevver)
(86)
where R is the number of rounds, Telient is the per-round client training time, and Teerver is the per-round server aggregation time.
Since FedDUAL achieves convergence in 440% fewer rounds (refer Fig. 4 while incurring 35% server-side overhead per round:
TEaDDUAL 0.6. (Telient + 1.35 Tserver) 0.6
0.35 T server
(87)
TedAvg
(Telient + T servvr))
Telient + Terverr
In typical FT. scenarios where client training dominates (Telient >> T seerver) we have Tamaa
R 0, yielding:
Tollens
FedDUAL 0.6 (1 + 0) 0.6.
(88)
Tedlwg
Thus, the 35% server overhead becomes negligible, and FedDUAL achieves approximately 40% reduction in total wall-clock time due to faster convergence.
To empirically validate the assumption that Tatient >> T serverr, we measured the average time required for single client update and the total server update time per round. The results, summarized in Table 3. corroborate the assumption, demonstrating that client-side updates dominate the overall computational cost which further validate the general assumption made above.
Table 3: Wall clock time (in seconds) of the proposed method and the Fedavg algorithm for client and server side update.
CIFAR-10
FMNIST
T client
T server
T client
T server
FedAvg
17.71 27.49
0.44
2.14
0.05
Proposed
1.94
2.62
1.95
D Runtime Analysis
Table presents the wall-clock training time (in hours) of the proposed method compared to several widely used federated learning baselines the CIFAR-10 and FMNIST datasets. This time is recorded for fixed number of rounds for all the algorithms (80 rounds for CIFAR10 dataset and 180 rounds for FMNIST dataset). As expected, the runtime varies across algorithms due to differences in their commuicattio strate- gies, local computation overhead, and auxiliary regularization terms. Among the baselines, FedAvg achieves the lowest runtime, since it performs simple model averaging without any additional constraints or control variates. FedProx and FedBN incur a slight increase in runtime due to the introduction of proximal terms and batch normalization handling the client side, respectively. FedNova and SCAFFOLD exhibit moderately higher runtimes, attributed to gradient normalization and control variate updates. MOON further increases
28

---

Published in Transactions on Machine Learning Research (11/2025)
computational cost because of its contrastive representation alignment loss, while shows additional overhead due to maintaining both personalized and shared components during optimization. FedDyn records the highest runtime among the existing methods, as its dynamic regularization term requires per-round model adjustment and additional global parameter updates. The proposed method, though computationally more expensive than the baselines, remains efficient considering the performance benefits (Higher accuracy and faster convergence) achieves. Specifically, requires 5.33 hours on CIFAR-10 and 3.01 holurs on FMNIST, which is competitive with other advanced methods such FedDyn and FedPVR, while providing superior convergence and generalization performance. Overall, the proposed method achieves favorable balance be- tween computational cost and performance gain, validating its practicality in real-world federated learning scenarios.
We additionally report the wall-clock time (in hours) required to reach target accuracies of 0.40 for CIFAR- and 0.70 for FMNIST across all baselines and the proposed method, as summarized in Table 5 The results demonstrate that the proposed method consistently achieves the target accuracy in less time than the baselines, thereby confirming that the faster convergence observed in rounds (see Fig. 4 effectively translates into faster overall wall-clock convergence.
Table 4: Wall clock performance (in hours) of the proposed method and the baselines. shows the respective method failed to converge.
CIFAR-10
FMNIST
FedAvg
2.66
1.75
FedProx
2.90
1.90
FodNova
3.10
*
FedBN
*
1.80
FedDyn
4.89
2.91
MOON
3.80
2.40
SCAFFOLD
FedPVR
4.50
2.90
Proposed
5.33
3.01
Table 5: Wall-clock time to reach larget accuracy (in hours). to converge.
iidicates that the respective method failed
CIFAR-10
FMNIST
FedAvg
0.76
1.46
FedProx rox
1.61
1.90
FedNova
1.24
FedBN
1.80
FedDyn
4.89
2.91
MOON
1.27
1.47
SCAFFOLD
*
FodPVR
1.31
1.21
Proposed
0.67
1.02
E Limitation and Future Work
While the proposed FedDUAL framework achieves superior performance and faster convergence under severe data heterogencity, where several statc-of-tho-art methods such as FodNova, FodBN, and SCAFFOLD fail converge, it introduces higher computational cost at the server due to the iterative calculation of the Wasser- stein Barycenter. On the client. side, FedDUAL remains lightweight. and is also communication-effiien
29

---

Published in Transactions on Machine Learning Research (11/2025)
17.5
15.0
10.0
0.75
0.50
-10% 2..5557...00 (01 25p PP.
0.00
0.25
-1.4% 25 28 18.0000 20 20.0% 2..5%.7..
0,00
Pertubetion
6.75%
-0.5 50
Portodation
-0.50
75
-1.00
M4
Figure 6: Visualization of the loss surface for the global model trained on the FMNIST dataset with non-IID data (a 0.01): (a) shows the loss surface for the global model trained using FedAvg, while (b) depicts the loss surface for the global model trained with the proposed method FedDUAL.
0.9
0.8
Test Accuracy
0.7
0.6
0.5
Adaptive Loss
Dynamic Aggregation Proposed
0.4
0
50
100
150
200
250
300
350
400
FL rounds
Figure 7: Learning curves of the individual modules and the proposed method.
than methods like SCAFFOLD and FedPVR, it only requires transmitting model updates without addi- tional control variates. Although this extra computational overhead is limited to the server, which typically has sufficient resources, future research will focus reducing the computational burden of the Wasser- stein Barycenter calculation by developing more efficient algorithms, aiming to maintain or even improve the performance the proposed framework. Moreover, FedDUAL is inherently compatible with standard privacy-preserving techniques. Tn particular, its adaptive loss and dynamic aggregation can be seamlessly integrated with Differential Privacy by adding noise client updates prior to aggregation, without modi- lying the design. Investigating this integration with formal privacy guarantees remains important direction for future research.
30

---

Published in Transactions on Machine Learning Research (11/2025)
91.9%%
72.59%
71.97%
70
66.3%
Test Accuracy (%)
60
50
30
20
10
A10%
10-5
10
100
10a
Epsilon value (log scale)
Figure 8: Performance of the proposed method with different epsilon values.
0.8
0.7
0.6
Test Accuracy
0.5
0.4
0.2
epsilon=0.11
epsilon=1
epsilon=10
epsilon=0.5
0.1
epsilon-0.0000 0000001 (Proposed)
100
200
300
400
500
rounds
Figure 9: Learning curve the proposed method with different epsilon values.
31

---

Published in Transactions on Machine Learning Research (11/2025)
Motheds
Nathods Fedkig PadProx
Methods Fedlag
Fediing
Touss
Fashrox
Fodilieo
Fedllli
PeeDyy
Moon FeePAR
FodPVR
FodPuR
Propesed
Proposed
FLRoands
FL RtRsndds
Rounds
Figure 10: Learning curves of the proposed method and baselines with error bars.
82
81.99%
81
80
Test Accuracy (%)
79.15%
79
78
77
76
74.48%
150
500
1000
Iterations for WB calculation
Figure 11: Performance across different number of Iterations for WB calculation.
32

---

Published in Transactions on Machine Learning Research (11/2025)
0.85
=iteration==500
#iterations-150 (Propesed)
0.80
Test Accuracy
0.60
0.55
0.50
100
200
300
400
500
FL rounds
Figure 12: Learning curves for different number of Tterations for WB calculation.
30.0
07.5
04.03
02.5
$1.70
$1.06
03.0
Test
72.5
760
Fodlog
Peuthoa
Femmi
Moodi
Fadtoyt
Prepormend
Methods
Figure 13: Illustrates the accuracy of the proposed method and baselines across different levels of data heterogeneity on the FMNIST dataset.
33

---

Published in Transactions on Machine Learning Research (11/2025)
0.86
0.84
Test Accuracy
0.82
Methods
FedAvg
0.80
FedProx
FedBN
0.78
FedDyn
Moon
FedPVR
0.76
Proposed
1000
Rounds
Figure 14: Learning curve of the proposed method and other baselines heterogeneity level a=0.1.
FMNIST dataset with data
34


Reviews:
[{"review_id": "Dite3e2hze", "reviewer": "Reviewer_bSgD", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes FedDUAL for non-iid federated learning (FL) when the data is highly skewed. The algorithm can be decomposed into two parts: from the client side, an adaptive loss that combines the cross-entropy loss of the local classification and the \"distance\" of the local parameters from the global (here the distance is made up of KL divergence between soft-maxed parameters); from the server side, a dynamic aggregation that applies Wasserstein barycenter (via Sinkhorn-Knopp) to final layers while keeping FedAvg for earlier layers. Experiments on CIFAR-10/100 and FMNIST with severe skewness ($\\alpha = 0.01$) show accuracy gains and faster convergence. Theoretical convergence guarantees are also given under assumptions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "1. Algorithm is clearly explained.\n\n2. Experiment setup is provided with details.\n\n3. Ablation study on each part of FedDUAL is conducted.\n\n4. Theory is proved in appendix.", "audience_yn": "Yes", "audience_explanation": "1. Data skewness is a possible scene in real-world applications of FL.\n\n2. Observation that final layers are more affected by large gradient norms than earlier layers is valuable.\n\n3. Algorithm focusing on gradient rather than parameter is interesting.", "requested_changes": "1. I recommend the authors to define widely accepted mathematical concepts (such as KL divergence and Wasserstein distance) in a separate subsection called \"Notations and definitions\" in the introductioon section. This would probably make the flow more fluent.\n\n2. I recommend the authors to explain more on why they choose to use the KL divergence of the soft-maxed parameters (as a distribution) rather than some other distances/divergences between other information (such as logits or layer-wise representations). Also, the authors could add more discussions of $\\beta$; now it still seems confusing to me that why one should force local to behave similarly to global when global is worse than local.\n\n3. Equation 1 is confusing. I think the authors should write \"$\\theta^\\ast \\coloneqq \\arg \\min F(\\theta), F(\\theta) \\coloneqq \\frac{1}{K} \\sum_{k} f_k(\\theta)$\" or something like that.\n\n4. I recommend the authors to generalize their observation on gradient norms to broader class of models, which should be valuable to the community.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "q40m99MORB", "reviewer": "Reviewer_ndLg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work begins with motivation experiments investigating non-IID loss landscape, and the instability emerging at the later layers in traditional FedAvg algorithm. To address these issue, the work proposes adaptive and dynamic aggregation strategy. This work performs detailed evaluations comparing state-of-the art algorithms with the proposed method, and shows consistent improvements. The empirical findings are detailed, and strong.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The empirical findings seem well-supported by well-structured experiment designs, and detailed ablations.", "audience_yn": "Yes", "audience_explanation": "The research topic seems relevant and its findings would be insightful for the TMLR research community.", "requested_changes": "All of my concerns raised in previous rebuttal has been addressed, authors have added more ablations on other datasets, added helpful reasoning where individual modules may underperform the FedAvg algorithm, and have also extended the motivating experiments to other models/datasets collecting more empirical evidence of the non-IID FL issues. That being said, I am not an expert in theoretical proofs and do not feel confident assessing the correctness of the theoretical contributions. I noticed that other reviewers have raised concerns regarding this aspect, and I would like to see their feedback on the theoretical soundness of the work.\n\nSome minor suggestions:\n1) in the abstract, authors mention \"we begin with comprehensive experiments to pinpoint the underlying issues in the FL training process.\" it would be helpful to the reader if those issues are mentioned. it would make the abstract self-contained, and also connect the proposal to that. \n2) May be Figure 2 and 6 should be in the main paper as a side-to-side case. \n3) Likewise, may be Table 3, 4 and 5 can be combined, with Accuracy + dataset name.", "broader_impact_concerns": "No ethical review is required for this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3gSx5HfHQH", "reviewer": "Reviewer_f21q", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Three main contributions:\n\n - A dual-strategy framework for federated learning under non IID data. One the client side, it introduces an adaptive loss function that balances local optimization with alignment to the server model. On the server side, it combines FedAvg with a Wasserstein barycenter for the final, unstable layers.\n - Theoretical convergence guarantees.\n - Empirical results showing accuracy and convergence speed over eight well-known baselines across CIFAR-10, CIFAR-100, and FMNIST.\n\nStrengths:\n - Empirical results show the method achieves higher accuracy and faster convergence than baselines under non-IID settings;\n - Theoretical analysis provides convergence guarantees with standard rates in both convex and non-convex settings.\n\nWeaknesses:\n - Novelty is unclear: the use of KL divergence and Wasserstein barycenter is not positioned relative to prior work.\n - No computational complexity or runtime analysis is provided, despite server-side cost being the only acknowledged limitation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The empirical results support the claims of improved accuracy and faster convergence (measured in rounds-to-accuracy). However, since the main limitation of the method is increased server-side overhead, the lack of runtime and computational complexity analysis prevents a complete assessment of the method’s practicality.", "audience_yn": "Yes", "audience_explanation": "Federated learning is certainly an area of interest to TMLR's audience and this paper addresses a common challenge—data heterogeneity.", "requested_changes": "- **[critical]** Clarify the novelty of the proposed dual strategy by situating each prposed update (client-side and server-side) within existing literature.\n\n  - **[critical]** Report performance as a function of runtime in addition to communication rounds. While rounds-to-accuracy is the standard FL metric, the proposed method introduces nontrivial server-side overhead. Including runtime comparisons would clarify the practical trade-off between improved convergence and increased computational cost, and directly address the only limitation acknowledged by the authors. In addition to wall-clock performance, it would also be relevant to provide a computational complexity analysis of the proposed updates and how they compare to existing methods.\n\n  - Strengthen the empirical validation by including either experients across different values of the Dirichlet $\\alpha$ parameter or experiments on real non-IID datasets (e.g. FEMNIST).\n\n  - Expand Section 3.0.2 with one or two sentences that restate the problem of unstable final layers under non-IID data and explain why the Wasserstein barycenter provides a more robust aggregation than simple averaging. (This would ensure the intuition is clear where the method is introduced and not only left to the introduction.)", "broader_impact_concerns": "I do not see major ethical concerns beyond the standard issues in federated learning (e.g. data privacy, fairness across clients), which are not specific to this paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)