ICLR 2025

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, we present
RAPID1, 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 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, suggest-
ing that retrieval-augmented DP training represents a promising direction for de-
veloping future privacy-preserving generative models. The code is available at:
https://github.com/TanqiuJiang/RAPID.
1
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 et al., 2024). However, training
performant diffusion models often requires massive amounts of training data, raising severe privacy
concerns in domains wherein data is sensitive. For instance, Carlini et al. (2023) show that compared
with other generative models,...