Confusion-Resistant Federated Learning via Diffusion-Based Data Harmonization on Non-IID Data

Authors: xiaohong  chen
canran  xiao
yongmei  liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:10:56

Abstract: Federated learning has become a pivotal distributed learning paradigm, involving collaborative model updates across multiple nodes with private data. However, handling non-i.i.d. (not identically and independently distributed) data and ensuring model consistency across heterogeneous environments present significant challenges. These challenges often lead to model performance degradation and increased difficulty in achieving effective communication among participant models. In this work, we propose Confusion-Resistant Federated Learning via Consistent Diffusion (CRFed), a novel framework designed to address these issues. Our approach introduces a new diffusion-based data harmonization mechanism that includes data augmentation, noise injection, and iterative denoising to ensure consistent model updates across non-i.i.d. data distributions. This mechanism aims to reduce data distribution disparities among participating nodes, enhancing the coordination and consistency of model updates. Moreover, we design a confusion-resistant strategy leveraging an indicator function and adaptive learning rate adjustment to mitigate the adverse effects of data heterogeneity and model inconsistency. Specifically, we calculate importance sampling weights based on the optimal sampling probability, which guides the selection of clients and the sampling of their data, ensuring that model updates are robust and aligned across different nodes. Extensive experiments on benchmark datasets, including MNIST, FashionMNIST, CIFAR-10, CIFAR-100, and NIPD, demonstrate the effectiveness of CRFed in improving accuracy, convergence speed, and overall robustness in federated learning scenarios with severe data heterogeneity.

Main paper: 
 introduction: 
 Federated Learning (FL) [McMahan et al., 2017b] has emerged as a powerful paradigm for distributed machine learning, enabling multiple clients to collaboratively train a shared model without exchanging raw data. This approach addresses critical concerns around data privacy and security, which are increasingly significant in various sectors such as healthcare [Antunes et al., 2022] , finance [Chatterjee et al., 2023] , and IoT [Li et al., 2020a , Pan et al., 2023 , Yao et al., 2024] . However, one of the fundamental challenges in FL is dealing with non-independent and identically distributed (non-IID) data, which can significantly impair the performance and convergence of the global model [Zhu et al., 2021] . As illustrated in Figure 1 , FL on non-IID data often suffers from issues like divergent model updates and inconsistent global models. Client models trained on heterogeneous data distributions tend to diverge [Ye et al., 2023] , making it difficult for the server to aggregate them into a coherent global model. This divergence is due to inconsistencies in data sources [Xiao and Liu, 2024] and distributions across clients [Duan et al., 2021] . This problem leads to reduced accuracy and slower convergence rates, highlighting the need for effective solutions to handle data heterogeneity. Existing research has made significant strides in improving the robustness and efficiency of FL in non-IID settings. Notable methods include FedProx [Li et al., 2020b] , which adds a proximal term to handle heterogeneity. Techniques such as MOON [Li et al., 2021b] and FedGen [Nguyen et al., 2021] introduce sophisticated strategies like contrastive learning and data generation to mitigate the effects of data heterogeneity. Despite these advancements, issues related to data distribution disparities and model inconsistency persist, limiting the scalability and effectiveness of FL in real-world scenarios.Given the critical gaps in existing FL approaches, particularly their limited robustness to severe non-IID data distributions, there is a pressing need for more resilient and adaptive solutions. This research is motivated by the necessity to enhance FL's capability to handle heterogeneous data efficiently. The primary objective of this study is to develop a novel FL framework, Confusion-Resistant Federated Learning via Consistent Diffusion (CRFed), which integrates advanced mechanisms to address data distribution disparities and enhance model consistency across clients.Our contributions can be summarized as follows:1. We propose a novel indicator function that dynamically adjusts sample weighting based on loss values and uncertainties, facilitating a self-paced learning approach that prioritizes more difficult samples over time. 2. Our framework employs a diffusion-based mechanism to harmonize data distributions, involving iterative noise injection and denoising processes that align local data with the desired distribution. 3. We implement a strategic client selection method based on the indicator function, ensuring the inclusion of the most reliable clients, which enhances the robustness and consistency of model updates. 4. Our extensive experimental evaluations demonstrate that CRFed achieves state-of-the-art performance on benchmark datasets. CRFed outperforms existing methods significantly in terms of accuracy and convergence speed under various non-IID settings.2 Related Work 
non-iid challenge in federated learning: 
 The issue of non-IID data in FL was initially highlighted by FedAVG [McMahan et al., 2017a] , and it has since been demonstrated that this challenge can significantly hinder the convergence and overall performance of the global model [Zhao et al., 2018 , Li et al., 2019] . Numerous studies, categorized as client-centric methods, have been proposed to tackle this problem by adjusting the local training objectives using insights from the global model and the local models of other clients [Wang et al., 2021] . For instance, FedProx [Li et al., 2020b] introduced a proximal term to constrain local updates by leveraging the global model. SCAFFOLD [Karimireddy et al., 2020] utilized control variates to correct for local training drift, while FedDyn [Acar et al., 2021] introduced a dynamic regularizer for parallelizing gradients among clients. MOON [Li et al., 2021b] applied contrastive learning to minimize the discrepancy between model representations, thereby correcting local training.Despite their contributions, these methods fall short of fully resolving the core of the non-IID issue and may experience performance limitations in scenarios with highly skewed data distributions [Li et al., 2022] . Beyond client-side adjustments, the server also plays a role in mitigating the adverse effects of non-IID data by calibrating the biased global model post-aggregation. For example, CCVR [Luo et al., 2021] uses virtual representations from an approximated Gaussian mixture model to correct the classifier. FedFTG [Zhang et al., 2022] employs data-free knowledge distillation to refine the global model with the knowledge derived from local models. Additionally, strategies such as client clustering [Ghosh et al., 2020 , Long et al., 2023] and client selection [Zhang et al., 2021 , Wang et al., 2020] , can be implemented by the server to alleviate the non-IID problem. IFCA [Ghosh et al., 2020] iteratively estimates client cluster identities based on local empirical loss and updates model parameters for each cluster via gradient descent.In federated learning (FL), data sampling strategies are vital for enhancing distributed training efficiency. [Tuor et al., 2020] proposed selecting local training data based on user-end data correlation analysis. This led to dynamic sampling strategies like [Li et al., 2021a] , where training sample importance is determined by model gradient magnitudes. Similarly, [Rizk et al., 2022] used gradient norms to derive sampling weights, minimizing theoretical convergence bounds. However, these methods require immediate gradient computations, increasing local overhead, and assume convex loss functions, which may not apply to deep learning models [Rizk et al., 2021] . Therefore, developing importance sampling methods suitable for deep learning-based FL remains an open challenge.FL convergence can be theoretically analyzed due to the model aggregation mechanism [98, [101] [102] 150] , with experimental validation for deep learning tasks [Wan et al., 2021] . Most studies rely on theoretical derivations, limiting practical application. This study aims to use a diffusion model to automate the modeling of optimal sampling strategies in FL.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: The paper discusses the limitation in Section 5, addressing potential biases from large language models. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
overview: 
 The CRFed framework, shown in Figure 2 , addresses challenges posed by non-i.i.d. data in FL. Our approach integrates a diffusion mechanism and a confusion-resistant strategy to ensure consistent and robust model updates across heterogeneous data distributions. The core idea of CRFed is that the performance of client i's data on the global model reflects its contribution to the training process. By using an optimal indicator function, we determine the optimal data sampling probability for each client, enhancing training efficiency and model performance.The framework comprises several key components: the current global model downloaded by clients at time t; the Model Encoder and Meta-model, which process the global model and client-specific data for the diffusion process; an indicator function, computed using client i's data on the global model; the Diffusion-based Data Harmonization Mechanism, which uses data augmentation, noise injection, and probabilistic modeling to mitigate data distribution disparities; and the Distribution Decoder, which aligns the denoised data distribution with the desired distribution.The Indicator Function I λ (l i , σ i ) is designed to measure the reliability of the i-th sample's loss value l i and its associated uncertainty σ i . The design is inspired by self-paced learning [Fan et al., 2017 , Castells et al., 2020] , which adjusts the weights of samples based on their loss values, allowing for a gradual learning process from easy to difficult samples. In the context of federated learning, this means that each client should adopt a self-paced learning paradigm, sampling its data in a way that allows the global model to learn from simple to complex tasks. The Indicator Function captures this performance and is defined as follows: This meta-model is then projected into a higher-dimensional space and concatenated with the indicator function, forming the combined representation z i . The diffusion-based data harmonization mechanism adds noise to this representation and iteratively denoises it to achieve the desired distribution. The distribution decoder then aligns the denoised data distribution. Client i's data is sampled based on importance sampling weights w i , calculated as the ratio of the optimal sampling probability P * i to the original data distribution P 0 . This ensures that the sampled data aligns with the desired distribution, following a curriculum learning approach that progresses from easy to difficult samples, thus enhancing overall model performance.I λ (l i , σ i ) = (l i -τ )σ i + λ(log σ i ) 2 (1)where λ is a pre-set regularization coefficient. τ is a confidence threshold that determines the difficulty of the sample based on its loss value. It can either be a fixed constant or a dynamically adjusted weighted average during the training process.The Indicator Function can be further explained using the following steps: for each client i, the loss value l i of each sample is calculated on the current global model, and the uncertainty σ i of each sample is estimated based on its difficulty. The Indicator Function I λ (l i , σ i ) is then used to assign weights to the samples, with easier samples having lower weights and more difficult samples having higher weights. This adaptive weighting mechanism ensures that the model focuses more on difficult samples over time, leading to improved learning efficiency and robustness. Theorem 3.1. In the CRFed framework, using the indicator function I λ (l i , σ i ) ensures stable and convergent updates for heterogeneous federated learning. For an appropriately chosen learning rate η, the model update rule for client i at iteration t,EQUATION)guarantees a decreasing step size, promoting convergence. Moreover, CRFed achieves a tighter bound on update steps than FedAvg, indicating faster convergence under the same conditions.The proof of Theorem 3.1 can be found in Appendix A.2.Given the global model θ, the optimal uncertainty σ * i can be derived through the following theorem: Theorem 3.2. The optimal uncertainty σ * i for a given loss value l i is obtained by minimizing the Indicator Function I λ (l i , σ i ). The solution is given by:σ * i (l i ) = exp -W 1 2λ max - 2 e , l i -τ (3)where W (•) is the Lambert W function.Proof. To find the optimal uncertainty σ * i , we first perform a variable transformation on the Indicator Function. Let c i = li-τ λ and x i = log σ i . Then, the problem can be rewritten as:σ * i (l i ) = arg min xi c i e xi + x 2 i (4)Taking the derivative with respect to x i and setting it to zero, we get:∂ ∂x i c i e xi + x 2 i = 0 (5)Simplifying this, we obtain:EQUATIONLet y = x i e xi , then the above equation can be expressed as:x i = -W c i 2 (7)where W (•) is the Lambert W function. Transforming back, we get:EQUATIONGiven c i ≥ -2 e , we can conclude that:σ * i (l i ) = exp -W 1 2λ max - 2 e , l i -τ (9)This optimal uncertainty σ * i dynamically adjusts sample weighting based on loss values, aligning with the self-paced learning paradigm. When the loss value l i is large, σ * i decreases, assigning a higher weight to the sample, indicating it is more difficult and needs more attention. Conversely, when l i is small, σ * i increases, assigning a lower weight, indicating it is easier and needs less focus. When we choose σ * i , the resulting indicator function is optimal, and the data distribution corresponding to this indicator function is the most optimal for training the current global model.Since the indicator function can be regarded as a measure of the client's data from the perspective of the global model, we need to consider the global model information when modeling the mapping between data distribution and the indicator function using the diffusion model. In this paper, we train a model encoder E using an autoregressive method to compress the global model θ t into a meta-model ϕ t (ϕ t = E(θ t )). The meta-model ϕ t is then projected into a higher-dimensional space P and concatenated with the indicator function I λ (l i , σ i ) to form z i = concat(P (ϕ t ), I λ (l i , σ i )). This combined representation z i is used as the input for the diffusion model (DiffusionModel(z i )). Refer to A.3 for more details of model encoder.The diffusion-based data harmonization mechanism is a critical component of the CRFed framework, responsible for mitigating data distribution disparities and ensuring consistent model updates across heterogeneous environments. The harmonization process is shown in Figure 3 , which involves adding noise to the data distribution and then iteratively denoising it to achieve the desired distribution. The workflow of this mechanism can be divided into two main processes: the forward diffusion process and the reverse denoising process.Forward Diffusion Process Suppose that a training sample x 0 is of a certain distribution, denoted as q(x 0 ). In the forward diffusion process, Gaussian noise with variance β t ∈ (0, 1) is added gradually to the sample x 0 for T steps, resulting in a latent sample x T ∼ N (0, I). The process is defined as follows:EQUATIONUsing notations α t = 1 -β t and ᾱt = t s=1 α s , the sample z t can be defined directly as:EQUATION)Reverse Denoising Process The reverse denoising process aims to sample reversely from z T through transition probabilities q(z t-1 |z t ) for timesteps T -1 through 1, yielding a sample drawn from q(z i ). The transition q(z t-1 |z t ) is a Gaussian distribution, tractable when conditioned on z i :EQUATIONwhere the mean μt and variance βt are calculated as:EQUATIONEQUATION)The reverse transition probability p θ (z t-1 |z t ) relies on the entire data distribution and is approximated through a neural network:EQUATIONThe detailed derivations and computations for these processes can refer to A.1. In the context of federated learning, the reverse denoising process starts from the optimal indicator function I * obtained in the forward diffusion process. By progressively denoising, we obtain the optimal sampling probability for client i, ensuring that the final data distribution aligns with the desired distribution. This minimizes the impact of data heterogeneity and ensures robust model updates across all clients.The confusion-resistant strategy is designed to address the challenges posed by data heterogeneity and model inconsistency in federated learning. It consists of three key components: client selection based on the indicator function, data sampling using the diffusion-based harmonization mechanism, and adaptive learning rate adjustment.To mitigate the adverse effects of data heterogeneity, we select clients based on their indicator function I λ (l i , σ i ), which quantifies the reliability of their data. Clients with the lowest indicator values, reflecting higher data reliability, are chosen for training. This approach follows the curriculum learning paradigm, where lower values indicate better data:EQUATION)where γ is a dynamically adjusted threshold ensuring the selection of the most suitable clients.Data Sampling Strategy For each selected client, the optimal sampling probability P * i is determined through the reverse denoising process, starting from the optimal indicator function I * . This ensures that the sampled data aligns with the desired distribution, enhancing the robustness of model updates. The importance sampling weight w i is calculated as follows:EQUATION)where P 0 is the original data distribution. Using w i , we sample the local training data (D sampled i = Sample(D i , w i )). This sampling ensures the effective sampling probability aligns with P * i . The distribution decoder, which is implemented as an autoencoder, is then used to decode the denoised data distribution. The autoencoder is trained to map the denoised samples back to the desired distribution, further ensuring that the data used for training is aligned with the ideal distribution. Refer to the A.3 for more details of distribution decoder.Adaptive Learning Rate Adjustment The learning rate η i for each client is adjusted based on the indicator function value, enhancing the influence of more reliable data:EQUATION)where η 0 is the base learning rate.The complete computational process(pseudocode) of CRFed is provided in the A.4.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: The paper includes a detailed proof for the convergence theorem in Section 3, providing all necessary assumptions and a complete proof. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems. • The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiments 4.1 experiment setup: 
 Figure 3 : The diffusion-based data harmonization mechanism in CRFed framework. The process involves a forward diffusion process where Gaussian noise is added to the initial data distribution, transforming it into a latent representation. This is followed by a reverse denoising process that iteratively removes the noise, aligning the data distribution with the desired target distribution.Datasets Our experiments are conducted on four widely used benchmark datasets: MNIST [LeCun et al., 1998 ], Fashion-MNIST [Xiao et al., 2017] , CIFAR-10 [ Krizhevsky et al., 2009] , and CIFAR-100 [Krizhevsky et al., 2009] . To simulate Non-IID data scenarios, we utilize the Dirichlet distribution [Yurochkin et al., 2019] Additionally, we use the NIPD dataset [Yin et al., 2023] , a benchmark specifically designed for federated learning in person detection tasks with Non-IID data. This dataset provides a realworld non-IID scenario to test the generalization of CRFed.Competing Methods Apart from FedAvg [McMahan et al., 2017a] , we compare the proposed algorithm with several benchmarking FL algorithms specialized for solving the non-IID problem, including FedProx [Li et al., 2020b] , MOON [Li et al., 2021b] , and FedGen [Nguyen et al., 2021] . We also compare our method against HFMDS-FL [ Li et al., 2024] , FRAug [Chen et al., 2023] , G-FML [Yang et al., 2023] , FedCD [Long et al., 2023] , FedNP [Wu et al., 2023] , and FedDPMS [Chen and Vikalo, 2023] , which are recent state-of-the-art approaches addressing non-IID data issues in federated learning. Effect of Data Heterogeneity We analyze the impact of data heterogeneity on the performance of the top 5 models by varying the Dirichlet concentration parameter β. Table 2 shows the performance of these models on CIFAR-100 and NIPD datasets for β values ranging from 0.1 to 0.5. As expected, the performance generally decreases with smaller β values due to increased data heterogeneity. Table 2 indicates that as β decreases, representing higher data heterogeneity, the performance of all models declines. CRFed consistently Convergence Rate The convergence performance of the top five models on FMNIST, CIFAR-10, and CIFAR-100 datasets is depicted in Figure 5 . As observed, CRFed demonstrates significantly faster and more stable convergence compared to the competing methods across all three datasets. This superior performance is attributed to the diffusion-based data harmonization mechanism, which effectively aligns data distributions, and the confusionresistant strategy that selects reliable clients and adaptively adjusts learning rates, ensuring efficient and robust training even in highly heterogeneous environments.To Previous importance sampling methods typically require prior analysis of the data relevance at each client-side [Hsu et al., 2020 , Tian et al., 2022] or necessitate deriving optimal sampling weights based on assumptions such as the convexity of the loss function [Rizk et al., 2022 , Zhu et al., 2024] . While these methods offer strong theoretical guarantees, they are somewhat limited in their adaptability to real-world federated learning (FL) scenarios. For instance, both FedIR [Hsu et al., 2020] and Harmony [Tian et al., 2022] assume that the server has knowledge of the local distributions of all clients. Although this assumption does not violate the privacy-preserving principles of FL, it can be challenging to obtain in real-world applications.In contrast, our CRFed does not depend on these assumptions. Instead, it iteratively adjusts the data distributions during the FL process itself, enabling the model to dynamically harmonize the diverse, non-IID data across clients without requiring explicit distributional assumptions or centralized access to all client data distributions. Guided by the indicator function, our CRFed can derive the optimal sampling strategy for each local node.Moreover, as shown in Table 3 , empirical experiments demonstrate that the diffusion model achieves superior performance, outperforming other benchmark methods. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: The paper specifies the datasets used, the training and testing details, the hyperparameters, and the evaluation metrics in Section 4, ensuring that the experiments can be reproduced. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: The paper discusses the potential positive impacts of improving clustering techniques for various applications and mentions possible negative impacts such as biases introduced by large language models and the risk of misuse in surveillance, along with mitigation strategies in Section 5.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
conclusion: 
 In conclusion, this study tackles the pressing challenge of handling non-i.i.d. data in federated learning environments. We propose the Confusion-Resistant Federated Learning via Consistent Diffusion (CRFed) framework. This framework introduces a novel Indicator Function that dynamically adjusts sample weighting, facilitating a self-paced learning paradigm that prioritizes more difficult samples over time. Additionally, our diffusion-based data harmonization mechanism ensures consistent and aligned data distributions through iterative noise injection and denoising processes, mitigating the adverse effects of data heterogeneity. Our strategic client selection method, guided by the Indicator Function, ensures that the most reliable clients are chosen for training, thus improving the robustness and consistency of global model updates.Despite the promising results, our approach has certain limitations. The reliance on complex diffusion mechanisms and adaptive strategies may introduce computational overhead, which could be a concern for resource-constrained environments. Future work should focus on optimizing the computational efficiency of the CRFed framework and exploring its applicability to a broader range of real-world federated learning scenarios [Zhang et al., 2023] . 
acknowledgement: 
 This research was funded by the Basic Science Center Project for National Natural Science Foundation of China (Grant No: 72088101), the Xiangjiang Laboratory Major Project(Grant No: 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: The paper does not involve the release of data or models that have a high risk for misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: The paper properly credits the creators of the datasets and models used, and mentions the licenses and terms of use, as detailed in Section 4. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Supplementary Explanation of the Diffusion-based Data Harmonization Mechanism Forward Diffusion Process In the forward diffusion process, Gaussian noise with variance β t ∈ (0, 1) is added gradually to the sample x 0 for T steps. The process is defined as:EQUATIONEQUATIONUsing notations α t = 1 -β t and ᾱt = t s=1 α s , the sample z t can be defined directly as:EQUATION)Reverse Denoising Process The reverse denoising process aims to sample reversely from z T through transition probabilities q(z t-1 |z t ) for timesteps T -1 through 1 to obtain a sample drawn from q(z i ). The transition q(z t-1 |z t ) is a Gaussian distribution, tractable when conditioned on z i :EQUATIONwhere the mean μt and variance βt are calculated as:EQUATIONEQUATION)The reverse transition probability p θ (z t-1 |z t ) relies on the entire data distribution and is approximated through a neural network:EQUATION)where Σ θ (z t , t) = βt I and the mean µ θ (z t , t) depends on a noise sample ϵ θ (z t , t) learned by a neural network. The learning process is guided by the objective function:EQUATION)while the output sample is obtained as:EQUATION)where z ∼ N (0, I) if t > 1 and z = 0 otherwise.Proof. To prove the convergence, we need to show how the indicator function promotes stable updates of the global model. Consider a simplified federated learning framework where the global model parameters θ are updated at iteration t as follows:EQUATION)where η is the learning rate and n is the number of clients. For simplicity, we consider a single client and expand ∇ θ I λ (l i , σ i ):EQUATION)By the chain rule, we have:EQUATION)Based on the definition of the optimal σ * i , ∇ θ σ * i can be approximated as proportional to ∇ θ l i , i.e.,EQUATION)where k is a constant. Thus, ∇ θ I λ (l i , σ * i ) simplifies to:EQUATION)Since σ * i is obtained by minimizing the indicator function, we have:EQUATION)Finally, the model update rule can be expressed as:EQUATION)To prove convergence, we note that the step size of parameter updates is finite:EQUATION)where C is a constant. Thus, as long as the learning rate η is appropriately chosen, the update step size will gradually decrease, ensuring the convergence of the model.Convergence speed Next, we compare the convergence speed of our method with the FedAvg algorithm.For FedAvg, the update rule is:EQUATION)where the update step size is:EQUATION)For our method with the indicator function, the update step size is:EQUATION)To demonstrate that our method converges faster or has a tighter bound, we analyze the total update step sizes over all clients.For FedAvg:EQUATION)For our method:EQUATION)Whenσ * i = exp W -(li-τ ) 2λ , λ ≥ -(li-τ ) 2e, and τ = l max -2λ ln 1 k , we can ensureEQUATION)Therefore, the update step size for our method is less than or equal to that of FedAvg.In the CRFed framework, both the model encoder and distribution decoder play crucial roles in ensuring effective data harmonization and robust model updates. These components are implemented using autoencoder architectures, designed to compress and reconstruct data representations efficiently.The model encoder E is responsible for compressing the global model parameters θ t into a lowerdimensional meta-model representation ϕ t . The encoder architecture comprises several fully connected layers activated by ReLU functions, followed by a linear transformation layer to produce the final compressed representation.Mathematically, given the input global model parameters θ t ∈ R d , the encoder outputs a compressed representationϕ t = E(θ t ) ∈ R d ′, where d ′ < d. The transformation is defined as follows:h 1 = ReLU(W 1 θ t + b 1 )h 2 = ReLU(W 2 h 1 + b 2 ) . . .h k = ReLU(W k h k-1 + b k )ϕ t = W out h k + b out (44)The training of the model encoder involves minimizing the mean squared error (MSE) between the original model parameters and their reconstructions. The loss function is given by:L E = 1 N N i=1 ∥θ ti -θti ∥ 2 (45)where θti = E -1 (E(θ ti )) and N is the number of samples. The optimization is performed using the Adam optimizer with a learning rate η. The detailed architecture includes an input layer of size d, hidden layers of sizes [128, 64, 32] , and an output layer of size d ′ = 16.The distribution decoder D aims to transform the denoised latent representations z t back into the desired data distribution. Like the encoder, the decoder uses a series of fully connected layers with ReLU activations, culminating in a linear layer to reconstruct the data.Given the input latent representation z t ∈ R d ′ , the decoder outputs the reconstructed data xt = D(z t ) ∈ R d . The transformations are defined as:h 1 = ReLU(W ′ 1 z t + b ′ 1 )h 2 = ReLU(W ′ 2 h 1 + b ′ 2 ) . . .h k = ReLU(W ′ k h k-1 + b ′ k )x t = W ′ out h k + b ′ out (46)The training process for the distribution decoder also minimizes the MSE, defined as:EQUATIONwhere xti = D(z ti ) and N is the number of samples. The optimization employs the Adam optimizer with a learning rate η. The architecture details include an input layer of size d ′ = 16, hidden layers of sizes [32, 64, 128] , and an output layer of size d.The model encoder and distribution decoder share a similar architectural approach, emphasizing efficient compression and reconstruction through deep learning techniques. The key parameters for both autoencoders are summarized as follows:• Model Encoder:- InputThe complete computational process of CRFed is illustrated in Algorithm 1.The detailed configuration of the models used in our experiments is provided below. Each table outlines the layers and parameters for the respective datasets. Adjust learning rate: Justification: The main claims in the abstract and introduction are supported by detailed descriptions, empirical evaluations, and theoretical analysis provided in the body of the paper (Sections 1, 3, and 4). Guidelines:η i = η 0 • I λ (li,σi) maxj I λ (lj ,• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The paper proposes a method to address data heterogeneity in FL by using a diffusion to harmonize clients data heterogeneity, with data noise injection and iterative denoising. The reviewers are generally positive, and all improved their ratings after engaging in the rebuttal process (including the lowest score).

Reviews: review: paper_summary: paper_summary: The work proposes a new FL approach, called CRFed, for addressing data heterogeneity in FL settings. CRFed relies on a diffusion based approach for harmonizing clients data heterogeneity  by performing data noise injection and iterative denoising, followed by a curriculum learning approach, which employs an indicator function to assign weights to training samples and indicate samples selection. The authors conduct of number of experiments across various domains and FL environments to showcase the promise of their approach.
strength_weakness: strength_weakness: - Some of the applied methods seem very ad-hoc and further elaboration is needed on their selection.
- More information is needed on the update frequency of the harmonization approach and respective curriculum sequence of the clients' training samples.
- Lack of theoretical framework limits the contribution of this work.
questions: questions: Comments, textual corrections and typos:
- Is the definition of the indicator function and the use of Lambert function based on previous work? Please cite accordingly and explain why these formulas were used. Moreover, why there is a substraction of confidence from the loss value?, also variables' reported value range in the indicator function is not clear (i.e., range of $\sigma_i, \tau$).
- Why the adaptive learning rate is adjusted based on the clients' indicator function? this is not clear at all.
- How often is the noise/denoise operation performed? At the beginning of training, at every federation round or after r-rounds? If it is done post-initialization then have you measured the effect of the approach in terms of wall-clock time? How expensive is the proposed approach in terms of time convergence (not round)? It would be great if you could include these results as well.
- Figure 4's is not readable. What is the $\beta$ value used for the two domains? Moreover, have you created 100 partitions and sub-sampling clients at each round or consider all available clients at every round?
- Please fix your citations style and add missing space between the text and the reference throughout the paper.
limitations: limitations: No.
paper_summary: paper_summary: This paper introduces CRFed, a framework designed to handle the challenges of non-i.i.d. data in federated learning. By using a diffusion-based data harmonization mechanism and a confusion-resistant strategy, CRFed aims to reduce data distribution differences among participating nodes and improve model consistency. Extensive experiments show that CRFed significantly enhances accuracy and convergence speed compared to existing methods.
In my opinion, the diffusion-based data harmonization mechanism is an innovative approach to dealing with data distribution differences. This paper not only introduces new theoretical concepts but also validates them through comprehensive experiments, making it a significant contribution to the field.
strength_weakness: strength_weakness: There are some areas where the paper could be improved to enhance its clarity, robustness, and applicability:
1.The method involves several hyperparameters, like the variance of Gaussian noise, the regularization coefficient, and the dynamically adjusted confidence threshold. A more thorough sensitivity analysis of these parameters would make the paper more complete.
2.In cases where data distributions are extremely diverse or have intrinsic properties that are difficult to capture through these processes, the effectiveness of the harmonization might be limited. 
3.While the method aims to enhance model consistency and reduce data disparities, the additional communication overhead introduced by the diffusion mechanism and importance sampling is not discussed.
4.Theorem 3.1 is valuable. However, during the derivation process, the purpose of equation (3) needs to be clearly explained.
5."I suggest highlighting the data showing performance advantages in Table 2 in bold.
questions: questions: Please refer to the weaknesses.
limitations: limitations: The limitations are discussed in the paper by the authors. There is no potential negative societal impact.
paper_summary: paper_summary: This paper presents a framework called CRFed to address the significant challenges posed by non-i.i.d. data in federated learning environments. This work introduces a diffusion-based data harmonization mechanism that effectively reduces disparities in data distributions across different nodes. Additionally, the paper proposes a confusion-resistant strategy that leverages an adaptive indicator function based on importance sampling.
Overall, this paper's writing is clear and easy to follow. The figures are well-drawn, allowing for a quick understanding of the research motivation and methodological design. The core contribution of this paper is the introduction of a diffusion-based data harmonization method to obtain valuable data distribution for global model aggregation, which is a very brave and innovative idea. Additionally, how to use synthetic data to enhance training effectiveness has long been an open question in the field of FL, and this work clearly provides a very promising approach. Therefore, I recommend that this paper be accepted.
strength_weakness: strength_weakness: 1. Some notations could be more clearly defined. For instance, in the definition of the Indicator Function $ I_\lambda (l_i, \sigma_i) $, it's not immediately clear how $\tau$ (the confidence threshold) is dynamically adjusted or chosen. A more detailed explanation of how $\tau$ impacts the learning process and its optimal selection criteria would be beneficial.
2. The iterative denoising process might face scalability issues with very large datasets or a high number of iterations. The paper should evaluate the performance and efficiency of the denoising steps in such scenarios to understand the method’s scalability better.
3. It would be beneficial to include a brief discussion on future work in the paper.
questions: questions: See above.
limitations: limitations: The authors have thoroughly outlined the limitations of their work as well as the potential negative societal impacts.
paper_summary: paper_summary: This paper proposes an importance sampling method with a diffusion model to achieve data harmonization in federated learning with non-i.i.d. data. The proposed method utilizes the indicator function from self-paced learning to measure the reliability of loss on each client and calculates the optimal data distribution according to the indicator function. The proposed method is evaluated in five different datasets empirically.
strength_weakness: strength_weakness: Despite the sufficient experimental results, the poor clarity of this paper weakens its soundness. There are several confusing points in the context of this paper:

1.	It is unclear what the indicator function means and how it is derived. As the core of this paper, the indicator function is only provided with a formulation and never explained how it is derived. Additionally, it is not straightforward to understand the relationship between the indicator function and the optimal data distribution.

2.	How the diffusion model is trained and used is not discussed in the paper. What we know from the paper is that the diffusion model takes the model embedding and the indicator function together to generate the optimal data distribution. It is weird why the diffusion model is necessary here. It seems that what we need is only a method that can estimate the optimal data distribution according to the indicator function.

3.	The definitions of the indicator function in Eq. (1) and Eq. (16) seem to be inconsistent. The indicator function is calculated for each sample in Eq. (1), while it is calculated for each client in Eq. (16). 

After all, the core idea of this paper is simple: re-assigning sampling probability to achieve overall balanced training data distribution such that the model update divergence can be mitigated. It is doubtful whether such a complicated framework is necessary to achieve this goal, while the improvement seems to be limited in the final converged accuracy.
questions: questions: 1.	Why is the optimal uncertainty obtained by minimizing the indicator function? 

2.	How is the distribution decoder trained? Since the decoder should output an unknown distribution from the denoised latent representation, there should not be a ground truth $x_{t_i}$ when training the decoder.

3.	Are the diffusion model, the model encoder, and the distribution decoder trained locally or globally? What is the training overhead of these components?
limitations: limitations: See the weaknesses and questions above.

Comments: title: Further discussions on motivation and computational costs
comment: Thanks again for your time and effort on reviewing our work!
I truly appreciate the authors' thorough feedback and their efforts in addressing my concerns, particularly their detailed explanation of the design principles for the indicator function and the additional analysis regarding the update step size when comparing CRFed to FedAvg. Overall, I am satisfied with the authors' response and their attention to my concerns, and I am pleased to increase my rating to 6 (Weak Accept).
We appreciate your recognition of our work's significance. Thank you again for your precious time and valuable suggestions.
Thank you for your timely feedback and for raising the score. We’ve provided the following response to address the remaining questions you raised:

### 1. Motivation for Diffusion Model
Previous importance sampling methods typically require prior analysis of the data relevance at each client-side [3,4] or necessitate deriving optimal sampling weights based on assumptions such as the convexity of the loss function [1,2]. While these methods offer strong theoretical guarantees, they are somewhat limited in their adaptability to real-world FL scenarios. For instance, both FedIR[3] and Harmony[4] assume that the server has knowledge of the local distributions of all clients. Although this assumption does not violate the privacy-preserving principles of FL, it can be challenging to obtain in real-world applications.

In contrast, the diffusion model based method we proposed does not depend on these assumptions. Instead, it iteratively adjusts the data distributions during the FL process itself. This enables the model to dynamically harmonize the diverse, non-IID data across clients **without requiring explicit distributional assumptions or centralized access to all client data distributions**. Guided by the indicator function, our CRFed can derive the optimal sampling strategy for each local node.

Moreover, as shown in the table below, empirical experiments demonstrate that **the diffusion model achieves superior performance, outperforming other benchmark methods**.

#### The performance of different importance sampling methods on CIFAR-100 under various β values.

| Method   | β=0.1 | β=0.3 | β=0.5 |
|----------|-------|-------|-------|
| ISFedAvg | 0.232 | 0.285 | 0.305 |
| ISFL     | 0.237 | 0.296 | 0.314 |
| FedIR    | 0.258 | 0.311 | 0.352 |
| Harmony  | 0.246 | 0.313 | 0.354 |
| CRFed    | **0.280** | **0.345** | **0.389** |

It’s worth noting that this comparison isn’t entirely fair, as each importance sampling method operates under different assumptions. For example, ISFL requires a validation set to update the empirical gradient Lipschitz constants for each local model, while FedIR requires all clients to upload the conditional distribution of images given class labels that matches the target distribution. Nevertheless, our **CRFed outperforms the others even under less restrictive conditions**—unlike ISFedAvg and ISFL, it doesn’t require assumptions about the loss function or gradient variance, and unlike FedIR and Harmony, it doesn’t require centralized access to all client data distributions before calculating the importance sampling weights.

### 2. Additional Computation
We acknowledge that CRFed requires additional computational resources, primarily on the server side. However, in FL, the main computational burden typically lies in the local model training and communication. In practical FL systems, **increasing server-side computation to achieve performance gains is often desirable**, as it does not compromise the real-time operation of the FL system and tends to offer stronger economic benefits. For example, in CLIP2FL[5], to mitigate data heterogeneity and class imbalance, the server generates federated features based on client-uploaded gradients and uses CLIP's text encoder for prototype contrastive learning. Similarly, in FedMRUR[6], the server must compute compressed data received from clients, calculate the corresponding logits, and perform global knowledge matching, which involves substantial computational intensity.

Lastly, we would like to express our sincere gratitude to Reviewer CpsC. Your feedback has inspired us to delve deeper into discussions regarding CRFed, and the additional experiments have further strengthened its contributions. We plan to include these results and discussions in future versions. We hope this response addresses some of your concerns.^_^

### References
[1] **ISFedAvg**. Federated learning under importance sampling. IEEE Transactions on Signal Processing. 2022.

[2] **ISFL**. Federated Learning for Non-iid Data with Local Importance Sampling. IEEE Internet of Things Journal. 2024

[3] **FedIR**. Federated Visual Classification with Real-World Data Distribution. ECCV. 2020.

[4] **Harmony**. Heterogeneity-aware hierarchical management for federated learning system. MICRO. 2022.

[5] **CLIP2FL**. CLIP-Guided Federated Learning on Heterogeneity and Long-Tailed Data. AAAI. 2024.

[6] **FedMRUR**. Federated learning with manifold regularization and normalized update reaggregation. Neurips. 2023.
Thank you for offering such a thorough rebuttal!
Considering the performance improvements and the importance of the problem tackled in this paper, the current overhead is acceptable. Future research can further enhance this aspect.
Overall, I am pleased with this work and would like to raise my score to 7.
I appreciate the response from the author, which helps me better understand the methodology in the paper. However, I am still concerned about the motivation for using the diffusion model for data importance sampling. The additional computation and data required for training the diffusion model offline are still unclear. And the comparison with other importance sampling methods for FL, e.g., FedIR (Federated Visual Classification with Real-World Data Distribution) and Harmony (HARMONY: Heterogeneity-Aware Hierarchical Management for Federated Learning System), is absent. Accordingly, I would increase my score to 4.