PRISM: Privacy-Preserving Improved Stochastic Masking for Federated Generative Models
Authors: Kyeongkook Seo
Dong-Jun Han
Jaejun Yoo
Venue: ICLR 2025 Poster (2025)
Keywords: Generative models
Federated learning
Primary area: generative models
TL;DR: PRISM is a new federated framework for generative models that addresses the challenges in communication efficiency, privacy, stability, and generation performance altogether, yielding a lightweight generator.

Abstract: Despite recent advancements in federated learning (FL), the integration of generative models into FL has been limited due to challenges such as high communication costs and unstable training in heterogeneous data environments. To address these issues, we propose PRISM, a FL framework tailored for generative models that ensures (i) stable performance in heterogeneous data distributions and (ii) resource efficiency in terms of communication cost and final model size. The key of our method is to search for an optimal stochastic binary mask for a random network rather than updating the model weights, identifying a sparse subnetwork with high generative performance; i.e., a ``strong lottery ticket''. By communicating binary masks in a stochastic manner, PRISM minimizes communication overhead. This approach, combined with the utilization of maximum mean discrepancy (MMD) loss and a mask-aware dynamic moving average aggregation method (MADA) on the server side, facilitates stable and strong generative capabilities by mitigating local divergence in FL scenarios. Moreover, thanks to its sparsifying characteristic, PRISM yields a lightweight model without extra pruning or quantization, making it ideal for environments such as edge devices. Experiments on MNIST, FMNIST, CelebA, and CIFAR10 demonstrate that PRISM outperforms existing methods, while maintaining privacy with minimal communication costs. PRISM is the first to successfully generate images under challenging non-IID and privacy-preserving FL environments on complex datasets, where previous methods have struggled.

Published as a conference paper at ICLR 2025
PRISM: PRIVACY-PRESERVING IMPROVED STOCHASTIC MASKING FOR FEDERATED GENERATIVE MODELS
Kyeongkook Seo', Dong-Jun Han"*, Jaejun You
1 Ulsan National Institute of Science and Technology (UNIST)
2 Department of Computer Science and Engineering, Yonsei University (kyeogkooksses, jae y @unist. .t..a.kk..k... djheyonsei .cc .cc.rr
ABSTRACT
Despite recent advancements federated learning (FL), the integration gener- ative models into FL. has been limited due challenges such high communi- cation costs and unstable training in heterogeneous data environments. To address these issues, we propose PRISM, FL framework tailored for generative models that ensures (i) stable performance in heterogeneous data distributions and (ii) re- source efficiency in terms of communication cost and final model size. The key of our method is to search for an optimal stochastic binary mask for a random network rather than updating the model weights, identifying sparse subnetwork with high generative performance; i.e., "strong lottery ticket".. By communi- cating binary masks in a stochastic manner, PRISM minimizes communication overhead. Combined with the utilization of maximum mean discrepancy (MMD) loss and a mask-aware dynamic moving average aggregation method (MADA) the server side, PRISM facilitates stable and strong generative capabilities by mitigating local divergence in FL scenarios. Moreover, thanks to its sparsifying characteristic, PRISM yields an lightweight model without extra pruning quan- tization, making it ideal for environments such as edge devices. Experiments on MNIST, FMNIST, CelebA, and CIFAR10 demonstrate that PRISM outperforms existing methods, while maintaining privacy with minimal communication costs. PRISM the first to successfully generate images under challenging non-IID and srivacy-preservingg FI. environments on complex datasets, where previous meth- ods have struggled. Our code available at atISM
INTRODUCTION
Recent generative models have demonstrated remarkable advancements in image quality and have been widely extended to various domains, including image-to-image translation (Choi et al., 20020] Saharia al. 2022)] layout generation (Seol al. 2024)]. text-to-image generation (Rombach
et al., 2022 Ramesh et al. 2.22)1. and video generation (Skorokhodov et al.] 2022 Kim et al.]
) Achieving high-quality generation with current generative models demands increasingly
large datasets, leading to concerns that publicly available data will soon exhausted (Villalobos et al.) 2024]]. Leveraging the vast amount of data stored on edge devices becomes a potential solution, but this poses significant challenges: not only does the private nature of the data make centralized train- ing impractical, but edge computing itself faces hurdles, including limited resources and prohibitive communication costs.
Federated learning (FL) (McMahan et al. 2017 is a promising paradigm tailored to this setup, enabling clients to collaboratively train global model without sharing their local datasets with third party. However, high communication costs, performance degradation due to data heterogeneity, and the need preserve privacy remain significant challenges in FL. These challenges are further intensified in the context of generative models. Unlike classification tasks, generative tasks lack well-defined objective function and focus on learning the data sample distribution, making the inte-
"Co-corresponding author.

---

Published as a conference paper at ICLR 2025
gration of FL and generative models even more difficult. A few recent works have made efforts to train generative models over distributed clients (Hardy et al.) 2019 Rasouli et al.) 2020 Li et al,
2022 Zhang et al., [021 Amalan et al. ] These methods
generally built upon generative adversarial networks (GANs) (Goodfellow et al., 2020)], which have shown impressive results in
the field of image generation. Fe- (Augenstein et al. (019)), GS-WGAN (Chen al. 2020 and Private-FGGAN (Xin et al. [ apply differential privacy (DP) (Dwork et al. 2006
Mironov 2017 mitigate the potential privacy risk in FL setups. However, existing works still
face several challenges: 1) Previous approaches (Farnia & Ozdaglar, 2020a b Wang et al., (022
underperform, especially in non-IlI (independent, identically distributed) data distribution scenar-
ios with strong data heterogeneity across FL clients due to the notorious instability of GANs. 2) Performance evaluations are limited to relatively simple datasets such as MNIST, Fashion MNIST. and EMNIST. 3) They suffer from significant communication overhead during model exchanges between the server and clients.
To overcome these challenges, we propose PRivacy-preserving Improved Stochastic-Masking for generative models (PRISM), a new strategy for training generative models in FL settings with the following key features: First, at the heart of PRISM is the strong lottery ticket (SLT) hypothe- sis (Frankle & Carbin| 2018)], suggesting the existence of a highly effective subnetwork within a randomly initialized network. PRISM shifts the focus towards identifying an optimal global binary mask, rather than updating the weights directly. This approach enables each client to transfer the binary mask the server instead of the full model, significantly reducing the overload in each com- munication round. Moreover, when training is finished. PRISM produces a lightweight final model, as each weight is already quantized, thanks to our initialization strategy. This feature provides signif- icant advantages for resource- resource-constrained edge devices. Second, PRISM incorporates the maximum mean discrepancy (MMD) loss (Gretton et al. 2006 2012)] during client-side updates, ensuring stable training for generative models. Third, a mask-aware dynamic moving average aggregation (MADA) introduced to alleviate local model divergence. This allows PRISM to maintain the previous global mask information and alleviate client drift under non- and DP oo-guarant sce- narios. By automatically adjusting the moving average parameter based on mask correlations, this approach requires neither regularization term nor hyperparameter tuning. Finally, PRISM offers a hybrid strategy that can flexibly trade-off between image quality and communication cost. Taken together, these features enable PRISM to consistently deliver robust performance in challenging non-lID and DP-guaranteeing FL settings, while maintaining minimal communication overhead.
Our experimental results reveal that PRISM sets a new standard in generative model performance, significantly outperforming GAN-based methods in both IID and non-IID scenarios. It acheives oftte-of-thh-arr image generation on complex datasets such as CelebA and CIFAR1O. whereas pre- vious methods were limited to simpler datasets like MNIST and FMNIST. This highlights PRISM's potential for scalable and resource-efficiet generative model learning in distributed environments. Overall, our main contributions can be summarized as follows:
We propose PRISM, an effective FL framework that achieves state-of-the-arr performance on various benchmark datasets. It is the first method to successfully generate images on complex datasets such as CelebA in FL scenarios that involves data heterogeneity and privacy preservation.
PRISM offers an efficient solution for federated generative models with minimal com- munication overhead by incorporating SLT with a stochastic binary mask. Even more, in conjunction with the weight initialization strategy, the final model acquired from PRISM becomes significantly lightweight, reducing to less than half the size of the initial model.
We further enhance the stability of federated learning for generative models by introducing MMD loss and a mask-aware dynamic moving average aggregation method (MADA).
To the best of our knowledge, this is the first work address the challenges in communication efficiency, privacy, stability, and generation performance altogether for federated generative models. We provide new directions to this area based on several unique characteristics, including SLT with stochastic binary mask, MMD loss, awsk- dynamic moving average aggregation strategy, and hybrid score/mask communications.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
Federated learning for classification models. FL has achieved significant success in training global model in a distributed setup, eliminating the necessity of sharing individual client's 's local datasets with either the server other clients. Research has been conducted various aspects of
FL, such as data heterogeneity (Zhao et al.) 2018 Li et al. [02[b)] communication efficiency (Isik)
ct al., 2022; Li et al., 2021a| Mitchell al., 2022 Basat et al. ]. privacy (Wei et al. ].
with most locusing on image classification tasks. Related to our approach, FedPM (1sik et al. (022
and FedMask (Li al., 2021aa] adopted binary mask communication to reduce the communication
costs in FL in classification tasks. FedMask (Li et al. [02 introduces binary mask commu- nication, focusing on communication efficiency and personalization in decentralized environments, while FedPM (Isik et al. [022)] utilizes stochastic masks to minimize uplink overhead and proposes a bayesian aggregation method to robustly manage scenarios with partial client participation. While the concept of incorporating SLT into FL paradigm has been studied in FedMask (Li et al. [02]aa]] and FedPM (Isik et al.] 2022)] for classification tasks, PRISM stands as an independent strategy tailored for the training of generative models across distributed clients: PRISM incorporates MMD loss for more robust performance compared to GAN-based approaches and introduces MADA to
maintain a stable image generation performance in heterogeneous and -p-gaaranteeing FL settings.
Federated learning for generative models. Generative models, like classification models, also offer potential benefits in various federated settings (See Appendix PP. Several recent works have aimed
to incorporate generative models into distributed settings (Hardy al. 2019 Amalan et al. 2022]
let al. 2022 Zhang et al., 2021; Rasouli et al., 2020: Augenstein et al. 2019 Chen et al. [020 Xin
al.
2020) MD-GAN (aady al. [ was the first attempt to apply generative models the
FL framework using GANs (Goodfellow et al. [020)]] extensively studied in image generation tasks.
In MD-GAN, each client holds discriminator, and the server aggregates the discriminator feedback from each client to train the global generator. To prevent overfitting of local discriminators, clients exchange discriminators, incurring additional communication costs. Multi-FLGAN (Amalan et al.) 2 proposed all vs. all game approach by employing multiple generators and multiple discrimi- nators and then selecting the most powerful network to enhance model performance. IFL-GAN (Li et al. ]022]] improves both performance and stability by weighting each client's 's feedback based on the MMD between the images generated by the global model and the local generator. This approach maintains balance between the generator and the discriminator, leading to Nash Equilibrium. Other works such as (Zhang et al. 2021 I |asouli et al., 2020p] have also explored the utilization of GANs in FL.. However, these works do not consider the challenge of privacy preservation in the context of FL and also suffer from resource issues during training and inference.
Federated learning for generative models with privacy consideration. Only a few prior works have focused on the privacy preservation in federated generative models (Augenstein et al. 2019
Chen al. 2020 DP-FedAvgGAN (Augenstein et al., [019)] introduces to combine federated
generative models and differential privacy (DP) Dwork et al. 2006 Mironov] 2017]] to ensure pri- vacy preservation. GS-WGAN (Chen et al., 20200] adopts Wasserstein GAN (Gulrajani et al.) [017)]
to bypass the cumbersome searching for an appropriate valuu, leveraging the Lipschitz property. While these approaches have successfully integrated FL. and generative models, they inherit draw- backs such as the notorious instability of GANs (Farnia & Ozdaglar) [2020a] bj Wang et al., [022)] and significant performance drop under data heterogeneity. Moreover, all existing approaches suffer from huge communication and storage costs during and after training, respectively.
3 BACKGROUND
3.1 STRONG LOTTERY TICKETS
Strong Lottery Ticket (SLT) hypothesis (Frankle & Carbin 2018 Malach et al. 2020 Orseau et al., 2020j] suggest the existence of sparse subnetwork within an initially random network that achieves superior performance. Pdgup (EP) algorithm (Ramanujan et al. 2020 is one of the most popular methods to discover SLT within the dense network, which introduces a scoring mechanism to select potentially important weights among the widespread initialized weight values. More specif- ically, given a randomly initialized dense network W Wintt a learnable score is trained while keeping the weight values frozen. These scores are designed to encapsulate the importance of each weight for
3

---

Published as a conference paper at ICLR 2025
(4) S.+1 Sigmoid ' (0s+1)
Local score update
(2-1)
Neural network W
Clients
Score s
(c, 8)-dp
Mask M
Winiz frozen
P (1 1)6 x
2
VinanD
Sigmoid
Client
Bern
V VIuD
Sigmoid
Probability Mask
Client 2
(2-2)
Bern
Upload
Sigmoid
Client k
Bern
Server
Figure 1: Overview of PRISM. PRISM finds the supermask for generative models in a FL scenario. At every round t, each client & updates local score of via MMD loss (Step and generates the privacy-preserving binary mask M₁ (Step 2-1) which is sent to the server. The server aggregates the masks obtain the global probability 0.+1 (Step 3), which is converted score 81+1 and broad- casted to the clients for the next round (Step 4). The global probability 0.+1 is gradually updated based on mask correlation À between M² and M
the objective function. As the scores get iteratively updated, the EP algorithm progressively shrinks the model by applying binary masks weights with higher scores, indicating their potentials to be included in the winning lottery ticket. The obtained SLT can be expressed as W Wine @ M, where M is the obtained binary mask and 3 denotes element-wise multiplication. SLT has been primarily explored within the context of classification tasks, while Yeo al. (Yeo al. 2023) have recently shown that SLT can also be found in generative models.
3.2 DIFFERENTIAL PRIVACY
Sharing each client's model or gradient can potentially lead to privacy risk. (c, (odd))ddffeeenill pri- vacy (DP) (Dwork et al. [2006]], (a, (o,,))-Reee J--éényi-Differnntil privacy (RDP) (Mironov 2017 are com- monly employed when tackling the privacy concerns in FL.
Definition ((e, Difffffrrtttial Privacy (Dwork al. 2006 randomized mechanism M X' + R is (e, 8) dddfffrrrtttial privacy, for any two adjacent datasets DD 'D' and for any measurable sets S: Pr(M(D) E S] s e' Pri S] + 5.
The above definition is designed to limit the impact of individual data points by introducing random- ness into M. The Gaussian mechanism (Mironov al., 2019) offers differential privacy guarantees by injecting Gaussian noise (0, M, where o² 21n 21m(1.25/8)24 and A3 is L2 sensitivity.
22
4 METHOD
We consider a FL setup with K clients, where each client i has its own local dataset D Starting from a randomly initialized model W inste the clients aim to collaboratively obtain a global generative model W that well-reflects all data samples in the system, i.e., in UK Dk
Overview of our approach. Figure I shows the overview of our PRISM. PRISM finds subnet- work with strong generative performance from the randomly initialized generative model W.nt- Rather than updating Wint, its focus is to find optimal binary mask M that has either 1 or 0 in its element and construct the final global model as W Wint @ M At a high-level, each client k generates binary mask M₁ based on its local dataset every communication round which is aggregated at the server. After repeating the process for multiple rounds / = 1, 2,2....TT. T, PRISM produces the final supermask M My. Our approach, which aims to find the SLT in federated generative model setting, is fundamentally different from existing FedGAN methods. In the following, we describe the detailed training procedure of PRISM along with advantages.
4

---

Published as a conference paper at ICLR 2025
4.1 PRISM: PRIVACY-PRESSRVING IMPROVED STOCHASTIC MASKING
Local score updates with MMD loss. Before training starts, the server randomly initializes the model Winit and broadcasts it to all clients, which remains fixed throughout the training process. At the start of each round f. all clients download the score vector s; from the server, representing the importance of each parameter in W Wint. Intuitively, if the score value of a specific parameter is high, the corresponding weight more likely be included in the final SLT. PRISM allows each client k to update the score vector 8, based on its local dataset to obtain k which is used to generate the local mask. In this local score update procedure, we leverage maximum mean discrepancy (MMD)
loss Gretton et al. 2006 2012 providing stable convergence for training generative models (Li al.) 2017a bj Binkowski al. 2018 Santos al., 2019 Ramanujan al. 2020 Yeo al. 2023 J.
The MMD loss measures the distance between two distributions by comparing their respective mean embeddings in a reproducing kernel hilbert space (RKHS) (Gretton et al. [2006] 2012) As in (San-)
tos et al.) 2019 Ramanujan et al.) 2020)], we take VGGNet pretrained on ImageNet as powerful characteristic kernel. Specifically, given the local dataset De = {사 |i=1 of client k and the fake image set D4 Dake = {u/) )==1 produced by its own generator, the local objective function at each client & is written as follows:
+ C)oo(((DD)) ( -
k MMD = E EENDE (x)] By
Cov(v (akee))
(I)
where (-) is a function that maps each sample to the VGG embedding space, Each client aims to match the mean and covariance between real and fake samples after mapping them to the VGG em- bedding space using kernel e.(.). Based on Eq. [ each client locally updates the scores to minimize the MMD loss according to s St v/ 2h M. Here, we note that the VGG network is utilized only for computing the MMD loss and is discarded when training is finished.
Binary mask generation and aggregation. After the local score update process, each client maps the score sf to probability value o5 E [0, 1] as of = Sigmoid(s$) where Sigmoid(-)) is the sigmoid function. The obtained 0f is then used as the parameter of the Bernoulli distribution to generate the stochastic binary mask Mt, according to M, ~ Bern(O") ). Each client l uploads only this binary mask M to the server, significantly reducing the communication overhead. At the server side, the received masks are aggregated to estimate the global Bernoulli parameter as 0++1= - Ek=1 Mtr which can be interpreted the probabilistic score reflecting the importance of the overall client's 's weights. d.+1 is then converted to the score through the inverse of the sigmoid function according to St+1 Sigmoid (0c+1)
(0c+1). ) which broadcasted the clients at the beginning of the next round. We provide the training details in Appendix C
Model initialization for storage efficiency. When training is finished after T' rounds of FL, Or is obtained the server. The supermask is then generated following M ~ Bern(Or), which is used to obtain the final global model as W = Winit @ M This final model W can be stored efficiently even in resource-constrained edge devices, thanks to the model initialization strategy. When initial- izing Wint in PRISM, we employ the standard deviation of Kaiming Normal distribution (He et al.) ], where the weight value in layer / is sampled from the set Hence, by storing the scaling factor 22/n,, each parameter in the initial model Winit is already quantized to a 1-bit value (See Appendix [E). This makes the final model exceptionally lightweight without extra pruning of quantization, which will be also showed via comparison in Section 5.5
{-v 12/n1-11 1/
4.2 PRIVACY
To consider the situation of potential privacy treats, we incorporate (e. -difffrrrrnt privacy (DP) (Dwork et al.] [006]] into our framework. We provide a more detailed description related to the privacy preservation of PRISM and a specific scenario where PRISM can achieve additional privacy benefits in Appendix A and Appendix respectively.
4.3 MASK-AWARE DYNAMIC MOVING AVERAGE AGGREGATION
Data heterogeneity and privacy preservation pose significant challenges in accurately estimating the correct update direction. Both the previous and current global models contain valuable information about whether the local models are diverging Praneeth Karimireddy et al. (2019): : Li et al. (0220);
5

---

Published as a conference paper at ICLR 2025
Mendieta et al. (2022). To address this, we propose p mask-aware dynamic moving average aggre- gation (MADA) that leverages information from previous aggregation rounds in a mask-aware man- ner. To the best of our knowledge, MADA is proposed for the first time, introducing the mask-aware similarity to determine the moving average ratio on the server side. Specifically, when a client mask Mc deviates significantly from the global model, the newly updated global mask M² will also differ considerably from the previous round's s global mask M However, such updates tend to favor the dominant clients. To mitigate this bias, the server calculates the mask correlation X, which measures the similarity between the current and previous global masks (M2 and M2- respectively). The server then interpolates between the two global masks using À to adjust how much of the current mask should be incorporated. While various metries can be used to compute the distance between masks, we employ the Hamming distance. We observe that using alternative distance metrics like cosine similarity yields similar performance (See Appendix L. The server-side aggregation process is defined as follows:
+
Mr. A= dist Mf),
K
This prevents excessive deviation by interpolating the aggregated mask with the previous Bernoulli parameter. As the global rounds progress, À gradually decreases, promoting stable convergence.
4.4 TRADING-OFF BETWEEN PERFORMANCE AND COMMUNICATION COST
While PRISM optimizes for minimal communication overhead and delivers satisfactory perfor- mance, it may not always meet the demand for higher-quality image generation. To address this, we introduces& a flexible solution, PRISM* (See Appendix F
5 EXPERIMENTS
In this section, we validate the effectiveness of PRISM on MNIST, FMNIST, CelebA, and CIFARIO datasets. The training set of each dataset is distributed across 10 clients following either IID or non-IID data distributions" [] We provide the detailed partitioning strategies for IID and non-IID sim- ulation in Appendix D. In addition, for a fair comparison, we set (9.8, 10 -0-D for all methods.
Baselines. We compare our method with several previous approaches for federated generative mod- els under both privacy-preserving (with DP) and privacy-free (without DP) scenarios: In the privacy- preserving scenario, we consider De-FedAvgGAN (Augenstein et al.) [019)] and GS-WGAN (Chen) cct al. 2020y] while in the privacy-free case, we adopt MD-GAN (Hardy et al. 2019)] and Multi- FLGAN (Amalan et al. 2022)] ]. In the case of Multi-FLGAN, the number of sync servers increases the communication cost quadratically, so we consider a 2 x multi generator and discriminator setup.
Performance metries. We evaluate the generative performance of each scheme using the commonly adopted metrics, including Fréchet Inception Distance (FID) (Heusel et al. 2017)), Precision & Re- call (Kynkäänniemi et al., 2019)), Density & Coverage (Naeem et al. 2020 We further demon- strate the efficiency of PRISM by comparing the uplink cost (MB) (See Appendix G at each FL round, the storage (MB) for the final models, and FLOPS (See Appendix
5.1 IID CASE
In this subsection, we examine an IID scenario where the training set of each dataset is uniformly distributed among clients. We compare various evaluation metrics under (c, 5)-DP guaranteed set- ting (Table D. Notably, PRISM outperforms current GAN-based models by large margin. Figure 2 reveals that existing in CelebA, while our method tends to generate high-quality results. The above results support that PRISM can effectively find a SLT, achieving performance gains despite charging 48% less commu- nication costs per round. It is worth noting that PRISM can further reduce the cost by applying extra techniques such as universal coding. Additionally, when comparing PRISM to PRISMÉ (MADA removed), PRISM demonstrates significant improvement in overall performance.
methods often produce distorted images, particularly evident
See Appendix H for additional results where large number of clients and partial of clients participate.

---

Published as a conference paper at ICLR 2025
Quantitative comparison in IID scenario with privacy budget (e, (9.8, 10 10-5). We compare FID, P&R, D&C, communication cost, and storage. Communication cost is the number of bytes exchanged between clients and server. i indicates that MADA is removed.
Table I:
Method (comm.cost)
Metric
MNIST
FMNIST
CelebA
Storage
FIDJ
71,1016
119.2589
230.7874
GS-WGAN
P&R
0.0975 0.1505
0,3694 0.0015
0,7951 0.0
15MB
(15MB)
D&C1
0.0257 / 0.0367
0.1264 0.0347
0.165 0.00021
FID
111.8855
118,5067
221.34
DP-FedAvgGAN (14MB)
P&R
0.2586 0.0047
0.5318 0.0163
0,1008 0.0
14MB
D&C1
0.0803 / 0.0014
0.2028
0.0341
0.0211 /0.0013
FID
48,5636
54.722
57.0573
PRISM
P&R
0.3343 0.4265
0,5836 0.1574
0.4998 / 0.1221
7.25MB
(5.75MB)
D&C1 1
0.1211 / 0,1151
0,2156 / 0.2432
0.2572 / 0.2189
FID
27.3017
46.1652
48.9983
PRISM
P&R
0.4377
0.5576
0.6355 / 0.211
0.6435 0.076
7.25MB
(5.75MB)
D&C1
0.1738 0.1982
0.4002/ 0.2971
0.4089 / 0.2415
9
3
9
/
0
8
8
7
0
7
4
7
6
3
8
1
5
/
0
3
9
3
3
7
9
9 8
5 9 3 J 2
9
1 7
7
4 5 7
7
5
6
0
9
8
8
1
7.
1
0
9
9
MNIST
4
8
9
7
6 6
4 3 7 /
3
7
7
4
9
7
7
9
9
1
3
3
0
7
3 0
5
9
2
4
7
4
9
2
4
3
1
1
9
9
5
9
7
5
J
7
4
9
4
7
S
4
9
9
9
5
5
4
9
9
J
3
5
9
7
6 2
7
7
9
3
9
3
D
5
[
9
9
5
7
4
A
FMNIST
CelebA
(a) GS-WGAN
(b) DP- DevveAAvAGAN
(c) PRISM
(d) PRISM
Figure 2: Qualitative results in IID scenario with a privacy budget (c. 6) (9.8, 10 10-5). We compare generated images from the models in Table MNIST, FMNIST, and CelebA. indicates that MADA is removed.
5.2 NON-IID CASE
We investigate more practical yet challenging non-IlI scenario, where clients exhibit diverse local data distributions, posing a significant challenge to train generative models. To establish such an en- vironment, we split the entire trainset based on mnards-partitioniss titative comparison of baselines and our methods when DP is guaranteed. We observe that PRISM achieves SOTA performance under the non-IID scenario, even with complex CelebA datasets Fig-
strategy Tabseepsesents quan-
We also provide the results of pirrttoce Non-IID in Appendix D
Since existing works the field of federated generative models struggles with MNIST-level generation, refer CelebA and CIFAR 10 complex dataset.
7

---

Published as a conference paper at ICLR 2025
Table 2:
Quantitative comparison in non-lID scenario with (c, (9.8, 10- We compare FID, P&R, D&C, communication cost, and storage. Communication cost refers to the number of bytes exchanged between clients and server. indicates that MADA is removed.
Method (comm.cost)
Metric
MNIST
FMNIST
CelebA
Storage
FID
338.6659
131.6166
228,9705
GS-WGAN
P&R
0.0/0.0
0.4186 0.0001
0.1363 0.0
15MB
(15MB)
D&C1 1
0.0
0.1569 0.0297
0.0307 / 0.0025
FID
153.9325
146.632
222.8257
DP-FedAvgGAN (14MB)
P&R
0.4371
0.0336
0.7207 0.0043
0.2331 0.0004
14MB
D&CT
0.1049/ 0.004
0.2589 0.01664
0.0668 0.0016
FID
49.6273
83.0481
59.4877
PRISM
P&R
0.3283
0,3844
0,4513
0.0775
0.4789 0.0898
7.25MB
(5.75MB)
D&C1
0.1101 0,1022
0.2355
0.1428
0.2392 / 0.2058
FID
34.2038
67.1648
39.7997
PRISM
P&R
0.4386
0.4236
0.4967
0.1231
0.6294 / 0.0713
7.25MB
(5.75MB)
D&C1
0.1734 0.1597
0.2748
0.1681
0.4565
0.2967
9
7
6
3
0
1
9
9
9
3
6
0
9
8
3
8
3
2
5
5
5
MNIST
2
9
7
7
9
9
7
3
J
5
4
2
4
7
4
9.
2
9
5
9
9
9
9
9
9
4
7
S
9
0
9
/
6
4
7
7
9
3
9
5
8
9
FMNIST
CelebA
(a) GS-WGAN
(b) DP-FedAvgGAN
(c) PRISM
(d) PRISM
Figure 3: Qualitative results in Non-IID scenario with privacy budget (e, 6) (9.8, 10-5), We compare generated images from the models in Table MNIST, FMNIST, and CelebA. indicates that MADA is removed.
illustrates that despite the heterogeneity of the data, our methods successfully generate high quality images, while traditional methods exhibit subpar quality. Additionally, Figure shows the overview of FID scores, final model parameters, and communication costs for each method.
ure
5.3 PERFORMANCE WITHOUT DIFFERENTIAL PRIVACY
We further explore the performance of PRISM without applying differential privacy. Quantitative and qualitative comparisons with MD-GAN (Hardy al. 2019 and Multi-FLGAN (Amalan al. 1, the current state-of-the-rr art under this condition, are shown in Table 3 and Figure 5 PRISM not only matches but also occasionally surpasses the performance of MD-GAN and Mult-FLGGN, all while significantly reducing communication overhead. Again, the fact that PRISM outperforms PRISM clearly demonstrates the effectiveness of MADA. Here, Bulti-FL.GAR takes 10 times more
8

---

Published as a conference paper at ICLR 2025
DP- oLAvgGGNN
AGAN
MNIST
350
GS-WIANN
MNIST FMNIST
200
FMNIST
300
CelebA
CelebA
230
GS WGAN
Cummusication cont
150
DF- UP-FedNggGA
200
100
DP FedAggGAA
DF-FedANGGAN
GS-GGAN
20D 4AE2 HeuD
150
DP-FodAN
DP-FedArgGAN
WGANN
GSWGAN
100
Ideal
PRISM
PRISM
PRISM
50
Edenl
PRISM
PRISM
PRISM
1
of final model parameters (M)
# final model parameters (M)
(a) IID case
(b) Non-IID case
Figure 4: The performance of baselines and our PRISM with privacy budget (c, (9.8, 10 50-5). X-axis represents the number of parameters of final generator, while Y-axis represents FID. The diameter of each circle denotes the required communication cost at every round. The ideal case is the bottom-left corner.
Table 3: Quantitative comparison in non-III scenario without DP. We compare FID, P&R, D&C. Communication cost refers to the number of bytes exchanged between clients and server. indicates that MADA is removed. We set a 80 for PRISM".
Method
Metric
MNIST
FMNIST
CelebA
CIFAR10
Storage
(comm.ccom cost)
FID
37,7971
55,5094
18,907
52.7159
MD-GAN
P&R f
0.3366 0,5435
0.5635 00005
0.7612
0.6425
0.827 0.1968
14MB
(14MB)
D&C t
0.1192 0.1405
0.3145 0.2033
0.7238
0.4267
1.2201/ / 0.3822
FID
32,1014
125.9276
314.8386
163.0540
Multi-FLGAN (52MB)
P&R
0.5659 / 0.3353
0,4781 / 0.0035
0,0
0.9345 / 0.0
14MB
D&CT 1
0.3171 0,2779
0.24 / 0.0595
0.0
0.3638/ / 0,0668
FID
15,2329
35.1448
24.2591
68.4238
PRISM
P&R
0.7128 0,5289 0.5106
0.7239 0.1949
0.7988 / 0.1888
0.65 0.1732 7.25MB
(5.75MB)
D&C
0.4851
0.645 0,3768
1.0746 0.5809
0.5575 / 0,3031
FID
9,698
32.7517
21,8567
61.1198
PRISM
P&R 1
0.7665 / 0.6253
0.7614 / 0,1281
0.7835 / 0.1615
0.5924 / 0.2323 7.25MB
(5.75MB)
D&CT
0.6088/ / 0.6003
0.8361 0.4383
1.047 0.6079
0.4334/ 0.3171
FID
6.9568
29.0081
13.0209
35,5326
PRISM
P&R 1
0.7717
0.7992
0.697 0.1572
0.7893 / 0.392
0.6662 / 0.3642 7,25MB
(15MB)
D&C1 1
0.6082 0.6499
0.7002 0.4056
1.072 / 0.7396
0.5764 / 00.4481
communication cost than PRISM due to multi GAN strategy. However, in FL setups, users may want trade-off between the communication cost and generative performance. As shown in Table and Figure 5 PRISM" outperforms across all benchmarks, including CIFAR1O, while maintaining a communication cost comparable to MD-GAN and far lower than Multi-FLGAN Further details on PRISM can be found in Appendix F
5.4 EFFECT OF DYNAMIC MOVING AVERAGE AGGREGATION
In this section, we empirically demonstrate the effectiveness of dynamic moving average aggrega- tion using the MNIST dataset. (igurr/ovvssuliizze local model updates over communication rounds At each global round f, clients receive the aggregated global mask M and continue local training for several epochs. Afterward, the trained local mask Mt is uploaded for the next communication round. We introduce the local divergence metric A, == hdd M MM²), which is defined as the Ham- ming distance between the received mask and the trained local mask to track the discrepancy of local model updates. For simplicity, we visualize the results for the first client, but similar trends were observed across other clients. In Figure 6 PRISM exhibits impressive FID scores and reduced local updates than PRISM clearly demonstrating that MADA not only restricts the local model divergence but also achieves significant performance gain across various challenging FL settings. This is accomplished by automatically obtaining À based on the current mask, without requiring additional regularization terms or hyperparameter tuning
*The convergence of X in MADA and its relationship with client drift discussed in Appendix

---

Published as a conference paper at ICLR 2025
4
4
3
2
5
7
3
4
7
4
7
9
5
8
5
1
5
0
8
6
1
9
)
1
7
0
3
9
3
9
2
3
0
2
S
5
3
7
7
4
MNIST
/
6
8 1
5
4 /
Q
4
4
8
3
3
/
8
8
8
7
4
3
5
4
5
2
K
7
2
9
9 7
5
1
9
9
8
9
3
4
4
9
5
8
8
6
2
7
2
L
2
3
5
9
6
f
8
3
5
4
5
9
4
4
4
9
4
3
5
7
8
4
R
5
7
J
9
0
5
4
8
7
7
8
3
5
7
3
4
6
CelebA
CIFARIO
(a) MD-GAN
(b) Multi-FLGAN
(c) PRISM
(d) PRISM
Figure 5: Qualitative results in non-IID scenario without considering privacy budget. Generated images from the models in Table 3 on MNIST, FMNIST, CelebA, and CIFARIO. Here, we set 80 for PRISM
PROSM
303
PRISEC ED om
PROSHI 300
0.55
0.35
19052-1ED
0:35
FRISM-I1
FRD
PSISKC EID
FROSHT F1D
F1155-710
(lobel nound
Globel munnd
Global round
MNIST
(b) FMNIST
(c) CelebA
Figure 6: Local divergence A, and FID values in the non-lID and DP-considering scenario.
5.5 RESOURCE RFFICIENCY AT INFERENCE TIME
The final model sizes of the baselines and our method are reported in Table Table 2 and Table Note that in addition our method, applying various lossless compression techniques (e.g., arith- metic coding (Rissanen & Langdon 1979h]) can further reduce the required resources of PRISM.
CONCLUSION
While image generation has emerged as promising area in deep learning, the fusion of FL with generative models remains relatively unexplored. In this paper, we proposed PRISM, an efficient and stable federated generative framework that capitalizes on stochastic binary masks and MMD loss. To further enhance stability under non-lll) and privacy-preserving scenario, we introduced a mask-aware dynamic moving average aggregation strategy (MADA) that mitigates client drift. Additionally, PRISM offers hybrid mask/score aggregation method, allowing for flexible and controllable trade-off between performance and efficiency. Our extensive experiments, including scenarios involving differential privacy and non-IID setups, demonstrated that PRISM is robust in unstable environments. To the best our knowledge, PRISM is the first framework to consistently generate high-quality images with significantly reduced communication overhead in FL settings.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) NNo.2022RICICIO08-440 cations Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT); Nko.2022---009555 No.RS- No.RS-2002-11 2022-11220959 (Part 2) Few-Shot Learning of Causal Inference in Vision and Language for Decision Making), INo.R--2022, 25-2022--1122264, Comprehensive Video Understanding and Generation with Knowledge-based Deep Logic Neural Network), (No. Ra.RR2220--12000 tificial Intelligence Graduate School Program (UNIST)), NN.RS-2021-112120 Artificial Intelli- gence Innovation Hub).
and Institute of Information & communi-
Ar-
REFERENCES
Akash Amalan, Rui Wang, Yanqi Qiao, Emmanouil Panaousis, and Kaitai Liang. Multi-flgans: Multi-distributed adversarial networks for non-iid distribution. arXiv preprint arXiv:2206. /2/78, 2022.
Sean Augenstein, H Brendan McMahan, Daniel Ramage, Swaroop Ramaswamy, Peter Kairouz, Mingqing Chen, Rajiv Mathews, et al. Generative models for effective ml on private, decentralized datasets. arXiv preprint arXiv:1911.06677 2019.
Ran Ben Basat, Shay Vargaftik, Amit Portnoy, Gil Einziger, Yaniv Ben-Itzhak, and Michael Mitzenmacher. Quick-fl: Quick unbiased compression for federated learning. arXiv preprint arX1v:2205.13341, 2022.
Mikotaj Binkowski, Danica Sutherland, Michael Arbel, and Arthur Gretton, Demystifying mmd gans. arXiv preprint arXiv:1801.014000 2018.
Dingfan Chen, Tribhuvanesh Orekondy, and Mario Fritz. Gs-wgan: gradient-sanitized approach for learning differentially private generators. Advances in Neural Information Processing Systems, 33:12673-12633 2020.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Wor Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8188-8199 2020.
Tim Dockhorn, Tianshi Cao, Arash Vahdat, and Karsten Kreis. Differentially private diffusion mod- els. arXiv preprint arXiv:2210.09929, 2022.
Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise to sensitivity in private data analysis. In Theory of Cryptography: Third Theory of Cryptography Conference, TCC 2006, New York, NY, USA, March 4-7, 2006. Proceedings 3, pp. 265 2244. Springer, 2006.
Farzan Farnia and Asuman Ozdaglar. Do gans always have nash equilibria? In International Con- ference on Machine Learning, pp. 3029-3039 PMLR, 2020a.
Farzan Farnia and Asuman Ozdaglar. Gans may have no nash equilibria. arXiv preprint arXiv:2002.091124 2020b.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: arXiv:/803.03. /803003635 2018.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-144 2020.
Arthur Gretton, Karsten Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alex Smola. A kernel method for the two-sample-prohlem. Advances in neural information processing systems, 19, 2006.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.

---

Published as a conference paper at ICLR 2025
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Corentin Hardy, Erwan Le Merrer, and Bruno Sericola. Md-gan: Multi-discriminator generative adversarial networks for distributed datasets. In 2019 IEEE international parallel and distributed processing symposium (IPDPS), pp. 866-877. IEEE, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings the IEEE international conference on computer vision, pp. 1026-002600000 1034, 2015.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840- 600166811110 2020.
Rui Hu, Yuanxiong Guo, and Yanmin Gong. Federated learning with sparsified model perturba- tion: Improving accuracy under client-level differential privacy. IEEE Transactions on Mobile Computing, 2348):8242-8255, 2023.
Jacob Imola and Kamalika Chaudhuri. Privacy amplification via bernoulli sampling, arXiv preprint arXiv:2105.10594, 2021.
Berivan Isik, Francesco Pase, Deniz Gunduz, Tsachy Weissman, and Michele Zorzi. Sparse random networks for communication--ffiiinne federated learning, arXiv preprint arXiv:2209, 15328, 2022.
Zepeng Jiang, Weiwei Ni, and Yifan Zhang. Pate-triplegan: Privacy-preserving preserving image synthesis with gaussian differential privacy. arXiv preprint wwXiv:2404.12740 2024.
Ye-Y Kim, Dong-Jur Jun Han, Jun Seo, and Jackyun Moon. Achieving lossless gradient sparsifica- tion via mapping to alternative space in federated learning. Forty-first International Conference on Machine Learning. 2024a.
Kihong Kim, Haneol Lee, Jihye Park, Seyeon Kim, Kwanghee Lee, Seungryong Kim, and Jaejun Yoo. Hybrid video diffusion models with 2d triplane and 3d wavelet representation. arXiv preprint arXiv:2402. 13729, 2024b.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in Neural Information Processing Systems, 32, 2019.
Ang Li. Jingwei Sun, Xiao Zeng. Mi Zhang, Hai Li, and Yiran Chen. Fedmask: Joint computa- tion and coamuuniattionefffceee personalized federated learning via heterogeneous masking. In Proceedings of the 19th ACM Conference on Embedded Networked Sensor Systems, pp. 42-55, 2021a.
Chun-Liang Li, Wei-Cheng Chang. Yu Cheng, Yiming Yang, and Barnabás Póczos. Mmd gan: Towards deeper understanding of moment matching network. Advances in neural information processing systems, 30, 2017a.
Chun-Liang Li, Wei-Cheng Chang. Yu Cheng, Yiming Yang, and Barnabás Póczos. Mmd gan: Towards deeper understanding of moment matching network. Advances in neural information processing systems, 30, 2017b.
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, 454, 2020.
12

---

Published as a conference paper at ICLR 2025
Wei Li, Jinlin Chen, Zhenyu Wang, Zhidong Shen, Chao Ma, and Xiaohui Cui. Ift -ft- Improved federated learning generative adversarial network with maximum mean discrepancy model aggre- gation. IEEE Transactions Neural Networks and Learning Systems, 2022.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-id features via local batch normalization. arXiv preprint arXiv:2102.07622 202 lb.
Eran Malach, Gilad Yehudai, Shai Shalev-Schwartz, and Ohad Shamir. Proving the lottery ticket hypothesis: Pruning is all you need. In International Conference on Machine Learning, pp. 6682- 6691. PMLR, 2020.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-cfficcee learning of deep networks from decentralized data. Artificial in- telligence and statistics, pp. 1273-1282 PMLR, 2017.
Matias Mendieta, Taojiannan Yang. Pu Wang, Minwoo Lee, Zhengming Ding, and Chen Chen. Local learning matters: Rethinking data heterogeneity federated learning. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8397-8406, 2022.
Ilya Mironov. Rényi differential privacy. In 2017 IEEE 30th computer security foundations sympo- sium (CSF), pp. 263-275. IEEF, 2017.
Iyya Mironov, Kunal Talwar, and Li Zhang. R 'eny differential privacy of the sampled gaussian mechanism. arXiv preprint arXiv: 1908. 10530, 2019.
Nicole Mitchell, Johannes Ballé, Zachary Charles, and Jakub Konečný. Optimizing the communication-acccure trade-off in federated learning with rate-distortion theory. arXiv preprint arXiv:2201.02664, 2022.
Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jaejun Yoo. Reliable fidelity and diversity metries for generative models. In International Conference on Machine Learning. pp. 7176-7185 PMLR, 2020.
Laurent Orseau, Marcus Hutter, and Omar Rivasplata. Logarithmic pruning is all you need. Ad- vances Neural Information Processing Systems, 33:2925-2934, 2020.
Sai Pranceth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Scbastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning, arXiv e-prints, pp. aiXiv-1910, 2019.
Vivek Ramanujan, Mitchell Wortsman, Aniruddha Kembhavi, Ali Farhadi, and Mohammad Raste- gari. What's hidden in a randomly weighted neural network? In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1 1199-1 1902, 2020.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. urXiv preprint 200.:22044.01125,, 2022.
Mohammad Rasouli, Tao Sun, and Ram Rajagopal. Fedgan: Federated generative adversarial net- works for distributed data. arXiv preprint arXiv:2006.072. 072228 2020.
Jorma Rissanen and Glen G Langdon. Arithmetic coding. IBM Journal of research and development, 1392):149-162, 1979.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings, pp. 1-10, 2022.
Cicero Nogueira dos Santos, Youssef Mroueh, Inkit Padhi, and Pierre Dognin. Learning implicit generative models by matching perceptual features. In Proceedings of the IEEE/CVF Interna- fional Conference on Computer Vision, pp. 4461-4470) 2019.
13

---

Published as a conference paper at ICLR 2025
Jaejung Seol, Seojun Kim, and Jaejun Yoo. Posterllama: Bridging design ability of langauge model to contents- aware layout generation. arXiv preprint arXiv:2404400995, 2024.
Ivan Skorokhodov, Sergey Tulyakov, and Mohamed Elhoseiny. Stylegan-v: continuous video generator with the price, image quality and perks of stylegan2. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3626-36333 2022.
Pablo Villalobos, Anson Ho, Jaime Sevilla, Tamay Besiroglu, Lennart Heim, and Marius Hobbhahn. Will we run out of data? limits of Ul scaling based on human-generated data. arXiv preprint, 2024.
Jianyuan Wang, Ceyuan Yang, Yinghao Xu, Yujun Shen, Hongdong Li, and Bolei Zhou. Improving gan equilibrium by raising spatial awareness. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11285-11293, 2022.
Yu-Xiang Wang, Borja Balle, and Shiva Prasad Kasiviswanathan. Subsampled rényi differential privacy and analytical moments accountant. The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1226-1235. PMLR, 2019.
Kang Wei, Jun Li, Ming Ding, Chuan Ma, Howard H Yang, Farhad Farokhi, Shi Jin, Tony QS Quek, and Vincent Poor. Federated learning with differential privacy: Algorithms and performance analysis. IEEE Transactions on Information Forensics and Security, 15:3454-3469, 2020.
Jun
Bangzhou Xin, Wei Yang. Yangyang Geng, Sheng Chen, Shaowei Wang, and Liusheng Huang. Pri- vate fl- fl-gan: Differential privacy synthetic data generation based on federated learning. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 2927-2931. IEBE, 2020.
Sangyeop Yeo, Yoojin Jang, Jy-yong Sohn, Dongyoon Han, and Jaejun Yoo. Can we find strong lottery tickets in generative models? In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 37, pp. 3267-3275, 2023.
Liping Yi, Wang Gang, and Liu Xiaoguang. Qsfl: A two-level uplink communication optimization framework for federated learning. International Conference on Machine Learning, pp. 25501 25513. PMLR, 2022.
Ashkan Yousefpour, Igor Shilov, Alexandre Sablayrolles, Davide Testuggine, Karthik Prasad, Mani Malek, John Nguyen, Sayan Ghosh, Akash Bharadwaj. Jessica Zhao, et al. Opacus: User-friendly differential privacy library in pytorch. arXiv preprint arXiv:2 /2298, 2021.
Yikai Zhang, Hui Qu, Qi Chang, Huidong Liu, Dimitris Metaxas, and Chao Chen. Training federated gans with theoretical guarantees: A universal aggregation approach. arXiv preprint arXiv:2102.04655, 2021.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint /80.008806.0088,, 2018.
Chenzhuo Zhu, Song Han, Huizi Mao, and William J Dally. Trained ternary quantization. arXiv preprint 1611: 2016.
A APPENDIX
A PRIVACY
Diffenrential privacy (DP) and Rényi Differential Privacy (RDP) (Dwork et al.) 2006 Mironov] 2017]] are the most popular definitions to analysis the privacy in FL environments. These help mitigate privacy concerns by limiting the contribution of individual data points. (c, 6DP and (a, e)- RDP basically calculates the distance of outcome for the algorithm of adjacent datasets.
Definition 2 ((e, Diffferrretia Privacy) A randomized mechanism M X + R is (c, differential privacy, for any two adjacent datasets D, D' and for any measurable sets S: Pr(M(D) E S) < Ee Pr(M(DD') E S +
14

---

Published as a conference paper at ICLR 2025
Definition 3 (la, c) Rényi Differential Privacy) For two probability distributions P and Q the Rényi divergence of order a > 1 defined as follows:
(al a-1 log R. P(x) Q(x)
(4)
then, randomized mechanism M X R. (a, Rényi differential privacy, for any two adjacent datasets 'D 'D' and for any measurable sets S:
Ra(MD))IM(')) SE
(5)
Theorem
Mironov
2017 | showed that if M is (ax, R)-DDP guarantee, is also (e + log -DP.
a-1
In this section, we provide more detailed explanation of privacy preserving in PRISM and also present updated results when DP is applied. To satisfy the (r,ō)-DP, our goal privatize the
probability vector E [0, 1] by adding gaussian noise (N0, o²), where x² = 26x(1.25/0)44 and e²
A₂ maxp,p' [IM ((2 When the local training is end, each client has scores R to choose which weight to prune. Recall that probability A [0, 1] can be obtained through sig- moid function. We inject gaussian noise and clip to E [c, cjd, where с is small value <<> 2. In out setup, we fix it at 0.1. Now, we ensure we use (c, 8) (9.8, 10-9) to PRISM and our baselines in all of our experiments. In addition, we regulate the global round ensure that the overall privacy budget does not exceed To track the overall privacy budget, we employ subsampled moments accountant (Wang et al. 2019)]. We refer to the Opacus library which is the user-friendly pytorch framework for differential privacy Yousefpour
(c. (DDP. For fair comparison,
et al. 2021
Imola & Chaudhuri 202 Isik et al. 2022) have shown that performing post processing to already
privatized vector such Bernoulli sampling enjoys privacy amplification under some conditions. By doing so, the overall privacy budget becomes smaller Camp I (c)) where %a is
the binary symmetry Rényi divergence expressed below:
1 (c) - (10)0(00000000)000000000
where a refers the order of the divergence. Note that d' limits the privacy amplification when the model size becomes large. Since PRISM assumes that the model size is large enough due to SLT, focus on communication efficiency rather than privacy amplification.
B PRIVACY AMPLIFIED SCENARIO
In this section, discuss about potential privacy benefit of PRISM under some conditions. Typical FL setting, a malicious third party can estimate W W from the communicated gradients. By analyzing these gradients, PRISM, only the binary masks M and M are exchanged, which hinders an attacker's efforts since they do not have access to the synchronized initial weight W int shared between the client and server.
attacker can extract information about the local data. However, with
C TRAINING DETAILS
In this section, we provide the detailed description of our implementations and experimental set- tings. In T we provide the model architectures used in our experiments. We use Nes generator and set the local epoch to 100 and learning rate 0.1. In addition, do not employ training schedulers or learning rate decay. Our code is based on (Santos et al.) 2019] Yeo et al.] ]. They employ the ImageNet-pretrained VGG19 network for feature matching by minimizing the Eq. M However, calculating the first and second moments require the large batch size to obtain the accurate statistics. To address this issue, (Santos et al., 2019)] introduces Adam moving average (AMA). With rate X, the update of AMA m is expressed as follows:
m XADAM (m A),
15

---

Published as a conference paper at ICLR 2025
Table 4: Generator architecture used in our experiments.
Layer
Type
Input Channels
Output Channels
Kernel Size
FirstConv
Conv
128
512
(4, 4)
Conv1
512
256
(3, 3)
Conv2
256
256
(3, 3)
ResblockO
BatchNorm22
512
512
ReLU
Bypass Conv
512
256
1)
Upsample
Convl
512
256
(3.3)
Bypass Conv
512
256
(1,
Resblock1
Upsample
Conv2
256
128
(3.3)
Conv3
128
128
(3,
Convl
128
64
(3, 3)
Conv2
64
64
(3,
Resblock2
BatchNorm2d
64
64
ReLU
Conv
64
(3,
LastConv
Tanh
where ADAM denotes Adam optimizer (Kingma & Ba 2014)] and A is the discrepancy of the means of the extracted features. Note that ADAM (m- - A can be interpreted as gradient descent by minimizing the L2 loss:
min mm AlI m 2
This means the difference of statisties (m A) passed through single MLP layer and updated
using the Adam optimizer to the direction of minimizing Eq. g By utilizing AMA, Eq. is for-
mulated MMDD
[xb=>>))((<<)] EywDY fode [vb(v)]
Cov(v(D*)) (D fake))
Algorithm [ the pseudocode for MADA and PRISM correspondingly. AMA is omitted simply express the flow of our framework. See our ccoee for pytorch implementation. We train the local generator for 100 local iterations with learning rate of 0.1. For the AMA layer, learning rate is set to 0.005. In addition, we use the Adam optimizer with 31 = 0.5, Bg = 0.999 to update the scores of the generators. After all clients complete their training, communication round is initiated. We set the global epoch to 150 for the MNIST dataset and 350 for the CclebA and CIFARIO datasets. As we do not adjust the parameters, note that there is room for performance improvements through hyperparameter tuning.
D DATASET PARTITIONING STRATEGIES ACROSS CLIENTS
Here, we elaborate on the dataset partitioning strategies used simulate the IID and non-IID setup. For non-lID setup, we used two strategies for partitioning datasets: 1) Shards-partitionns and 2)
Dirichlet.partitioned..
II setup. In the case of IID scenario, we assign equal-sized local datasets by uniformly sampling from the entire training set. This allows for balanced distribution of data across the clients.
Non-IID setup (Shards-partitionee) In Section 5.2 and Section 5.3 we partition the MNIST, FMNIST, and CIFARIO datasets into 40 segments based on the sorted class labels and randomly assign four segments to each client. For the CelebA dataset, which contains multiple attributes per image, defining a clear non-III distribution for splitting is inherently ambiguous. this work, we divide the dataset into tow partitions based number partitions corresponds to the number clients, with each client being assigned either the positive negative subset of images for the pivotal attribute. However, the remaining 39 attributes
pivotal attribute (gender, in case). The total
16

---

Published as a conference paper at ICLR 2025
Algorithm MADA
Parameter: learning rate 7, communication rounds T. local iterations / Input: local datasets U k= D* ImageNet pretrained VGGNet yú, random noise
Server execute:
Initialize random weight W and score vector s, then broadcasts to all clients. for round = 1,... T do
Client side:
for each client k E [1, K] do 2t St
b Download score vector
for local iteration i = ,,,, do
0f + Sigmoid stt)
M ~ Bern (
W < Wintt 3 M
fake < (2)
D. Generate fake images
Extract real and fake features (DD)), (akk)) 8 < + n/ (MMDMDD(( (D) fake)) end for A t Sigmoid(s$)
D Update local score vector
0, = 0f +N(0, Ia²)
Clip to [c. 1-c]
M ~ Bern(Of))
Upload binary mask M the server.
end for
Server side;
Be+1 +- Σx= M
Aggregate the received binary masks
8t+1 + Sigmoid (9c+1)
+ Me(M Bern (0.+1))
> Compute the hamming distance
81+1 + (1 A)st A8
end for
Sample the supermask M Bern(Or) (fr)
Obtain the final model W <- W Wint ( M
Algorithm PRISM
Input: ratio of score layer
Output: probability (100 a) and binary mask M (a:
Client side: for I L do
if Is))or)Laaer((a..)) then Return probability (1)
else
Return binary mask M (1)
end if end for
are still shared among clients, resulting relatively weak heterogeneity. This explains why the performance drop CelebA is not significant compared to the IID case. For more detailed discussion, please refer to Appendix D
Non-IID setup (Dirichlet-partiin partitioned). We further explore more realistic and challenging non- IID scenario, where datasets are partitioned using Dirichlet distribution. Specifically, we set Dirichlet parameter 0.005 to create more label-skewed distribution. For CelebA, assign data to clients such that each client possesses the pivotal attribute in different proportions using Dirichlet distribution. For example, client has 60% male and 40% female, while client 2 has 20% male and
17

---

Published as a conference paper at ICLR 2025
Table 5: Quantitative comparison in Dirichlet non-Ilf scenario with (c, (9.8, 10- We compare FID, P&R, D&C. indicates that MADA is removed.
Method
Metric
MNIST
FMNIST
CelebA
FID 4
128.4401
134,4054
108.8792
GS-WGAN
P&R
0.0851 / 0.0633
0.3927 / 0.0001
0.8891 0.0014
D&C1 f
0.0196/
0.0071
0.1252 / 0.0348
0.2165 0.052
FID
175.3729
209.1346
323.3559
DP-FedAvgGAN
P&R
0.0408
0.1982
0.0814 / 0.0421
0.1357 / 0.0
D&C1 f
0.0102 0.0048
0.0214 / 0.0123
0.0271 0.0001
FID
74.739
137.7545
57,4051
PRISM
P&R
0.2209 0.1156
0.3041 / 0.0274
0.482 / 0.1058
D&C1 1
0.0619 0.0415
0.1086 0.0312
0.2317 / 0.1934
FID 1
70.9146
137.7359
38.0283
PRISM
P&R
0.4119 0,1563
0.3394 0.0378
0.6305 0.0746
D&C1 f
0.1513/ 0.0759
0.1 / 0.0394
0.4493 0.3144
Table 6: Quantitative comparison in Dirichlet non-III scenario without DP. We compare FID,
P&R, D&C. i indicates that MADA is removed. We set 80 for PRISM
Method
Metric
MNIST
FMNIST
CelebA
CIFAR10
FID
106.3468
86.0443
30.7787
54.38
MD-GAN
P&R
0.7109 0.2431
0.5586 0.0359
0.7612 0.6425
0.8941 / 0903
D&C1 f
0.6884/ 0.4895
0.2137 / 0.0917
0.7238 0.4267
1.4087 / 0.355
FID
Multi-FLGAN
P&R f
diverge
D&C1
FID
43.3057
76.1413
21.0935
89.0257
PRISM
P&R
0.6186 0.216
0.6093 / 0.023
0.797 0.22229
0.6916 0.0854
D&C1 1
0.3219 / 0.1799
0.3443 0.1163
0.991 0.6533
0.5885 / 0.2157
FID
31.6191
73.8701
20.0883
74, 1502
PRISM
P&R
0.5871 / 0.36
0.5239 0.0378
0.782/ 0.2024
0.572 0.1601
D&C1 T
0.2828
0.2328
0.2919 1054
1.1148 0.6893
0,3838 0.2535
FID
38.8984
72.8571
13.6906
$7.9337
PRISM
P&R 1
0.5671 / 0.5423
0.4926 / 0.045
0.784/ 0.3266
0.6395 / 0.2645
D&CT
0.27 / 0.2251
0.2628 0.1177
0.9627 0.7352
0.567 0.327
80% female. To the best of our knowledge, methods for splitting multi-label datasets into extreme non-lID scenarios have not been extensively explored. Table 5 and Table 6 present quantitative comparison of baselines and our methods under both with and without DP. PRISM demonstrates robust performance even under highly heterogeneity distributions. Note that Multi-FLGGAN fails to successfully train generative model under highly heterogeneity.
E EFFICIENT STORAGE AND QUANTIZATION STRATEGY IN PRISM
In this section, provide further detail description regarding storage efficiency of PRISM. Afore- mentioned in Section 4.1 the final model can be expressed as W = Winit @ M Since Mfinnt is lightweight binary masks, need to focus the storage of Winit- Note that Winst ini- tialized as signed constant using Kaiming Normal distribution. In other words, the initialized pa- ... /2/mm11 E Rd This design makes rameter of layer becomes { 22/m=11 be enough to represent each parameter, represented as storing the scaling factor
Wini @ Msign, where Msign is defined (+1.-1,-. +1} E Rd. Our memory-efficient approach is closely related to ternary quantization (Zhu et al., (016)], which typically compresses neural network weights into 0, through thresholding and projection. However, PRISM dis-
18

---

Published as a conference paper at ICLR 2025
60
FTD forward
MDD Jerward
FTD backward
FID backeward
forrand
20
55
Spe_Jorwarl
20
50
backwant
Cost (MB)
Oost backward
13
50
45
Cosl.
45
40-
35
Sweeet apot
Swept
0
20
40 60
80
100
0
20
40
60
80
100
Portion of score layer (%)
Portion of score layer
(a) IID case
(b) Non-IID case
Figure 7: Analysis of PRISM using CIFAR10 dataset. The backward path selects a% of score layers from deeper layers, closer to the output, while the forward path chooses from the opposite end, nearer to the input. Solid-line demonstrates FID following each direction while dash-line shows communication cost (MB) of each path.
Less cost quality
High High quality
(a) a
(b) 20
(c) 50
(d) 70
(c) a 100
Figure 8: The effect of adjusting of PRISM". Qualitative comparison according a. Note that a = is identical to PRISM.
tinguishes itself by directly initialized frozen weights as signed constants using the Kaiming Normal distribution. Since storing signed arra is negligible in terms of memory overhead, this design makes the final model exceptionally lightweight without extra pruning or quantization.
F ANALYSIS OF HYBRID AGGREGATION
In this section, we further analyze PRISM". which leverages both binary mask and score communi- cations. To explore the trade-off between communication cost and generative capability, consider two strategies: the backward path and forward path. The backward path progressively increases the number of score layers from deeper layers earlier layers. Conversely, in the forward path, select score layers from earlier layers to deeper layers. Figure 7 visually demonstrates the trade-off' between communication cost and FID of both strategies. The FID gradually improves as we increase values in both cases. Note that the additional communication cost of the backward path tends to increase more smoothly. Based on these observations, we adopt the backward path Table and Figure 5 Figure 8 provides comprehensive comparison across wide range of values, showing that PRISM" consistently produces high quality images. Note that all experiments are conducted in privacy-free scenario.
G COMMUNICATION COST DURING DOWNLINK PROCESS
In typical FL setup, the server possesses powerful computational capabilities, whereas the clients do not. As a result, several prior studies primarily focus on addressing the limited bandwidth of clients (Kim et al. 2024a Hu et al. 2023 Yi al. 2022)]. Accordingly, the reducing communication cost during the uplink process, transmitting a float-type score during downlink to mitigate the information loss. Last but not least, binary mask communication can also
have concentrated
19

---

Published as a conference paper at ICLR 2025
be employed in the downlink to further enhance communication efficiency, despite the occurrence of information loss during the conversion of the global mask into local scores.
Table Quantitative results on cross-device environment. We set the number of clients to 50 and participant ratio to 20%.
Case
Metric
MD-GAN
DP-FedAvgGAN
GS-WGAN
PRISM
FID
N/A
118.3975
98.6553
34.7157
Non-IID w/ DP
P&R
N/A
0.1095 / 0.3723
0.8477 / 0.0355
0.4344 / 0.3401
D&C1 1
N/A
0.0301 /0.0289
0.2621 /0.0105
0.1692 0.1476
FID
15.4119
N/A
N/A
14.3168
Non-IID w/o DP
P&R 1
0.7305 0.359
N/A
N/A N/A
0.7533 0.4757
D&C 1
0.5266 / 0,3803
N/A
N/A
0.5804 0.5293
H CROSS-DEVICE ENVIRONMENTS
PRISM is designed to provide efficient communication cost, making it ideal for cross-device en- vironments with numerous clients and limited bandwidth. Typically, cross device setting involves large number of clients and an unreliable network, such as those encountered in iPhone and wearable devices. This complexity makes it difficult ensure that all clients in every communication round. To simulate this cross-device environment, we set the number of clients to 50 and the participant ratio as 20%. Table|sshww the quantitative results using MNIST dataset, both with and without DP. PRISM demonstrates a clear performance advantage over other baselines and is well-suited for the cross-device environment.
I ANALYSIS OF CONVERGENCE OF IN MADA
In this section, we provide detailed discussion on how the reduction in X, which measures the similarity between two consecutive global masks, helps mitigate client drift.
One possible explanation is that MADA implicitly modulates the learning rate for global updates on the server-side without requiring an explicit learning rate scheduler. This behavior reduces the need for client-side learning rate scheduling regularize local objectives. To illustrate this, consider Ai. where 1) is constant learning rate, and M is local model update from client k. In standard FedAvg, the global update be expressed WaudA / n A, which aggregates all client updates equally. In contrast, MADA introduces mask-aware aggregation mechanism, and the global model is formulated as WMAGA =
local mask update represented as w w
(1 - A + wf., where wf = w--1 7) AG. Substituting the local mask update into this equation, MMAGA = (1 - A)wY + (w 17 Σx = w 17 AΣ indicating that effectively scales the global updates step size. When is large, the global model closely follows the client updates, while small reduces the influence client updates, acting as a dampening factor.
The intuition behind this behavior is rooted in the nature early-stage training. During the initial rounds, the model is far from convergence, resulting in large gradient norms and substantial param- eter changes. Consequently, consecutive global masks exhibit significant differences, leading large X. As training progresses and the model stabilizes, mask updates become more similar, de- creasing This gradual reduction in prevents excessive global updates, stabilizing convergence and mitigating client drift. We elaborate on the behavior of client drift. Even when client-drift is biased toward a specific dominant client, the update difference for the dominant client naturally de- creases from the perspective of global support, whereas the difference from the local solutions of other clients increases. This pattern reflected the aggregated model, and the difference between consecutive global models can partially capture this behavior (e.g., indicated by a large A). As train- ing progresses and the model approaches the global optimum across local objectives, the differences between consecutive global models gradually decrease, leading to a reduction in X. Conversely, in the absence of client drift, such as in IID scenarios, the model learns fairly across all clients and naturally converges without notable fluctuations in global model differences.
20

---

Published as a conference paper at ICLR 2025
In summary, MADA updates the global model in a way that avoids local divergence, gradually ducing global model differences even without learning rate decay. As shown in Figure decreases sharply in the early stages of training due significant local divergence but gradually converges over time. This result suggests that our understanding aligos with the behavior of X.
Table 8: Comparison of FLOPS between PRISM and other baselines.
DP-FedAvgGAN
GS-WGAN
PRISM
FLOPS
0.002 B
1.94
0.34 B
J COMPUTATION COST ANALYSIS
One might say that SLT process within PRISM increases computational complexity, particularly in large networks. As the number of model parameters grows, the required score parameters also scale linearly. However, in SLT, weight parameters not optimized via gradient descent. The only additional computations stem from the sigmoid and Bernoulli processes introduced during the bi- nary masking procedure. While these operations introduce some additional computational overhead, they remain negligible in the context of GPU-accelerated gradient descent. To demonstrate that the additional computational cost is not substantial, we compare Doattin-point-operee (FLOPS) in Table g GS-WGAN exhibits increased FLOPS due to spectral normalization, which applied stabilize GAN training. In contrast, PRISM, despite incorporating the SLT process, maintains sufficiently low FLOPS, ensuring its computational efficiency.
per second
3
9
2
6
7
Figure 9: Generated images of PRISM where ddpm is used instead of MMD loss.
K SCALABILITY OF PRISM
Due to the challenges associated with training generative models in FL setups, research on federated generative models has primarily focused on simple architectures (e.g., GANs) and datasets, such as MNIST and FMNIST level generation. However, it is important to note that the field of federated generative models remains far behind the current trend of using more complex and powerful models, such as diffusion models, and large-scale datasets. This gap arises from the inherent difficulties in adapting these advanced models FL setups. Existing works have struggled to achieve stable cases, even with MNIST-
performance in challenging setups, such non-lID and privacy-preservin level datasets. Even for centralized differential privacy generative models (Dockhorn et al. 2022, Jiang et al. 2024)) J.raiing generative models is highly challenging and often unstable. Due to these obstacles, prior studies have primarily focused on simpler datasets, such as MNIST, and have been largely restricted to GANs.
K.1 EXTEND TO DIFFUSION MODEL
We further investigate the applicability of PRISM to larger and mode powerful generative models, such diffusion models. Specifically, retain the core concept of identifying SLT but replaced
21

---

Published as a conference paper at ICLR 2025
the MMD loss with DDPM loss (Ho et al. 2020 to guide score updates. Notably, all other figurations remain unchanged. We conduct experiments on the MNIST dataset under a non-llD, privacy-free scenario. TigurePpresent the generated images the server-side, which exhibit slight noisiness. We anticipate that by improving the approach for identifying SLT within diffusion models and further curating the experimental setup, can achieve more promising results. On the other hand, directly applying the MMD loss-based approach used in PRISM, which requires generating set of samples for each update, poses practical challenges due to the well-known limitations of DDPM, including slow sampling and convergence rates. These limitations lead to significant com- putational costs, making the approach practically infeasible. Since our primary focus is lightweight communication cost and stable performance, we leave this future work, efficiently incorporating the powerful and diffusion model federated settings.
K.2 EVALUATION ON LARGE-SCALE DATASETS
Figure 10: Generated images of CelebA 128x dataset and CIFAR100 dataset under non-III and privacy-free scenario.
Table 9: Performance of PRISM using CelebA 128x128 and CIAFR100 datasets under non-IlI and privacy-free scenario.
Dataset
FID
P& R
D& C
CelebA 128
40.2927
0.7738 / 0.006
1.0027 / 0.348
CIFAR100
74.1609
0.6655 0.0719
0.602 / 0.3121
Note that the field of federated generative models have struggled achieve stable performance in challenging setups, such as non-lID and criaacypreserving cases, even with MNIST-level datasets. However, even under these conditions, PRISM the first to consistently demonstrate stable per- formance under these conditions. Despite the poor conditions of federated generative models, it is important to address whether PRISM can be extended to large-scale datasets under limited resources. We provide quantitative and qualitative results on large-scale datasets in Table Figure such as CelebA 128x128 and CIFAR 100 datasets, under non-lID and privacy-free scenario. Both experi- ments are conducted under exactly same settings in Section 5 As result, PRISM demonstrates acceptable performance on high-resolution images; however, it reveals certain limitations when ap- plied to large-scale datasets with diverse classes, such CIFAR 100.
22

---

Published as a conference paper at ICLR 2025
Table 10: FID comparison where distance metric is cosine similarity for MADA.
Distance
Case
MNIST
FMNIST
CelebA
PRISM w/ hd
27.3017 27.7895
46.1652
48,9983
IID, DP
PRISM w/ cos
44.4084
49.0747
PRISM w/ hd
34.2038
67.1648
39.799
Non-IID, DP
PRISM w/ cos
34.9577
69.2994
51.1734
L SELECTION OF OTHER DISTANCE METRIC
Table TO provide the FID performance of PRISM when using cosine similarity instead of hamming distance for MADA. PRISM with cosine similarity achieves competitive performance compared to the scheme using Hamming distance most cases. Although the performance degrades for CelebA in the non-IlD scenario, it still outperforms all baseline methods, further demonstrating the robust- ness of the approach.
M THE EFFECT OF MASKING MECHANISM
Table 11: Effect mask selection on MNIST.
Dataset
Metric
Random
Top-k %
Bernoulli (ours)
Weight
FID
391.8257
20.3616
12.7373
5.9895
MNIST
P&R
0.0/0.0006
0.5232/ 0.3782
0.7323 / 0.5904 0.5556 / 0.5313
0.6783 0.8414
D&C
0.0/0.2511
0.2854 0.0
0.446 0.678
In Table 11, we compare the effect of the mask extraction algorithm on generative model training- For this, relied on prior studies on the existence of SLT and the convergence of MMD with characteristic kernel [46, 14]. Our results empirically show SLT's convergence under mask aver- aging generative FL setting. We also explore the Random and Top-k% algorithms. The Bernoulli method employed in PRISM outperformed top-k% while Random and Weight represent the lower and upper bounds of performance achievable by SLT, respectively.
N EVALUATION ON THE CENTRALIZED SETTING
Table 12: Comparison of PRISM with centralized setting.
Method
Metric
MNIST
FMNIST
CelebA
FID
34.2038
67.1648
39.7997
PRISM
P&R 1
0.4386 / 0.4236
0.4967 0.1231
0.6294 0.0713
D&C
0.1734
0.1597
0.2748 0.1681
0,4565
0,2967
5.8238
5.5004
19. 19.1512
PRISM (vanilla)
P&R
0.6913 0.851
0.6985 / 0.8534
0.6621 0.3895
D&C
0.4689 0.679
0.4864 / 0.6965
0.5348 0.5947
In Table 12 we report the results of models trained under FL vanilla (centralized data) setups. PRISM shows performance degradation due to privacy, data heterogeneity, and communication over- head in FL settings. However, centralized training is not applicable in scenarios where distributed samples canmot be shared, which is the primary focus of our work. Therefore, our original submis- sion did not include this comparison, consistent with existing FL research.
23

---

Published as a conference paper at ICLR 2025
O ADDITIONAL RESULTS
O.1 DIFFERENT RANDOM INITIALIZATION OR SEED
Table 13: Quantitative performance of PRISM across different random seeds on the MNIST dataset.
IID, DP
seed
FID
P& R
D& C
30
27.3017
0,4377 0.5576
0.1738 / 00.1982
PRISM
123
26.8707
0,4342 0.5399
0.174 / 0.2047
1234
26.9664
0.4417 0.5842
0.1785 0.1997
Non-IID, DP
seed
FID
P& R
D& C
30
34.2048
0.4386 00.4336
0.1734 / 0.1597
PRISM
123
34.0586
0,4372 / 0.3306
0.1732 / 0. 1618
1234
32.9423
0.3921 0.4182
0.1496 0.1587
PRISM learns a binary mask over frozen, randomly initialized weights. One might argue that the way model parameters initialized can significantly impact performance. To validate the robustness of PRISM, we present quantitative performance across various random seeds using the MNIST dataset in Table 13 Note that the random seed 30 in the first row of each table corresponds to the default configurations described Section all of experimental setup is identical to Section 5.1 and Section 5.2
O.2 ABLATION STUDY ON THE LEARNING RATE FOR BASELINES
Table 14: Baseline's FID performance for adjusting learning rate using MNIST dataset.
Non-IID, DP
1e-4, le-4 (default)
1e-3, 1e-4
1e-5, le-4
Ie-4, 1e-3
1e-4, 1e-5
GS-WGAN
diverge
108.0657
diverge
diverge
96.1892
Non-lID, DP
5e-4 (default)
1e-3, 5e-3
1e-2, 5e-4
1e-4, 5e-4
1e-3, 1e-5
Fe-eeAAVGAAN
153.9325
206.759
diverge
177.3752
232.2336
In this section, we present the additional experiments to validate the fairness of our entire experi- ments. In Table 14 sweep learning rates for GS-W GS-WGAN and DP-FedAvgGAN and a FID perfor- mance on MNIST dataset with Non-IID and srivacy-preserving scenario. Since they utilize GANs, there separate learning rates for the generator and discriminator, and we examined these combi- nations. In the tables below, the top row represents the discriminator' s Ir generator's s Ir, while the subsequent row reports the FID scores. Due the notorious instability of GAN training caused by the coupled learning dynamics between the generator and discriminator, we observed divergence in some settings. For GS-WGAN, the best performance achieved with le-4 / 1e-5, but it still failed to generate MNIST images properly.
P CONTEXTUALIZE THE PRACTICAL SCENARIOS
Federated learning (FL) is commonly motivated by real-world scenarios such healthcare, finance, and mobile devices, where direct data sharing is infeasible due to privacy concerns. In this con- text, our work does not aim generate private data directly but instead focuses on enabling stable and effective generative modeling under these constraints. For instance, in healthcare, PRISM allows multiple institutions to collaboratively train generative models without sharing sensitive patient data. The resulting models can facilitate privacy-preservin synthetic data generation, imputation of miss- ing medical records, data augmentation for low-data scenarios, and solving inverse problems such denoising and reconstruction. Similarly, in mobile devices, PRISM enables decentralized train- ing of generative models while preserving user privacy, which can enhance personalized content generation and user-specific data augmentation.
24


Reviews:
[{"review_id": "9NWYlCrFNH", "reviewer": "Reviewer_Aeqb", "summary": "This paper introduces PRISM, a federated learning (FL) framework designed specifically for generative models to learn under heterogeneous data and lower the communication costs. The main idea leverages the strong lottery ticket (SLT) hypothesis by identifying an optimal sparse subnetwork with high generative performance through stochastic binary masking. The communication overhead is reduced by exchanging binary masks rather than full model weights. PRISM also includes a mask-aware dynamic moving average aggregation (MADA) to mitigate client drift, and maximum mean discrepancy (MMD) loss for stable training in generative tasks. Experiments demonstrate that PRISM outperforms existing federated generative models, such as GAN trained with DP-FedAvg on different image datasets.", "questions": "- What is the baseline performance of purely server-side SLT? E.g. without any FL setup. \n- Would the approach work for diffusion-based generative models?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The algorithm is a novel for learning the binary mask in SLT hypothesis under FL setting. The paper is well-written and the included MMD loss and MADA are well-justified approaches to handle data heterogeneity.\n- The experiments considered different settings (e.g. IID v.s. non-IID and privacy v.s. no privacy) and the proposed PRISM showed considerable gain compared to the baseline approaches.", "weaknesses": "There seems to be limited advances in GAN in the recent generative models literature so I am not sure if the experiment setting of GAN with simple image datasets (e.g. MNIST or CIFAR10) would still be a significant result for the community.", "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": "jNPLNf4UsO", "reviewer": "Reviewer_1FdJ", "summary": "The core idea of this manuscript is to search for the optimal random binary mask of a random network to identify sparse sub-networks with high-performance generative capabilities. By communicating the binary mask randomly, PRISM minimizes communication overhead. Combining maximum mean discrepancy (MMD) loss and a mask-aware dynamic moving average aggregation method (MADA) on the server side, PRISM achieves stable and robust generative capabilities by mitigating local divergence in federated learning scenarios. Moreover, due to its sparse nature, models generated by PRISM are lightweight and well-suited for environments like edge devices without requiring additional pruning or quantization.", "questions": "None", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The PRISM framework demonstrates a degree of innovation; although techniques such as the strong lottery ticket hypothesis and MMD are existing methods, the authors have effectively incorporated them into the federated learning setting.\n\n2. In terms of privacy, PRISM alleviates privacy leakage to some extent by introducing an (ϵ, δ)-differential privacy mechanism.", "weaknesses": "1. Although PRISM performs well on small-scale datasets, the manuscript does not adequately discuss its scalability and performance on large-scale datasets and larger models.\n\n2. The search process for random binary masks may increase computational complexity, especially in large networks. It is recommended that the authors analyze the specific computational costs of this process.", "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": "46XKECH6dr", "reviewer": "Reviewer_NWco", "summary": "The authors propose a federated masking method for generative models that reduces communication and memory costs as well as improves performance in non-iid settings. A binary mask over the entire model is selected via the Strong Lottery Ticket (SLT) hypothesis, by which a maximum mean discrepancy (MMD) loss is used to provide gradient feedback to update weight-importance scores. Devices use these scores to generate (via Bernoulli distribution) a binary mask that is passed to the server for aggregation (thereby reducing communication costs). Finally, the server performs mask-aware dynamic moving average aggregation (MADA) whereby the server determines the drift of the global model and uses the value to determine how much of the new mask update should be incorporated. Experimental results show large increases in performance metrics while reducing communication and memory costs.", "questions": "Is MADA derived from other literature, or is this proposed here for the first time?\n\nCould the authors further detail the communication savings of PRISM? It seems that there is no downlink communication savings?\n\nI was a bit confused about the model memory savings (Line 242). Could the authors clarify this? What I took away is that since model weights are frozen, and since they are initialized using the Kaiming Normal distribution, only the -/+ bit needs to be saved? Is this where the memory savings stem from (as shown in Tables 1/2 and Figure 4)?\n\nThe authors mention in Line 276:\"As the global rounds progress, λ gradually decreases, promoting stable convergence.\" Is this due to model convergence arising from learning rate decay? What I mean is that, especially in non-iid settings, the global model generally converges once the learning rate decays and thus this should in turn decrease λ. Is that the correct characterization?\n\nCould the authors provide new non-iid experiments that showcase a more realistic setting?\n\nOverall, I feel that the application of sparse binary mask training to GANs is powerful and effective as shown in this paper. While each component is not especially novel, the combination suited for GANs is a novelty (albeit not a major one). However, the empirical performance is impressive and I believe in conjunction leads to a nice paper. The only disclaimer is that I am not especially well-versed within GAN or FL sparse binary mask literature.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The strongest aspect are the empirical results of the proposed method. PRISM outperforms all of the baselines by a good margin while also being more efficient (in communication and memory). This is impressive.\n\nThe proposed MADA method is an interesting idea to alleviate model drift and improves performance empirically. As detailed in the questions, it seems to improve performance overall but not alleviate the effects of non-iid settings.\n\nThe paper writing and presentation is great. I could follow easily, the diagram is very helpful and the empirical results (both tables and figures) are presented cleanly. Great job!", "weaknesses": "I do not believe that the PRISM method itself is especially novel or specifically tailored for generative models as the authors market it as. The major performance driver behind PRISM seems to be the idea of training a sparse binary mask. This is shown to be quite effective in FedMask (Li et a. 2021), and it makes sense why it would be effective on GANs as well. However, this seems to be a simple application of FedMask coupled with the MMD loss (used in GAN training) and the WADA update to GANs. In fact, outside of using MMD loss, I do not see how PRISM is tailored towards generative models. The entire process could be applied to any other ML models if the MMD loss, which has already been proposed before, is substituted out.\n\nCommunication costs only seem to be saved during the uplink process, when only the binary mask is sent to the server. However, the downlink communication cost remains the same since the learnable weight-importance scores must be sent down to all devices. This is not mentioned by the authors and thus does not tackle all the communication issues. \n\nThe non-iid empirical setting seems slightly odd. Namely the authors partition \"datasets into 40 segments based on class labels and randomly assign four segments to each client\". Usually a Dirichlet split is the best way to simulate a non-iid split of labels amongst classes. As a result, it is odd that the performance of all algorithms stays around the same or *improves* for certain metrics. Precision and Recall improves for some of the other baselines which do not account for non-iidness. Generally non-iid settings should degrade performance in FL settings and that is not the case here for all baselines. This makes me feel that the partitioning was not non-iid enough.", "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": "MsCKJbPMQ1", "reviewer": "Reviewer_EU7U", "summary": "This paper uses the strong lottery ticket hypothesis to search for valuable weights within initializations. They optimize for utilitarian weight masking structures, which provides communication efficiency and memory benefits. They apply their approach to federated learning, specifically within the context of generative models. The authors argue that PRISM is robust for non-IID data and promotes resource efficiency.", "questions": "(1) Given that the initialization is frozen, how robust is PRISM to the random seed? Does this robustness hold across identical hyperparameters?\n\n(2) Could the authors help to contextualize an industrial or academic setting in which PRISM may be deployed in practice?\n\n(3) Are 10 clients really enough to evaluate cross-device \"non-IID\" performance? For example, the authors note that they split MNIST into 40 disjoint partitions based on class labels. I'm assuming this means that each digit has 4 partitions. Then, the 40 portions are sampled uniformly without replacement and 4 each are assigned to a single client. Am I correct to assume the participation percentage is 100% (which would imply cross-silo)? This does imitate label imbalance, but the number of the clients seem far too low. I would have assumed 2 participating clients out of, say, 200 mobile devices to evaluate for non-IID (e.g., [5]). Also, to my knowledge, a more robust and realistic non-IID partitioning strategy in converting centralized datasets to FL datasets is to use LDA partitioning (widely used in topic modeling).\n\n[1] Sparse Random Networks for Communication-Efficient Federated Learning (Isik et al., ICLR 2023) \n\n[2] Can We Find Strong Lottery Tickets in Generative Models? (Yeo et al., AAAI 2023)\n\n[3] DiLoCo: Distributed Low-Communication Training of Language Models (Douillard et al., Arxiv 2023)\n\n[4] Asynchronous Local-SGD Training for Language Modeling (Liu et al., ICML workshop 2024)\n\n[5] Efficient Adaptive Federated Optimization (Lee et al., ICML workshop 2024)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "(1) Figure 1 is quite effective in compactly communicating the main ideas. \n\n(2) Originality and significance: This paper is an interesting read. I will reserve my opinion on these two factors until after further clarification from the authors for (1) in weaknesses. \n\n(3) Clarity and quality: It is clear that much effort has been put into composing the related works.", "weaknesses": "(1) Can the authors help me understand the differences between PRISM, and the works [1-2]? The central methodology of this paper appears to just be a combination of the two papers, up to some minor changes. \n\n(2) The MNIST digits generated in Figures 2/3 for PRISM and baselines look problematic. As these are tasks that we would expect even a simple VAE to be able to do, I am concerned about the significance of the contributions given in this paper. The fact that Figure 2 is in the IID scenario is also concerning, as one might expect contributable training paradigms to be able to generate MNIST digits properly. I am beginning to wonder if I am missing the point of the paper. Moreover, I do not agree with the assertion that CIFAR10 and CelebA are \"complex datasets\" which allows evaluation of \"state-of-the-art image generation\" in line 083. Figure 3 gives some sensible-looking digits. However, this is an empirical paper, without theoretical results. Additionally, the central idea of training for lottery subnetworks has already been examined in prior works, e.g., cited above. Given this, MNIST, CelebA, and CIFAR-10 are not difficult enough to yield publishable empirical results by themselves, and there are far more informative datasets out there (e.g., ImageNet subsets). \n\n(3) In Appendix G, the authors state that they do not include comparisons with centralized settings as they study settings in which \"distributed samples cannot be shared\". I think the concern here is that the evaluated datasets are so mature, and generative central training on them easily practicable on even modern Laptops, such that if the evaluated methods cannot even beat the central model, it is difficult to see the contribution being made. Again, my assessment would differ significantly if this was a theoretical paper, but this paper is mainly empirical. Are there any reasonable, practical settings in which we would be forced to train in such restrictive scenarios? Training generative image models on iPhones with user data? \n\n(4) Could the authors point me to where they discuss (or add in a discussion of):\n\n(a) What is the architecture specification used in the paper?\n\n(b) How are the hyperparameters tuned?\n\nRegarding (a), it would be very helpful to explicitly discuss the architecture details in the text so that it is self-contained (c.f., Table 2 in [1]). That way, readers do not have to go line-by-line through the code link to get this information. \n\n(5) The empirical setup and discussions seem to be mismatched. The prior discussions seem to focus on communication costs, memory, privacy, etc. This strongly suggests the resource-restricted cross-device setting in FL, which is especially the case given the simplicity of model architectures as well as datasets being used for generation. However, the experiments use 10 clients, which is clearly in the cross-silo setting; such datacenters are capable of training far more advanced models (e.g., [3-4]). In this case, that would be diffusion models.", "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)