Fake It Till Make It: Federated Learning with Consensus-Oriented Generation
Authors: Rui Ye
Yaxin Du
Zhenyang Ni
Yanfeng Wang
Siheng Chen
Venue: ICLR 2024 poster (2024)
Keywords: Federated learning
data heterogeneity
Primary area: general machine learning (i.e., none of the above)

Abstract: In federated learning (FL), data heterogeneity is one key bottleneck that causes model divergence and limits performance. Addressing this, existing methods often regard data heterogeneity as an inherent property and propose to mitigate its adverse effects by correcting models. In this paper, we seek to break this inherent property by generating data to complement the original dataset to fundamentally mitigate heterogeneity level. 
As a novel attempt from the perspective of data, we propose federated learning with consensus-oriented generation (FedCOG). FedCOG consists of two key components at the client side: complementary data generation, which generates data extracted from the shared global model to complement the original dataset, and knowledge-distillation-based model training, which distills knowledge from global model to local model based on the generated data to mitigate over-fitting the original heterogeneous dataset.
FedCOG has two critical advantages: 1) it can be a plug-and-play module to further improve the performance of most existing FL methods, and 2) it is naturally compatible with standard FL protocols such as Secure Aggregation since it makes no modification in communication process.
Extensive experiments on classical and real-world FL datasets show that FedCOG consistently outperforms state-of-the-art methods. Code is available at https://github.com/rui-ye/FedCOG.

Published as a conference paper at ICLR 2024
FAKE IT TILL MAKE IT: FEDERATED LEARNING WITH CONSENSUS-
GENERATION
Rui Yel.,, Yaxin D Zhenyang Nit.3, Yanfeng Wang"1, Siheng Chen
"Shanghai Jiao Tong University. 'Shanghai AI Laboratory
$ "Melti-Agent Governance & Intelligence Crew (MAGIC) fyr991129, dorothydu, 0107nzy, wangyanfeng, sihenge) cn
ABSTRACT
In federated learning (FL), data heterogeneity one key bottleneck that causes model divergence and limits performance. Addressing this, existing methods of- ten regard data heterogeneity as an inherent property and propose to mitigate its adverse effects by correcting models. In this paper, we seek to break this inherent property by generating data to complement the original dataset to fundamentally mitigate heterogeneity level. As novel attempt from the perspective of data, we propose federated learning with cansensus-orientee FedCOG consists of two key components at the client side: complementary data generation, which generates data extracted from the shared global model to com- plement the original dataset, and www.lededdiitillattonnn-- which distills knowledge from global model to local model based on the generated data mitigate over-fitting the original heterogeneous dataset. FedCOG has two critical advantages: 1) it can be a plug-and-play module to further improve the performance of most existing FL methods, and 2) naturally compatible with standard FL protocols such as Secure Aggregation since it makes no modification in communication process. Extensive experiments on classical and real-worke FL. datasets show that F FedCOG consistently outperforms tfte-ooftthe-ar methods. Code is available at httpss//gittuu..ooo/uuuyyyffCCCt
generation (F (FedCoG).
model training,
INTRODUCTION
Federated learning (FL.) an emerging privacy-preserving training paradigm that enables multiple clients to train global model collaboratively without directly accessing their data (McMahan ctt al. 2017). With the increasing privacy concerns and legislation, FL has attracted much attention from industry and academia (Kairouz al. 2019 Yang et al. 2019). FL be widely used in diverse of real-world applications such as key board prediction (Hard et al.) 2018 Yang et al. 2018).
As one of the most common and critical issues in FL, data heterogeneity fundamentally limits FL's practical performance (McMahan et al. 2017) as clients' datasets could be distinctly different from each other (Li et al. 2020a). To tackle this issue, enormous previous works focus model-level correction at either the client or server side. At the client side, many methods propose to enhance the consistency among local models through regularization in model space (Li et al. 2020b) or feature space (Li et al. 2021a), or local gradient correction (Karimireddy et al. 2020). At the server side, some methods introduce update momentum (Hsu et al., (019), apply knowledge distillation based on public dataset (Lin et al., 2020), or adjust aggregation manner (Wang et al., (022b)). However, these previous works regard the data heterogeneity as an inherent property and attempt to mitigate its negative effects via model correction, leaving the training adversely affected throughout the FL process as the inherent heterogeneity persists to cause client drift (Karimireddy et al. 2020).
In this paper, seek an orthogonal approach address the data heterogeneity issue: correcting the data itself to mitigate the heterogeneity level. To correct the data, our core idea is to generate data from the shared global model as consensus to complement the original data, which contributes to mitigate the effects of data heterogeneity by making all local datasets more homogeneous (e.g., local categorical distributions are more balanced).
Following this spirit of data correction, we propose a new FL algorithm, Federated Learning with Consensus-Oriees Generation, denoted as FedCOG. FedCOG includes novel components:

---

Published as a conference paper at ICLR 2024
complementary data generation and plementary dataset generation, each client generates data that is accurately predicted by global model but incorrectly predicted by local model. In this case, the generated data not only contains consen- sual knowledge in the global model but also serves as an informative dataset complement, mitigating During local model training, beside minimizing the conventional
model training. 1) During com-
the level of data heterogeneity. task-driven loss on the original dataset, each client distills the knowledge from global model to cur- rent local model based on the generated dataset. Such knowledge distillation contributes to enhance the consensus among local models, which further alleviates the impact of data heterogeneity. Over- all, FedCOG improves the performance by mitigating both data heterogeneity level and its effects.
FedCOG has two critical properties: 1) Blug-and-play property and 2) compatibility with conven- tional real-world FI. protocol. First, FedCOG focuses on data correction, which orthogonal to most existing FL. methods (e.g., those that focus model correction (Li et al., 20206, Karimireddy al. 2020)) and thus can be easily combined with them to further enhance the performance. Second, FedCOG is compatible with practical standard FL protocols such Secure Aggregation (Bonawitz ct al. 2017) since it does not modify the communication process on standard FL, making it conve- nient to deploy in real-world application.
With extensive experiments across different representative datasets including real-worle multi- FL dataset FLAIR (Song et al. (022). heterogeneity types and heterogeneity levels, we show that our proposed FedCOG consistently outperforms 11 representative baselines and is of plug-and-play property. Besides, conduct further empirical analysis provide more insights in FedCOG from the perspective of alleviated model dissimilarity and enhanced local model performance.
The main contribution is as follows:
We propose novel attempt to tackle data heterogeneity in FL from the perspective of data correction, which seek to mitigate its adverse effects from the dataset itself.
We propose novel algorithm, FedCOG, which mitigates the issue of data heterogeneity by generating data to complement original data and distilling knowledge from global model to local model through the generated data,
With extensive experiments, we show FedCOG consistently achieves the best and helps relieve model dissimilarity and forgetting.
2 RELATED WORK
Data Heterogeneity in Federated Learning. Data heterogeneity is one of the most fundamental
and essential challenges in FL (McMahan et al. 2017 Kairouz et al., 2019), where data distribu- tions of clients could distinctly vary from each other, and shown to affect the performance of FL empirically (Zhao et al., 2018 Li et al. 2020b) and theoretically (Li et al. 2019 Yu et al., 2019). Addressing this, many methods are proposed from the aspects of local model correction and global model adjustment. 1) Local model correction aims for more similar local models at the client side. FedProx (Li al.) 2020b) and FedDyn (Acar et al. (2020) propose regularizing the distance between local and global model. SCAFFOLD (Karimiredyy et al. 2020) introduces control variate to cor- rect local gradients, MOON (Li al. 2021a), FedFM Ye et 2023b) and FedDecorr (Shi al., 2022) regularize the local models from the feature space via feature alignment or feature correlation regularization. 2) Global model adjustment aims for better- perttor-pe global model the server side. FedAvgM (Hsu et al.) 2019) and FedOPT (Reddi et al.] (021) introduce momentum to stabilize global model updating. FedNova (Wang et al. 2020b) and FedDisco (Ye et al. (202cc) modify the aggregation rules by adjusting aggregation weights. Others may apply knowledge distillation based on a public dataset (Lin et al. 2020) or design client sampling strategies (Cho et al. 2020).
Different from these two aspects, our method addresses this issue in novel way by directly ducing heterogeneity level from the aspect of data. We propose to generate additional data samples to complement the original heterogeneous dataset for each client, which is achieved by inversely optimizing inputs given the global and local model.
Data Generation and Augmentation in FL. (1) Recently, some methods (Zhang et al. 2022a b Zhu et al. 202T Hao et al. (02T) apply data generation at the server side to tune the global model for better performance, which generate data based on all local models. However, this could be un-
2

---

Published as a conference paper at ICLR 2024
acceptable in practice since Secure Aggregation technique (Bonawitz al. 2017 7: So al. 2022 is
often required so that the server can only receive the aggregated global model, making these methods less practical and privacy-preserving In contrast, our proposed FedCOG is compatible with Secure Aggregation and fundamentally mitigates the data heterogeneity by generating complementary data. (2) Some methods (Tang et al.) 2022. Chen et al. [022) Yuan et al.] 2022) apply data augmentation, though, they all introduce additional communication objects other than model parameters (McMa- han et al. 2017), which increases communication cost and/or risks of privacy leakage. Others (Xu et 2022] Yoon et al. (202T) may generate data to address the problem of forgetting, Unlike these, FedCOG makes no compromise on privacy and communication cost, and is a novel method that adopts task-specific data generation to tackle data heterogeneity.
Data Generation from Model. There are a line of works that focus on data generation from a single model (Fredrikson et al., 2015 Mahendran & Vedaldi, 2015 2016) to understand the repre- sentations of neural networks. Some works extend data generation to achieve data-free knowledge distillation (Fang et al. 2019; Nayak et al. [019 Mordvintsev et al. 2015 Yin et al. 2020). For ex- ample, DrepDream (Mordvintsev et al. (2135) generates data by penalizing classification loss, total variance and l2 norm while DeepInversion (Yin et al. 2020) additionally adds feature regulariza- tion and divergence term. These works are orthogonal to ours as they do not consider FL scenarios and our method can be further enhanced by combining with more advanced generation techniques.
PROBLEM FORMULATION
The objective of federated learning (FL) is to enable multiple clients collaboratively train a global model under the coordination of server without directly accessing data (Wang et al.) 2021) Specifically, suppose that there are K clients, where each client k holds a private local original dataset Dr, the global objective of FL is: ming Pe Dz), where pk= TP: is the
|De
relative dataset size, L.(-) is the task-driven loss function. Conventionally, this optimization problem is solved by the two iterative steps, including 1) client step: client downloads global model 0 from the server and conducts training on local dataset D& to obtain local model 0z. 2) server step: server receives local models }0x} from clients and aggregates them to obtain global model 0 := Σ. PRO.
Since clients' datasets } } can be heterogeneous, the trained local models at client step can be significantly different and eventually results in slow convergence and limited performance of global model. Some methods are proposed to mitigate the difference between local models {0x} at client step (Li et al.) [2020b] Karimireddy al. (2220) while some improve the global model O the server step (Reddi et al. [202]] Jhunjhunwal et al. 2023).
However, they do not consider modification on the original datasets (Dx), which is the root cause of heterogeneity issue. In this paper, we propose a novel attempt for mitigating the effects of data heterogeneity by generating data complement the original dataset for each client.
4 METHODOLOOGY
We propose federated learning with consensus-oriented generation (FedCOG), which mitigates data heterogeneity by generating complementary data and refining training at the client side. FedCOG follows the standard FL protocol that consists of client and server steps (McMahan et al., 2017). The novel designs of FedCOG lie in the client step, which is composed of two key modules: generating data to complement the original dataset and conducting knowledge distillation on the local model, which are elaborated in the following; also see details Algorithm and illustration in Figure [
4.1 COMPLEMENTARY DATASET GENERATION
As a novel attempt to mitigate the adverse effects of data heterogeneity from the perspective of data, the key step in FedCOG is generating data to complement the original dataset of each client, reducing the heterogeneity levels. We propose to generate task-specific and client-specific data by learning to produce samples that can be accurately predicted by the current global model and falsely predicted by the previous local model. Such generated data can contain consensual knowledge for the targeted task and complement the client dataset to increase data diversity.
3

---

Published as a conference paper at ICLR 2024
FedCOG
Step Complementary Generation
Step -D-ssed Local Model Training
Client
Gen. Data
Global Model
Task-driven Loss
Gan. Data
Local Model 6i share weights
Knowledge Distilaation
Hill
Server
learnable
flxed
Local Model 0%
Disagreement Loss
Ori.
Local Model 0
Task-driven Loss
0x
At Round
Client
Figure 1: Overview of FedCOG, which consists of two steps at the client side. 1) Complemen- tary data generation by optimizing the inputs, supervised by task-driven and disagreement loss. 2) wwwlledgeddisiilation-boe
local model training by distilling knowledge from global model to lo- cal model through generated data and minimizing task-driven loss through original data.
To obtain the desired samples, we set the inputs of the model as learnable parameters and fix the model parameters of global model and local model, where the learnable inputs are optimized by the guidance of task-driven loss and disagreement loss. Specifically, let 0 be the global model at round after r SGD iterations of local model
I and 0% 1." be client k's local model at previous round training. We use gelr represent that the model parameters are fixed and O represent that the model parameters are learnable. Denote Dj = {(2,, y: Ui)li N) as the generated dataset client where N is the number of samples generate, number of channels) is learnable tensor and y; is an arbitrary pre-defined target label. Each client generates data by solving the following optimization problem:
IR (H, W, denotes height, width and
D = arg min Cyen (0-1,, 0 De
1
Co(0ri: Adii 00 17, +),, (D
Σ
DR
(1)
where (-)n()) the overall generating objective. Lo(.) the task-driven loss, Adis is a hyper- parameter, and (-) (-) is the defined disagreement loss. Taking single-label classification task as concrete example, we define the task-driven loss as a cross-entropy loss and the disagreement loss
based on Sansen-Shanno divergence (Lin 1991 Yin et al. 2020), which is formulated as:
Lais (o 0 1. :)) :
KL (KL(
where p(0; àt) denotes the output of model 0, p = 1 (p(0 a) p(0-1111. +)) is the averaged output, and K L.) is the Kullback-Leibler divergence (Kullback & Leibler 1951).
The proposed two terms contribute to generating consensual and reasonable data to diversify and complement the local client dataset. 1) The first task-driven loss encourages generating data that is accurately recognized by the global model. Since the global model is the knowledge aggregation of multiple local models, such generated data is task-specific and can be of all clients. 2) The second disagreement term encourages the generated data to cause global-local model disagreement, serving as better complement for the original client dataset the generated data generally 'unfamiliar' for the client's local model. The generated dataset Dj. then utilized in local model training. See discussions on the applicability to other modalities in Appendix A.1.2
consensual knowledge
Details of designing target labels for generated data. One basic generating strategy uniformly generating data for all possible target labels. Taking CIFAR- 10 as concrete task example, the target label list can be [0,... 0.,...,,,.......
I,...,
9, ..., ...,9) resulting in 10 n samples to generate. Though, there
can also be more specific designs for particular tasks; see Appendix A.I.I
4.2 LOCAL MODEL TRAINING WITH KNOWLEDGE DISTILLATION
After data generation, the pivotal next step involves incorporating the generated data into local model training. A simple and straightforward solution is merge the generated dataset with the original dataset, creating a new dataset for local model training. However, since there are inherent discrep- ancies between generated and real- woaldw data, direct training the generated data in hard-labeled
4

---

Published as a conference paper at ICLR 2024
Algorithm FedCOG: federated learning with consensus-oriented generation.
Initialization: number of clients K, number of rounds T, number of iterations T. 2: for t = .,...,, I' -1 do
D> FL rounds in sequence
3: Send global model Ø to each client
Model communication
4: for k 0, 1, K 1 in parallel do
> Clients in parallel
5:
Dx - arg min, ggen (0 0% ;D))
Complementary data generation
6:
@ == 0t
b Model synchronization
7:
for r 0, .....
do
Iterations in sequence
8:
0x
0%
n (F.( + dRa (())
KD DD-based model training
9:
Send local model 0% to server M
Model communication
=EK kk=1 kk
Model aggregation
11:
Return: g"
format may result overly stringent supervision. Thus, we advocate using the global model's s soft labels for guidance, striking balance between learning from valuable insights from the generated data while cireumventing the potential biases from the generation process.
Specifically, client & initially employs the global model to reinitialize local model: 0% := 0' and subsequently launches local model training both of the original dataset Dj and the generated dataset DR. with the optimization objective defined
min 1 [Dal
Σ (m..))EDD
L.(0;
+
Akd
Σ (t))
(3)
[Del
(#::
where p(0-1*; the output of global model, p(0; a)) is the output currently optimizing local model, and Aud is a hyper-parameter. Note that the outputs of the global model given generated data can be extracted from the previous generation step without the need recomputation. Solving the optimization problem for steps of standard SGD in total, each client obtains the trained local model 0% which is then uploaded to the server for further model aggregation.
In this objective function, the optimization of the local model is balanced between task-driven loss (optimizing on the original real-world dataset) and knowledge distillation loss (optimizing on the generated dataset). This approach serves two purposes: first, to prevent over-fitting on the locally heterogeneous original dataset; and second, to preserve the knowledge of the global model and prevent excessive loss of consensus knowledge during local training.
4.3 DISCUSSIONS
Convergence analysis. We provide a theoretical convergence analysis in Appendix A.5 A.5
Compatibility, FedCOG is compatible with standard FL communication protocols including Se- cure Aggregation (SA) (Bonawitz et al.) 2017) since it does not modify the communication process compared with standard FL. (McMahan et al. 2017). However, many previous works are not com- patible with SA since they assume that the server can directly access individual local models (Lin) et al. 2020 Zhang et al. 2022a by; while SA requires that the server can only obtain the summation of local models enhance security and privacy (Zhu al. 2019). Besides, FedCOG has the plug- and-play property that can be combined with many existing FL works (Hsu et al. 2019) and can be potentially applied future works performance improvement; see empirical evidence in Table
Communication, privacy, and computation. As pointed in Kairouz et al. 2019). privacy and communication efficiency are two first-order concerns in FI., our proposed FedCOG does not com- promise either of these two aspects since it does not introduce any additional communication objects other than model parameters as FedAvg (McMahan al.) 2017). As there is no free lunch for privacy, utility, and efficiency in FL (Zhang et | (022c), we search enhance utility by slightly sacrificing computation efficiency, which is relatively minor concern FL. Though, the introduced computation overhead is acceptable as 1) the introduced generator is lightweight; see the comparison of training cost in Table and 2) generating data for only a few rounds is adequate to bring evident performance improvement; experiments in Table We provide further discussions in A.5.2
5

---

Published as a conference paper at ICLR 2024
Table I: Accuracy (%) comparisons across three datasets and two heterogeneous scenarios. The last column records the averaged accuracy across settings and FedCOG achieves the highest accuracy.
Dataset
Fashion-MNIST
CIFAR-10
CIFAR-100
Avg.
Heterogeneity
NIID-1
NIID-2
NIID-1
NIID-2
NIID-1
NIID-2
FedAvg
73.07 20.08
64.11
=0.78
64.36
#0.11
50.55
20.45
39.07
20.36
30,93
20.23
53.68
FedAvgM
76.81 10.48
67.73
20.65
64.36
20.33
51.32
70.74
39.21
10.19
30.97
10.11
55.07
FedProx
73.35 #0.07
64.44
#1.36
63.91
=0.36
53.51
+0.32
38.81
10.03
30.76
10.21
54.13
SCAFFOLD
77.02 #0.63
58.47
#3.05
64.36
+0.43
48.35
+1.86
42.25
+0.40
34.27
+0.22
54.12
MOON
73.21 +0.22
63.11
#1.08
63.65
+0.04
51.01
+1.34
39,34
+0.18
31.17
#0.13
53.58
FedSAM
76.44 +:0.06
68.14
+0.00
61.38
+0,23
50.76
+0.36
37.97
+0026
30.75
+0.11
54.24
VHL
73.08 10.89
70.26
1177
60.81
11.33
51.60
1:098
37.55
11006
28.41
10.46
53.62
FedExP
65.03 11114
59.49
1.60
48.96
12.11
37.97
11.75
35.14
11.25
22.47
10.29
44.84
FedDecorr
72.77 10.70
60.70
#4.59
63.40
10.29
47.53
11.99
37.68
10.40
27.08
10029
51.53
FedReg
70.75 #0.23
69.89
#0.74
55.40
#0.33
46.64
#0.42
30.95
+0.30
23.70
#0.31
49.52
FedGen
72.64 11.27
61.34
#2.4
62.45
#0.40
49.12
#0.51
37.99
+0.23
26.93
10.15
51.75
FedCOG (Ours)
77.34 +0.07
73.68
#0.38
64.83
=0.12
54.00
+0.84
42.88
+0.02
34.80
10.42
57.92
5 EXPERIMENTS
We conduct comprehensive comparisons with many FL baselines multiple datasets, including real-world FL dataset FLAIR (Song et al. 2022). See more details and results in Section A
5.1 IMPLEMENTATION DETAILS
Heterogeneous datasets. Overall, consider three classical datasets and one real-world FL. dataset. 1) For the three classical datasets, including Fashion-MNIST (Xiao et al.] (2017). CIFAR- 10/100 (Krizhevsky et al. 2009). we consider two types of heterogeneity, namely NIID-1 and NIID- 2. NIID-1 follows Dirichlet distribution Dira (default / = 0.1), which is a common setting in FL Yurochkin et al., 2019 Wang et al., 2020a). NIID-2 makes each client hold data of only par- tial labels (McMahan et al.) 2017 Li et al.] 20206) (2 for Fashion-MNIST and CIFAR-10, 10 for CIFAR-100). 2) FLAIR (Song et al. (022) is a recently released real-world FL multi-label dataset, where each client is user from Flickr and we adopt the task with 17 labels. It is a challenging task since each user has different data distribution and few data samples. See more in Section A.3
Training setups. For the three classical datasets, we consider 70 communication rounds, K = 10 clients, 7 = 400 iterations of model training, and apply a simple CNN (Li et al. 2021a).. For the more challenging FLAIR (Song et al. 2022) dataset, we consider 400 communication rounds and apply ResNet18 (He et al. 2016 model. Besides, we sample = 200 clients in total and 10 clients in each FL round, and set datasets, we use SGD as the optimizer with learning rate 0.01. For FedCOG, we apply the most compatation-efficient way by setting the inputs learnable, supervised by 256 pre-defined target labels. The default settings of Auta and Akd are 0.1 and 0.01, respectively.
10 because each client holds relatively few data samples. For all
Baselines. We compare with 11 baselines. FedAvg (MeMahan et al. 2017) the basical baseline. FedProx (Li et al. (2220b). SCAFFOLD (Karimireddy et al. 2020)). MOON eLi al. 2021a), Fed- SAM (Qu al. 2022). FedDecorr (Shi et al.) (022) focus on local model correction, VHL (Tang et al., 2022). FedReg (Xu et al. 2022) and FedGen (Zhu et al., 202T) focus on data augmenta- tion. FedAvgM (Hsu et al., 2019) and FedExP (Jhunjhunwala et al. (2023) focus on global model adjustment. Among these, SCAFFOLD requires double communication cost.
5.2 COMPARISONS WITH STATE-OF-THE-ME METHODS
Applicability on standard datasets. Here, we conduct experiments under two types of hetero- geneous scenarios on three standard datasets used in FL. literature (Li et al.; 2020b 2021al Ta Shii ctt al. 2022 Jhunjhunwala et al., 2023), including Fashion-MNIST (Xiao et al. 2017). CIFAR- 10/100 (Krizhevsky et al., 2009). We run 50 rounds of FedAvg before running 20 rounds of Fed- COG for higher computation efficiency. We report accuracy comparisons in Table I From the table, see that 1) FedCOG consistently performs the best across different datasets and heterogeneous scenarios, indicating the effectiveness of tackling data heterogeneity from the perspective of data. 2) For the relatively more challenging scenario, NIID~-2, some methods (e.g., SCAFFOLD)

---

Published as a conference paper at ICLR 2024
Table 3: Efficient plug-and-pe and-play property. FedCOG consistently brings performance improvement across different scenarios, with only round for CIFAR-10 and rounds for FLAIR.
Method + FedCOG?
FedAvg
FedAvgM X
FedProx
SCAFFOLD
MOON
X
X
x
V
V
1IFAR-1--1
63.34
64.47
61.71
64.63
64.15
64.87
63.67
65.36
62.56
63.24
CIFAR-10-2
49.30
49.50
44.73
45.14
50.35
50.89
46.47
51.26
49.64
50.72
FLAIR O-FI
47.91
51.35
48.77
51.60
45.71
54.57
42.50
45.22
44.99
52.78
FLAIR C-F1
20.63
20.89
20.11
21.98
18.66
19.08
16.74
18.96
18.83
19.92
perform worse than FedAvg; while in contrast, our proposed FedCOG brings significantly larger performance gain. This indicates that complementary data generation is more effective in cases where local dataset misses data from several categories (i.e., more heterogeneous).
Applicability on the real-world FL multi-label dataset, FLAIR (Song al.) 2022). We run ad- ditional 5 rounds of FedCOG based on the global model trained by 400 rounds of FedAvg, for the sake of computation efficiency. We compared with other baselines with 405 rounds for fair compari- son by evaluating per-class (denoted by C-) preci- sion, recall, F1 score, and per-sample (denoted by O-) Fl score in Table From the table, we see that 1) FedCOG consistently achieves the best per- formance under different evaluation metrics, indi- cating the effectiveness of data generation in Fed- COG for multi-label task. While some methods perform worse than the vanilla FL method FedAvg or even inapplicable to this task (e.g., VHL (Tang al.] 2022) and FedReg (Xu et al. (0222) our proposed FedCOG even brings larger performance improvement for such real-world multi-label task.
Experiments on real-world FL multi- label dataset, FLAIR. C- and O- denote aver- aged metric across 17 labels and samples, re- spectively, P: precision, R: recall.
Table 2:
Method
C-P
C-R
C-F1
O-F1
FedAvg
30.76
14.23
18.30
44.53
FedAvgM
29.76
16.76
20.44
48.43
FedProx
32.27
14.95
18.66
45.71
SCAFFOLD
31.10
12.37
16.74
42.50
MOON
31.41
14.71
18.83
44.99
FedSAM
32.33
14.39
17.94
46.18
VHL
N/A
N/A
N/A
N/A
FedDecorr
33.44
13.36
17.51
41.89
FedExP
29.14
15.70
19.64
46.57
FedReg
N/A
N/A
N/A
N/A
FedCOG
37.87
18.78
20.89
51.35
Specifically, it outperforms FedAvg by 14% and 15% relatively with regard to C-FI and O-F1.
5.3 EFFICIENT BUIG-AND-PLAY PROPERTY
One valuable advantage of FedCOG is its plug-and-play nature since FedCOG focuses on data- level modification, which is orthogonal to most existing methods. In this section, we select five representative baseline methods, including FedAvg (McMahan et al. 2017), FedAvgM (Hsu al. 2019) FedProx (Li et al.] (020b)), SCAFFOLD (Karimireddy et al. 2020) and MOON (Li et al. 202 to investigate the compatibility of FedCOG with existing methods.
Here, we first conduct 50 rounds of training using the baseline method X. Subsequently, on one hand, we continue training method X for round obtain Model A, and on the other hand, we conduct 1 round of training of method X+FedCOG to obtain Model B. Finally, we compare the performances of Model A and B in Table 3 From the table, we see that 1) FedCOG consistently brings performance gain when combined with these baseline methods, indicating its plug-and-pla property. Specifically, can bring 19% O-FI improvement FedProx (Li al. (020b) on FLAIR dataset. 2) there is already evident performance improvement through only I round of FedCOG, indicating that FedCOG can take effects in a computation-effiiien way.
5.4 EMPIRICAL ANALYSIS OF FEDCOG
We dive into the training process by evaluating model difference, generalization, and personaliza- tion of local models to provide more deeper insights. Experiments are conducted on a common setting (Shi et al. 2022: Jhunjhunwala et al. 2023), NIID-1 on CIFAR-10. We start from the same global model initaaizzaion (trained by 50 rounds of FedAvg (McMahan et al. ()) then launches 1 round of FedAvg, FedProx (Li et al. 2020b), and FedCOG, respectively. See more in Section A.5
Model difference of local models. Here, we measure the &, difference between each local model and the aggregated global model, which is a reasonable indicator of the effects of data heterogene-
7

---

Published as a conference paper at ICLR 2024
between
Germafooniann local modils
Persormtiaaionn socuresy koal models
PedProx
FedDOG
FedAug
WdCOG
wdcOG
60
85
3,0
55
50
75
45
1.5
70
Cienn
Di
Clem
(a) Model difference
(b) Generalization analysis
(c) Personalization analysis
Figure 2: Empirical analysis of FedAvg. FedProx, and FedCOG. a) FedCOG has the lowest model difference between local model (LM) and global model (GM). b) FedCOG consistently and signif- icantly achieves the best preservation of the generalization ability of local models. c) Interestingly, FedCOG also consistently achieves best personalization performance.
ity (Li et al. 2020b). From Figure 2(a) we see that FedCOG achieves the lowest model difference between local and global model, indicating that FedCOG can significantly reduce the effects of data heterogeneity on the consistency of local models.
Generalization ability of local models. Here, we evaluate each local model on the uniform test dataset (used in Table [ to examine the preservation of general knowledge during local training. From Figure 2(b) we see that FedCOG consistently and significantly achieves the highest accuracy of local models. indicating that FedCOG preserves the generalization ability best since it generates data from global model to complement the original dataset. Additionally, the achieved global model accuracies of FedAvg, FedProx and FedCOG are 63.34, 64.57 and 65.59, respectively, indicating that FedCOG achieves best generalization performance of both local models and global model.
Personalization ability of local models. Here evaluate each local model on the personalized test dataset, which follows the data distribution of its training dataset. Interestingly, from Figure 22c) we find that FedCOG also consistently achieves the highest personalization performance, indicating that FedCOG can also enhances the fitting ability on local dataset. Such a performance gain of FedCOG could be attributed to aspects: 1) From global view, FedCOG leverages data generation and knowledge distillation to achieve well-performing aggregated global model. This global model then used as the initialization for local model training, contributing improved performance. 2) From local perspective, the generated data serves two purposes: it enhances model learning by providing additional knowledge, and it mitigates the risk of over-fitting. See the further explanations and comparisons of FedCOG against several representative personalized FL methods in A.4
Computational cost. We compare the total local time re- quired by each client throughout: 51 rounds and the final accuracy in Table Note that we have include the gener- ation time (which takes only 0.81 seconds for each round) required for FedCOG for fair comparison. 1) Results show that FedCOG introduces moderate training and gen- erating time while achieving the highest accuracy. 2) This acceptable as communication and privacy more im- portant in FI. (Kairouz et al., (219)) while FedCOG makes no compromise on these two aspects, and the utility is improved with evident gain. Specifically, compared with FedProx and MOON, FedCOG achieves much higher performance with the least computation time.
Table 4: Local time of clients and test accuracy of global model. FedCOG introduces minor training cost while achieves the highest performance.
Method Local Time (s)
Acc (%)
FedAvg
983 100
63.34
FedProx
1051 + 87
64.57
MOON
1276 + 147
64.44
FedCOG
1001 + 98
65.59
Applicability on different heterogeneity levels. Here, we tune the argument / E (0.1, 0.5, 5.0, for Dirichlet distribution NIID-1 and compare with two representative base- lines in Figure ((a] Note that smaller 3 denotes a more heterogeneous setting. From the figure, we see that 1) generally, FI. methods have worse performance under the more heterogeneous sce- narios, indicating that data heterogeneity is critical issue that limits FL performance, 2) FedCOG consistently performs the best across different heterogeneity levels, indicating its effectiveness on tackling data heterogeneity.
8

---

Published as a conference paper at ICLR 2024
Effects Heterogenelty Levels
Effects
66
FadAug
FedProx
64
64
FedCOG
NIID-1
62
53
NIID-2
52
NIID-
63
0.01
5.,001
(a) Heterogeneity Level
(b) Disagreement Adsa
(c) Knowledge Distillation Apd
Figure 3: Ablation study. (a) FedCOG is effective across different heterogeneity levels. (b) Effects
of disagreement in data generation. (c) Effects of knowledge distillation in model training.
5.5 ABLATION STUDY
To evaluate the impact of each module on the final performance of global model, we conduct exper- iments CIFAR-10 for ablation study, See effects of FL arguments in Section A.6
Complementary data from noise or gener- ation. To validate the effectiveness of task- specific data generation from global model and local model, we also run experiments on pure Gaussian noise. From @ & @ in Table 5 we see that data generation from global model and local model leads to sig- nificantly better performance than Gaussian noise, indicating the effectiivneess of data generation through optimizing the inputs.
Table 5: Ablation study of FedCOG. @: Comple-
mentary data is from Gaussian noise. @: Data is gen- erated from global model. @: Hard-label supervision on generated data. @ Full version of FedCOG.
No.
Gen. Data Source
Super.
NIID-1
NIID-2
1
Gaussian Noise
Soft
60.98
49.67
Gen. from GM
Soft
64.72
51.61
3
Gen. from GM&LM
Hard
63.59
51.64
Gen. from GM&LM
Soft
64.88
53.20
Generating data from global model
global & local models. By comparing @ with in Table 5 we can see that generating data from merely global model achieves moderate performance while adding the disagreement term between global and local model contributes to better performance, indicating the effectiveness of disagree- ment loss in FedCOG. More detailed analysis is in Figure 3(b
we see that soft- label format supervision leads to better performance, indicating the effectiveness of knowledge- distillation-based model training in FedCOG. More detailed analysis is in Figure 33c)
Supervision in hard- or soft-label format. By comparing 3 and in Table 5
Effects of disagreement in generation and knowledge distillation in model training. Un- der two heterogeneous settings on 10, we tune Adis E {0.01, 0.1, 1.0, 10.0} and Apd E (0.001, 0.01, 0.1, 1.0). showing results in Figure 3(b) and 3(c) respectively. From the figure, we see that generally Adis 0.1 and A = 0.01 can lead to better performance for both settings.
6 CONCLUSION AND FUTURE WORKS
In this paper, we seek to tackle the issue of data heterogeneity in FL from the novel perspective of modifying local dataset. To achieve this, we propose a novel FL algorithm, federated learning with consensus-oriente generation (FedCOG), to mitigate the heterogeneity level. FedCOG consists of two key components, complementary data generation to reduce heterogeneity level and knowledge- distillation-based model training mitigate the effects of heterogeneity. FedCOG is plug-and-play in most existing FL methods, is compatible with standard FL protocol such as Secure Aggregation, and makes no compromise experiments on classical and real-world FI. datasets show FedCOG consistently outperforms state- of-the-art methods. We believe that FedCOG inspire more future subsequent works along this line via adding regularization terms during generation or introducing advanced generative models.
communication cost and privacy while improves utility. Extensive

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
This research is supported by the National Key R&D Program of China under Grant 2021ZD01 12801, NSFC under Grant 62171276 and the Science and Technology Commission of Shanghai Municipal under Grant 2151 100900 and 22DZ2229005.
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Con- ference Leurning Representations, 2020.
Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for privacy- preserving machine learning. proceedings of the 2017ACM, SIGSAC Conference on Computer and Communications Security, pp. 1175-1191, 2017.
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-33 311, 2018.
Haokun Chen, Ahmed Frikha, Denis Krompass, and Volker Tresp. Fraug: Tackling federated learn- ing with non-iid features via representation augmentation. arXiv preprint arXiv:2205. 14900, 2022.
Yac Jee Cho, Jianyu Wang, and Gauri Joshi. Client selection in federated learning: Convergence analysis and power-of-choice selection strategies. arXiv preprint arXiv:2010.01243, 2020.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared repre- sentations for personalized federated learning. In International Conference on Machine Learning, pp. 1008922099 PMLR, 2021.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Fedavg with fine tuning: Local updates lead to representation learning. Advances in Neural Information Processing Sys- tems, 35:10572-10588 2022.
Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. In International Conference on Learning Represen- tations, 2020.
Ziqing Fan, Jiangchao Yao, Bo Han, Ya Zhang, Yanfeng Wang, et al. Federated learning with bilateral curation for partially class-disjoint data. Advances in Neural Information Processing Systems, 36, 2024.
Gongfan Fang, Jie Song, Chengchao Shen, Xinchao Wang, Da Chen, and Mingli Song. Data-free adversarial distillation. arXiv preprint arXiv:1912.11006. 2019.
Matt Fredrikson, Somesh Jha, and Thomas Ristenpart. Model inversion attacks that exploit confi- dence information and basic countermeasures. In Proceedings of the 22nd ACM SIGSAC confer- ence on computer and communications security, pp. 1322-1333, 2015.
Weituo Hao, Mostafa El-Khamy, Jungwon Lee, Jianyi Zhang, Kevin Liang, Changyou Chen, and Lawrence Carin Duke. Towards fair federated learning with zero-shot data augmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 3310-3319, 2021.
Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: /8/ I.0 1.03604, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
10

---

Published as a conference paper at ICLR 2024
Feng Hong, Jiangchao Yao, Zhihan Zhou, Ya Zhang, and Yanfeng Wang. Long-tailed partial label learning via dynamic rebalancing. In International Conference on Learning Representations. OpenReview.net, 2023.
Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout. Advances Neural Information Processing Systems, 34: 1876-11188, 2021.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06333 2019.
Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Fedexp: Speeding up federated averaging via extrapolation. arXiv preprint arXiv:2301.0900 2023.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurelien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. arXiv preprint xrXiv:1912.009977 2019.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference Machine Learning, pp. 5132-5143, PMLR, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathe- matical statistics, 22(1):79-86, 86, 1951.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10713-10722, 2021a.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-60, 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Sys- tems, 2:429 2020b.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pp. 6357- 6368. PMLR, 202
Xiang Li. Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations, 2019.
Jianhua Lin. Divergence measures based on the shannon entropy. IEEE Transactions on Information theory, 3751):145-151, 1991.
Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363,, 2020.
Aravindh Mahendran and Andrea Vedaldi. Understanding deep image representations by inverting them. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 5188-5196, 2015.
Aravindh Mahendran and Andrea Vedaldi. Visualizing deep convolutional neural networks using natural pre-images. International Journal of Computer Vision, 120:233-255. 2016.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. caanuuniiaaiooefffccioiooo learning deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.

---

Published as a conference paper at ICLR 2024
Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. In International Conference on Learning Representations. OpenReview.net, 2021.
Alexander Mordvintsev, Christopher Olah, and Mike Tyka. Inceptionism: Going deeper into neural networks. 2015.
Gaurav Kumar Nayak, Konda Reddy Mopuri, Vaisakh Shaj, Venkatesh Babu Radhakrishnan, and Anirban Chakraborty. Zero-shot knowledge distillation in deep networks. In International Con- ference Machine Learning, pp. 4743-4751. PMLR, 2019.
Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. In International Conference on Machine Learning, pp. 18250- 18280. PMLR, 2022.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021.
fforum?
1d=LkFG31B13U5
Yujun Shi, Jian Liang, Wenqing Zhang, Vincent YF Tan, and Song Bai. Towards understand- ing and mitigating dimensional collapse in heterogeneous federated learning. arXiv preprint arXiv:2210.00226, 2022.
Jinhyun So, Chaoyang He, Chien-Sheng Yang, Songze Li, Qian Yu, Ramy Ali, Basak Guler, and Salman Avestimehr. Lightsecagg: lightweight and versatile design for secure aggregation in federated learning. Proceedings of Machine Learning and Systems, 4:694-720, 2022.
Congzheng Song, Filip Granqvist, and Kunal Talwar. Flair: Federated learning annotated image repository. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
Zhenheng Tang, Yonggang Zhang, Shaohuai Shi, Xin He, Bo Han, and Xiaowen Chu. Virtual ho- mogeneity learning: Defending against data heterogeneity federated learning, In International Conference Machine Learning, pp. 11111-21133 PMLR, 2022.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conference Learning Represen- tations, 2020a.
net /to//oooo/m////m Biauumidd--kkuussF.
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:7611-7623 2020b.
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide to feder- ated optimization. arXiv preprint arXiv:2107.06917. 2021.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint arXiv: 2017.
Chencheng Xu, Zhiwei Hong, Minlie Huang, and Tao Jiang. Acceleration of federated learning with alleviated forgetting in local training, 2022.
Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
Timothy Yang, Galen Andrew, Hubert Eichner, Haicheng Sun, Wei Li, Nicholas Kong, Daniel Ra- mage, and Françoise Beaufays. Applied federated learning: Improving google keyboard query suggestions, arXiv preprint arXiv: 18X2:28//200903: 2018
12

---

Published as a conference paper at ICLR 2024
Rui Ye, Zhenyang Ni, Fangzhao Wu, Siheng Chen, and Yanfeng Wang. Personalized federated learning with inferred collaboration graphs. In International Conference on Machine Learning, pp. 39801-398117 17. PMLR, 2023a.
Rui Ye, Zhenyang Ni, Chenxin Xu, Jianyu Wang, Siheng Chen, and Yonina C Eldar. Fedfm: Anchor- based feature matching for data heterogeneity in federated learning. IEEE Transactions on Signal Processing, 2023b.
Rui Ye, Mingkai Xu, Jianyu Wang. Chenxin Xu, Siheng Chen, and Yanfeng Wang. Feddisco: Fed- erated learning with discrepancy-aware collaboration. In International Conference Machine Learning, pp. 39879-39902.. PMLR, 2023c.
Hongxu Yin, Pavlo Molchanov, Jose M Alvarez, Zhizhong Li, Arun Mallya, Derek Hoiem, Ni- raj K Jha, and Jan Kautz. Dreaming to distill: Data-free knowledge transfer via deepinversion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8715-8724, 2020.
Jaehong Yoon, Wonyong Jeong, Giwoong Lee, Eunho Yang, and Sung Ju Hwang. Federated contin- ual learning with weighted inter-client transfer. In International Conference on Machine Learn- ing, pp. 12073-12086. PMLR, 2021.
Hao Yu, Sen Yang. and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5693-5700, 2019.
Haolin Yuan, Bo Hui, Yuchen Yang, Philippe Burlina, Neil Zhenqiang Gong, and Yinzhi Cao. Addressing heterogeneity federated learning via distributional transformation. In Computer Mision-ECCy 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceed- ings, Part XXXVIII. pp. 179-195. Springer, 2022.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In Interna- tional Conference on Machine Learning, pp. 7252-7261. PMLR, 2019.
Jie Zhang, Chen Chen, Bo Li, Lingjuan Lyu, Shuang Wu, Shouhong Ding, Chunhua Shen, and Chao Wu. Dense: Data-free one-shot federated learning. Advances in Neural Information Processing Systems, 35:21414-21428, 2022a.
Lin Zhang, Li Shen, Liang Ding, Dacheng Tao, and Ling- Yu Duan. Fine-tuning global model via data-free knowledge distillation for non-iid federated learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 10174-10183, 2022b.
Xiaojin Zhang, Yan Kang, Kai Chen, Lixin Fan, and Qiang Yang. Trading off privacy, utility and efficiency in federated learning. arXiv preprint arXiv:2209.00233. 2022c.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv:1806.00530 2018.
Zhihan Zhou, Jiangehao Yao, Yan-Feng Wang, Bo Han, and Ya Zhang. Contrastive learning with boosted memorization. In International Conference Machine Learning, pp. 27367-27377. PMLR, 2022.
Zhihan Zhou, Jiangchao Yao, Feng Hong, Ya Zhang, Bo Han, and Yanfeng Wang. Combating rep- resentation learning disparity with geometric harmonization. In Advances in Neural Information Processing Systems, 2023.
Ligeng Zhu, Zhijian Liu, and Song Han. Deep leakage from gradients. Advances in neural infor- mation processing systems, 32, 2019.
Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogeneous federated learning. In International Conference on Machine Learning, pp. 12878-12888 PMLR, 2021.
13

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 DATA GENERATION
A.1.1 DETAILS OF DESIGNING TARGET LABELS FOR GENERATING DATA
One basic generating strategy is uniformly generating data for all possible target labels. Taking CIFAR- 10 as a concrete task example, the target label list can be (.
9
result-
ing 10 X n samples to generate.
Though, there can also be more specific designs for particular tasks such that the generated dataset can better complement the original heterogeneous dataset. Here, we take 5-way classification task as example, where the ultimate objective is training global model that works well fairly across categories while different clients have different preference on categories. Suppose the categorical distribution of client k is d2 [500, O, 400, 200, 400), where the i-th element du,; denotes the number samples that belongs to category i. Comparing with generating equal number of samples across categories, categories with fewer samples, better complementing the original dataset and empowering the client perform better those minor categories.
reasonable and targeted strategy is to generate more samples for those
Specifically, in order to make the complimented dataset more uniform, we generate data according to the following complementary categorical data distribution dę. = marr(d)) - d2, that is, dg. = 500 d2 [0, 500, 100, 300, 100). Then, given the budget of generating samples N, we will allocate N de label i's in the target label list, e.g., generating N x 1000
500 samples for category 1.
dul
A.1.2 APPLICABILITY TO OTHER MODALITIES
Though we primarily focus on vision tasks to verify our idea of tackling data heterogeneity from data perspective, our method can also extend other input modalities. Here, we elaborate this point on two types of modalities.
(1) For modalities such continuous signals and images where the raw inputs can be optimized by gradient-base optimization, the pipeline is exactly the same as illustrated in this paper, since we can optimize the inputs in an end-to-end manner.
(2) For modalities such as text, where the raw inputs are harder to be directly optimized, it can be achieved with slight difference. Since the raw inputs can optimized via end-to-end gradient- based optimization, we do not optimize/generate raw inputs (which is discrete) but optimize/generate the intermediate embeddings (which is continuous). For example, suppose the overall pipeline of the model is: [a] discrete inputs are transformed into continuous embeddings after going through to- kenizing and embedding layers, [b] continuous embeddings go through the downstream task model, [c] the model gives outputs. Then, data generation process is generating the continuous embed- dings at step [b], which are regarded as consensus data.
A.2 THEORETICAL CONVERGENCE ANALYSIS
For simplicity, assume that all the clients share the same generated consensus data (this can be achieved by generating data based on only global model and sharing the same random seed). We denote the global objective function
(0;
Σ 1=1
Pk (O;D) = Σ Pk (0) ( Q))0:)) k=
(4)
where F3 is task loss, Q is KD loss on generated dataset D. and >> Pk = 1. Then, we show four conventional assumptions in FL theoretical literature Wang et al. (2020b).
Assumption A.1 (Smoothness). Each loss function 4 D) is Lipschitz
Assumption A.2 (Bounded Scalar). (0;(D) is bounded below by Pinf-
14

---

Published as a conference paper at ICLR 2024
Assumption A.3 (Unbiased Gradient and Bounded Variance). For each client, the stochastic gradient unbiased: Eg(g)(01E)] = (0,(0; and has bounded variance: Eg gu (0 (0;(0))))11 S o².
Assumption A.4 (Bounded Dissimilarity). For any set of weights {pk 2
0 k= subject to
ΣX 1= 1 Ph = 1, there exists constants 3 > 1 and x² 2 such that Σx=1 [O < ( +) 2
Additionally, we show a lemma that is critical for the derivation of our theoretical analysis.
Lemma A.5 -occeancgg Global Loss).
( ((0(+)) D(t)).
(5)
Proof. We first expand both two sides of equation 5
((0(+11), D(++))= ( = F(g((+1)) + 1(0((+1)) p((+1))
D((+1)),
((0(+))) ((t)) = F(0((+)) + ++((Q(((()))) D(O)) = F(o(+)))) + K KL(O ) D(()),
Since K (((t+))) (((+1=), D((+1)) = 0 and K ( D(0) 2 0, we naturally have:
((0(H))) & 0(0(++)); D()))
(8)
Based on these assumptions and this lemma, can finally prove the following theorem:
Theorem A.6 (Optimization bound of the global objective function), Under these assumptions, if we set I S min 1
(2r(r-1)(2 ), the optimization error will be bounded as follows:
27:
min V+ 2()))
( dissl
+ An 2
Pk 3(T - 11)m2022 Gr(r 12
TrJ
where 17 is learning rate for local model training and T is the number of local model updates.
Theorem A.6 establishes that when the number of communication rounds, T, tends towards infinity (T -> 00), the expectation of the optimization error remains bounded by constant number for fixed 1J; see detailed proof in Appendix [A.2.2] Further, based on Theorem[A.A we can deduce the following corollary when suitable learning rate, denoted as 1), established:
Corollary A.7 (Convergence of the global objective function). By setting V our method can converge to a stationary poins given an infinite number of communication rounds Specifically, the bound could be reformulated as follows:
1
min (1) (D()))
4 (1)) D(
+ Pnnsl
+ 4LO² EK
Ek-iP
3(7 - 1
6(T - 11)T2
rT +
(10)
VT
T
=0( VFT + 0
1
(
+ ) +
1
(11)
This corollary indicates that as T - oc, the error's upper bound approaches 0. Also, given a finite T. there exists a best that minimizes the emorrs upper bound. These analyses show that our method can achieve the same convergence rate as most methods, such as FedAvg et al. (219); Wang et al. 2020b).. Please detailed proof in Appendix A.2.2
15

---

Published as a conference paper at ICLR 2024
A.2.1 PRELIMINARIES
For ease of writing, we use g to denote mini-batch gradient gl (0(E) and (0) denote full-batch gradient. We further define the following two notions:
Averaged Mini-bate Gradient:
d
Σ 9k (0
(12)
7
Averaged Full-batch Gradient:
h
(4)
Σ
V (0!
(13)
r=0
Then, the update of the global model between two rounds is as follows:
K
( -((0 = -TV)
P
(14)
A.2.2 PROOF OF THEOREM
According to the Lipschitz-smonn smooth property in Assumption A.1 al. (2018):
we have its equivalent form Bottou
E ((((+1); [(1)]
((0(0);(())
4 ( D(t)), ((++1) - 00()) }} <
L
1 000011
(15)
2
K
LT217²
K Σ
B
p !!)
+
E
Ped
(1)
(16)
2
A=1
T2
where the expectation taken over mini-batches e((rr) wk: E 1,2, K, E 0, ...,. 1.
For ease of writing, we use ((0(t) and (x(((()) denote and (0(0)) ((t)) when it causes no ambiguity for the rest of paper.
Bounding Ti in (6)
K
- (1 ) +
T
V V0(0 Σ
Ps (d
Σ
ps
h
(t)
(17)
K
E
(2(0
p
(18)
k=
K
/|
+ E
1
h
(t)
1 E
(
p k
(19)
2
2
k=1
where 18 uses the unbiased gradient assumption in Assumption
A.3
such that E(d -h
=
h (0)
(99 uses the fact that 2 (a,b) b) ||a]|
+ ||6||
||a
b1
Bounding T3 in (16)
16

---

Published as a conference paper at ICLR 2024
K
K
T2 = E
(1 (t)
-h(²
+
Ph: h
(20)
A=1
k=1
K
S 2E
Σ
Pk
(2
((()))
+ 2E
Σ pah (8)
(21)
2
E
d (4)
|
+ 2B
Σ Ph: h (e)
(22)
K
2
Σ
LE
(0
(0(
+ 2E
Σ p (0)
(23)
72
K
2
Pk
E
(on(0((
V0 (cr))
+ 2E
Σ P&b
(e)
(24)
72
r=0
20²
K
Pk + 2E
2
Σ
P
(t)
(25)
k=1
A=1
where 21 uses |/a + 6112 < 2 | + 1|b1
22 uses the fact that clients are independent to each
other so that E (a h _ - hm O. vk + n. (44 uses Lemma 2 in Wang et al. (2020b)
and 25 uses bounded variance assumption in Assumption A.3
Plug 19 and
25 back into [6 we have
E
(00++)))
D(o))
=0((());
Tn (1 2TTL)E
Σ h (4)
2
K
+ LTnja 0²
Σ
Pk
B E
Σ
pk (1)
h
(26)
k=1
When 2rTT. 2 0, we have
(1)
220(0))
K
S
LTT 2 0²
Pk
+ TV 2
E
V
Σ h (1)
(27)
2
k=1
K
K
(
LTn 02
2 Pk
+
TM
J
Pk V t ((())) -b(2)
(28)
2
k=1
K
K
<
Tq
2
+ LTm 2 0²
pk
+
T1 2
Pk
(0()) - (1)
(29)
2
A=1
Ta
1
where (9 uses Jensen's Inequality
Σx Pr-
< E Pk |
17

---

Published as a conference paper at ICLR 2024
Bounding T3 in (29)
E
(0(0)
-
V
(
& (0 (0(trr))
(30)
E
(0()))
(0()))
(31)
(0 Von (0 (o()))))
EE
(32)
122
g
00 1
(33)
T4
where 32
uses Jensen's 's Inequality and
33 follows Lipschitz-smooth property.
Bounding Ti in (33)
|//r)) | T/
Σ Ih (e
0
(34)
<2n E
go (((())))) (0(()(())))
+ 2
V V (0ftmt))
(35)
2n Σ E
v0. (1.0) [1]
+21/BB
Σ
V ( (1.0)))
(36)
sm0
<2m1 0² + 2
(0 (1:00)
(37)
=2m1120 2² + +2rmj
E
V V ( ((((.)))
(38)
<2r1/ 0² + 2rm2
E
(0
(39)
where uses la + b1| 2 ||||
+
|bb13
36
uses Lemma in Wang et al. 2020b
(37 uses
the bounded variance assumption in Assumption
38 uses Jensen's 's Inequality.
Plug 33 back into (33) and use this equation
((--1)
have
V ( - -b//
I
E (0()) - 0//.r)/
(40)
0=0
<(r - ( + ( 1)T?,
(0
(41)
T3
where T3 in (11 can be further bounded.
18

---

Published as a conference paper at ICLR 2024
Bounding T3 in
41
V (0 /
E
4 (0 (0()) | 21E
||
22
(42)
|100 / 11
$2L'E
+ 2
(43)
where 42
uses l|a + 6112 - 2| 12 2 1661. (43 uses Lipschitz-smooth property. Plug
43 back
we have
1,2
Σ
E
7
S(T - 1)L337002 + 2(T 1)n7 L
E
+2( - + 0 - 00.0011
V
(010)
==0
(44)
After rearranging, we have
- -n []
E
V (0
[ of(r)) 0 19]
E
(45)
(r 002
27(T 1),
|
(46)
1 2T(T - 1)
1 1 2T(T - 1)
(T 12?
A E
(0(+)))
(47)
1- A
where we define A = 2T(T 1)m12 2 < 1. Then
1 - -h 11(011]
P&L
2 k=1
( 12
A
|1
Pk
E
(
v
(48)
1-
1 A
7T
L?n
Arn
||]
ATRR2 2(1
E
(49)
2(1 - A)
2(1 - A)
A)
where 44
follows bounded dissimilarity assumption in Assumption
A.4
Plug
49
back to
29
we have
() D(t>]] D()))
E
K
s TT
V (010 D(.)) |²
+ LTij² o²
pk
2
1=1
T(r - 1)02
Arr
( 1(0)
ATTJK²
IE
+
(50)
2(1 A)
2(1 A)
2(1 A)
=- T1) (1 1 )
A
(T(0(():(0)) |
+ L
pk
2
7(T 1)0² 1,233 1,2133
ATRI2
2
2(1 A)
2(1 A)
R=1
(51)
19

---

Published as a conference paper at ICLR 2024
A8 2. then 1 I
1 2
and have
E (1)0
((0(+t)) D(())] (( D(O)
s F b[oia L 0²
k
T(T
- 1)02
1
Arnk²
1
Pk
2
+
2
(1 +
232
+
2
(1+ +
2
(52)
K
< V |(0))] LTn o²
Pk
+ 4
3
T( - 2 22 ATMM?
3
(53)
where 53 follows B 2 in Assumption A.4
Then
((0(+)); ((((tt))((tt)) -0000;
B
(0(0(+t1) |(o)) - (((()) D(O))
E
(54)
T!/ ( I 0²
4
Σ
Pk
- - 1)0 (7(--1)02171 123 + ATTR?
3
4
3
(55)
where 54 follows our key lemma in Lemma
A.5
Finally, by taking the average expectation across all rounds, we finish the proof of Theorem A.6
min E (D000)
E
(56)
t
4 (0000000) - Pansl
K
An 2
Σ
ph +
3(T 1)n200222 Gr - 1
(57)
TIT
k=1
Constraints on local learning rate. Here, summarize the constraints on local learning rate ry:
2T 20,
(58)
S1
(59)
2T(T - 1)1
2,32
that is,
1
nL S min
(60)
2r
2 1) +1)
A.2.3 COROLLARY
By setting 7) V+z the above bound can be written
min E V :(
(000.00) DO)
Dins] + 41.22 >
VTT
3(7
1)02 TT
+ 67(T - 1)L22²
+
(61)
7T
VTT - + 0
(62)
A.3 EXPERIMENTAL DETAILS
A.3.1 IMPLEMENTATION DETAILS
Datasets and Heterogeneity. Overall, we consider three classical datasets and one real-world FL
dataset.
20

---

Published as a conference paper at ICLR 2024
Data Distribution Clients
Data Distribution Clients
Distribution over Clients
510
Cilent ID
Client ID
Client
-ax
100
SISARRR 2- 2:
Class Fashion-MNIST
Class
(b) CIFAR-10
(c) CIFAR-100
Figure 4: Data distribution of NIID-1 1 for Fashion-MNIST, CIFAR-10 10 and CIFAR-100
1. Fashion-MNIST (Xiao et al. 2017 is a 10-category classification dataset that consists of 10 types of clothes;
2. CIFAR-I0 and CIFAR-100 (Krizhevsky et al. 2009)
classification datasets that
consists of 10 and 100 types of objects, respectively;
3. FLAIR (Song et al. 2022 is recently released real-world FL. multi-label dataset, where each client is a user from Flickr with multiple images and corresponding tags. We adopt the task of predicting the existence of 17 categories of objects in each image.
For the three classical datasets, we consider two types of data heterogeneity, namely NIID-1 and NIID-2.
NIID-1 follows Dirichlet distribution Dirg (default 8 0.1), which a common setting in FL Yurochkin et al. 2019) Wang et al.) 2020a b Li et al. 2021a al Shi et al., [022] Fan et al. 2024). The argument controls the level of data heterogeneity, where a smaller 3 corresponds to a more heterogeneous level and / is set to 0.1 for default setting. The distribution for these datasets are shown in Figure.
2. NIID-2 makes each client hold data of only partial labels (McMahan et al. 2017 Li et al. 2020b) (2 for Fashion-MNISS and CIFAR-10, 10 for CIFAR-100).
For FLAIR (Song al. 2022). the data heterogeneity is inherent different local dataset is drawn from different user while different user has different preference and thus different data distribution.
Models. 1) For MASII and CIFAR-10, we use a simple convolutional neural network (CNN) that frequently used in FL literature (Li al. 2021a). The data flow in this model is: x 5 convolution, max-pooling, 5 x 5 convolution, three fully-comnected layer with hidden size For CIFAR-100, 100, use slightly different simple CNN. The data flow in this model is: 3 x 3 convolution, max-pooling, 3 x 3 convolution, max-poot max-pooling, 3 x 3 convolution, two fully-connected layer with hidden size of 128 and 100 respectively. 3) For FLAIR, we use ResNet18 (He et al. (0166) from Py Torch API.
of 120, 84 and 10 respectively.
A.3.2 HYPER-PARAMETTE
KNOWLEDGE DISTILLATION FOR IMAGE CLASSIFICATION.
Following the example in A.I.I for classification task, also propose Syyper-prramtte-freee knowledge distillation manner. Specifically, denote the number of samples from the original real dataset as Nreat [dg] and the number samples required complement a balanced dataset Nuwal Ngen = [d we set the coefficient of task-driven loss L. as Nreat and the coefficient of
knowledge-distillate loss Led as NreattN,e. In this case, samples from all categories are super-
Ngen
vised with similar intensity, relieving the issue of imbalanced learning (Menon et al. 2021 Hong
et al. 2023 Zhou et al. 2022 2023) and contributing to enhanced overall performance.
Baselines and hyper-parameters. Here, we list the applied spperr-arameters for the compared baselines.
21

---

Published as a conference paper at ICLR 2024
1. FedAvgM (Hsu et al.
2019): the momentum coefficient is set to 0.1.
2. FedProx (Li et al.
2020b); the model regularization strength is set to 0.01.
3. MOON (Li et al.
a0 the feature regularization strength is set to 0.01.
4. FedSAM (Qu et al. 0.5.
2022): the constant to control the radius of the perturbation p is set to
5. VHL Tang et al.
2022): the coefficient of conditional distribution mismatch penalty is set to 1.0 and the number of epochs for feature alignment is set to 5.
6. FedExP (Jhunjhunwala et al. 2023): the is set to le 1.., 3. Note that in our experiments, the performance of FedExP is unstable. Through careful checking, we find that the experiments
in the original paper of FedExP are conducted by setting local iteration т 20. While in our paper, we are more interested in larger number of iterations (c.g.. T = 400) which is common in FL. We conjecture that such large iteration number causes the model update more divergent and thus leads to unstable performance of FedExP.
7. FedDecorr (Shi et al. 2022): the feature regularization strength is set to 0.01.
A.4 PERSONALIZATION ANALYSIS
Methodology. FedCOG provides a better global model, which further provides a better initializa- tion for local, personalized models. Specifically, with data generation and knowledge distillation, FedCOG achieves a retterpperforee aggregated global model, which subsequently used as ini- tialization for local model training. Such better-performed global model provides a better initial point for local model training and thus benefits the model training (improves the personalization). This phenomenon actually common as FedAvg with Fine-tuning generally a strong baseline in personalized FL (Collins ef al. 2021 2022).
Ye et al. 2023a), which is also explained in (Collins et al.)
Experiments. (1) In Figure c we have shown the advantages of FedCOG in the personalized setting by comparing it with two baselines personalized FL: FedAvg/FedProx with Fine-tuning (local models before aggregation in FedAvg). Both baseline methods are strong and have been intensively verified in many personalized FL papers (Collins al., [22T] Ye al., 2023a] Collins et al. 2022). Through the figure, we see that FedCOG consistently achieves the highest accuracy across clients. (2) Besides, we conduct further evaluations on personalization by comparing with several representative personalized FI. methods, namely, FedAvg with Finetuning, FedProx with Finetuning, FedRep (Collins et al. 202I), and Ditto (Li et al.) 2021b). We conduct experiments on CIFAR-10 NIID-1 and the results are shown in Table 6 From the table, we see that FedCOG outperforms all these strong and representative baselines.
Table 6: Accuracy(%e ()) comparisons with four representative personalized FL methods across CIFAR-10 dataset NIID-I heterogeneity level.
Method FedAvg+FT FedProx+FT
FedRep
Ditto
FedCOG+FT
Acc %
76.16
76.38
74.90
76.57
78.34
A.5 FURTHER ANALYSIS OF FEDCOG
A.5.1 GENERATED DATA ANALYSIS
Here, we provide analysis on the generated data.
Visualization of generated data. We visualize the generated data after various generation steps. CIFAR-I0 consists of 32x32 colour images in 10 classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship and truck. From Figure 5 we see that 1) when the generation step is set to 0, the synthetic data produced is purely random noise. As the number of iterations increases, the generated images start contain more specific patterns related to their corresponding classes. For instance, when generating images of ships and horses, the abstract shape of hull and a horse are included in the output. 2) In this experiment, validate the federated learning framework is privacy preserving.
22

---

Published as a conference paper at ICLR 2024
Table 7:
Effects of generation steps. A moderate size of generation step (e.g.. 100) contributes to higher accuracy.
Generation step 0
50
100
500
1000
Accuracy
60.98
65.40
65.44
64.80
64.92
clients could not recover detailed private data from either client side or server side. The patterns in each image are highly abstract and contain certain category features, but overall, we cannot invert a trained network to generate class-specific input images from random noise.
Effect of generated data. (I) The synthetic data is optimized to be correctly predicted by the global model, thus, it can be seen as a representative of the knowledge of global model. Then, during training,
the synthetic data to distill knowledge from global model to local model, which basically forces the local model to not forget what the global model has already known. This would benefit the updating process of global model and will not be affected even if the synthetic data is not visually the same as real data. (2) We do not directly use the generated data to train the local model (i.e., using cross entropy loss to train local model without guidance from global model), which prevents the local model from overly fitting the generated data that could contain some unrelated properties.
Number of generation steps. Here we experiment on how the number of generation steps affects performance on CIFAR-10 and report the results in Table From the table, we see that 1) generally, a moderate size (c.g.. 100) of generation step contributes to better performance. 2) The performance is degraded when the generation step is too small too large. This is reasonable generation step is small, the image is similar to pure noise that may fail to well complement the original dataset; when the generation step is large, the optimizing process leads to over-fitiing too many details and thus could even bring negative impacts.
when the
Effects of introducing generator. In this experiment, we analyze the effects of introducing learn- able generator into the process of data generation. We insert learnable convolutional neural net- work between the learnable noisy inputs and the taks-specific models. The introduced generator consists of linear layer with hidden size 8192, upsampling and three 3 x 3 convolution lay- ers. After introducing the generator, we keep using the same hyper-parameters as before. We see that introducing a generator brings performance gain to NIID-2 setting (1.5% absolute accuracy im- provement), but degrades performance to NIID-1 setting (0.93% absolute accuracy deterioration). This indicates that using the same set of hyper-parameters, FedCOG without generator and FedCOG with learnable generator perform on par. Such results are reasonable since do not carefully tune hyper-parameter and the generator is training from scratch. However, we believe that with careful hyper-parameter tuning and per-trained generator, the performance can be further improved, though, we leave this to future works.
A.5.2 FEDCOG UNDER RESOURCE HETEROGENEITY
In practice, FL clients may vary in their available resources. Several methods can be employed to fit this scenario, including:
First, we can use small batch size during data generation for resource-constrainee clients, thereby mitigating memory constraints on these clients. By adjusting the batch size to match the limited memory capacity, we enable efficient data generation without necessitating any alterations to the underlying algorithm. An illustrative example from our CIFAR- experiment underscores the feasibility of this approach. Using simple convolutional neural network, the required number of learnable parameters for local model training is 69380. While for instance, when employing a batch size of 1, the number of learnable parameters during local model training dramatically reduces to 3x32x32-3072, thereby rendering it practical to generate data for clients with restricted resources.
Second, another interesting angle is use different model sizes for different clients. In system where different clients have different computing resources, we can tailor model sizes to match indi- vidual capabilities (Dioo al. 2020 Horvath at. 2021). By embracing this dynamic model-sizing
23

---

Published as a conference paper at ICLR 2024
airplane automobile
bird
cat
deer
dog
frog
horse
ship
truck
(a) Original CIFAR-10 dataset
airplane automobile
bird
cat
deer
dog
frog
horse
ship
truck
(b) Generation step:
airplane automobile
bird
cat
deer
dog
frog
horse
ship
truck
(c) Generation step: 100
airplane automobile
bird
cat
deer
dog
frog
horse
ship
truck
(d) Generation step:
airplane automobile
bird
cat
deer
dog
frog
horse
ship
truck
(c) Generation step: 500
airplane automobile
bird
cat
deer
dog
frog
horse
ship
truck
(f) Generation step: 1000
Figure 5: Visualization of generated data and CIFAR-10. The generated data contains some abstract patterns for the corresponding category, but will not reveal private details.
24

---

Published as a conference paper at ICLR 2024
strategy, clients constrained by limited resources harness models with reduced dimensions, such as employing sub-sampled CNNs with fewer channels (Diao et al. 2020
Third, we introduce new avenue for future exploration under our framework; that is, the adoption of variable generator sizes based on the clients' resource profiles. Specifically, resource-constrained clients solely focus learning the input tensors (no generator, done in our paper). While clients with more resources could learn to generate data with the help of a pre-trained generative model. We believe that this approach would be an interesting future direction.
A.5.3 GENERATION TIME V.S. NUMBER OF GENERATED SAMPLES
The process of data generation in FedCOG is efficient as the only learnable parameters are the inputs (which is relatively small compared to model parameters) and the required optimization steps is only 100. To further verify this, we adjust the number of generated samples and record the required optimization time for each client at each round. We vary the number of generated samples in the range of [16, 32, 64, 128, 256, 512, 1024] and report the corresponding generation time in g From the table, we see that generation process is quite efficient. Throughout the ppaer, set the number of generated samples 256, which costs only 0.81 seconds. As reference, the conventional SGD-based local model training takes 14.81 seconds. Also, even when the number of generated samples is 1024, the required generation time is still only 1.27 seconds, indicating that our method is efficient.
Table
Table 8: Required generation time under different numbers of generated samples.
Num.
16
32
64
128
256
512
1024
Time
0.72
0.78
0.79
0.80
0.81
0.83
1.27
A.6 EFFECTS FL ARGUMENTS
In order validate the robustness of method for various federated learning settings, we con- ducted experiments using the CIFAR-10 dataset with NIID- distribution.
Iteration number of local model training. Here evaluate the effect of the number of iteration for local training on CIFAR-I0 with NIID-1 data distribution. The results are shown in Table g From the table, we see that 1) when J is relatively small, the accuracy of all methods rise with the number of local iteration increasing; while when the number of local iteration gets too large (e.g., 800), the accuracy of all approaches drops due the drift of local updates. 2) Nevertheless, our method FedCOG consistently outperforms the other methods in all scenarios.
Table 9: Effects of iteration number of local model training т. FedCOG consistently performs the best for different T.
T
200
400
600
800
FedAvg
61.46
64.68
63.45
62.39
FedProx
61.39
63.90
62.13
61.48
MOON
61.12
63.59
63.27
62.81
FedCOG
61.73
65.61
64.89
64.98
Partial client participation scenarios. For this experiment, we consider K = 50 clients in total and sample partial clients to be available at each round. We compare FedCOG with FedAvg (McMa- han et al. 2017). FedProx (Li et al. 2020b), MOON (Li et al. 202 here and report the results in Table 10 From the table, we see that 1) as the participation rate increases, the performance of available each round. 2) FedCOG consistently out- performs the other methods across all participation rates. Specifically, FedCOG outperforms the second-best method (FedProx (Li et al. 2020b)) by 14% relatively under 0.2 participation rate.
all methods increase since more clients
25

---

Published as a conference paper at ICLR 2024
Table 10: Effects of partial client participation rate. We consider K 50 clients in total and change the sample rate for each round. FedCOG consistently performs the best for different participation rates.
K = 50
0.1
0.2
0.3
0.4
FedAvg
18.64
37.44
41.47
46.73
FedProx
20.42
38.36
42.38
47.42
MOON
20.84
35.55
41.03
46.57
FedCOG
21.36
43.44
48.28
48.37
FASHION-NISS NID-1
CIFAR NIDA NID-4
CIFAR- 10 NIDD2
70
80
60
FedAvg
20
FedAug
Feddug
FodProx
FedProx
20
FedProx
NOON
MOON
MOON
FecCOG
FedCOG
10
FedCOG
10
30
50
70
20
10
50
10
40
50
Communaatinn Round
Communication Round
Communication Round
(a) Bushion-MNIST-NIIE
(b)
(c) CIFAR-IO-NIID
Figure 6: Visualization convergence curves. see that performance saturation after 70 rounds for most methods. Despite only launching 20 rounds of FedCOG (FedAvg for initial 50 rounds), our method still performs the best.
Client number. We investigate the effect of the number of clients for participation for CIFAR-10 see results in Table II Our method consistently outperforms all other approaches across different numbers of clients.
Table 11: Effects of client number K. FedCOG consistently performs the best for different client numbers.
K
5
10
20
30
FedAvg
61.00
61.77
57.51
58.03
FedProx
61.28
61.90
57.08
58.98
MOON
60.64
60.41
56.25
58.35
FedCOG
61.72
61.98
57.67
59.14
A.7 VISUALIZATION OF CONVERGENCE CURVES
Here, we visualize curves different methods, where x-axis denotes the round index and the y-axis denotes the test accuracy of global model. In this experiment, our method is based on intial 50 rounds of FedAvg and rounds of FedCOG to see a clearer improvement brought by FedCOG. From Figure 6 of all methods saturate after rounds of communication. 2) see that our method consistently performs the best.
that I) after 70 rounds of communication, the performance of global models
Please note that this case, we are more interested at performance improvement brought by FedCOG over FedAvg, which is significant through the figures. FedCOG can be further integrated with other methods such as FedProx Li et al. (2020b) and MOON Li et al. (2021a); please refer to Table 3
26


Reviews:
[{"review_id": "FrM9o8so01", "reviewer": "Reviewer_iXGo", "summary": "This paper focuses on federated learning (FL) in the presence of data heterogeneity. Different from the existing methods which usually consider this data heterogeneity as an inherent property and attempt to mitigate the adverse effects, this paper proposes to handle the heterogeneity by generating new data, called FedCOG. There are two key components in FedCOG, including complementary data generation and knowledge-distillation-based model training. It can be plug-and-play, and naturally compatible with standard federated learning protocols.  Extensive experiments on classical and real-world datasets proved the effectiveness of the proposed method.", "questions": "(see above)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is novel, handling the data heterogeneity from the perspective of the data generation, instead of correcting the model.\n2. The proposed method can be a plug-and-play model, and it is naturally compatible with standard FL protocols.", "weaknesses": "1. The motivation for why we need to use data generation, instead of recent popular methods based on model correction, is somewhat not clear. When the training dataset is very large, the proposed method therefore needs to generate a large amount of data in order to achieve alignment, which is costly, then it seems like model correction is a better choice in such a scenario.  \n2. As the paper mentioned, the proposed method FedCOG has two advantages, i.e., plug-and-play and compatibility with standard FL protocols, however, the unique advantages of this data generation method, compared to previous model correction methods, are still ambiguous. Could you please elaborate more regarding them?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "cG2gpTFzko", "reviewer": "Reviewer_wN99", "summary": "This paper proposed FedCOG, a Federated Learning (FL) algorithm that facilitates learning via augmented data generated from the global model, which is later used for knowledge distillation between the global and client models. This scheme is compatible with most existing FL algorithms. Its effects have been empirically verified on real-world datasets.", "questions": "\\- Since tackling data heterogeneity is the key of this paper, I suggest authors conduct more experiments on data with Dirichlet distribution by varying the hyper-parameter $\\beta$. More results with changing heterogeneities would further validate the effects of the proposed methods.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "\\+ This paper tackles a crucial challenge in FL which is Data heterogeneity. Their core idea of data correction for achieving global data consensus is well-motivated.\n\n\\+ Data generation by capturing the residual knowledge between the global and the client model is novel.\n\n\\+ Sensitivity analysis is well designed and conducted.\n\n\\+ This paper is clearly written. Related work is comprehensive.", "weaknesses": "\\- Data generation on the client step brings extra computation workload compared with classic FL or FL with data generation on the server side.\n\n\\- The method seems to be designed purely for vision tasks. Authors could discuss if the proposed method can be extended to scenarios with other input modalities, such as text inputs.\n\n\\- This paper would further benefit from theoretical derivations to interpret why generated data on the client side helps in improving global model performance.\n\n\\- Concerns on Experiments: All methods in Table 1 achieve notably lower accuracies than SOTA FL methods. I am concerned that the model arch, communication round, or optimizer setting is not well set up for appropriate comparison.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "kZ2jlsK9Ui", "reviewer": "Reviewer_SDr1", "summary": "This paper introduced a novel consensus scheme based on data generation to solve data heterogeneity problems in federated learning. It achieved a relatively higher accuracy on four public datasets with different degrees of heterogeneity. In the context of each individual client, the present study implemented a methodology wherein the global mode and local model which is extracted from the previous epoch were frozen. The objective was to train the generated data with the aim of optimizing the disparity between predictions made by the global model and those made by the local model, all the while mitigating any potential impact on the overall accuracy of the global model. All goals are evaluated on the generated dataset. Unlike current works focusing on the model, this paper provides a novel perspective from the local dataset. By enhancing the distribution of the local dataset, it claims to achieve better convergence. It achieved relatively higher accuracy on public datasets (FLAIR, Fashion-MNIST, CIFAR-10, and CIFAR-100) with different degrees of heterogeneity compared with federated learning (FL) algorithms like FedProx.", "questions": "1. This paper introduced data distribution from other clients, will this cause privacy leakage, making it easier for the attacker to learn data information from the clients? \n\n2. Will generating new data for each client be identical to amplifying the global weight update direction collected in the last epoch?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The novel proposed method has a low overhead on each client, making it easy to apply in general FL tasks.\n- Extensive experiments have been done to prove the advantages of the applied method.\n- The paper is well organized and it's easy to follow.", "weaknesses": "- As far as I know, an enormous amount of work has proved that in the vision task, introducing unbalanced label distribution will influence the performance of the global model, and according to the results of the experiments, it's possible that this empirical idea is true. For details, please refer to the detailed comment C1.\n- Evaluation is not strong enough; For details, please refer to the detailed comment C1~C3.\n- No analysis of convergence is provided. For details, please refer to the detailed comment C3.\n\nDetailed comments:\n- C1 In the experiment results, the final accuracy on CIFAR-10 is relatively low, please try some more complicated networks other than the 5-layer CNN. \n- C2 It's possible the network is not converged. To eliminate such a possibility, please provide a graph depicting the trend of convergence with the number of rounds on the server side. \n- C3 What's more, the proposed method only achieved a little improvement in accuracy, it's not sure whether it's caused by insufficient experiments, please repeat and provide mean and standard error for all results.\n- C4 We kindly request further experimentations involving the generation of datasets of varying sizes, with corresponding meticulous documentation of the associated overhead. Furthermore, if feasible, we encourage experimentation on datasets comprising high-resolution images uniformly, to enhance the comprehensiveness of the analysis.\n- C5 Please add proofs for the convergence analysis. If possible, please add a formal security analysis to your method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "bJ53da3Cw7", "reviewer": "Reviewer_3E8Z", "summary": "This paper proposed FedCOG, a synthetic data-assisted federated learning system, to mitigate the data heterogeneity in the training. The design mainly focused on the local training part. In the local training part, FedCOG first generates task-specific and client-specific data, and then uses knowledge distillation to train the local model. The experiment on computer vision benchmark datasets demonstrates that FedCOG performs well compared to existing FL baselines.", "questions": "1. How would the FedCOG be harmonized with FedProx? I am curious about how the FedCOG does the local proximal term in the KD-based model training?\n\n2. In Table 4, why FedProx took longer local training time compared to the FedCOG?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The proposed method is easy to follow. Different from other synthetic-data based methods, the FedCOG proposed task-specific and client-specific data for generation, which is novel and practical.\n\n2. The paper is well-structured. The experiments include several existing FL baselines. The usage of the real-world FL multilabel dataset FLAIR is very rare in the FL literature.", "weaknesses": "1. I am confused about the data generation part of the reading. To my understanding, FedCOG took a learnable parameter for the data generator. What is the structure of the data generator? Does FedCOG update the weight of the data generator during the training as well? Could the author address more about how the data is generated locally?\n\n2. In the experiment part, what are the sample numbers of the synthetic data in your setup? \n\n3. The client number is so limited for the experiment related to standard datasets.\n\n4. I am concerned that none of the selected baselines is a synthetic data-based method. I see the paper cites FedGen in the related work section. Why does the author not compare with the recent synthetic data-based methods such as FedGen[1] and DynaFed[2]?\n\n[1]. Zhu, Zhuangdi et al. “Data-Free Knowledge Distillation for Heterogeneous Federated Learning.” Proceedings of machine learning research 139 (2021): 12878-12889 .\n\n[2]. Pi, Renjie et al. “DYNAFED: Tackling Client Data Heterogeneity with Global Dynamics.” 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) (2022): 12177-12186.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)