PRISM: Privacy-Preserving Improved Stochastic Masking for Federated Generative Models

Authors: Kyeongkook Seo, Dong-Jun Han, Jaejun Yoo
Primary area: generative models
Keywords: Generative models, Federated learning

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.

Decision: accept (poster)
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729567174524

thread_type:
official_review

<thread_object>
creator: reviewer_eu7u

<content>
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.

strengths:
(1) Figure 1 is quite effective in compactly communicating the main ideas. 

(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. 

(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. 

(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). 

(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? 

(4) Could the authors point me to where they discuss (or add in a discussion of):

(a) What is the architecture specification used in the paper?

(b) How are the hyperparameters tuned?

Regarding (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. 

(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.

questions:
(1) Given that the initialization is frozen, how robust is PRISM to the random seed? Does this robustness hold across identical hyperparameters?

(2) Could the authors help to contextualize an industrial or academic setting in which PRISM may be deployed in practice?

(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).

[1] Sparse Random Networks for Communication-Efficient Federated Learning (Isik et al., ICLR 2023) 

[2] Can We Find Strong Lottery Tickets in Generative Models? (Yeo et al., AAAI 2023)

[3] DiLoCo: Distributed Low-Communication Training of Language Models (Douillard et al., Arxiv 2023)

[4] Asynchronous Local-SGD Training for Language Modeling (Liu et al., ICML workshop 2024)

[5] Efficient Adaptive Federated Optimization (Lee et al., ICML workshop 2024)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730384503453

thread_type:
official_review

<thread_object>
creator: reviewer_nwco

<content>
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.

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.

The 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.

The 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.

Communication 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. 

The 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.

questions:
Is MADA derived from other literature, or is this proposed here for the first time?

Could the authors further detail the communication savings of PRISM? It seems that there is no downlink communication savings?

I 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)?

The 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?

Could the authors provide new non-iid experiments that showcase a more realistic setting?

Overall, 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730603386465

thread_type:
official_review

<thread_object>
creator: reviewer_1fdj

<content>
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.

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.

2. 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.

2. 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.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730675395631

thread_type:
official_review

<thread_object>
creator: reviewer_aeqb

<content>
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.

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.
- 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.

questions:
- What is the baseline performance of purely server-side SLT? E.g. without any FL setup. 
- Would the approach work for diffusion-based generative models?
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a federated masking method for generative models that reduces communication and memory costs as well as improves performance in non-iid settings. The authors have addressed the comments and questions from the reviewers. None of the reviewers against the acceptance of the paper. Please revise and incorporate all the discussions into the final version.