Title: Federated Virtual Learning on Heterogeneous Data with Local-global Distillation
Venue: NeurIPS 2023
Accepted: no
Keywords: Federated Learning

Dataset Distillation

Data Heterogeneity
Abstract: Despite Federated Learning (FL)'s trend for learning machine learning models in a distributed manner, it is susceptible to performance drops when training on heterogeneous data. In addition, FL inevitability faces the challenges of synchronization, efficiency and privacy. Recently, dataset distillation has been explored in order to improve the efficiency and scalability of FL by creating a smaller, synthetic dataset that retains the performance of a model trained on the local private datasets. We discover that using distilled local datasets can amplify the heterogeneity issue in FL. To address this, we propose a new method, called Federated Virtual Learning on Heterogeneous Data with Local-Global Distillation (FedLGD) which trains FL using a smaller synthetic dataset (referred as virtual data) created through a combination of local and global dataset distillation. Specifically, to handle synchronization and class imbalance, we propose iterative distribution matching to allow clients to have the same amount of balanced local virtual data; to harmonize the domain shifts, we use federated gradient matching to distill global virtual data that are shared with clients without hindering data privacy to rectify heterogeneous local training via enforcing local-global feature similarity. We experiment on both benchmark and real-world datasets that contain heterogeneous data from different sources, and further scale up to an FL scenario that contains large number of clients with heterogeneous and class imbalance data. Our method outperforms state-of-the-art heterogeneous FL algorithms under various settings with a very limited amount of distilled virtual data.

1

Abstract

Despite Federated Learning (FL)'s trend for learning machine learning models in a distributed manner, it is susceptible to performance drops when training on heterogeneous data. In addition, FL inevitability faces the challenges of synchronization, efficiency, and privacy. Recently, dataset distillation has been explored in order to improve the efficiency and scalability of FL by creating a smaller, synthetic dataset that retains the performance of a model trained on the local private datasets. We discover that using distilled local datasets can amplify the heterogeneity issue in FL. To address this, we propose a new method, called Fed erated Virtual Learning on Heterogeneous Data with L ocalG lobal D istillation (FEDLGD), which trains FL using a smaller synthetic dataset (referred as virtual data ) created through a combination of local and global dataset distillation. Specifically, to handle synchronization and class imbalance, we propose iterative distribution matching to allow clients to have the same amount of balanced local virtual data ; to harmonize the domain shifts, we use federated gradient matching to distill global virtual data that are shared with clients without hindering data privacy to rectify heterogeneous local training via enforcing local-global feature similarity. We experiment on both benchmark and real-world datasets that contain heterogeneous data from different sources, and further scale up to an FL scenario that contains large number of clients with heterogeneous and class imbalance data. Our method outperforms state-of-the-art heterogeneous FL algorithms under various settings with a very limited amount of distilled virtual data.

1 Introduction

Federated Learning (FL) [29] has become a popular solution for different institutions to collaboratively train machine learning models without pooling private data together. Typically, it involves a central server and multiple local clients; then the model is trained via aggregation of local network parameter updates on the server side iteratively. FL is widely accepted in many areas, such as computer vision, natural language processing, and medical image analysis [25, 12, 41].

Onthe one hand, clients with different amounts of data cause asynchronization and affect the efficiency of FL systems. Dataset distillation [39, 5, 46, 44, 45] addresses the issue by only summarizing smaller synthetic datasets from the private local datasets to ensure each client owns the same amount of data. We refer this underexplored strategy as federated virtual learning , as the models are trained from synthetic data [40, 10, 16]. These methods have been found to perform better than modelsynchronization-based FL approaches while requiring fewer server-client interactions.

On the other hand, due to different data collection protocols, data from different clients inevitably 34

face heterogeneity problems with domain shift, which means data may not be independent and 35

identically distributed (iid) among clients. Heterogeneous data distribution among clients becomes a 36

key challenge in FL, as aggregating model parameters from non-iid feature distributions suffers from 37 client drift [18] and diverges the global model update[26]. 38

We observe that using locally distilled datasets can amplify the heterogeneity issue. Figure 1 shows 39 the tSNE plots of two different datasets, USPS [31] and SynthDigits [9], each considered as a client. 40 tSNE takes the original and distilled virtual images as input and embeds them into 2D planes. One 41 can observe that the distribution becomes diverse after distillation. 42

To alleviate the problem of data heterogeneity 43 in classical FL settings, two main orthogonal 44 approaches can be taken. Approach 1 aims 45 to minimize the difference between the local 46 and global model parameters to improve conver47 gence [25, 18, 38]. Approach 2 enforces con48 sistency in local embedded features using an49 chors and regularization loss [37, 47, 42]. The 50 first approach can be easily applied to distilled 51 local datasets, while the second approach has 52 limitations when adapting to federated virtual 53 learning. Specifically, VHL [37] samples global 54 anchors from untrained StyleGAN [19] suffers 55 performance drop when handling amplified het56 erogeneity after dataset distillation. Other meth57 ods, such as those that rely on external global 58 data [47], or feature sharing from clients [42], 59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

are less practical, as they pose greater data privacy risks compared to classical FL settings 1 . Without hindering data privacy, developing strategies following approach 2 for federated virtual learning on heterogeneous data remains open questions on 1) how to set up global anchors for locally distilled datasets and 2) how to select the proper regularization loss(es).

To this end, we propose FEDLGD, a federated virtual learning method with local and global distillation. We propose iterative distribution matching in local distillation by comparing the feature distribution of real and synthetic data using an evolving feature extractor. The local distillation results in smaller sets with balanced class distributions, achieving efficiency and synchronization while avoiding class imbalance. FEDLGD updates the local model on local distilled synthetic datasets (named local virtual data ). We found that training FL with local virtual data can exacerbate heterogeneity in feature space if clients' data has domain shift (Figure. 1). Therefore, unlike previously proposed federated virtual learning methods that rely solely on local distillation [10, 40, 16], we also propose a novel and efficient method, federated gradient matching , that integrated well with FL to distill global virtual data as anchors on the server side. This approach aims to alleviate domain shifts among clients by promoting similarity between local and global features. Note that we only share local model parameters w.r.t. distilled data. Thus, the privacy of local original data is preserved. We conclude our contributions as follows:

This paper focuses on an important but underexplored FL setting in which local models are trained on small distilled datasets, which we refer to as federated virtual learning . We design two effective and efficient dataset distillation methods for FL.

We are the first to reveal that when datasets are distilled from clients' data with domain shift, the heterogeneity problem can be exacerbated in the federated virtual learning setting.

We propose to address the heterogeneity problem by mapping clients to similar features regularized by gradually updated global virtual data using averaged client gradients.

Through comprehensive experiments on benchmark and real-world datasets, we show that FEDLGD outperforms existing state-of-the-art FL algorithms.

2 Related Work 86

2.1 Dataset Distillation 87

Data distillation aims to improve data efficiency by distilling the most essential feature in a large88 scale dataset (e.g., datasets comprising billions of data points) into a certain terse and high-fidelity 89 dataset. For example, Gradient Matching [46] is proposed to make the deep neural network produce 90 similar gradients for both the terse synthetic images and the original large-scale dataset. Besides, 91 [5] proposes matching the model training trajectory between real and synthetic data to guide the 92 update for distillation. Another popular way of conducting data distillation is through Distribution 93 Matching [45]. This strategy instead, attempts to match the distribution of the smaller synthetic 94 dataset with the original large-scale dataset. It significantly improves the distillation efficiency. 95 Moreover, recent studies have justified that data distillation also preserves privacy [7, 4], which is 96 critical in federated learning. In practice, dataset distillation is used in healthcare for medical data 97 sharing for privacy protection [22]. Other modern data distillation strategies can be found here [33]. 98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

2.2 Heterogeneous Federated Learning

FL performance downgrading on non-iid data is a critical challenge. A variety of FL algorithms have been proposed ranging from global aggregation to local optimization to handle this heterogeneous issue. Global aggregation improves the global model exchange process for better unitizing the updated client models to create a powerful server model. FedNova [38] notices an imbalance among different local models caused by different levels of training stage (e.g., certain clients train more epochs than others) and tackles such imbalance by normalizing and scaling the local updates accordingly. Meanwhile, FedAvgM [15] applies the momentum to server model aggregation to stabilize the optimization. Furthermore, there are strategies to refine the server model from learning client models such as FedDF [27] and FedFTG [43]. Local training optimization aims to explore the local objective to tackle the non-iid issue in FL system. FedProx [25] straightly adds L 2 norm to regularize the client model and previous server model. Scaffold [18] adds the variance reduction term to mitigate the 'clients-drift". Also, MOON [24] brings mode-level contrastive learning to maximize the similarity between model representations to stable the local training. There is another line of works [42, 37] proposed to use a global anchor to regularize local training. Global anchor can be either a set of virtual global data or global virtual representations in feature space. However, in [37], the empirical global anchor selection may not be suitable for data from every distribution as they don't update the anchor according to the training datasets.

2.3 Datasets Distillation for FL

Dataset distillation for FL is an emerging topic that has attracted attention due to its benefit for 118 efficient FL systems. It trains model on distilled synthetic datasets, thus we refer it as federated 119 virtual learning. It can help with FL synchronization and improve training efficiency by condensing 120 every client's data into a small set. To the best of our knowledge, there are few published works on 121 distillation in FL. Concurrently with our work, some studies [10, 40, 16] distill datasets locally and 122 share the distilled datasets with other clients/servers. Although privacy is protected against currently 123 existing attack models, we consider sharing local distilled data a dangerous move. Furthermore, none 124 of the existing work has addressed the heterogeneity issue. 125

126

127

128

129

3 Method

In this section, we will describe the problem setup, introduce the key technical contributions and rationale of the design for FEDLGD, and explain the overall training pipeline.

3.1 Setup for Federated Virtual Learning

We start with describing the classical FL setting. Suppose there are N parties who own local datasets 130 ( D 1 , . . . , D N ), and the goal of a classical FL system, such as FedAvg [29], is to train a global model 131

Serve

r

Client N

Client 1

…

Local-global distillation

Local-global model update

! "

!

"

!

E

"

"

!

E

C

Label

Dist.

Match

(c) Local virtual data update

#

##$

$

!%&

$

!%&

If

% &' (

:

!

"

!

'

Global

virtual

E

C

Label

#

$()!

Grad.

Match

∇

*

"

# !

#

+,

! "

!

'

,

*

*

"

$

(a) Global virtual data update

#

+,

! "

-

"

!

E

C

Label

+#

+./

#

+,

+

! "

!

'

E

C

(d) Local model update with virtual data

(b) Global model update

!

!

"

!#$

"

+

1

% &∆!

!

%

!

! "

!

'

*

*

"

$

!

"

!

Repeat:

., 0

∆!

!

%

!

=

!

!

%

!

-!

!&$

%

!

Local

virtual

Local

private

Local virtual

Global virtual

#

+,

If

% &' (

:

Repeat: c

, 1

forward

gradient

update

E

C

feature

extractor

classification

head

server

to client

client

to server

random

image

#

loss

function

Else

:

Repeat:

0

Else:

Repeat:

1

with parameters θ on the distributed datasets ( D ≡ ⋃ i ∈ [ N ] D i )). The objective function is written as: 132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

where L i ( w ) is the empirical loss of client i .

In practice, different clients in FL may have variant amounts of training samples, leading to asynchronized updates. In this work, we focus on a new type of FL training method - federated virtual learning, that trains on distilled datasets for efficiency and synchronization (discussed in Sec.2.3.) Federated virtual learning synthesizes local virtual data ˜ D i for client i for i ∈ [ N ] and form ˜ D ≡ ⋃ i ∈ [ N ] ˜ D i . Typically, | ˜ D i | /lessmuch | D i | and | ˜ D i | = | ˜ D j | . A basic setup for federated virtual learning is to replace D i with ˜ D i in Eq (1), namely FL model is trained on the virtual datasets. As suggested in FedDM [40], the clients should not share gradients w.r.t. the original data for privacy concern.

3.2 Overall Pipeline

The overall pipeline of our proposed method contains three stages, including 1) initialization, 2) iterative local-global distillation, and 3) federated virtual learning. We depict the overview of FEDLGD pipeline in Figure 2. However, FL is inevitability affected by several challenges, including synchronization, efficiency, privacy, and heterogeneity. Specifically, we outline FEDLGD as follows:

We begin with the initialization of the clients' local virtual data ˜ D c by performing initial rounds of distribution matching (DM) [45]. Meanwhile, the server will initialize global virtual data ˜ D g and network parameters θ g 0 . In this stage, we generate the same amount of class-balanced virtual data for each client and server.

Then, we will refine our local and global virtual data using our proposed local-global distillation 150 strategies in Sec. 3.3.1 and 3.3.2. This step is performed for a few selected iterations (e.g. τ = 151 { 0 , 5 , 10 } ) to update θ using L CE (Eq 3), ˜ D g using L Dist (Eq 5), and ˜ D c using L MMD (Eq 2) in early 152 training epochs. For each selected iterations, the server and clients will update their virtual data for a 153 few distillation steps. 154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

and where L CE is the cross-entropy measured on the virtual data, and L Con is the supervised contrastive 197 loss where I is the collection of all indices, A ( i ) indicates all the local and global virtual data indices 198 without i (i.e. A ( i ) ≡ I \{ i } ), z = ψ ( x ) is the output of feature extractor, P ( i ) represents the set of 199

Finally, after refining local and global virtual data ˜ D g and ˜ D c , we continue federated virtual learning in stage 3 on local virtual data ˜ D c using L total (Eq 3), with ˜ D g as regularization anchor to calculate L Con (Eq. 4). We provide implementation details, an algorithm box, and an anonymous link to our code in the Appendix.

3.3 FL with Local-Global Dataset Distillation

3.3.1 Local Data Distillation

Our purpose is to decrease the number of local data to achieve efficient training to meet the following goals. First of all, we hope to synthesize virtual data conditional on class labels to achieve classbalanced virtual datasets. Second, we hope to distill local data that is best suited for the classification task. Last but not least, the process should be efficient due to the limited computational resource locally. To this end, we design Iterative Distribution Matching to fulfill our purpose.

Iterative distribution matching. We aim to gradually improve distillation quality during FL training. To begin with, we split a model into two parts, feature extractor ψ (shown as E in figure 2) and classification head h (shown as C in figure 2). The whole classification model is defined as f θ = h ◦ ψ . The high-level idea of distribution matching can be described as follows. Given a feature extractor ψ : R d → R d ′ , we want to generate ˜ D so that P ψ ( D ) ≈ P ψ ( ˜ D ) where P is the distribution in feature space. To distill local data during FL efficiently that best fits our task, we intend to use the up-to-date server model's feature extractor as our kernel function to distill better virtual data. Since we can't obtain ground truth distribution of local data, we utilize empirical maximum mean discrepancy (MMD) [11] as our loss function for local virtual distillation:

where ψ t and ˜ D c,t are the server feature extractor and local virtual data from the latest global iteration t . Following [46, 45], we apply the differentiable Siamese augmentation on virtual data ˜ D c . K is the total number of classes, and we sum over MMD loss calculated per class k ∈ [ K ] . In such a way, we can generate balanced local virtual data by optimizing the same number of virtual data per class.

Although such an efficient distillation strategy is inspired by DM [45], we highlight the key difference that DM uses randomly initialized deep neural networks to extract features, whereas we use trained FL models with task-specific supervised loss. We believe iterative updating on the clients' data using the up-to-date network parameters can generate better task-specific local virtual data. Our intuition comes from the recent success of the empirical neural tangent kernel for data distribution learning and matching [30, 8]. Especially, the feature extractor of the model trained with FEDLGD could obtain feature information from other clients, which further harmonize the domain shift between clients. We apply DM [45] to the baseline FL methods and demonstrate the effectiveness of our proposed iterative strategy in Sec. 4. Furthermore, note that FEDLGD only requires a few hundreds of local distillations steps using the local model's feature distribution, which is more computationally efficient than other bi-level dataset distillation methods [46, 5].

Harmonizing local heterogeneity with global anchors. Data collected in different sites may have different distributions due to different collecting protocols and populations. Such heterogeneity will degrade the performance of FL. Worse yet, we found increased data heterogeneity among clients when federatively training with distilled local virtual data (see Figure 1). We aim to alleviate the dataset shift by adding a regularization term in feature space to our total loss function for local model updating, which is inspired by [37, 20]:

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

images belonging to the same class y i without data i , and τ temp is a scalar temperature parameter. In such a way, global virtual data can be served for calibration, where z g is from ˜ D g as an anchor, and z p and z a are from ˜ D c . At this point, a critical problem arises: What global virtual data shall we use?

3.3.2 Global Data Distillation

Here, we provide an affirmative solution to the question of generating global virtual data that can be naturally incorporated into FL pipeline. Although distribution-based matching is efficient, local clients may not share their features due to privacy concerns. Therefore, we propose to leverage local clients' averaged gradients to distill global virtual data and utilize it in Eq. (4). We term our global data distillation method as Federated Gradient Matching .

Federated gradient matching. The concept of gradient-based dataset distillation is to minimize the distance between gradients from model parameters trained by original data and distilled data. It is usually considered as a learning-to-learn problem because the procedure consists of model updates and distilled data updates. Zhao et al. [46] studies gradient matching in the centralized setting via bi-level optimization that iteratively optimizes the virtual data and model parameters. However, the implementation in [46] is not appropriate for our specific context because there are two fundamental differences in our settings: 1) for model updating, the gradient-distilled dataset is on the server and will not directly optimize the targeted task; 2) for virtual data update, the 'optimal' model comes from the optimized local model aggregation. These two steps can naturally be embedded in local model updating and global virtual data distillation from the aggregated local gradients. First, we utilize the distance loss L Dist [46] for gradient matching:

where ˜ D c and ˜ D g denote local and global virtual data, /triangleinv θ L ˜ D c CE is the average client gradient. Then, our proposed federated gradient matching optimize as follows:

where θ c i ∗ = arg min θ L i ( ˜ D c ) is the optimal local model weights of client i at a certain round t .

Noting that compared with FedAvg [29], there is no additional client information shared for global distillation. We also note the approach seems similar to the gradient inversion attack [49] but we consider averaged gradients w.r.t. local virtual data, and the method potentially defenses inference attack better (Appendix D.8), which is also implied by [40, 7]. Privacy preservation can be further improved by employing differential privacy [1], but this is not the main focus of our work.

4 Experiment

To evaluate FEDLGD, we consider the FL setting in which clients obtain data from different domains while performing the same task. Specifically, we compare with multiple baselines on benchmark datasets DIGITS (Sec. 4.2), where each client has data from completely different open-sourced datasets. The experiment is designed to show that FEDLGD can effectively mitigate large domain shifts. Additionally, we evaluate the performance of FEDLGD on another benchmark dataset, CIFAR10C [14], which collects data from different corrupts yielding data distribution shift and contains a large number of clients, so that we can investigate varied client sampling in FL. The experiment aims to show FEDLGD's feasibility on large-scale FL environments. We also validate the performance under medical datasets, RETINA , in Appendix. B.

4.1 Training and Evaluation Setup

Model architecture. We conduct the ablation study to explore the effect of different deep neural networks' performance under FEDLGD. Specifically, we adapt ResNet18 [13] and ConvNet [46] in our study. To achieve the optimal performance, we apply the same architecture to perform both the local distillation task and the classification task, as this combination is justified to have the best output [46, 45]. The detailed model architectures are presented in Appendix D.4.

Comparison methods. We compare the performance of downsteam classification tasks using state-of244 the-art (SOTA) FL algorithms, FedAvg [29], FedProx [26], FedNova [38], Scaffold [18], MOON [24], 245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

and VHL [37] 2 . We directly use local virtual data from our initialization stage for FL methods other than ours. We perform classification on client's testing set and report the test accuracies.

FLtraining setup. Weuse the SGD optimizer with a learning rate of 10 -2 for DIGITS and CIFAR10C . If not specified, our default setting for local model update epochs is 1, total update rounds is 100, the batch size for local training is 32, and the number of virtual data update iterations ( | τ | ) is 10. The numbers of default virtual data distillation steps for clients and server are set to 100 and 500, respectively. Since we only have a few clients for DIGITS and RETINA experiments, we will select all the clients for each iteration, while the client selection for CIFAR10C experiments will be specified in Sec. 4.3. The experiments are run on NVIDIA GeForce RTX 3090 Graphics cards with PyTorch.

Proper Initialization for Distillation. We propose to initialize the distilled data using statistics from local data to take care of both privacy concerns and model performance. Specifically, each client calculates the statistics of its own data for each class, denoted as µ c i , σ c i , and then initializes the distillation images per class, x ∼ N ( µ c i , σ c i ) , where c and i represent each client and categorical label. The server only needs to aggregate the statistics and initializes the virtual data as x ∼ N ( µ g i , σ g i ) . In this way, no real data is shared with any participant in the FL system. The comparison results using different initialization methods proposed in previous works [46, 45] can be found in Appendix C.

4.2 DIGITS Experiment

Datasets. We use the following datasets for our benchmark experiments: DIGITS = {MNIST [21], SVHN [31], USPS [17], SynthDigits [9], MNIST-M [9]}. Each dataset in DIGITS contains handwritten, real street and synthetic digit images of 0 , 1 , · · · , 9 . As a result, we have 5 clients in the experiments, and image size is 28 × 28 .

Comparison with baselines under various conditions. To validate the effectiveness of FEDLGD, we first compare it with the alternative FL methods varying on two important factors: Image-per-class (IPC) and different deep neural network architectures (arch). We use IPC ∈ { 10 , 50 } and arch ∈ { ResNet18(R), ConvNet(C)} to examine the performance of SOTA models and FEDLGD using distilled DIGITS . Note that we fix IPC = 10 for global virtual data and vary IPC for local virtual data. Table 1 shows the test accuracies of DIGITS experiments. In addition to testing with original test sets, we also show the unweighted averaged test accuracy. One can observe that for each FL algorithm, ConvNet(C) always has the best performance under all IPCs. The observation is consistent with [45] as more complex architectures may cause over-fitting in training virtual data. It is also shown that using IPC = 50 always outperforms IPC = 10 as expected since more data are available for training. Overall, FEDLGD outperforms other SOTA methods, where on average accuracy, FEDLGD increases the best test accuracy results among the baseline methods of 2.1% (IPC =10, arch = C), 10.4% (IPC

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

(a) Comparison of Reg. loss

(b) Vary

τ

(c) Vary

steps

(d) Vary

steps

|

|

Figure 3: (a) Comparison between different regularization losses and their weighting in total loss ( λ ). One can observe that supervised contrastive loss gives us better and more stable performance with different coefficient choices. (b) The trade-off between | τ | and computation cost. One can observe that the model performance improves with the increasing | τ | , which is a trade-off between computation cost and model performance. Vary data updating steps for (c) DIGITS and (d) CIFAR10C . One can observe that FEDLGD yields consistent performance, and the accuracy tends to improve with an increasing number of local and global steps.

=10, arch = R), 2.2% (IPC = 50, arch = C) and 3.9% (IPC =50, arch = R). VHL [37] is the closest strategy to FEDLGD and achieves the best performance among the baseline methods, indicating that the feature alignment solutions are promising for handling heterogeneity in federated virtual learning. However, VHL is still worse than FEDLGD, and the performance may result from the differences in synthesizing global virtual data. VHL [37] uses untrained StyleGAN [19] to generate global virtual data without further updating. On the contrary, we update our global virtual data during FL training.

4.3 CIFAR10C Experiment

Datasets. We conduct real-world FL experiments on CIFAR10C 3 , where, like previous studies [24], we apply Dirichlet distribution with α = 0 . 5 to generate 3 partitions on each distorted Cifar10-C [14], resulting in 57 clients each with class imbalanced non-IID datasets. In addition, we apply random client selection with ratio = 0.2, 0.5, and 1 and set image size as 28 × 28 .

Comparison with baselines under different client sampling ratios. The objective of the experiment is to test FEDLGD under popular FL questions: class imbalance, large number of clients, different client sample ratios, and data heterogeneity. One benefit of federated virtual learning is that we can easily handle class imbalance by distilling the same number (IPC) of virtual data. We will vary IPC and fix the model architecture to ConvNet since it is validated that ConvNet yields better performance in virtual training [46, 45]. One can observe from Table 2 that FEDLGD consistently achieves the best performance under different IPC and client sampling ratios. We would like to point out that when IPC=10, the performance boosts are significant, which indicates that FEDLGD is well-suited for FL conditions when there is a large group of clients and each of them has a limited number of data.

4.4 Ablation studies for FEDLGD

The success of FEDLGD relies on the novel design of local-global data distillation, where the selection of regularization loss and the number of iterations for data distillation plays a key role. In this section, we study the choice of regularization loss and its weighting ( λ ) in the total loss function. Recall that among the total FL training epochs, we perform local-global distillation on the selected τ iterations , and within each selected iteration , the server and clients will perform data updating

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

label

0

Featureshift

D

8

9

data source

Featureshift

Alignedfeature

cliento

(a)FedAvg

(b)FedLGDwithoutregularization

(c)FedLGD

for some pre-defined steps . The effect of local-global distillation iterations and data updating steps will also be discussed. We also perform additional ablation studies such as computation cost and communication overhead in Appendix C.

Effect of regularization loss. FEDLGD uses supervised contrastive loss L Con as a regularization term to encourage local and global virtual data embedding into a similar feature space. To demonstrate the effectiveness of the regularization term in FEDLGD, we perform ablation studies to replace L Con with an alternative distribution similarity measurement, MMD loss, with different λ 's ranging from 0.1 to 20. Figure 3a shows the average test accuracy. Using supervised contrastive loss gives us better and more stable performance with different coefficient choices.

To explain the effect of our proposed regularization loss on feature representations, we embed the latent features before fully-connected layers to a 2D space using tSNE [28] shown in Figure 4. For the model trained with FedAvg (Figure 4(a)), features from two clients ( × and ◦ ) are closer to their own distribution regardless of the labels (colors). In Figure 4(b), we perform virtual FL training but without the regularization term (Eq. 4). Figure 4(c) shows FEDLGD, and one can observe that data from different clients with the same label are grouped together. This indicates that our regularization with global virtual data is useful for learning homogeneous feature representations.

Analysis of distillation iterations ( | τ | ). Figure 3b shows the improved averaged test accuracy if we increase the number of distillation iterations with FEDLGD. The base accuracy for DIGITS and CIFAR10C are 85.8 and 55.2, respectively. We fix local and global update steps to 100 and 500, and the selected iterations ( τ ) are defined as arithmetic sequences with d = 5 (i.e., τ = { 0 , 5 , ... } ). One can observe that the model performance improves with the increasing | τ | . This is because we obtain better virtual data with more local-global distillation iterations, which is a trade-off between computation cost and model performance. We select | τ | = 10 for efficiency trade-off.

Robustness on virtual data update steps . In Figure 3c and Figure 3d, we fix | τ | = 10 , and vary (local, global) data updating steps. One can observe that FEDLGD yields stable performance, and the accuracy slightly improves with an increasing number of local and global steps. Nevertheless, the results are all the best when comparing with the baselines. It is also worth noting that there is still trade-off between steps and computation cost (See Appendix).

5 Conclusion

In this paper, we introduce a new approach for FL, called FEDLGD. It utilizes virtual data on both client and server sides to train FL models. We are the first to reveal that FL on local virtual data can increase heterogeneity. Furthermore, we propose iterative distribution matching and federated gradient matching to iteratively update local and global virtual data, and apply global virtual regularization to effectively harmonize domain shift. Our experiments on benchmark and real medical datasets show that FEDLGD outperforms current state-of-the-art methods in heterogeneous settings. Furthermore, FEDLGD can be combined with other heterogenous FL methods such as FedProx [26] and Scaffold [18] to further improve its performance. The potential limitation lies in the additional communication and computation cost in data distillation, but we show that the trade-off is acceptable and can be mitigated by decreasing distillation iterations and steps . Our future direction will be investigating privacy-preserving data generation. We believe that this work sheds light on how to effectively mitigate data heterogeneity from a dataset distillation perspective and will inspire future work to enhance FL performance, privacy, and efficiency.

348

References

Abadi, M., Chu, A., Goodfellow, I., McMahan, H.B., Mironov, I., Talwar, K., Zhang, L.: Deep 349 learning with differential privacy. In: Proceedings of the 2016 ACM SIGSAC conference on 350 computer and communications security. pp. 308-318 (2016) 351

Batista, F.J.F., Diaz-Aleman, T., Sigut, J., Alayon, S., Arnay, R., Angel-Pereira, D.: Rim-one dl: 352 A unified retinal image database for assessing glaucoma using deep learning. Image Analysis & 353 Stereology 39 (3), 161-167 (2020) 354

Carlini, N., Chien, S., Nasr, M., Song, S., Terzis, A., Tramer, F.: Membership inference attacks 355 from first principles. In: 2022 IEEE Symposium on Security and Privacy (SP). pp. 1897-1914. 356 IEEE (2022) 357

Carlini, N., Feldman, V., Nasr, M.: No free lunch in" privacy for free: How does dataset 358 condensation help privacy". arXiv preprint arXiv:2209.14987 (2022) 359

360

361

362

Cazenavette, G., Wang, T., Torralba, A., Efros, A.A., Zhu, J.Y.: Dataset distillation by matching training trajectories. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 4750-4759 (2022)

363

364

365

366

367

368

369

370

Diaz-Pinto, A., Morales, S., Naranjo, V., Köhler, T., Mossi, J.M., Navea, A.: Cnns for automatic glaucoma assessment using fundus images: an extensive validation. Biomedical engineering online 18 (1), 1-19 (2019)

Dong, T., Zhao, B., Lyu, L.: Privacy for free: How does dataset condensation help privacy? arXiv preprint arXiv:2206.00240 (2022)

Franceschi, J.Y., De Bézenac, E., Ayed, I., Chen, M., Lamprier, S., Gallinari, P.: A neural tangent kernel perspective of gans. In: International Conference on Machine Learning. pp. 6660-6704. PMLR (2022)

Ganin, Y., Lempitsky, V.: Unsupervised domain adaptation by backpropagation. In: International 371 conference on machine learning. pp. 1180-1189. PMLR (2015) 372

Goetz, J., Tewari, A.: Federated learning via synthetic data. arXiv preprint arXiv:2008.04489 373 (2020) 374

Gretton, A., Borgwardt, K.M., Rasch, M.J., Schölkopf, B., Smola, A.: A kernel two-sample test. 375 The Journal of Machine Learning Research 13 (1), 723-773 (2012) 376

Hard, A., Rao, K., Mathews, R., Ramaswamy, S., Beaufays, F., Augenstein, S., Eichner, H., 377 Kiddon, C., Ramage, D.: Federated learning for mobile keyboard prediction. arXiv preprint 378 arXiv:1811.03604 (2018) 379

He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: Pro380 ceedings of the IEEE conference on computer vision and pattern recognition. pp. 770-778 381 (2016) 382

Hendrycks, D., Dietterich, T.: Benchmarking neural network robustness to common corruptions 383 and perturbations. arXiv preprint arXiv:1903.12261 (2019) 384

Hsu, T.M.H., Qi, H., Brown, M.: Measuring the effects of non-identical data distribution for 385 federated visual classification. arXiv preprint arXiv:1909.06335 (2019) 386

Hu, S., Goetz, J., Malik, K., Zhan, H., Liu, Z., Liu, Y.: Fedsynth: Gradient compression via 387 synthetic data in federated learning. arXiv preprint arXiv:2204.01273 (2022) 388

Hull, J.J.: A database for handwritten text recognition research. IEEE Transactions on pattern 389 analysis and machine intelligence 16 (5), 550-554 (1994) 390

Karimireddy, S.P., Kale, S., Mohri, M., Reddi, S., Stich, S., Suresh, A.T.: Scaffold: Stochastic 391 controlled averaging for federated learning. In: International Conference on Machine Learning. 392 pp. 5132-5143. PMLR (2020) 393

Karras, T., Laine, S., Aila, T.: A style-based generator architecture for generative adversar394 ial networks. In: Proceedings of the IEEE/CVF conference on computer vision and pattern 395 recognition. pp. 4401-4410 (2019) 396

Khosla, P., Teterwak, P., Wang, C., Sarna, A., Tian, Y., Isola, P., Maschinot, A., Liu, C., 397 Krishnan, D.: Supervised contrastive learning. Advances in Neural Information Processing 398 Systems 33 , 18661-18673 (2020) 399

LeCun, Y., Bottou, L., Bengio, Y., Haffner, P.: Gradient-based learning applied to document 400 recognition. Proceedings of the IEEE 86 (11), 2278-2324 (1998) 401

402

403

Li, G., Togo, R., Ogawa, T., Haseyama, M.: Dataset distillation for medical dataset sharing. arXiv preprint arXiv:2209.14603 (2022)

Li, Q., Diao, Y., Chen, Q., He, B.: Federated learning on non-iid data silos: An experimental 404 study. arXiv preprint arXiv:2102.02079 (2021) 405

406

407

Li, Q., He, B., Song, D.: Model-contrastive federated learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 10713-10722 (2021)

408

409

410

411

412

413

414

415

Li, T., Sahu, A.K., Zaheer, M., Sanjabi, M., Talwalkar, A., Smith, V.: Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems 2 , 429-450 (2020)

Li, X., Huang, K., Yang, W., Wang, S., Zhang, Z.: On the convergence of fedavg on non-iid data. International Conference on Learning Representations (2020)

Lin, T., Kong, L., Stich, S.U., Jaggi, M.: Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems 33 , 2351-2363 (2020)

Van der Maaten, L., Hinton, G.: Visualizing data using t-sne. Journal of machine learning research 9 (11) (2008)

McMahan, B., Moore, E., Ramage, D., Hampson, S., y Arcas, B.A.: Communication-efficient 416 learning of deep networks from decentralized data. In: Artificial intelligence and statistics. pp. 417 1273-1282. PMLR (2017) 418

Mohamadi, M.A., Sutherland, D.J.: A fast, well-founded approximation to the empirical neural 419 tangent kernel. arXiv preprint arXiv:2206.12543 (2022) 420

421

422

Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu, B., Ng, A.Y.: Reading digits in natural images with unsupervised feature learning (2011)

Orlando, J.I., Fu, H., Breda, J.B., van Keer, K., Bathula, D.R., Diaz-Pinto, A., Fang, R., Heng, 423 P.A., Kim, J., Lee, J., et al.: Refuge challenge: A unified framework for evaluating automated 424 methods for glaucoma assessment from fundus photographs. Medical image analysis 59 , 101570 425 (2020) 426

Sachdeva, N., McAuley, J.: Data distillation: A survey. arXiv preprint arXiv:2301.04272 (2023) 427

428

429

Schuster, A.K., Erb, C., Hoffmann, E.M., Dietlein, T., Pfeiffer, N.: The diagnosis and treatment of glaucoma. Deutsches Ärzteblatt International 117 (13), 225 (2020)

430

431

432

433

434

435

Shokri, R., Stronati, M., Song, C., Shmatikov, V.: Membership inference attacks against machine learning models. In: 2017 IEEE symposium on security and privacy (SP). pp. 3-18. IEEE (2017)

Sivaswamy, J., Krishnadas, S., Joshi, G.D., Jain, M., Tabish, A.U.S.: Drishti-gs: Retinal image dataset for optic nerve head (onh) segmentation. In: 2014 IEEE 11th international symposium on biomedical imaging (ISBI). pp. 53-56. IEEE (2014)

Tang, Z., Zhang, Y., Shi, S., He, X., Han, B., Chu, X.: Virtual homogeneity learning: Defending 436 against data heterogeneity in federated learning. arXiv preprint arXiv:2206.02465 (2022) 437

Wang, J., Liu, Q., Liang, H., Joshi, G., Poor, H.V .: Tackling the objective inconsistency problem 438 in heterogeneous federated optimization. Advances in neural information processing systems 439 33 , 7611-7623 (2020) 440

Wang, T., Zhu, J.Y., Torralba, A., Efros, A.A.: Dataset distillation. arXiv preprint 441 arXiv:1811.10959 (2018) 442

Xiong, Y., Wang, R., Cheng, M., Yu, F., Hsieh, C.J.: Feddm: Iterative distribution matching for 443 communication-efficient federated learning. arXiv preprint arXiv:2207.09653 (2022) 444

Xu, J., Glicksberg, B.S., Su, C., Walker, P., Bian, J., Wang, F.: Federated learning for healthcare 445 informatics. Journal of Healthcare Informatics Research 5 , 1-19 (2021) 446

Ye, R., Ni, Z., Xu, C., Wang, J., Chen, S., Eldar, Y .C.: Fedfm: Anchor-based feature matching 447 for data heterogeneity in federated learning. arXiv preprint arXiv:2210.07615 (2022) 448

Zhang, L., Shen, L., Ding, L., Tao, D., Duan, L.Y .: Fine-tuning global model via data-free knowl449 edge distillation for non-iid federated learning. In: Proceedings of the IEEE/CVF Conference 450 on Computer Vision and Pattern Recognition. pp. 10174-10183 (2022) 451

Zhao, B., Bilen, H.: Dataset condensation with differentiable siamese augmentation. In: Inter452 national Conference on Machine Learning. pp. 12674-12685. PMLR (2021) 453

Zhao, B., Bilen, H.: Dataset condensation with distribution matching. In: Proceedings of the 454 IEEE/CVF Winter Conference on Applications of Computer Vision. pp. 6514-6523 (2023) 455

Zhao, B., Mopuri, K.R., Bilen, H.: Dataset condensation with gradient matching. ICLR 1 (2), 3 456 (2021) 457

Zhou, T., Zhang, J., Tsang, D.: Fedfa: Federated learning with feature anchors to align feature 458 and classifier for heterogeneous data. arXiv preprint arXiv:2211.09299 (2022) 459

Zhu, H., Xu, J., Liu, S., Jin, Y .: Federated learning on non-iid data: A survey. Neurocomputing 460 465 , 371-390 (2021) 461

Zhu, L., Liu, Z., Han, S.: Deep leakage from gradients. Advances in neural information 462 processing systems 32 (2019) 463

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method called FedLGD that utilizes distilled virtual data on both clients and the server to train FL models. To address the synchronization issue and class imbalance, the authors use iterative distribution matching to distill the same amount of local virtual data on the clients for local model training, thereby improving the efficiency and scalability of FL. The authors also reveal that training on local virtual data exacerbates the heterogeneity issue in FL. To address this problem, they use federated gradient matching to distill global data on the server and add a regularization term to the local loss function to promote the similarity between local and global features. They evaluate the proposed FedLGD method on benchmark and real datasets and show that FedLGD outperforms existing heterogeneous FL algorithms.\n\nSTRENGTHS:\n1. The authors use visualization to reveal the limitation of local data distillation in federated virtual learning, which makes the motivation of the proposed method clear.\n2. The proposed method preserves local data privacy by leveraging averaged local gradients to distil global virtual data.\n3. The experiment results validate performance improvement and privacy protection.\n\nWEAKNESSES:\n1. The initialization for data distillation requires each client to calculate the data statistics and the server to aggregate these statistics, which still raises privacy concerns since the statistics contain some private information. How about using random initialization or other strategies? The authors need to justify it. \n2. Compared with VHL that uses untrained StyleGAN without further updates, the proposed FedLGD method needs to update the global virtual data iteratively.  Therefore, it is not surprising that FedLGD outperforms VHL. If the StyleGAN can be updated the same number of times as FedLGD, does FedLGD still outperform it? This requires justification or experimental validation.\n3. The structure of the proposed method is not clear enough, which makes it difficult to follow. The authors first present the overall pipeline and then describe each component. However, the connection between the components and the overall pipeline is not clear. This requires significant revision.\n4. The presentation quality is not satisfactory. There are too many typos and grammatical errors. Some notations are unclear, e.g., $i$ represents both the data index and client index; the subscript $t$ disappears in many places; $\\tau$ is a set but denoted as a scalar in the caption of Figure 2.\n\nQUESTIONS:\n1. For the selected rounds, is the regularization term $L_{con}$ needed for the local model update? As shown in Algorithm 1, the term disappears in these rounds.\n2. In Table 1, why is VHL better than the proposed method on SVHN with IPC=10?\n\nLIMITATIONS:\n1. As pointed out by the authors, data distillation incurs additional communication and computation cost. Further investigation is required to enhance the efficiency.\n2. The proposed method performs well but lacks of theoretical analysis to support its performance improvement.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a method to address data heterogeneity from the perspective of dataset distillation, named FedLGD. Specifically, the proposed iterative distribution matching and federated gradient matching strategies are used to iteratively update the local balanced data and the global shared virtual data, and the global virtual regularization is applied to coordinate the data domain drift between clients effectively. This method can effectively solve the problem of client data imbalance and domain drift in heterogeneous data scenarios. Extensive experimental results show the effectiveness of the proposed method.\n\nSTRENGTHS:\n1. The problem of data heterogeneity studied in this paper is important for applying federated learning in real-world scenarios.\n2. The idea of this paper to solve the problem of data heterogeneity in federated learning through the dataset distillation method is novel.\n3. The authors perform various experiments to analyze the proposed method\n\nWEAKNESSES:\n1. Some symbols are written differently. The authors should unify these symbols. In section 3.1, local virtual data is written as $\\widetilde{D}_{I}$, but in section 3.2 and 3.3 is written as $\\widetilde^{D}{c}$, in Figure 2 is $\\widetilde{D}_{t}^{c_{I}}$. In Eq. (3), $L_ {Con} $is about \\ widetilde ^ {D} {g} and \\ widetilde {D} ^ {c} function. But they do not appear in Eq. (4). The authors should give more details about Eq. (3) and (4). \n2. other approaches to address heterogeneity are personalized federated learning, e.g. FedAMP[1] Ditto[2] KT-pFL[3], etc. However, it is not mentioned in related work, and there is no comparison in experimental methods.\n3. The legend and curve in Figure 3a do not match. In Table 1, ResNet18 generally performs worse than CNN. The authors mention that overfitting may be happening (at line 275). The authors should increase the dataset size or use a smaller model to make the results more convincing.\n4. According to the results in Figure 4, the visualization results of FedLGD do not define the boundaries of each class well. Although FIG. 4 can prove that FedLGD solves the data drift of both clients, the degree of clustering of each class looks reduced. The authors should analyze it further. \n\n\n[1] Huang, Yutao, et al. \"Personalized cross-silo federated learning on non-iid data.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 35. No. 9. 2021.\n[2] Li, Tian, et al. \"Ditto: Fair and robust federated learning through personalization.\" International Conference on Machine Learning. PMLR, 2021.\n[3] Zhang, Jie, et al. \"Parameterized knowledge transfer for personalized federated learning.\" Advances in Neural Information Processing Systems 34 (2021): 10092-10104.\n\nQUESTIONS:\n1. Can the proposed approach be generalized to other data types such as text and corresponding text models? \n2. Can the proposed method be generalized to more client scenarios which is more practical in reality? Is the proposed method theoretically guaranteed?\n\nLIMITATIONS:\nThe authors discuss the limitations in section 5.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
To solve the challenges of synchronization, efficiency, and privacy, this paper presents a local-global distillation mechanism for FL (FedLGD). In FedLGD, an iterative distribution matching scheme is proposed to distill global virtual data to alleviate the heterogeneous problem. Experiments have shown superiority of FedLGD compared with existing FL methods.

STRENGTHS:
1. The whole pipeline of FedLGD is well depicted in Figure 2. Each component involved in the pipeline is carefully designed.

2. It is an interesting idea to solve the existing FL challenges from the virtual learning perspective. This can inspire future studies in this direction. 

3. Experimental results look solid with sufficient implementation details.

WEAKNESSES:
1. It seems that only feature heterogeneity is considered in this work. How the proposed method performs on different heterogeneous cases should be discussed. 

2. The definition of small distilled dataset is not very clear, which can affect the readers’ understanding towards the motivation and detailed technical parts. 

3. Privacy concern. Since there are image-level data transferred between the server and clients, it is better to further discuss the potential privacy-preserving risks.

QUESTIONS:
What is the detailed setting of Figure 1? Is it reasonable to use the size of the overlapped area to represent the level of heterogeneity? Is there any theoretical basis? Also, will the results be different when using different synthetic schemes?

LIMITATIONS:
Yes. The authors have addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-07

comments: SUMMARY:
This paper introduces an approach on Federated learning using dataset distillation techniques

STRENGTHS:
1. The idea of using dataset distillation for FL is interesting
2. The solution is reasonable 
3. The experimental results show the effectiveness of the proposed approach

WEAKNESSES:
1. In a few equations, the details is not provided. For instance $L_CE$ in 3, $Dist$ in 5. The paper should be self-contained 
2. The technical contribution is low
3. In the experimental results, Table 1, can you highligh both first and second place? In MNIST-M, the winner should be VHL/R, 85.7> FedLGD 85.2.

QUESTIONS:
1. Can you try to provide details on those equations, see above?
2. I didn't see how the sizes of different clients play a role in the global data distillation. Can you provide some details?
    Also, what is the overall objective functions for both local training and global training? The current equations  are not very clear on this.

LIMITATIONS:
The limitation on using virtual data should be discussed. Is there any drawbacks on using fake data instead of real data?
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-21