RAPID: Retrieval Augmented Training of Differentially Private Diffusion Models

We thank the reviewer for the valuable feedback on improving this paper! Please find below our response to the reviewer’s questions.

Could the authors further explain the motivation for integrating RAG into DPDM? For example, consider the "public" model (or pre-trained model) which uses public data to train the diffusion model without adding differential privacy noise. Notice that this model does not include sensitive private data. Is it possible that the differential privacy noise introduced in the fine-tuning step will degrade utility, causing RAPID to underperform this public model? What are the possible scenarios where fine-tuning with differential privacy noise can indeed improve utility?

Thanks for the insightful question. In the pre-training/fine-tuning paradigm, public data typically captures diverse patterns while private data is more task-specific and may exhibit distributional differences. Under reasonable privacy constraints, DP fine-tuning using private data tends to improve model performance. For example, under the setting of ImageNet$\rightarrow$CIFAR10, the model pre-trained on the public data (ImageNet) achieves an FID score of 78.63 with respect to the private data (CIFAR10). After DP fine-tuning with RAPID, the score improves to 25.4 at $\epsilon = 10$, although the gain marginally diminishes under stricter privacy budgets (63.2 at $\epsilon = 1$).

What is the possible reason that DPDM improves sharply in the "EMNIST→MNIST" setting as the privacy budget increases, i.e., for FID 125.7 to 12.9 (the best among all methods), as $\epsilon$ goes from 0.2 to 10?

Unlike pre-training-based approaches, DPDM trains directly on private data with DP, making it more sensitive to privacy constraints. This sensitivity is evident in Table 4, where DPDM's FID scores on CelebA32 deteriorate from 33.0 to 153.1 as $\epsilon$ decreases from 10 to 1. DPDM's strong performance on MNIST at $\epsilon = 10$ can be attributed to the dataset's simplicity (single-channel, grayscale images). For this specific case, using autoencoders and latent diffusion models (as in DP-LDM and RAPID) may be suboptimal. However, DPDM's performance degrades rapidly with stricter privacy budgets or more complex datasets like CIFAR10.

Could the authors further analyze the time complexity of the end-to-end algorithm? Though this method improves inference efficiency by skipping intermediate sampling steps via RAG, is it time-consuming to build the trajectory knowledge base? How does the runtime depend on the knowledge base size? It would be helpful if the authors could provide rough runtime comparisons.

One key feature of RAPID is its efficient knowledge base generation (detailed in Appendix C5). Unlike prior work on retrieval augmented generation (e.g., ReDi), which requires all the trajectories to share the same latent and builds the knowledge base by iteratively sampling tens of thousands of trajectories from a pre-trained diffusion model, RAPID directly computes the trajectory for each sample in the public dataset in a forward pass. The table below shows RAPID's runtime efficiency for various knowledge base sizes on a workstation running one Nvidia RTX 6000 GPU. For comparison, ReDi requires over 8 hours to build a 10K-sample knowledge base.

We thank the reviewer for the valuable feedback on improving this paper! Please find below our response to the reviewer’s questions.

As previously mentioned, some published works on DP diffusion models are not included in this paper's experimental comparisons, which makes the current evaluation less convincing. Notable examples include: “Differentially Private Synthetic Data via Foundation Model APIs 1: Images,” ICLR 2024 “PrivImage: Differentially Private Synthetic Image Generation using Diffusion Models with Semantic-Aware Pretraining,” USENIX Security 2024

We thank the reviewer for highlighting the recent work. We have added a discussion of these papers in Section 2 and conducted comparative experiments with them (details in Appendix C3).

DPSDA (Lin et al., 2024) synthesizes a dataset similar to the private dataset by iteratively querying commercial image generation APIs (e.g., DALL-E 2 and Stable Diffusion) in a DP manner. For a fair comparison with RAPID, instead of using commercial APIs trained on vast datasets (hundreds of millions of images), we train a latent diffusion model (architecture in Appendix B) on ImageNet32 as the query API for DPSDA, using CIFAR10 as the private dataset. The table below shows that RAPID outperforms DPSDA across various $\epsilon$ values in terms of FID scores. Intuitively, RAPID represents a more effective approach for leveraging the public data than querying a generative model trained on such data.

PrivImage (Li et al., 2024) uses a pretraining-then-finetuning approach, querying the private data distribution to select semantically similar public samples for pretraining, followed by DP-SGD finetuning on the private data. The table below compares RAPID and PrivImage's performance across different $\epsilon$ values on CIFAR10 and CelebA64.

CIFAR10

CelebA64

RAPID outperforms PrivImage in most scenarios, with one exception: CIFAR10 at $\epsilon = 1$. This likely occurs because PrivImage selects public data similar to the private data for pretraining. With clearly structured private data (for instance, CIFAR10 contains 10 distinct classes), using a targeted subset rather than all the public data tends to improve DP finetuning, especially under strict privacy budgets. However, this advantage may diminish with less structured private data (e.g., CelebA64). We consider leveraging the PrivImage’s selective data approach to enhance RAPID as our ongoing research.

Are the timesteps $k$ and $v$ fixed during training? If so, how are these values selected, and what is the robustness to different choices? If not, how is the value of the retrieved timestep $v$ determined, given that the key-value pair only stores the embeddings $(h(x_k), x_v)?

The parameters $k$ and $v$ remain fixed during training, defining three key phases: initial sampling $(T-k)$ steps, skipping $(k-v)$ steps, and privatization $(v)$ steps. A larger $k$ allows more extensive knowledge base search through longer initial trajectories, enhancing search accuracy at the expense of computational cost. Meanwhile, our ablation study (see Figure 7) shows that increasing $v$ leads to worse FID scores despite better coverage, highlighting a quality-diversity trade-off. Based on empirical testing, we set default values of $k/T = 0.8$ and $v/T = 0.2$ to balance these factors.

For the comparisons, it would improve clarity to present performance across varying $\epsilon$ in a figure, with each method represented as a curve.

Following the reviewer’s suggestion, we have included figures to compare the performance of different methods under varying settings of $\epsilon$ (details in Appendix C.6).

We thank the reviewer for the valuable feedback on improving this paper! Please find below our response to the reviewer’s questions.

Why is the cost of DP hyperparameter tuning small (Line 307)?

The general DP hyperparameters follow standard practices, with $\delta$ set smaller than the reciprocal of private data samples. RAPID's unique hyper-parameters are $k$ and $v$, which determine the number of initial sampling steps $(T-k)$, skipped steps $(k-v)$, and privatized steps $(v)$. We explore limited combinations of $k,v$ (e.g., $0.3T$ - $0.7T$). Empirically, RAPID shows robustness to $k/v$ settings, minimizing hyperparameter tuning overhead.

Does using public data in this way ever reduce the generative quality of the model? For instance, what happens when the public data is wholly unrelated to the private data?

Thanks for the insightful question! To answer this question, we conduct further experiments to evaluate RAPID’s performance when using dissimilar public/private datasets. We use ImageNet32 as the public dataset. For the private dataset, we use VOC2005 (Everingham, 2005), a dataset used for object detection challenges in 2005, which significantly differs from ImageNet32 and contains only about 1K images. We apply RAPID in this challenging setting, with results shown in the table below. Notably, RAPID outperforms baseline methods (e.g., DP-LDM) across various privacy budgets $\epsilon$. While the performance gains are less substantial compared to scenarios with more similar public/private datasets (e.g., ImageNet32$\rightarrow$CIFAR10), the results still demonstrate RAPID's robustness even when working with dissimilar public/private data.

Moreover, we compare RAPID's performance (without DP) to direct training on the VOC2005 dataset. RAPID improves the FID score from 77.83 to 54.60, showing its ability to effectively leverage the public data even when it differs substantially from the private data.

We also explore the possible explanations (details in Appendix D2). Existing studies (Meng et al., 2022; Zhang et al., 2023) establish that in diffusion models, early stages only determine image layouts that can be shared across many generation trajectories, while later steps define specific details. Wu et al. (2023) further discover diffusion models' disentanglement capability, allowing the generation of images with different styles and attributes from the same intermediate sampling stage (as shown in Appendix D2). Therefore, this disentanglement property enables RAPID to maintain robust performance even when public and private dataset distributions differ significantly, provided their high-level layouts remain similar.

Again, we thank the reviewer for the valuable feedback. Please let us know if there are any other questions or suggestions.

Best,

Authors

We thank the reviewer for the valuable feedback on improving this paper! Please find below our response to the reviewer’s questions.

Over-Reliance on the Quality and Relevance of Public Data. The method’s effectiveness depends heavily on the quality and relevance of the public data used in the early stages of training. If the public data does not adequately represent the private data or is of low quality, the retrieval-augmented generation process will suffer from poor performance.

Like other DP diffusion model approaches (e.g., DPDM, DP-LDM, PrivImage) and the broader pre-training/fine-tuning paradigm, RAPID assumes access to a diverse public dataset that captures a range of patterns. However, RAPID is more flexible: the public and private datasets need not closely match in distribution, as long as the public dataset contains similar high-level layouts (more details in Appendix D2).

The explanation is as follows. Existing studies (Meng et al., 2022; Zhang et al., 2023) establish that in diffusion models, early stages only determine image layouts that can be shared across many generation trajectories, while later steps define specific details. Wu et al. (2023) further discover diffusion models' disentanglement capability, allowing generation of images with different styles and attributes from the same intermediate sampling stage (as demonstrated in Appendix D2). Therefore, this disentanglement property enables RAPID to maintain robust performance even when public and private dataset distributions differ significantly, provided their high-level layouts remain similar.

The integration of RAG to enhance the privacy-utility tradeoff could introduce a critical flaw in privacy protection. The retrieval mechanism in RAPID depends on accessing external knowledge bases or data repositories to inform the generative process. This retrieval is based on similar data points or trajectories, which risks the accidental exposure of private data that was used to build the retrieval database.

We would like to clarify that private data is only accessed during model training, not during the sampling process. Meanwhile, the private data's influence on the trained diffusion model is strictly bounded through i) gradient clipping and ii) noise injection (see Algorithm 2), ensuring differential privacy (DP) guarantees. By the post-processing property of DP, since the trained model is differentially private, all its subsequent uses preserve the DP guarantees.

Given that the method heavily relies on public data in the early stages, how do you ensure that the public data sufficiently represents or aligns with the private dataset? Have you tested the model’s robustness when using diverse or low-quality public datasets that may not closely match the private data distribution?

Following the reviewer’s suggestion, we conduct further experiments to evaluate RAPID’s performance when using degraded and dissimilar public/private datasets. We use ImageNet32 as the public dataset with added Gaussian noise $\mathcal{N}(0, 0.1)$ to degrade its quality. For the private dataset, we use VOC2005 (Everingham, 2005), a dataset used for object detection challenges in 2005, which significantly differs from ImageNet32 and contains only about 1K images. We apply RAPID in this challenging setting, with results shown in the table below. Notably, RAPID outperforms baselines (e.g., DP-LDM) across varying $\epsilon$, indicating its robustness to dissimilar public/private datasets.

Moreover, we compare RAPID's performance (without DP) to direct training on the VOC2005 dataset. RAPID improves the FID score from 77.83 to 54.60, showing its ability to effectively leverage the public data even when it differs substantially from the private data.

We also explore the possible explanations (details in Appendix D2). Existing studies (Meng et al., 2022; Zhang et al., 2023) establish that in diffusion models, early stages only determine image layouts that can be shared across many generation trajectories, while later steps define specific details. Wu et al. (2023) further discover diffusion models' disentanglement capability, allowing the generation of images with different styles and attributes from the same intermediate sampling stage (as shown in Appendix D2). Therefore, this disentanglement property enables RAPID to maintain robust performance even when public and private dataset distributions differ significantly, provided their high-level layouts remain similar.

Again, we thank the reviewer for the valuable feedback. Please let us know if there are any other questions or suggestions.

Best,

Authors