RAPID: Retrieval Augmented Training of Differentially Private Diffusion Models
Authors: Tanqiu Jiang
Changjiang Li
Fenglong Ma
Ting Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion model
retrieval augmented generation
differential privacy
Primary area: generative models
TL;DR: We integrate differentially private training with retrieval augmented generation to significantly improve the utility of diffusion models.

Abstract: Differentially private diffusion models (DPDMs) harness the remarkable generative capabilities of diffusion models while enforcing differential privacy (DP) for sensitive data. However, existing DPDM training approaches often suffer from significant utility loss, large memory footprint, and expensive inference cost, impeding their practical uses.  
To overcome such limitations, we present RAPID: Retrieval Augmented PrIvate Diffusion model, a novel approach that integrates retrieval augmented generation (RAG) into DPDM training. Specifically, RAPID leverages available public data to build a knowledge base of sample trajectories; when training the diffusion model on private data, RAPID computes the early sampling steps as queries, retrieves similar trajectories from the knowledge base as surrogates, and focuses on training the later sampling steps in a differentially private manner. Extensive evaluation using benchmark datasets and models demonstrates that, with the same privacy guarantee, RAPID significantly outperforms state-of-the-art approaches by large margins in generative quality, memory footprint, and inference cost, suggesting that retrieval-augmented DP training represents a promising direction for developing future privacy-preserving generative models. The code is available at:  https://github.com/TanqiuJiang/RAPID

Published as a conference paper at ICLR 2025
RAPID: RETRIEVAL AUGMENTED TRAINING OF DIF- FERENTIALLY PRIVATE DIFFUSION MODELS
Tanqiu Jiang'
Changjiang Li
Fenglong Ma"
Ting Wang
Stony Brook University *Pennsylvania State University
ABSTRACT
Differentially private diffusion models (DPDMs) harness the remarkable gener- ative capabilities of diffusion models while enforcing differential privacy (DP) for sensitive data. However, existing DPDM training approaches often suffer from significant utility loss, large memory footprint, and expensive inference cost, impeding their practical uses. To overcome such limitations, present RAPIE" a novel approach that integrates retrieval augmented generation (RAG) into DPDM training. Specifically, RAPID leverages available public data to build a knowledge base of sample trajectories; when training the diffusion model on private data, RAPID computes the early sampling steps as queries, retrieves sim- ilar trajectories from the knowledge base as surrogates, and focuses on training the later sampling steps in differentially private manner. Extensive evaluation using benchmark datasets and models demonstrates that, with the same privacy guarantee, RAPID significantly outperforms state-of-the-art approaches by large margins in generative quality, memory footprint, and inference cost, suggest- ing that -uumenel DP training represents a promising direction for de- veloping future privacy-preserving generative models. The code is available at: https::/ com Taaoo//aanuuaaaaa
INTRODUCTION
The recent advances in diffusion models have led to unprecedented capabilities of generating high- quality, multi-modal data (Ho et al. 2020 Kong et al., 2020 Bar-Tal al. 2024). However, training performant diffusion models often requires massive amounts of training data, raising severe privacy concerns in domains wherein data sensitive. For instance, Carlini et al. (2023) show that compared with other generative models, diffusion models are especially vulnerable to membership inference attacks, due to their remarkable modeling capabilities; meanwhile, Wen et al. (2023) show that text- conditional diffusion models trained with text-image pairs produce images almost identical to certain training samples with proper prompting.
This pressing need has spurred intensive research enforcing privacy protection in diffusion model training. Notably, [Dockhom et al.] (2023) proposed the concept of differentially private diffusion models (DPDMs), which incorporate DP-SGD (Abadi et al. 2016) into diffusion model training, providing guaranteed privacy: Ghalebikesabi privacy:Chalebikeaahi et al. (2023) further applied DP during diffusion model fine-tuning, first pre-training denoising diffusion probabilistic model (Ho al.! 2020) public data and then fine-tuning the model on private data under DP constraints. Similarly, Lyu et al. (2023) employed the same strategy but extended it latent diffusion models (Rombach al. 2022). ) How-
ever, existing methods suffer from major limitations. (i) Significant utility loss The quality of gen- erated samples often drops sharply under tightened privacy budgets. For instance, DPDM (Dockhorn et al. 2023) fails to synthesize recognizable images on CIFARIC 10 under a DP budget of 1. (ii) Large memory footprint- - To reduce the noise magnitude applied at each iteration, most approaches adopt excessive batch sizes (e.g., B = 8,192 samples per batch). As common DP frameworks (e.g.. ll ) have peak memory requirements of O(B²), this severely limits the size of usable diffusion models. (iii) Expensive inference cost Similar to non non-private diffusion models, existing methods require expensive, iterative sampling at inference, impeding them from synthesizing massive amounts of data.
OPACUS (Yousefpour et al. 202
R APID: Retrieval- Augmented Prívate Diffusion model.

---

Published as a conference paper at ICLR 2025
To address such challenges, present RAPID, novel approach that integrates retrieval augmented generation (RAG) (Lewis et al. 2020 Blattmann et al. 2022) into DPDM training. Our approach based on the key observation that small perturbations to the early sampling steps of diffusion models have a limited impact on the overall sampling trajectory (Khalil 2008). This allows us to reuse previously generated trajectories, if similar to the current one, as effective surrogates to reduce both training and inference costs (Zhang et al. [202).. Leveraging this idea, RAPID utilizes available public data to pre-train diffusion model and build a knowledge base of sampling trajectories. It further refines the model using private data: it first computes the early sampling steps, retrieves similar trajectories from the knowledge base as surrogates, and focuses on training the later sampling steps under DP constraints. Compared to prior work, RAPID offers several major advantages:
It achieves a more favorable privacy-utility trade-off by fully utilizing public data and only train- ing the later sampling steps on private data;
It significantly reduces batch-size requirements by leveraging the retrieved sample trajectories, making the use of large diffusion models feasible;
It greatly improves inference efficiency by skipping intermediate sampling steps via RAG.
Extensive evaluation using benchmark datasets and models demonstrates that RAPID outperforms state-of-the-art methods by large margins in there key areas: (i) generative quality (e.g.. improving FID score to 63.2 on CIFARIO under a DP budget e = 1), (ii) memory footprint (c.g., reducing the required batch size to just 64 samples per batch), and (iii) inference efficiency (c.g., saving up to 50% of the inference cost). Our findings suggest that integrating RAG into DP training represents a promising direction for developing future ppeeeeeires generative models.
RELATED WORK
Differentially private data generation. As an important yet challenging problem, enforcing DP
into training variety of advanced generative models has attracted intensive research effort (Hu
et al. 2023). including generative adversarial networks (Goodfellow et al. | Chen et al.) 2020aa
Yoon et al. 2019), variational autoencoders (Jiang et al.) 2022). and customized architectures (Liew
et al. 2022; Vinaroz et al., 2022, Harder et al. 2021). For instance, |Harder et al. (2023) pre-train perceptual features using public data and fine-tune only data-dependent terms using maximum mean
discrepancy under the DP constraint.
Differentially private diffusion models. In contrast, the work on privatizing diffusion models is relatively limited. Notably, DPDM (Dockhorn et al. 2023) integrate DP-SGD (Abadi et al. 2016) with a score-based diffusion model (Song et al. (202 1): Ghalebikesabi et al. (2023) propose to pre- train diffusion model with public data and then fine-tune the model using DP-SGD private data; DP-LDM (Lyu et al. 2023) apply a similar fine-tuning strategy to a latent diffusion model (Rombach) al. 2022); Privlmage (Li et 2024) queries the private data distribution to select semantically similar public samples for pretraining, followed by DP-SGD fine-tuning on the private data. How- ever, all these methods share the drawbacks of significant utility loss, excessive batch sizes, or expensive inference costs. Beyond the pre-traiiing/ine-tuniin
paradigm, recent work also explores synthesizing DP datasets by querying commercial image generation APIs (e.g., Stable Diffusion and
DALL-E2) to approximate the distributions of private data (Wang et al. 2024 Lin et al. 2024).
Retrieval augmented generation. Initially proposed to enhance the generative quality of NLP models by retrieving related information from external sources Khandelwal et al.] (2019); Lewis al. 2020); Guu al. (2020), RAG has been extended to utilize local cohorts in the training data to facilitate image synthesis. For instance, rather than directly outputting the synthesized sample,
Casanova et al. (202I) I) compute the average of the sample's nearest neighbors in the training data. Blattmann et al. (2022) use external image dataset to provide enhanced conditional guidance, augmenting the text prompt during the text-to-image generation. Zhang et al. (2023) explore RAG to accelerate the inference process of a diffusion model by reusing cre-coeputed sample trajectories as surrogates for skipping intermediate sampling steps. However, existing work primarily focuses on employing RAG in the inference stage to improve generative quality efficiency.
2

---

Published as a conference paper at ICLR 2025
XT
XD
Diffusion Model Ep
Initial Sampling
Continual Sampling
Early Steps
Later Steps
Feature Extraction
Xk
Z
Nearest Neighbor Search
2,-
Value
Query
Key
Trajectory Knowledge Base
Latent Feature Extractor
Nearest Neighbor
Figure I: Overall framework of RAPID,
To our best knowledge, this presents the first work on integrating RAG in the DP training of diffusion models, aiming to improve the generative quality, memory footprint, and inference efficiency over the existing DPDM approaches.
3 RAPID
Next, present RAPID, novel approach for training differentially private diffusion models by leveraging retrieval-augmmnntd generation.
3.1 PRELIMINARIES
A diffusion model consists of forward diffusion process that converts original data x its latent x, (where denotes the timestep) via progressive noise addition and a reverse sampling process that starts from latent X; and generates data Xo via sequential denoising steps.
Take the denoising diffusion probabilistic model (DDPM) (Ho et al. 2020) as an example. Given Xo sampled from the real data distribution Qdata, the diffusion process formulated Markov chain:
g(x,|x;-1) = (x;; V1 (,
where {3 E (0,1) specifies the variance schedule. For sufficiently large T, the latent xT approaches isotropic Gaussian distribution. Starting from p(xr) = N(xx 0, I), the sampling process maps latent Xy to data Xg in qdata as a Markov chain with learned Gaussian transition:
po (xx-1 (xx)) (xx po (x,, t), Eg (x,, t))
To train the diffusion model eo t) that predicts the cumulative noise up timestep t for given latent x,, DDPM aligns the mean of the transition (x, 1 with the posterior q(x,-1 1 [x, Xo):
min Ex 2-a. -eepellleeeeeeaaee Eelva + vI - āge, 1) 0
where ā, III (1 3r) 7=1
Once trained, starting from Xy ~ NN I), the sampling process iteratively invokes eg:
X = Ex((x, t),
(3)
which generates the following trajectory (xr, ..
3.2 DESIGN RAPID
Prior work training DPDMs (Dockhorn et al. 2023 Ghalebikesabi al. 2023 Lyu et al. 2023 often applies SPDSGD (Abadi et al.) (016) to fine-tume tune the entire sampling process using private data, resulting in significant utility loss and inference cost. However, it is known that, within given sampling trajectory (xr. X7-1.... xo), the early steps only determine the high-level image layout shared by many latents, while the later steps determine the details (Khalil) 2008 Meng al. 2021
3

---

Published as a conference paper at ICLR 2025
Algorithm/ 1: Training latent feature extractor.
reference data 2D, pre-trained diffusion model EØ, timestep le Output: latent feature extractor h while not converged yet do
Input:
foreach x E D do
// generate positive and negative pairs
generate x, xt, and N&x with random augmentations:
generate latents timestep k
sample Šr. X and &g for x E following Eq
// compute contrastive loss compute (x) following Eq / update feature extractor update h to minimize (x);
return
Thus, instead of privatizing the end-to-en sampling process, by fully utilizing the public data, we may skip intermediate steps and focus on fine-tuning the later steps using private data.
Motivated by this idea, as illustrated in Figure I RAPID first pre-trains diffusion model Eş using the public data Drub Further, RAPID builds knowledge base KBB by calculating the diffusion tra- jectories of DDub, RAPID then fine-tunes using the private data DDrv as follows. Corresponding to each input x Dery, generates its initial steps XT:k in the sampling process, uses X (at timestep k) query to retrieve similar trajectory ST:0 from KCB, and resumes DP training the sampling process, starting from x, (at timestep v) of the retrieved trajectory, to reconstruct x. Intuitively, this RAG strategy skips the sampling process from timestep & to e, thereby improving privacy saving, generative quality, and inference efficiency. Next, we elaborate on the implementation of RAPID'S key components.
3.3 BUILDING TRAJECTORY KNOWLEDGE BASE
ref We divide the public data Deub into two parts pre and Drub to avoid overfitting, with Dpub to
pre pre-train the diffusion model Ep and DPub to construct the trajectory knowledge base K.B.
For each X E DDub we construct its sampling trajectory by iteratively applying Eq to generate sequence of latents (x1....; x J. and store (xx: as key-value pair (k: > v) in KB. During RAG, we may sample random latent Xy at timestep T and generate its early trajectory by iteratively invoking Eg (Eq.3) until timestep k: {xT, (XT,XT-1)) XT-1..... x&} and use X2 as query to search for its nearest neighbors (in terms of (2-norm) in K.B. For simple datasets (e.g.. MNIST), due to their distributional sparsity, this straightforward approach effective as possible enforce all the trajectories to share fixed initial latent xy (Zhang al., 2023). However, for more complex datasets (e.g., 10), their distributional density necessitates allowing different trajectories to have distinct initial latents, making this approach much less effective. More importantly, enforcing the same initial latent severely limits the model's generative quality and diversity.
Thus, instead of applying the similarity search the latents {x+} directly, we first extract their features (by applying a feature extrac- tor h) and perform the search in their feature space. To this end, we first project X to the input space by applying one-step denois- ing on Xy using the pre-trained diffusion model eg, and then apply the feature extractor on the denoised X to extract its feature: Z h(c,(Xe, k)).. For simplicity, we omit the one-step denoising in the following notations: Z h(xx).
7%
Festure Extraction h
Latents al Timestep À
X
Nake Augmented Sampling
X
We employ contrastive learning (Chen et 2020b c) train the feature extractor h. Intuitively, contrastive learning learns represen- tations by aligning the features of the same input under various aug- mentations (c.g., random cropping) while separating the features of different inputs. In our current implementation, we extend the SimCLR (Chen al. (0200b framework, as illustrated in Figure Specifically, for each input DDuh pair of its augmented views (xx*+) forms "positive" pair,
Augmestation
X
Figure 2: Contrastive learning of noise-augmented latents {xx
4

---

Published as a conference paper at ICLR 2025
Algorithm 2: RAPID.
Input: private data Dr pre-trained denoiser cø, fcature extractor h, trajectory knowledge base K.B, batch B, timestep k, number iterations I, gradient norm bound C, noise seale o Output: fine-tuned diffusion model Eq
for
[7] do
sample a batch B of size B from Drry via Poisson sampling:
foreach x E B do
// find nearest neighbor
sample X5 following Eq. find key 2g closest to h(xx) ) in XB;;
// skip intermediate steps
fetch value &. corresponding to 2n: compute gradient g(x) t Valnu x.) following Eq.5
// clip gradient
[(x) g(x)/ max(1, L
// apply DP noise
(BB) + Exee 8(x) + CN(
10 0 < Adam g(B));
11 return <s:
while a set of augmented views of other inputs N forms the "negative" samples. The contrastive loss defined by the InfoNCE loss (Oord et al. 2018). which aims maximize the similarity of positive pairs relative to that of negative pairs:
cxp(sim(N(xx))MM333)))
(xr(() = llgg Ex-ENt ENG. exp(sim((((x))) (xp(()))) h(())))
(4)
where X5 denotes the sampled latent timestep l corresponding x (similar for x and x ) sim is the similarity function (e.g., cosine similarity), and r is the hyper-parameter of temperature. Algorithm sketches the training of the latent feature extractor.
After training the latent feature extractor h. we build the trajectory knowledge base KCB. For each DPub sample its trajectory Xg.... XT); consider x 's feature, Z₁ h(xx)) as the key and x, as the value, and store the key-value pair (Zk, x.) into K.B.
3.4 TRAINING DIFFERENTIALLY PRIVATE DIFFUSION MODEL
Leveraging the trajectory knowledge base KBB, further train the denoiser eg on private data Dery Notably, we focus on training Eg from timestep to 0, leading to the advantages of fully utilizing limited private data and reducing overall privacy costs.
As outlined in Algorithm 2. at each iteration, we sample batch B from Dprv using Poisson sampling for privacy amplification (Mironov al. 2019). For each input x B, we (i) sample its early trajectory X3 up to timestep k. (fi) use its feature z = h((x;) as a query to find 2 's nearest neighbor 2. in KBB as the starting point at timestep v, and
2n K.B, (iii) reuse the value &, corresponding (iv) train eg to reconstruct x. In other words, ep is fine-tuned to predict the random noise (x x,) at timestep v. To make the training differentially private, we extend DP-SGD (Abadi et al. 2016 during updating Eg. For each input x in batch B, compute its diffusion loss as:
X, - Va,x VI a.
(&. eolVa. /1 C. Cs'
(x, x.) = Exr=)((),,))
vx)x),")) (5)
a.
where the random noise (x
xx) is scaled with a randomly sampled timestep v'. We compute the gradient g(x) = VolDM x.). To bound g(x)'s influence on cg. we clip g(x) using its & norm. We then sanitized the per- gradient as gB)) by applying random Gaussian noise before updating eg using the Adam optimizer (Kingma & Ba 2014).
We prove the privacy guarantee of Algorithm 2 under Rényi differential privacy (RDP) (Mironoy 2017, which can be converted to (e,8)-DP The following theorem formulates the guarantee pro- vided by RAPID (proof deferred to (A).
Theorem 1. Using the sanitized per-batch gradient g(B) to update eo satisfies (a, 20 R
5

---

Published as a conference paper at ICLR 2025
The overall privacy cost of RAPID is computed via RDP composition (Mironov 2017). which can be further improved using more advanced privacy accounting (Gopi et al. (021).
The inference of RAPID runs as follows. By sampling random Gaussian noise xy at timestep T, generate its early trajectory X up to timestep k; using the feature žy of x as the query, we search for 2 's nearest neighbor d in K.B; then the corresponding value &, as the starting point at timestep v to resume the sampling. Compared with prior work (Lyu et al.] [203]] Dockhorn et al.) 2023 this RAG-based inference also significantly improves inference efficiency.
EVALUATION
4.1 EXPERIMENTAL SETTING
Datasets. We focus on the image synthesis task. In each task, we use the public dataset Dpub to pre-train the diffusion model and build the trajectory knowledge base for the retrieval-augmented generation, and use the private dataset Drub to further fine-tune/train the diffusion model in a dif-
ferentially private manner. Specifically, we consider the following settings: i) EMNIST (Co- hen et al. 2017) (public) and MNIST (Deng, 2012) (private), ii) ImageNet32 Deng et al. (009)) (public) and CIFARIO (Krizhevsky et al. 2009) (private), iii) FFHQ32 (Karras et al. 2019 (pub- lic) and CelebA32 (Liu et al.) 2015) (private), and iv) FFHQ64 (Karras et al. 2019) (public) and CelebA64 (Liu et al. 2015) (private). More details of these datasets are deferred to Table
Diffusion models. We primarily use the latent diffusion model (Rombach et al. 2022) as the under- lying diffusion model and DDIM (Song et al. 2020 as the default sampler.
Baselines. We mainly consider two ofate--o-thheeee DPDM methods as baselines: differential private diffusion model (DPDM) (Dockhorn et al. 2023) and differential private latent diffusion model (Dp- LDM) Lyu et al. (2023).
Metrics. Following prior work (Dockhorn et al. 2023 Lyu et al. 2023), we use the Frechet Incep- tion Distance (FID) score to measure the generative quality of different methods. In addition, adopt the coverage metric (Nacem et al. 2020)) to measure the generative diversity. Intuitively, given reference real dataset Deeal the coverage is measured by the proportion of samples from Dreal that have at least one sample from the synthesized data Deyn in their neighborhood (with neighborhood size fixed as 5 (Lebensold et al. 20244). Formally,
Coverage
Σ
(6)
|Dreal
x
where 1 is the indicator function and N denotes x's neighborhood.
Privacy. We use OPACUS Yousefpour et al. (2021, ),, DP-SGD library, for DP training and privacy accounting. Following prior work (Dockhorn et al. 2023), we fix the setting ó as 10-5 for the CIFAR10 and MNIST datasets and 10 for CelebA dataset that 8 is smaller than the reciprocal of the number of training samples. Similar to existing work, we also do not account for the (small) privacy cost of hyper-parameter tuning.
To simulate settings with modest compute resources, all the experiments are performed on a work- station running one Nvidia RTX 6000 GPU.
4.2 MAIN RESULTS
We empirically evaluate RAPID and baselines. To make fair comparison, we fix the default batch size as 64 for RPPID and DP-LDM; we do not modify the batch size (i.e., 8,192) for DPDM because the impact of batch size on its performance is significant that it stops generating any recognizable images with smaller batch sizes. By default, we fix the sampling timesteps as 100 across all the methods. For MNIST, we train the diffusion model under three privacy settings c = {0.2, 1, 10), corresponding to the low, medium, and high privacy budgets; for the other datasets, we vary the (1,10).
privacy budget
Conds-con generation. We evaluate the quality of methods, in which, besides the input image guidance signal y (e.g., x's class label) is also pro-
generation by different

---

Published as a conference paper at ICLR 2025
Setting
Privacy (e)
DPDM
DP-LDM
RAPID
0.2
125.7
50.8
24.0
EMNIST-MMNST
50.5
34.9
18.5
12.9
27.2
79.1
63.2
ImageNer32 :CIFARIO
10
109.9
33.3
25.4
Table 1. FID scores of class-conditional generation by different methods.
vided for training (and inference). We consider the EMNIST-MNIST and ImageNet32--IIARI settings. Table compares the generative quality of different methods. Observe that, under the same privacy budget, RAPID considerably outperforms the baselines across most cases. For instance, un- der the BegeNet322~CIFARR
setting, with 1, RAPID attains an FID score of 63.2, while DPDM fails to produce any sensible outputs, highlighting the effectiveness of RAG in facilitating the DP training of diffusion models.
Setting
Privacy
DPDM
DP-L.DM
RAPID
0.2
85.77%
11.35%
96.43%
EMNIST-MMIST (CNN)
1
95.18%
74.62%
98.11%
10
98.06%
95.54%
99.04%
I
50.39%
63.61%
ImageNet32-<CIFARI (ResNet)
10
30.41%
66.02%
67.37%
Table 2. Downstream accuracy of classifiers trained on synthesized data.
We further evaluate the utility of the data synthesized by different methods to train downstream classifiers. For fair comparison, we use the same classifier architecture to measure the down- stream accuracy. For the synthesized MNIST data, train Convolutional Neural Network (CNN) (Krizhevsky et al.] (012) and test its performance on the MNIST testing set. For the synthesized CIFARTO data,
train ResNet-9 (He et al. (016) and evaluate its accuracy the CIFARIO testing set, with results summarized in Table 2 Observe that the classifier trained on RAPID'S synthesized data largely outperforms the other methods. For instance, under the AmageNet32-<CIFFE 10 setting with e = 1. RAPID attains 63.6% downstream accuracy with 10% higher than the baselines, indicating the high utility of the data synthesized by RAPID.
RAPID
DeDM
Airplane Automobile
Bird
Cat
Deer
Dog
Frog
Horse
Ship
Truck
Figure 3: Random samples synthesized by RAPID and baselines trained under the ImageNet32+CIFARRIO setting with = 10.
We also qualitatively compare the class-conditional samples generated by DPDM, DP-LDM, and RAPID trained under the M.NNeec2--CIFARIII setting (with = 10), as shown in Figure It observed that across different classes, RAPID tends to produce samples of higher visual quality, compared with the baselines.
Unconditional generation. We further evaluate the unconditional generation by different methods under the FFHQ32 -222-CCeebb33 and FFHQ64 settings. Table summarizes the FID and coverage scores of different methods. Observe that RAPID outperforms the baselines by large margins in terms of generative quality (measured by the FID score). For instance, in the case of FFHHQ32->CeleAA32 under = 10, RAPID achieves an FID score of 37.3, which is 19.1 % and 51.8% lower than DP-LDM and DPDM. respectively, highlighting its superior generative quality. Meanwhile, across all the cases, RAPID attains the highest (or the second highest) generative di-

---

Published as a conference paper at ICLR 2025
Setting
Privacy (e)
DPDM
DP-LDM
RAPID
135.9
0.087
65.3 0.74
52.8 0.96
FFHQ32 CCelebA32
10
29.8
0.55
38.0 0.98
28.0 0.98
72.2 0.59
60.5 0.90
FTHQ64-> CCelebbA4
10
80.8
0.094
45.2 0.94
37.3 0.93
Table 3. FID (left) and coverage (right) scores of unconditional generation by different methods.
versity (measured by the coverage score). Overall, RAPID strikes the optimal balance between generative quality and diversity among all three methods.
RAPID
Dp-LoM
DPDM
Figure Random samples synthesized by RAPID and baselines FFHQ64 (with = 10).
trained under the setting of
We also qualitatively compare the unconditional samples generated by different methods trained on the CelebA64 dataset (with = 10). 3 shows random samples synthesized by DPDM, DP-LDM, and RAPID (more samples in $C). Observe that in general RAPID tends produce un- conditional samples of higher visual quality, compared with the baselines.
4.3 ABLATION STUDIES
Next, we conduct ablation studies understand the impact of various key factors, such batch size and knowledge base size, on RAPID'S performance. We use the FFHQ32->CelebAB3 (with c = 10) as the default setting.
Retrieval accuracy. Recall RAPID relies retrieving the most similar trajectory from the knowl- edge base. crucial question is thus whether RAPID indeed retrieves semantically relevant neigh- bors. To answer this question, under the class-conditional generation, we evaluate the accuracy of RAPID in retrieving the neighbors from the class corresponding to the given label (c.g., "automo- bile"), We calculate the top-k (k 1)5) accuracy based on the true labels of the retrieved neighbors. Under the setting of ImageNet32- -octNNtt22-IIIAA -CIARROO, RAPID attains 81.2% top-1 accuracy and 94.4% top-5 accuracy, respectively, indicating its effectiveness.
Batch size. In contrast to existing methods (e.g., DPDM and DP-LDM) that typically quire excessively large batch sizes (e.g., 8,192 samples per batch), by fully utilizing the pub- lic data using its RAG design, RAPID can gen- erate high-quality samples under small batch sizes. Here, we evaluate the impact of batch- size setting on the performance of RAPID and DP-LDM, with results illustrated in Figure It can be noticed that RAPID attains an FID score of 29.5 under batch size of 16 while its score steadily improves to 26.67 as the batch size in- creases from 16 to 256, highlighting its superior performance under small batch sizes.
12.
EID Score
37.5
32.5
DP-1 DM
RAPID
Batch Size
Figure 5: Impact of batch size.
Sampler. By default, use DDIM (Song al. 2020 the underlying sampler. Here, evaluate the influence the sampler on RAPID. Specifically, we consider the state-of-the-e PNDM sam- pler (Liu al. 2022) and evaluate the performance of different methods in unconditional generation
8

---

Published as a conference paper at ICLR 2025
Setting
Privacy (e)
DPDM
DP-LDM
RAPID
153.1
72.2
56.6
33.0
42.6
30.3
78.7
68.9
FF11064 #CelebA64
10
86.2
50.3
41.1
Table 4. FID scores of unconditional generation with the PNDM sampler.
tasks. By comparing 3 and Table 4 it is observed that RAPID achieves similar FID scores in both cases, indicating insensitivity to the underlying sampler.
and
Knowledge base size. One key component of RAPID is its trajectory knowledge base that supports RAG. We now evaluate the impact of knowledge base size. As shown in Fig- ure 6 we evaluate how RAPID's performance varies as the knowledge-bast size grows from 100 to 50,000. As expected, both RAPID's FID and coverage scores improve greatly with the knowledge-base size. Meanwhile, even der a small knowledge base (e.g., of size 100), RAPID attains satisfactory performance (e.g., with an FID score of about 40 and a coverage of about 0.9).
1.02
40
0.98
EID Stom
-1.94
Coverage
HD
30
0.50
0.86
100
500
2000
10000
50000
Knowledge Base Size
Figure Impact knowledge-base size.
34
Fraction of privatized steps. Recall that, over the sample trajectory. RAPID samples the initial (T - k) steps, skips the intermediate (A: - steps, and privatizes the later e steps. By de- fault, set k/T 0.8 and v/ 0.2. We now evaluate the influence of the fraction pri- vatized steps v/T RAPID. Figure 7 shows how RAPID'S s FID score varies as v/T increases from 0.3 to 0.7 (with k//T fixed 0.8) under
32-
-.98
30
Coverage
nD
0.3
0.4
0.6
0.7
Fraction Privatized Steps
Figure 7: Impact the fraction privatized steps.
the FFHQ$2-<CelebA33 setting (with = 10). Notably, as more steps are privatized, R RAPID's FID) score deteriorates while its coverage score improves marginally, suggesting an interesting trade-off between the generative quality and diversity. Intuitively, larger fraction of v/l indicates less re- liance the retrieved trajectory, encouraging more diverse generations but negatively impacting the generative quality.
CONCLUSION AND FUTURE WORK
The work represents a pilot study of integrating vetrieval-augmented generation (RAG) into the pri- vate training of generative models. We present RAPID, novel approach for training differentially private (DP) diffusion models. Through extensive evaluation using benchmark datasets and mod- els, we demonstrate that RAPID largely outperforms state-of-the-an methods terms of generative quality, memory footprint, and inference efficiency. The findings suggest that integrating RAG with DP training represents promising direction for designing privacy-preserving generative models.
This work also opens up several avenues for future research. (i) Like other approaches for training DP diffusion models and the broader
traaning/fine-tuniinn paradigm. RAPID relies on access to diverse public dataset that captures a range of patterns and shares similar high-level layouts with the private data. It worth exploring scenarios with highly dissimilar public/private data (Liu al.) b Fuentes et al. 2024). (ii) In its current implementation, RAPID retrieves only the top- nearest trajectory RAG. Exploring ways to effectively aggregate multiple neighboring trajectories could improve generative quality and diversity. (iii) While R RAPID's 's privacy accounting focuses on privatizing the fine-tuning stage, is worth accounting for random noise introduced by the diffusion process further improve its privacy guarantee (Wang et al., 2024). (iv) Although this work pri- marily focuses on image synthesis tasks, given the increasingly widespread use of diffusion models, extending RAPID to other tasks (e.g., text-to-video generation) presents intriguing opportunity.

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
This work is partially supported by the National Science Foundation under Grant No. 2405136, 2406572, and 2212323.
REFERENCES
Martin Abadi, Andy Chu, Ian Goodfellow, H Brendan McMahan, llya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. Proceedings of the ACM Conference on Computer and Communications (CCS), 2016.
Omer Bar-Tal, Hila Chefer, Omer Tov, Charles Herrmann, Roni Paiss, Shiran Zada, Ariel Ephrat, Junhwa Hur, Yuanzhen Li, Tomer Michaeli, et al. Lumiere: A space-time diffusion model for video generation. ArXiv e-prints, 2024.
Andreas Blattmann, Robin Rombach, Kaan Oktay, Jonas Müller, and Bjöm Ommer. Retrieval- augmented diffusion models. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2022.
Nicholas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramèr, Borja Balle, Daphne Ippolito. and Eric Wallace. Extracting training data from diffusion models. In Proceedings of the USENIX Security Symposium (SEC), 2023.
Arantxa Casanova, Marlene Careil, Jakob Verbeek, Michal Drozdzal, and Adriana Romero Soriano. Instance-conditioned gan. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2021.
Dingfan Chen, Tribhuvanesh Orekondy, and Mario Fritz. Gs-wgan: gradient-sanitized approach for learning differentially private generators. In Proceedings of the Advances in Neural Informa- tion Processing Systems (NeurIPS), 2020a.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations, In Proceedings of the IEEE Conference on Machine Learning (ICML), 2020b.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. ArXiv e-prints, 2020c.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. Proceedings of International Joint Conference on Neural Networks (LJCNN, 2017.
Luke N Darlow, Elliot Crowley, Antreas Antoniou, and Amos Storkey. Cinic- is not imagenet or 10. ArXiv e-prints, 2018.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2009.
Li Deng. The mnist database of handwritten digit images for machine learning research. 1EEE Signal Processing Magazine, 29(6):141-142, 2012.
Tim Dockhorn, Tianshi Cao, Arash Vahdat, and Karsten Kreis. Differentially Private Diffusion Models. Transactions on Machine Learning Research, 2023.
Mark Everingham. The PASCAL Visual Object Classes Challenge 2005, 2005. URL http:// host robots......ooo ox. ac. uk wtt.kk/aaaaa//////o//00000
Miguel Fuentes, Brett C Mullins, Ryan McKenna, Gerome Miklau, and Daniel Sheldon. Joint selection: Adaptively incorporating public information for private synthetic data. In Proceedings of the International Conference Artificial Intelligence and Statistics (AISTATS), 2024.
10

---

Published as a conference paper at ICLR 2025
Sahra Ghalebikesabi, Leonard Berrada, Sven Gowal, Ira Ktena, Robert Stanforth, Jamie Hayes, Soham De, Samuel L Smith, Olivia Wiles, and Borja Balle. Differentially private diffusion models generate useful synthetic images. ArXiv e-prints, 2023.
lan J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2014.
Sivakanth Gopi, Yin Tat Lee, and Lukas Wutschitz. Numerical Composition of Differential Privacy. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2021.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training, In Proceedings of the IEEE Conference on Machine Learning (ICML). 2020.
Frederik Harder, Kamil Adamczewski, and Mijung Park. Dp-merf: Differentially private mean embeddings with random features for practical on Machine Learning Research, 2021.
data generation. Transactions
Frederik Harder, Milad Jalali, Danica Sutherland, and Mijung Park. Pre-trained perceptual fea- tures improve differentially private image generation. Transactions on Machine Learning Re- search, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Proceed- ings of the Advances in Neural Information Processing Systems (NeurIPS), 2020.
Yuzheng Hu, Fan Wu, Qinbin Li, Yunhui Long, Gonzalo Munilla Garrido, Chang Ge, Bolin Ding, David Forsyth, Bo Li, and Dawn Song. Sok: Privacy-preserving data synthesis. In Proceedings of the IEEE Symposium on Security and Privacy (S&P), 2023.
Dihong Jiang, Guojun Zhang, Mahdi Karami, Xi Chen, Yunfeng Shao, and Yaoliang Yu. Dpž-vae: Differentially private pre-trained variational autoencoders. ArXiv e-prints, 2022.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
Hassan K. Khalil. Nonlinear Systems Third Edition. Prentice Hall, 2008
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. ArXiv e-prints, 2019.
Diederik Kingma and Jimmy Ba. Adam: Method for Stochastic Optimization. ArXiv e-prints, 2014.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. ArXiv e-prints, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Muster's 's thesis, Department of Computer Science, University of Toronto, 2009.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep con- volutional neural networks. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2012.
Jonathan Lebensold, Maziar Sanjabi, Pietro Astolfi, Adriana Mike Rabbat, and Chuan Guo. Dp-rdm: Adapting diffusion models to private domains without fine-tuning. ArXiv e-prints, 2024.
Kamalika Chaudhuri,
11

---

Published as a conference paper at ICLR 2025
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmente gener- ation for knowledge-intensiie nlp tasks. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2020.
Kecen Li, Chen Gong, Zhixiang Li, Yuzhong Zhao, Xinwen Hou, and Tianhao Wang. PrivImage: Differentially Private Synthetic Image Generation using Diffusion Models with Semantic-Aware Pretraining. In Proceedings of the USENIX Security Symposium (SEC), 2024.
Seng Pei Liew, Tsubasa Takahashi, and Michihiko Ueno. Pearl: Data synthesis via private embed- dings and adversarial reconstruction learning. In Proceedings of the International Conference on Learning Representations (ICLR), 2022.
Zinan Lin, Sivakanth Gopi, Janardhan Kulkarni, Harsha Nori, and Sergey Yekhanin. Differentially Private Synthetic Data via Foundation Model APIs. In Proceedings of the International Confer- ence on Learning Representations (ICLR), 2024.
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In Proceedings of the International Conference on Learning Representations (ICLR), 2022.
Terrance Liu, Giuseppe Vietri, Thomas Steinke, Jonathan Ullman, and Steven Wu. Leveraging public data for practical private query release. In Proceedings the the IEEE Conference on Machine Learning (ICML), 2021a.
Terrance Liu, Giuseppe Vietri, and Steven Z Wu. Iterative methods for private synthetic data: Uni- fying framework and new methods. Proceedings of the Advances Neural Information Pro- cessing Systems (NeurIPS), 2021b.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE International Conference Computer Vision (ICCV), 2015.
Saiyue Lyu, Margarita Vinaroz, Michael F Liu, and Mijung Park. Differentially private latent diffu- sion models. ArXiv e-prints, 2023,
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Er- mon. SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.
Ilya Mironov. Renyi Differential Privacy. In Proceedings of the IEEE Computer Security Founda- fions Symposium (CSF), 2017.
Ilya Mironov, Kunal Talwar, and Li Zhang. Rényi Differential Privacy of the Sampled Gaussian Mechanism. ArXiv e-prints, 2019.
Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jacjun Yoo. Reliable fidelity and diversity metries for generative models. In Proceedings of the IEEE Conference on Machine Learning (ICML), 2020.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. ArXiv e-prints, 2018
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Jiaming Song, Chenlin Meng, and Stefano Ermon. e-prints, 2020.
Denoising diffusion implicit models. ArXiv
Yang Song, Jascha Sohl- Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Proceedings of the International Conference Learning Representations (ICLR), 2021.
12

---

Published as a conference paper at ICLR 2025
Margarita Vinaroz, Mohammad-Anne Amiin Charusaie, Frederik Harder, Kamil Adamczewski, and Mi- jung Park. Hermite polynomial features for private data generation. In Proceedings of the IEEE Conference Machine Learning (ICML), 2022.
Haichen Wang, Shuchao Pang, Zhigang Lu, Yihang Rao, Yongbin Zhou, and Minhui Xue. Dp- promise: Differentially private diffusion probabilistic models for image synthesis. In Proceedings of the USENIX Security Symposium (SEC), 2024.
Yuxin Wen, Yuchen Liu, Chen Chen, and Lingjuan Lyu. Detecting, explaining, and mitigating memorization in diffusion models. In Proceedings of the International Conference on Learning Representations (ICLR), 2023.
Qiucheng Wu, Yujian Liu, Handong Zhao, Ajinkya Kale, Trung Bui, Tong Yu, Zhe Lin, Yang Zhang, and Shiyu Chang. Uncovering the Disentanglement Capability in Text-to- Image Diffusion Mod- els. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. PATE-GAN: Generating synthetic data with differential privacy guarantees. In Proceedings of the International Conference on Learning Representations (ICLR), 2019.
Ashkan Yousefpour, Igor Shilov, Alexandre Sablayrolles, Davide Testuggine, Karthik Prasad, Mani Malek, John Nguyen, Sayan Ghosh, Akash Bharadwaj, Jessica Zhao, Graham Cormode, and Ilya Mironov. Opacus: User-Friendly DockhormDifferential Privacy Library in PyTorch. ArXiv e-prints, 2021.
Kexun Zhang, Xianjun Yang, William Yang Wang, and Lei Li. Redi: efficient learning- free diffusion inference via trajectory retrieval. In Proceedings of the IEEE Conference on Machine Learning (ICML), 2023.
13

---

Published as a conference paper at ICLR 2025
A PROOFS
Definition 1. (Rényi differential privacy (Mironov 2017)) A randomized mechanism M D -
R over domain D and range R satisfies (a, RDP for any two adjacent d, E
M((M()))M(d)) S where Da denotes the Rényi divergence of order
RDP can be converted to DP. If mechanism satisfies (a,e) (a,e)-RDP, it also satisfies (c + DP (Mironov 2017).
log 1/6 a-1 a-I, 8)-
Notably, Gaussian mechanism can provide RDP. Specifically, for function f with sensitivity Af maxa,1 [I] - [(d') releasing [(d) (O o²) satisfies (a, a )-DP (Mironov al. 2019
Now, we prove T'eorrm Recall that Algorithm 2 computes the per-sample gradient g(x) and clips bound its influence gxx g(x)/ max(1, RbOll) then computes the per-batch gradient
g(B) B Exer g(x) and applies Gaussian noise: g(B) - g(B) + (,
Proof. Consider two adjacent mini-batches B and H' that differ by one sample x /xt: B' B (x-) U {. We bound the difference of their batch-level gradients as follows:
llg(B) g(B')l22 11 Σ B(x) - B
Σ
g(x)ll2
XEB
XEB
1
B 5(x ) B 1((x+))1 ||2
V 11ő(x-) )11 + 110(x+) 28 (x"((x>> (x))=(x))
B
S B
VC2 + C2 +20
2C
B
where we use the fact that 0(x and ğ(x+) are bounded by C' and the Cauchy-Schware inequality.
Thus, the sensitivity of g(BB B Following the RDP Gaussian mechanism, releasing the sanitized
batch-level gradient O(B) provides (ax, 2/ )DDP, corresponding to [2 log 1/8
a-1 8)-DP.
B EXPERIMENTAL SETTING
Table summarizes the setting of public and private datasets our experiments.
Public dataset Doub
Private dataset DerV
Pre-training (DPub
Trajectory knowledge base Drns
EMNIST (50K)
EMNIST (10K)
MNIST
Image/Ne132 (1.2M)
ImageNet32 (70K) (Darlow al. 2018)
CIFARIO
FFHQ32 (60K)
FFHQ32 (1OK)
CelebA32
FFHQ64 (60K)
FFHQ64 (10K)
CelebA64
Table 5. Setting of public/private datasets experiments.
Table 6 lists the default parameter setting for training the autoencoder in the latent diffusion model under different settings, while Table 7 summarizes the default parameter setting for training the diffusion model.
C ADDITIONAL RESULTS
C.1 QUALITATIVE COMPARISON
Figure g illustrates random samples synthesized by RAPID and baselines under the setting of FFHQ32 CelebA32 (with 10).
14

---

Published as a conference paper at ICLR 2025
EMNIST--MNIST
ImageNet-sCIFARII -FO -
FFHQ64
Input size
32x32x3
32x32x3
32x32x3
64x64x3
z-shape
4 3
16 16 x
16 16 x
64 x 64 x 3
Channels
128
128
128
192
Channel multiplier
[1, 3,
[1,
11,
[1, 2]
Attention resolutions
[32, 16, 8]
[16, 8]
[16, 8]
[16, 8]
# ResBlocks
2
2
2
2
Batch size
64
64
64
64
Table 6. Hyper-parameters for training autoencoders under different settings.
EMNIST-MMNIST (mneANCC-CCCFFARRFFFFO------ CTC.G//AAAAA/LLEAH------CCCCAA
Input size
32x32x3
32x32x3
32x32x3
64x64x3
z-shape
4 x x
16 16 x
16 x 16 x
32 32 x
# Channels
64
128
192
192
Channel multiplier
[1,2]
[1, 2, 4]
[1,2, 4]
[1,2,4 4
Attention resolutions
[1,
[1, 2, 4]
[1,2, 4
[1,2, 4
ResBlocks
2
2
2
Heads
2
8
8
8
Batch size
64
64
64
32
Spatial transformer
True
True class_label
False
False
Cond stage key
class_label
class
cabes
Conditioning key
crossattn
crossatin
crossattn
crossatin
# Classes
26
1000
1000
1000
Embedding dimensions
512
512
512
Transformer depth
2
2
2
Table 7. Hyper-parameters for diffusion models under different settings.
DP-LOM
DeDM
Figure 8: Random samples synthesized by RAPID and basclines trained under the FFHQ32 -CelebA.3 setting (with 10).
15

---

Published as a conference paper at ICLR 2025
C.2 DISSIMILAR PUBLIC/PRIVATE DATA
To evaluate RAPID's robustness to the distributional shift between public and private data, we duct additional experiments to evaluate RAPID'S performance when using dissimilar public/private datasets. Specifically, we use ImageNet32 as the public dataset with added Gaussian noise N ((, 0.1) to degrade its quality. For the private dataset, we use VOC2005 (Everingham) 2005) (resized to 32x32). a dataset used for object detection challenges in 2005, which significantly differs from Im- ageNet32 and contains only about 1K images. We apply RAPID in this challenging setting, with results shown Table 8 Notably, RAPID outperforms baselines (e.g., DP-LDM) in terms of FID scores across varying e, indicating its robustness to dissimilar public/private datasets.
Privacy (c)
DP-LDM
DPSDA
PRIVIMAGE
RAPID
164.85
142.20
139.07
93.17
10
147.86
130.42
123.89
82.56
Table Performance RAPID and baselines (measured by FID scores) in ImageNec2- case.
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, highlighting its ability to effectively leverage the public data even when it differs substantially from the private data.
C.3 ADDITIONAL BASELINES
We further compare RAPID with more recent work on DP diffusion models. DPSDA (Ln et al. 2024) synthesizes dataset similar to the private data by iteratively querying commercial image generation APIs (e.g.. DALL-E 2) in a DP manner. For fair comparison with RAPID, instead of using commercial APIs trained vast datasets (hundreds of millions of images). following the setting of Li et al. (2024) that replicates DPSDA'S results, we use ImageNet32 for pre-training the public model (also the query API for DPSDA) and CIFARIO the private dataset.
Model Size 90M
Model Size 337M
Privacy (c)
DPSDA
RAPID
DPSDA
RAPID
I
113.6
63.2
89.1
66.5
10
60.9
25.4
43.8
29.0
Table 9. Performance DPSDA and RAPID (measured by FID scores) in the ImageNe:32-sCIFARRI case.
Note RAPID and DPSDA represent two distinct approaches to training DP diffusion models, with the pre-trained model size affecting their performance differently.
For DPSDA, which uses DP evolution rather than DP training to synthesize data, larger tre-trained models tend to lead to better performance. This is demonstrated in DPSDA'S 's ablation study (Lin et al. 2024), where increasing the model size from 100M to 270M parameters improves results by enhancing the quality of selected data. In contrast, methods involving DP training (such as DPDM, DP-LDM, PRIVIMAGE, and RAPID) may not benefit from heavily over-parameterized models, as shown in (Dockhorn et al. 2023). This is because the f--norm noise added in DP-SGD typically grows linearly with the number of parameters.
To empirically evaluate how model complexity affects different approaches, we conduct experiments varying the size of the pre-trained model from 90M 337M parameters (by increasing the latent diffusion model's architecture from 128 to 192 channels and expanding its residual blocks from 2 to 9 compares the performance (measured by FID scores) of DPSDA and RAPID across dif- ferent pre-trained model sizes. As model complexity increases, DPSDA achieves better FID scores, while RAPID shows only marginal performance degradation. Notably, when using the same public dataset and pre-trained model, RAPID consistently outperforms DPSDA, suggesting that it is more effective at leveraging public data under DP constraints.
4). Table
IMAE Li ct al. (024) uses the fine-tuning approach, querying the private data distribution to select semantically similar public samples for pretraining, followed by DP-SGD fine-tuning on the private data. The table below compares RAPID and PRIVIMAGE'S performance across different values on CIFAR10 and CelebA64.
Notably, RAPID outperforms PRIVIMAGE in most scenarios, with exception: CIFAR10 under == 1. This likely occurs because PRIV IMAGE selects public data similar to the private data for pre- training. With clearly structured private data (for instance, CIFAR10 contains 10 distinet classes),
16

---

Published as a conference paper at ICLR 2025
CIFARIO
CelebA64
Privacy (e)
PRIVIMAGE
RAPID
PRIVIMAGE
RAPID
29.8
63.2
71.4
60.5
10
27.6
25.4
49.3
37.3
Table 10. Performance comparison of PRIVIMAGE and RAPID (measured by FID scores).
using a targeted subset rather than all the public data tends to improve DP fine-tuning, 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.
C.4 IMPACT OF SEEEIEVAL-AGGEENTTE TRAINING
RAPID can integrate with existing methods for training DP diffusion models since it is agnostic to model training, though its neighbor retrieval operates on latents, making it compatible only with la- tent diffusion models. To measure the impact of RAPID, we use a latent diffusion model as the back- bone model for both DP-LDM and DPDM, evaluating their performance with and without RAPID. Table shows results on MNIST and CIFAR 10 at = 10. The substantial FID score improvement demonstrates the effectiveness of strieval--augmenter training.
DPDM
DD-
Dataset
w/o
w/
w/o
w/
MNIST
42.9
25.4
27.2
14.1
CIFAR1O
82.2
54.1
33.3
25.4
Table 11. Impact retrieval-augmentee training on existing methods (c = 10).
C.5 KNOWLEDGE BASE GENERATION
While prior work on retrieval augmented generation (e.g., REDI Zhang et al. (2023) requires all the trajectories to share the same latent, building the knowledge base needs to iteratively sample tens of thousands trajectories from a pre-trained diffusion model (e.g.. Stable Diffusion), which highly expensive. For instance, hours to build a 10Klla knowledge base.
workstation rumning one Nvidia RTX 6000 GPU, REDI requires over
Knowledge Base Size
10K
20K
30K
40K
50K
60K
70K
RAPID
2.10s
4.17s
6.12s
8.23s
10.33s
12.19s
14.48s
Table 12. Runtime of RAPID for knowledge base construction.
In comparison, RAPID eliminates this constraint, which allows to directly compute the trajectory for each sample in the public dataset in forward pass. Table 12 shows RAPID'S runtime efficiency for various knowledge base sizes, achieving orders of magnitude faster performance than prior work.
C.6 PERFORMANCE WITH VARYING 6
FID Scores for Different OP Diffusion Methods
Downstream Classification Accuracies for Different DP Diffusion Metheds
120
20-44
300
baitd
27
32
Rrvacy Budget
Figure 9: Performance of RAPID and baselines with varying EMNIST-MNIS (a) FID and (b)
downstream classification accuracy
Figure 9 compares the performance (measured by FID scores and downstream classification ac- curacy) of RAPID and baselines (DPDM, DP-LDM, DPSDA, and PRIVIMAGE) in the case of
EMNIST MNIST under varying settings. Observe that, under the same privacy budget, RAPID considerably outperforms the baselines across most cases.
17

---

Published as a conference paper at ICLR 2025
D DISCUSSION
D.1 COMPARISON OF RAPID AND REDI
REDI Zhang et 2023) also employs some strategies similar RAPID such as constructing tra- jectory knowledge bases at early stages to bypass intermediate steps in the generation process. How- ever, the two methods differ several fundamental aspects.
First, REDI employs RAG in the inference stage to improve generative efficiency, while RAPID inte- grates RAD into the DP training diffusion models. Second, unlike REDI that builds its knowledge base by iteratively sampling tens of thousands of diffusion trajectories from pre-trained latent dif- fusion model (c.g.. Stable Diffusion), RAPID constructs the knowledge base by directly computing the diffusion trajectories via adding scaled version of the initial latent to each sample in the pub- lic dataset, which greatly reduces the computational cost. Last, all the trajectories REDI share the same initill latent. In contrast, the initial latents in RAPID are randomly sampled, significantly improving the diversity of generated samples.
D.2 IMPACT OF PUBLIC/PRIVATE DATA SIMILARITY
Like other DP diffusion model approaches (e.g.. DPDM, DP-LDM, PRIVIMAGE) and the broader pre-training/fine-tmiee paradigm, RAPID assumes access to diverse public dataset that captures range of patterns. However, RAPID more flexible: the public and private datasets need not closely match in distribution, as long as the public dataset contains similar high-level layouts. Here, we explore the possible explanations.
Intermediate Step
Final Step
Figure 10: Disentanglement effects of diffusion models.
Existing studies (Meng et al., 2021 Zhang et al. 2023) establish that diffusion models, early stages determine image layouts that can be shared across many generation trajectories, while later steps define specific details. Wu et al. (2022) further discover diffusion models' disentanglement capability, allowing generation of images with different styles and attributes from the same inter- mediate sampling stage, as shown in Figure TO, 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.
18


Reviews:
[{"review_id": "qAQnjWdPEL", "reviewer": "Reviewer_2WH8", "summary": "The paper presents RAPID, a new approach to train differentially private diffusion models by integrating retrieval-augmented generation (RAG) techniques. RAPID utilizes public data to build a knowledge base of sample trajectories, allowing for efficient use of private data by focusing on training later sampling steps under privacy constraints. This approach improves generative quality, memory efficiency, and inference speed compared to existing methods, demonstrating RAPID's effectiveness in privacy-sensitive generative modeling.", "questions": "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?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Privacy-Utility Trade-off: RAPID effectively balances privacy and utility by using public data for initial steps and focusing on private data for detailed refinement. This selective approach enhances the quality of generated outputs.\nEfficiency: The retrieval mechanism reduces memory and computation demands by bypassing intermediate steps, which enhances scalability for larger models.\nGenerative Quality: Experimental results show that RAPID achieves higher FID and accuracy scores than existing differentially private diffusion models (DPDM), indicating superior output quality and utility.\nInnovative Use of Public Data: The approach of leveraging public data to construct a trajectory knowledge base is novel and significantly reduces the computational cost of training on private data.", "weaknesses": "1.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.\n2.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.", "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": "x5VwhV1tdW", "reviewer": "Reviewer_EUmh", "summary": "This paper introduces RAPID, an approach to training differentially private diffusion models that leverages retrieval augmented generation (RAG). RAPID works by using public data to train a knowledge base of sample trajectories. Then, when training the diffusion model on private data, it retrieves similar trajectories from the knowledge base and focuses on training the later sampling steps in a differentially private manner. \n\nRAPID is compared against baseline private diffusion methods for the task of image synthesis. RAPID outperforms the baselines in most settings in terms of generative quality (measured by FID score).", "questions": "1. Why is the cost of DP hyperparameter tuning small (Line 307)?\n2. 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "* The paper is clearly written, readable, and well-organized.\n* The authors provide a thorough evaluation of RAPID on a variety of benchmark datasets for the task of image synthesis.\n* The results show that RAPID outperforms state-of-the-art methods.\n* This work opens an avenue for future work incorporating RAG into training.", "weaknesses": "- No discussion of the limitation of this approach. See question 2.", "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": "blS9CC9pHY", "reviewer": "Reviewer_AmQE", "summary": "This paper investigates differentially private (DP) data synthesis through diffusion models, proposing an approach that leverages *public* pre-training data to: (1) pre-train diffusion models, (2) train a feature extractor, and (3) derive trajectories from the forward diffusion process to establish a knowledge base. For training on *private* (sensitive) data, this method queries the knowledge base with a noisy *private* sample to retrieve a partially denoised *public* counterpart, which is then used to update the diffusion model via DP-SGD. This process effectively uses public knowledge to bypass certain intermediate training steps, simplifying model fitting. Experimental results indicate promising improvements compared to standard DP diffusion model training.", "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:\n“Differentially Private Synthetic Data via Foundation Model APIs 1: Images,” ICLR 2024\n“PrivImage: Differentially Private Synthetic Image Generation using Diffusion Models with Semantic-Aware Pretraining,” USENIX Security 2024\n\n- 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)$?\n- For the comparisons, it would improve clarity to present performance across varying $\\epsilon$ in a figure, with each method represented as a curve.\n- As noted above, the experiments should more thoroughly investigate performance under strong distributional shifts between public and private data (e.g., by using private data with specific domain features, such as medical data).\n- As far as I understand, the \"Retrieval Augmented Training\" approach is orthogonal to model training and architecture. Therefore, it would be more convincing to compare the addition or omission of “Retrieval Augmented Training” on both DPDM and DP-LDM training, showing that adding the proposed  “Retrieval Augmented Training” consistently improves performance in both cases.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper addresses a relevant and important topic, presenting the material clearly. \n- The proposed method is intuitive, straightforward to implement, and readily applicable in practice.\n- Experimental results are promising, showing effective improvements over standard DP diffusion model training.", "weaknesses": "While the technical contribution is relatively modest (more of a training technique specifically designed for DP diffusion models, which somewhat narrows its scope), the paper’s main value seems to lie in demonstrating performance improvements with this approach.  \n\nHowever, I have concerns regarding the experimental evaluation: (1) some key DP diffusion model baselines are missing, making it difficult to thoroughly assess the claimed performance gains, and (2) the paper lacks adequate investigation into large domain shifts between public and private data, a critical consideration for methods that rely on public data. These gaps leave the reported performance improvements less convincing", "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": "JdyRol9Akf", "reviewer": "Reviewer_iaJi", "summary": "This paper introduces a novel method called RAPID that integrates retrieval-augmented generation (RAG) into the training of differentially private diffusion models (DPDMs). This approach leverages public data to create a knowledge base of sample trajectories, which are used as surrogates during the training on private data. By focusing differential privacy constraints on the later sampling steps and reusing similar public trajectories for the initial steps, RAPID significantly improves the trade-off between privacy and utility. The method improves existing DPDMs by reducing utility loss, memory footprint, and inference costs. Evaluations on benchmark datasets demonstrate that RAPID outperforms state-of-the-art methods in generative quality, memory efficiency, and inference speed. This work suggests that integrating RAG with differentially private training represents a promising direction for developing privacy-preserving generative models.", "questions": "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 $\\varepsilon$ goes from 0.2 to 10?\n\n2. 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.\n\n3. If the private data and the public data are too \"dissimilar,\" how will the trajectory knowledge base trained by the public data affect the later sampling steps?\n\nMinor: \n\n4. Line 10 of Algorithm 2: should it be $\\tilde{g}$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The experiment results demonstrate great improvements in generative quality, compared to existing differentially private diffusion models.\n\n2. The method reduces the memory footprint by minimizing the batch size requirements and enhances inference efficiency by skipping intermediate sampling steps. \n\n3. The paper is well-structured and clearly written, making it easy to follow the progression of ideas and understand the contributions.", "weaknesses": "1. 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?\n\n2. See question 2.", "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)