Generative Modeling with Explicit Memory

Yi Tang1,∗ Peng Sun1,2,∗ Zhenglin Cheng1,2 Tao Lin1,†

1
Westlake University 2 Zhejiang University
yitang86@[Link], sunpeng@[Link],
chengzhenglin@[Link], lintao@[Link]
arXiv:2412.08781v1 [[Link]] 11 Dec 2024

20 20
18 GMem SiT GMem
GMem REPA
18
REPA
Network Network Memory Bank 16 16 SiT
14
14
FID-50K
Generalization Decouple

FID-50K
Generalization Memorization
12
12
10× faster
10
8 46.7× faster 10
Memorization 6 8 5× faster
4
16× faster
6
100K 350K 1M 4M 7M 25 50 250
Training Iterations Number of Functional Evaluations
(a) Key idea of GMem (b) Training efficiency of GMem (c) Sampling efficiency of GMem

Figure 1. GM EM significantly enhances the training and sampling efficiency of diffusion models on ImageNet 256 × 256. We propose
an external memory bank that enhances the semantic information storage, improving both training stability and sampling efficiency in
diffusion models. Sub-figure (a) illustrates the key concept of GM EM, where ϵ represents the initial noise and x0 the generated samples.
In GM EM, we disentangle the generalization and memorization capabilities originally intertwined within diffusion models, assigning
memorization to an external memory bank M. This decoupling reduces computational and capacity overhead for memorizing internal
semantic information, thereby accelerating the overall process. Sub-figure (b) demonstrates the training efficiency of GM EM on ImageNet
256 × 256. At an FID of 5.8, GM EM achieves over 16× speedup relative to REPA [70]. At an FID of 7.3, GM EM achieves over 46.7×
speedup relative to SiT [44]. Sub-figure (c) illustrates the sampling efficiency of GM EM on ImageNet 256 × 256. At the same target level
of FID, GM EM requires 5× less NFEs compared to REPA and 10× less NFEs compared to SiT.

Abstract SiT training by over 46.7×, achieving the performance of a

SiT model trained for 7M steps in fewer than 150K steps.
Recent studies indicate that the denoising process in deep Compared to the most efficient existing method, REPA [70],
generative diffusion models implicitly learns and memorizes GM EM still offers a 16× speedup, attaining an FID score of
semantic information from the data distribution. These find- 5.75 within 250K steps, whereas REPA requires over 4M
ings suggest that capturing more complex data distributions steps. Additionally, our method achieves state-of-the-art gen-
requires larger neural networks, leading to a substantial in- eration quality, with an FID score of 3.56 without classifier-
crease in computational demands, which in turn become the free guidance on ImageNet 256 × 256. Our code is available
primary bottleneck in both training and inference of diffusion at [Link]
models. To this end, we introduce Generative Modeling with
Explicit Memory (GM EM), leveraging an external memory
bank in both training and sampling phases of diffusion mod-
els. This approach preserves semantic information from data
1. Introduction
distributions, reducing reliance on neural network capacity Deep generative models based on denoising [8, 9, 24, 37,
for learning and generalizing across diverse datasets. The 40, 47, 59, 68]—which prioritize the generation of high-
results are significant: our GM EM enhances both training, quality, realistic data—have achieved notable success within
sampling efficiency, and generation quality. For instance, the deep learning community. These methods demonstrate
on ImageNet at 256 × 256 resolution, GM EM accelerates exceptional performance in complex tasks such as zero-shot
† Corresponding author. text-to-image and video generation [4, 15, 50, 51, 56]. How-
∗ Equal contribution. Peng is the project leader. ever, training and sampling diffusion models suffer from

1
the issue of high computational burden [28, 29], due to the We outline the main contributions of this paper as follows:
necessity of using neural networks with larger capacities for (a) We propose that diffusion models operate through two
better empirical performance [28, 29]. fundamental perspectives: memorization and generaliza-
In the meanwhile, a line of research examines the dif- tion. Utilizing neural networks for memorization incurs
fusion models from the perspective of representation learn- significant computational and model capacity burdens.
ing [8, 36, 46, 66], where diffusion models can capture se- (b) We present GM EM, a simple yet extremely effective
mantic features within their hidden layers, with advanced approach to construct an external memory bank that
models producing even stronger representations [66]. Yang stores semantic information inherent in data distribu-
and Wang [68] illustrate that diffusion models strike a bal- tions, alleviating the memorization burden on neural
ance between learning meaningful features and regularizing networks within the diffusion models.
the model capacity. Kadkhodaie et al. [27] further interpret (c) Our framework significantly enhances efficiency and
the generalization behavior in diffusion models, indicating quality of the SoTA diffusion models, for both training
that diffusion models aim to learn and encapsulate semantic and inference (see Figure 1 and Figure 2 ).
information from the training data, subsequently generaliz-
ing to the true data distribution. 2. Related work
In this paper, we extend these insights [27, 68] and conjec-
ture that the functions of diffusion models can be implicitly Generative models. Generative models—including Gen-
divided into two components: (i) memorization of semantic erative Adversarial Networks (GANs) [19, 57, 67], Varia-
information and (ii) generalization to true data. The in- tional Autoencoders (VAEs) [21, 32], flow-based methods,
trinsic generation procedure of diffusion model then can be and diffusion-based methods [12, 24, 45]—aim to learn the
explained as transforming input noisy samples into the inter- data distribution p(x) and generate data through sampling,
nal semantic distribution of a model, followed by “decoding” demonstrating remarkable performance in producing realis-
it into the data distribution. Specifically, tic samples [38]. Recently, diffusion-based methods utilize
stochastic interpolation to model the forward process and
(a) in the training stage, diffusion models require significant
subsequently convert the Gaussian distribution back into the
neural network parameters and computational resources
original image distribution during the reverse process, yield-
to extract and memorize semantic information from
ing realistic samples. Such methods have achieved SoTA
noisy samples;
results in deep generative modeling and are the focus of this
(b) in the sampling/inference stage, models often remap
study [14, 45, 60].
Gaussian noise to the internal semantic distribution
Training diffusion-based models is computationally in-
before retrieving the data distribution, rather than
tensive and unstable [60, 70]. Additionally, the generative
directly transforming noise to data.
process necessitates multi-step sampling, incurring substan-
Building upon the aforementioned analyses, the effi- tial computational costs [42] Accelerating the training and
ciency challenges of training and sampling diffusion models sampling processes of diffusion models has emerged as a
stem from the need to learn and memorize the internal significant challenge, prompting various efforts in previous
semantic information, resulting in higher computational cost studies to address this issue.
and greater demands on model capacity. As a remedy, we
propose GM EM, a new paradigm that uses an external mem-
ory bank to accompany the diffusion models to approximate Diffusion modeling and representation learning. To
the true data distribution, therefore avoiding memorizing overcome the instability and computational inefficiency of
the data distributions in the neural network. Extensive diffusion models, recent studies [17, 45, 70] have turned
experiments demonstrate the effectiveness and efficiency to representation learning to enhance diffusion models. On
of our approach, on both diffusion transformers [44] and one hand, diffusion models are capable of learning high-
convolutional neural networks [24, 28, 61]. For example, quality representations [70]. For instance, Tang et al. [63]
integrating our GM EM with SiT on ImageNet 256 × 256 demonstrated that feature maps extracted from diffusion net-
only requires 150K training iterations to achieve an FID works can establish correspondences between real images,
of 7.3, which is 46.7× faster compared to the baseline indicating a strong alignment between the learned represen-
SiTs [44] and 16× than REPA [70]. Under the same tations with actual image. Furthermore, Yang and Wang [68]
number of training iterations as 400K, our model enhances conducted a detailed analysis of the trade-off between the
the sampling efficiency by 10×, where our model requires quality of learned representations and the penalization of the
only 25 sampling steps to achieve an FID = 12.3 while optimal parameter spectrum.
the vanilla SiT needs 250 sampling steps. Ultimately, our Conversely, well-trained representation models can im-
approach achieves SoTA performance with an FID of 3.56 prove performance and expedite the training of diffusion
after 2M training steps, without classifier-free guidance. models. [45] accomplish this by adjusting the weighting

2
Figure 2. Selected samples on ImageNet 256 × 256. This figure showcases a selection of images generated by GM EM. We utilized the
SiT-XL/2 backbone and trained GM EM for 2M steps without classifier-free guidance.

function in the denoising score matching objective to en- enabling the model to learn better representations. As a
hance representation learning. REPA [70] introduces an result, our method significantly enhances both the training
alignment loss for intermediate layer representations, signifi- speed and effectiveness of diffusion model.
cantly accelerating the training process by over 17.5 times.
In general, these methods enhance diffusion modeling by
employing regularization techniques, which implicitly inte- Text to image generation. Text-to-image generation seeks
grate semantic knowledge from pre-trained representation to synthesize realistic images from free-form textual prompts.
models. Despite advancements, current text-to-image models con-
tinue to experience challenges with text-image misalign-
ment [35, 53]. Recent works [6, 16] have attempted to ad-
Self-supervised learning in diffusion models. Self- dress this issue by intervening in cross-modal attention acti-
supervised learning, as a powerful approach within represen- vations that influence the internal states of Stable Diffusion
tation learning that leverages unlabeled data to learn useful models. However, this approach increases computational
representations, has garnered increasing attention in the dif- complexity and may impede the learning of optimal repre-
fusion modeling community. Many recent works [2, 26, 39] sentations. Indeed, these models often require several days
have adopted self-supervised learning techniques to enhance to weeks for training.
the representation learning capabilities of diffusion mod-
els, thereby improving their performance and efficiency. In 3. Preliminary
particular, Li et al. [39] proposed a novel diffusion training
paradigm that introduces an additional diffusion model to In this section, we provide a concise introduction to the
generate a 1D feature map as a condition to guide the train- training and sampling processes of flow-based and diffusion-
ing of the generative model. However, the inclusion of an based models. See Appendix G for more details.
additional diffusion model introduces extra sampling during Both diffusion-based and flow-based models [44] derive
inference, which increases computational overhead and can their training procedures from a deterministic T -step noising
introduce additional errors. In this work, differing from Li process applied to the original data, formalized as:
et al. [39], we directly utilize representation from ‘memory
bank’ as condition. This straightforward strategy eliminates xt = αt x0 + σt ϵ, ϵ ∼ N (0, I) , (1)
the need for an additional diffusion model, thereby enhanc-
ing representation learning while reducing computational where xt represents the noisy data at time t, x0 ∼ p(x) is
costs. It significantly improves both the training speed and a real data sample from the true distribution, αt and σt are
the effectiveness of the model. In this work, differing from Li time-dependent decreasing and increasing functions respec-
et al. [39], we directly utilize representation from ‘memory tively satisfying αt2 + σt2 = 1.
bank’ as condition. This straightforward strategy not only According to (1), each marginal probability density
eliminates the unnecessary computational burden but also pt (xt ) corresponds to the distribution of a Probability Flow
serves as an effective self-supervised learning method that Ordinary Differential Equation [61] (PF ODE) with velocity

3
Algorithm 1 Training GM EM using memory bank M in diffusion models that separates generalization and memo-
procedure T RAIN GM EM(vθ , D, M, T, αt , σt ) rization: neural networks are employed for generalization,
Initialize model parameters θ while an external memory bank is used to store semantic
for each training iteration do information of data.
Sample a batch of data x0 ∼ D Specifically, we find that conventional generative mod-
Sample timesteps t ∼ {0, . . . , T } uniformly els [24, 60] typically use a neural network to achieve
Generate noise ϵ ∼ N (0, I) both generalization and memorization of true data distri-
Compute noisy data xt = αt x0 + σt ϵ butions. However, the UNet and Transformer-based neural
Sample memory snippet s ∼ M using (9) networks [24, 49] used in modern diffusion models have
Predict velocity vθ (xt , t, s) limited parameters. This makes it challenging for these net-
Compute loss using (8) and (10) works to memorize distribution information, leading to two
Backpropagate and update θ using Optimizer main issues [27]: (a) limited parameters reduce the ability to
end for memorize extensive data distributions, and (b) memorization
end procedure via parameter optimization incurs high computational and
time costs.
To address these challenges and alleviate the memoriza-
field v(x, t), defined as: tion burden on neural networks in diffusion models, we
propose employing an external memory bank to serve as the
v(x, t) = α̇t E[x0 | xt = x] − σ̇t E[ϵ | xt = x] , (2) memorization function, therefore significantly improving
training efficiency and the models’ ability to capture
where α̇t = dα dσt
dt and σ̇t = dt . By solving this ODE
t
distributions. The framework is depicted in Figure 3 .
starting from xT = ϵ ∼ N (0, I), we obtain the probability
density function p0 (x0 ), which can be used to estimate the 4.1. External Memory Bank
ground-truth data distribution p(x).
Alternatively, the aforementioned noise-adding process Inspired by previous studies [27], which demonstrated that
can be formalized as a Stochastic Differential Equation [59, generalization in diffusion models arises from geometry-
61] (SDE): adaptive harmonic representations, we aim to design a mem-
ory bank that: (a) Provides sufficient semantic and represen-
dxt = m(xt , t) dt + g(t) dWt , (3) tational information to assist the diffusion model in gener-
ating high-quality, realistic images that align with the true
where Wt is a Wiener process [23], m(xt , t) is the drift
data distribution. (b) Avoids excessive detailed or textual
coefficient defined as m(xt , t) = − 12 β(t)xt , and g(t) is
p information to prevent overfitting to the training data and
the diffusion coefficient, set as g(t) = β(t) with β(t) ensure robust generalization.
being a time-dependent positive function controlling the We define the memory bank for diffusion models as a ma-
noise schedule. trix M ∈ Rn×m , encompassing n memory snippets. Each
The corresponding reverse process is represented by the snippet s contains m units and satisfies ∥s∥ = 1:
reverse-time SDE:
 ⊤
dxt = m(xt , t) − g(t)2 s(xt , t) dt + g(t) dW̄t , (4) M = s1 , s2 , . . . , sn . (6)
 

where W̄t is a reverse-time Wiener process, and s(xt , t) We employ a robust representation model f under the con-
is the score function, defined by the gradient of the log dition that for every input x sampled from data distribution
probability density: D, the produced normalized feature f (x)/∥f (x)∥ always cor-
1 responds to a vector s that exists in the set M. In conclusion,
s(xt , t) = ∇xt log pt (xt ) = − E [ϵ | xt = x] . (5) our goal is to ensure that all semantic information contained
σt
in D is captured by M.
By solving the reverse-time SDE (4), starting from the ini- To achieve this goal, we formalize the optimization objec-
tial state xT = ϵ ∼ N (0, I), we can obtain p0 (x0 ), thereby tive for M as follows (see our analysis in Appendix A.1 ):
estimating the true data distribution p(x).
(    2 )
f (x)
4. Methodology arg min Ex∼D max M −1 , (7)
M∈Rn×m ∥f (x)∥
The motivation for our method arises from the observation
that diffusion models inherently encode two distinct capa- where max(·) denotes the maximum value of a vector. Note
bilities: generalization and memorization of semantic infor- that n and m are predefined hyperparameters balancing the
mation within data [8, 66, 68, 70]. We propose a paradigm memorization capacity of M against storage constraints.

4
 (c) using cosine similarity:
 
(b) f (x0 )
(b) s = arg max · si . (9)
si ∈M ∥f (x0 )∥
(a) To ensure a fair comparison with the baselines, we use the
M
exact same parameter settings of αt and σt as in previous
studies [44, 70]. For a detailed description of the training and
Figure 3. Data generation via GM EM-enhanced diffusion mod- generation processes, please refer to the framework outlined
els. (a) Sampled noise ϵ is used to index a memory snippet from in Algorithm 1 .
the memory bank. (b) Both the sampled noise ϵ and the memory
snippet s are simultaneously fed into the neural network. (c) The Multiple usage of representation models for enhancing
neural network generates data using SDE or ODE solvers. training diffusion models. A very recent work REPA
[70] introduces a regularization term that leverages represen-
tation models within the diffusion model’s loss function to
Specifically, increasing n and m enhances the memoriza- enhance training efficiency. Therefore, we can integrate this
tion capacity of M but also increases storage requirements. regularization with our modified diffusion loss (8):
Detailed settings and analyses of n and m are provided in
Appendix A.2 . L(θ) ← L(θ) + LREPA (θ) (10)
with LREPA (θ) := −Ex0 ,t [sim (f (x0 ), ψ(h))] ,
Usage of representation models for building the memory where f (x0 ) denotes the pretrained representation, and ψ(h)
bank. Our memory bank framework is adaptable to various represents the projection of the intermediate output of diffu-
representation models f , each with distinct capabilities in sion transformer block h using a trainable projection head
capturing image information. This variability impacts the ψ, implemented as a multilayer perceptron (MLP).
performance of the trained generative models. Therefore, in
this study, we utilize self-supervised representation models 4.3. Sampling with Memory Bank
based on the following reasoning: The sampling stage of unconditional generation in diffusion
(a) Previous research [70] demonstrates that self-supervised models [24, 42] aims to transform Gaussian noise into the
models can expedite the training of diffusion models, true data distribution within an end-to-end framework. This
which can be seamlessly integrated with a memory bank. iterative process employs SDE or ODE solvers (refer to
(b) Prior work [3, 62, 74] points out that self-supervised Section 3) to progressively convert the noise into the desired
models capture semantic information more effectively data distribution. Our framework GM EM with the memory
than supervised alternatives. bank also follows this end-to-end generation process, namely,
We rigorously assess the effectiveness of different represen- feeding the noise into the neural network, together with the
tation models within our framework to build the memory corresponding memory snippet s.
bank like CLIP visual encoder [54] and BERT [11] (see Specifically, our GM EM transforms the noise distribution
Appendix A.3 and Appendix F ). into a corresponding bank snippet index distribution (a.k.a.,
noise-to-index transformation) and samples indices from
4.2. Training Objective this distribution to obtain memory snippets. We propose a
Our proposed memory bank is designed to supply semantic two-step transformation methodology below—namely (1)
information about the data distribution, supporting both the normal distribution to uniform distribution and (2) uniform
training and inference phases of diffusion models. Therefore, distribution to index distribution—that rigorously converts
we incorporate semantic information into diffusion models Gaussian-distributed noise into the index distribution. Such
during the training phase by slightly modifying the diffusion a strategy allows the index distribution exhibiting two key
model loss function as follows: characteristics: firstly, similar noise inputs should corre-
spond to similar indices, thereby maintaining coherence in
Z T
2 the selection process; secondly, each generated index should
L(θ) = E∥vθ (xt , s, t) − α̇t x0 − σ̇t ϵ∥ dt , (8) possess an equal probability of being selected during the
0
training phase, promoting unbiased learning.
where x0 ∼ D, ϵ ∼ N (0, I), α̇t = dα dσt
dt , σ̇t = dt , λ ∼
t

p(λ) and vθ is the velocity estimated by SiT backbone. This Step 1: normal distribution to uniform distribution.
approach trains diffusion models using a interpolant way The initial step leverages the Probability Integral Trans-
over various noise levels [44], incorporating an additional form [52] to convert Gaussian-distributed noise into a uni-
input snippet s, which was selected from memory bank M formly distributed variable. Specifically, we begin with a

5
high-dimensional noise vector ϵ ∈ Rd that follows a Gaus- Backbones. Similar to prior image generation works [70],
sian distribution. By computing and scaling the sum of we utilizes SiT [44] as backbone. This architecture rep-
the vector,
Pdwe obtain a one-dimensional Gaussian variable resent most recent advances in diffusion and flow-based
ϵ = √1d i=1 ϵi , which follows the standard normal distri- transformers [70], and are widely used in image generation
bution N (0, 1). The cumulative distribution function of this tasks [7, 20, 70].
Gaussian variable, denoted by Φ(ϵ), transforms the Gaussian
noise into a uniform distribution over the interval [0, 1] and
Baselines. For a fair comparison, we compare to the SoTA
can be mathematically expressed as:
image generation methods on both training efficiency and
Z ϵ    performance. Specifically, for pixel-space image genera-
1 2 1 ϵ
Φ(ϵ) = √ e−t /2 dt = 1 + erf √ , tion, we consider the following three categories of baselines:
2π −∞ 2 2
First, we compare GM EM with traditional generative models,
where erf(·) denotes the error function erf(z) = including Diffusion GAN [67], Diffusion StyleGAN [65],
R z −t2
√2 e dt. The resulting uniformly distributed variable DMD2 [69]. We also compare the SoTA diffusion mod-
π 0
serves as the foundation for the subsequent transformation. els with UNets, including DDPM [24], Score SDE [61],
EDM [28], DPM-Solver [43], ADM [12], EDMv2 [29],
Step 2: uniform distribution to index distribution. CTM [30], SiD [72], EMD [28]. Finally, we also com-
Given that the memory bank size may be much smaller pare to the SoTA flow-based transformer methods, including
than the total number of images in the dataset, a direct DiT [49], SiT [44], and the most recent yet contemporaneous
uniform distribution may not accurately reflect the true in- work REPA [70].
dex distribution. To address this, we first record the in-
dex distribution from the training set, denoted as Pindex = Metrics. In line with prior research [25, 44], all generation
[p1 , . . . , pi , . . . , pn ], where each pi represents the probabil- quality metrics are reported as FID-50K (FID) [22] scores.
ity of selecting the i-th snippet si during training. Utilizing For CIFAR-10 dataset, we use the training set as the refer-
inverse transform sampling [52] , we map the uniform vari- ence set, and for ImageNet 64 × 64 and ImageNet 256 × 256,
able to the index distribution by defining the variable k as we use the validation set as the reference set. Following Kar-
Pi smallest index i for which the cumulative probability
the ras et al. [28], we also report the Number of trained Images
j=1 pj exceeds the uniform variable: (NI) as a measure of training efficiency. With similar FID
  scores, the model with less NI is considered more efficient1 .
i
 X  Moreover, following [42], we evaluate the efficiency of the
k = min i ∈ {1, . . . , n} Φ(ϵ) ≤ pj .
  sampling by comparing the Number of Function Evaluations
j=1
(NFE) required to achieve a certain FID score. Under our
Upon completion of aforementioned noise-to-index transfor- setting, a lower NFE corresponds to reduced computational
mation, the original noise vector ϵ and the selected memory costs during sampling, thus making the model more efficient.
snippet sk are combined and fed into the diffusion trans-
former backbone. The backbone is then interfaces with ODE
Implementation details. We strictly adhere to the setup
or SDE solvers [24, 59] to generate images.
details outlined in [44] to ensure a fair comparison. For
latent-space generation, we utilize ImageNet [10], where
5. Experiments
images are pre-processed to a resolution of 256 × 256. We
5.1. Experimental Setting follow the ADM framework [12] for additional preprocess-
ing strategies. Followed by Yu et al. [70], the preprocessed
We list the experimental settings below (see more details in
images then been fed into a Stable Diffusion VAE [55] and
Appendix B ).
embedded into the latent space z ∈ R32×32×4 . For model
configurations, we employ the B/2, L/2, and XL/2 architec-
Datasets. For pixel space generation, we evaluate GM EM tures, as introduced by SiT [44], for ImageNet 256 × 256.
on CIFAR-10 [34] for its diversity in classes and widespread
We consistently use the batch size of 128 for CIFAR-10
use in benchmarking image generation models. We also
following [61] and 256 for ImageNet following [29] during
evaluate the scalability of GM EM on ImageNet 64 × 64 [10]
training. Additional experimental implementation details are
to assess the performance on larger resolution generation.
provided in Appendix C .
We also conduct experiments on ImageNet 256 × 256 [10]
to evaluate the ability of GM EM to model the latent space 1 100M NI corresponds to a total of ∼ 400K training steps with the
distribution, a key approach in current image generation batch size of 256, and baselines of different batch-sizes or training steps
research [28, 44, 49, 70]. can be converted to NI in a similar way.

6
Table 1. Sampling quality on various datasets. We report the performance of GM EM on CIFAR-10 (left), ImageNet 64 × 64 (right), and
ImageNet 256 × 256 (middle). GM EM achieves comparable FID with fewer sampling steps across multiple datasets, further highlighting
its advantages in accelerating training.

CIFAR-10 ImageNet 64×64 ImageNet 256×256

METHOD NI (↓) FID (↓) METHOD NI (↓) FID (↓) METHOD NI (↓) FID (↓)
Traditional generative models Traditional generative models Diffusion models (UNets)
Diffusion GAN [67] 51.2M 3.75 DMD2 [69] 98M 1.28 ADM [12] 512M 10.94
Diffusion StyleGAN [65] 50M 3.19 StyleGAN-XL [57] - 1.52 VDM++ [31] 1.43B 3.36
Diffusion models (UNets) MDTv2-XL/2 [18] 332.8M 9.60
Diffusion models (UNets)
Simple diffusion [25] 1B 3.71
DDPM [24] 50M 3.17
SiD [72] 409M 1.52
Score SDE (deep) [61] 166M 2.20 Diffusion models (Transformer)
CTM [30] 61M 1.73
EDM [28] 200M 2.01
EMD [28] 200M 2.20 MaskDiT [71] 1B 5.69
SiD [72] 800M 1.92
DPM-Solver [43] - 3.42 Simple diffusion [25] 1B 3.75
Diffusion (Transformer) EDM (Heun) [28] 2500.5M 2.44 SD-DiT [73] 512M 5.06
SiT-XL/2 [44] 25.6M 6.68 EDM2 (Heun) [29] 2147.5M 1.33 DiT-XL/2 [49] 1.8B 9.60
+ REPA [70] 10M 4.52 SiT-XL/2 [44] 1.8B 8.30
Diffusion models (Transformer) + REPA [70] 1B 5.90
+ GM EM (ours) 10M 1.59
+ GM EM (ours) 2.6M 4.08 SiT-XL [44] + GM EM (ours) 100M 4.98
+ GM EM (ours) 22.5M 1.22 + GM EM (ours) 22.5M 2.10 + GM EM (ours) 512M 3.56

Table 2. Ablation study and sensitive analysis. All models are ing efficiency of GM EM with other baselines on Ima-
based on SiT-L/2 and trained for 100M NI (400K iterations under geNet 256 × 256. We observe that GM EM achieves
the batch size of 256) without using classifier-free guidance on a FID of 4.98 with merely 100M NI, outperforming
ImageNet 256 × 256. ↓ indicates lower values are better.
the current SoTA methods, including SiT and REPA.
Dropout Bank size Bank epoch Solver FID↓ Moreover, we demonstrate that this result can be further
optimized to 3.56 when extends the NI to 512M.
0.2 640,000 5 epochs SDE 5.81
0.2 128,000 5 epochs SDE 9.72
0.2 640,000 1 epoch SDE 5.85
GM EM is a 46.7× training accelerator. We compare
0.2 1,280,000 (Full) 0 epoch SDE 4.21 the advantages of GM EM in terms of training efficiency.
Table 1 and Table 4 also presents the FID scores achieved
0.1 640,000 5 epochs SDE 5.93 by GM EM with limited NI.
0.3 640,000 5 epochs SDE 5.68
For the pixel-space generation, on CIFAR-10, GM EM
0.2 640,000 5 epochs ODE 6.56 achieves performance similar to REPA with only 2.6M NI,
offering more than 3.85× speedup over REPA and over 10×
speedup compared to traditional SiT.
Sampler configurations. Following SiT [44], we use Eu- In latent-space generation, GM EM attains FID of 5.75
ler Maruyama sampler (sde solver) and set the NFE as 50 with NI= 64M, surpassing the most efficient diffusion
for CIFAR-10 and 250 for ImageNet 64 × 64 and ImageNet transformer baseline, REPA, which requires more than
256 × 256 by default. NI= 1000M. This indicates a training speedup of over
16×. Furthermore, GM EM achieves an FID of 7.26 at
5.2. Main Results NI= 38.4M, providing a speedup exceeding 46.7× com-
GM EM surpasses SoTA methods in both training effi- pared to SiT, which achieves the similar FID = 8.3 with
ciency and effectiveness. NI= 1800M.
(a) Pixel-space generation. For pixel-space generation,
as shown in Table 1 , we benchmark GM EM against GM EM is a 10× sampling accelerator. As a benefit of
prior methods using FID and NI. GM EM outperforms substantially lower FID score, GM EM is capable of generat-
all diffusion transformer baselines and achieves com- ing competitive samples with significantly fewer NFEs, re-
petitive performance with traditional generative models, sulting in over a 5× in sampling efficiency. Table 3 presents
obtaining FID =1.22 on CIFAR-10 and FID =2.10 on the FID scores of GM EM at various NFEs.
ImageNet 64 × 64. Specifically, reducing the NFE from 250 to 50 leads to
(b) Latent-space generation. For latent-space generation, increase in FID of 1.7, while still outperforming the previous
in Table 1 , we compare the performance and the train- SoTA, REPA, thereby achieving a sampling speedup of more

7
Table 3. FID with respect to NFE on ImageNet 256×256. In impact on the optimal performance, GM EM still consistently
this table, we present the impact of variations in NFE on FID. generates high-quality images efficiently. The results are
All models are based on SiT-L/2, trained for 100M NI (400K reported in Table 2 and Table 4 .
iterations with a batch size of 256) without using classifier-free
guidance. GM EM achieves higher image quality within NFE = 25,
demonstrating over a 10× improvement in sampling efficiency. Dropout ratio of memory snippets. We first investigate
the impact of different dropout ratios on model perfor-
Model NI NFE↓ FID↓ mance. As shown in Table 2 , introducing a moderate level
of dropout proves advantageous, with a 20% rate applied to
SiT-L/2 25.6M 250 18.8
the memory snippets yielding the best results. We argue that
+ REPA 25.6M 250 8.4
this dropout mechanism implicitly reduces the model’s re-
+ GM EM 25.6M 250 5.8
liance on memory snippets, thus improving generalizability
+ GM EM 25.6M 50 7.5
and overall performance.
+ GM EM 25.6M 25 12.3

Configurations of building memory bank. We exam-

than 5×. Additionally, decreasing the NFE from 250 to 25
ine the impact of various memory bank configurations on
allows GM EM to generate samples that remain competitive
performance, we reported the results in Table 2 . Our find-
with those of SiT, attains an FID of 12.3, providing more
ings indicate that larger memory banks generally yield su-
acceleration of more than 10×.
perior results, with bank size = 1,280,000 leading to op-
timal result of FID = 4.21. We propose that excessively
GM EM is NOT a memory machine. Storing informa- small memory banks impede the retention of essential knowl-
tion related to source data often raises privacy and copy- edge, thereby degrading FID performance, while overly large
right concerns, particularly in the context of image gener- memory banks increase storage requirements and compu-
ation [5, 58, 64]. Although our approach incorporates a tational costs. Typically, allocating 50% of total images in
memory bank M, we clarify that GM EM does not function the dataset to the memory bank—approximately 1.8 GB for
as a system solely for storing pixel-level detailed information ImageNet 256 × 256—provides satisfactory outcomes. Ad-
from images. Instead, the memory bank stores high-level ditionally, we demonstrate that training the memory bank
semantic representations, enabling M to serve as an external for five epochs is sufficient to achieve the optimal results.
knowledge repository. This design facilitates faster training
and sampling processes while mitigating privacy considera-
tions. Choice of SDE/ODE solver. We observe that SDE solver
We demonstrate the semantic information within the consistently produces higher quality images than the ODE
memory bank M through experiments conducted on Im- solver, achieving a FID score reduction of 0.25, as shown
ageNet 256 × 256. Specifically, for one memory snippet s in in Table 2 . Consequently, we employ SDE solvers in all
M, we apply different noise during the sampling process and main experiments.
compare with the generated samples. As shown in Figure 4 ,
GM EM generates images that are semantically aligned with Neural network architecture. We delve into the scalabil-
the original input while differ in specific details. This out- ity of the neural network architecture of GM EM by varying
come confirm that GM EM does not merely memorize M or its size. Table 4 presents the FID-50K scores of GM EM
learn a direct mapping from M to Z, but instead effectively across different model sizes on ImageNet 256 × 256. Our
leverages the high-level semantic knowledge stored in M findings reveal that larger models not only converge more
during the image generation process. rapidly but also achieve lower FID scores, aligning with
Appendix D further demonstrates GM EM’s flexibility in the observations of diffusion transformers by Yu et al. [70]
capturing the knowledge from M, where we interpolate be- and Ma et al. [44]. Notably, compared to smaller models,
tween representations from two different classes to generate GM EM-XL consistently reaches the same FID levels more
a series of smoothly transitioning images. quickly and ultimately converges to a lower FID of 3.56 on
ImageNet 256 × 256.
5.3. Ablation Study This scalability trend extends to pixel-space generation. For
The effectiveness of GM EM depends on several key factors, example, on CIFAR-10, as depicted in Figure 5 , GM EM-
including dropout ratio, bank size, bank training settings, XL achieves an FID of 1.22 more efficiently than smaller
solver, and backbone architectures. In this subsection, we set model variants. These results underscore the scalability of
NI to 100M and conducted a series of ablation studies on GM EM, demonstrating its ability to maintain high genera-
ImageNet 256 × 256. We systematically examined these fac- tion quality while enhancing training and sampling efficiency
tors and concluded that, although each variable has a slight across diverse architectures and datasets.

8
Figure 4. Diverse image generation. The figure presents selected samples from ImageNet 256 × 256 generated by the SiT-XL/2 + GM EM
model using checkpoints with NI = 512M and without classifier-free guidance. The top row shows the original image corresponding to the
memory snippet, followed by four rows of images generated through different noises when sampling. For instance, in the sixth column,
GM EM adeptly captures the semantic essence of a castle nightscape from memory snippets, while presenting variations in elements such as
appearance, lighting, perspective, and other nuanced details.

FID50K-NI (CIFAR-10) Table 4. FID comparisons with various SiT backbones. In this
140 Architectures table, we report the FID of GM EM with different backbones on
GMem-XL GMeM-B ImageNet 256 × 256 at various NI values. For comparison, we
120
GMem-L also present the performance of the state-of-the-art baseline REPA
100 at similar NI or comparable FID values. Note that ↓ indicates
FID 50K

80 that lower values are preferred and all results reported are without
60 classifier-free guidance.
40 Model #Params NI↓ FID↓
20
SiT-B/2 130M 100M 33.0
0
1.50 1.75 2.00 2.25 2.50 2.75 3.00 3.25 3.50 + REPA 130M 100M 24.4
NI (Millions) + GM EM 130M 100M 15.4

Figure 5. Ablation study on the variation of FID-50K with re- SiT-L/2 458M 100M 18.8
spect to NI. We analyze the impact of NI during the training of + REPA 458M 180M 8.4
GM EM using different backbones on CIFAR-10 dataset. Specifi- + GM EM 458M 100M 5.8
cally, while GM EM-XL initially shows a higher FID than GM EM-L
at the early stages of training, after 1.75M NI, GM EM-XL consis- SiT-XL/2 675M 100M 17.2
tently outperforms both GM EM-L and GM EM-B in terms of FID + REPA 675M 100M 7.9
and maintains this lower value until convergence, where the final + GM EM 675M 100M 5.0
FID reaches 1.22.
SiT-XL/2 675M 1800M 8.3
+ GM EM 675M 38.4M 7.3
+ REPA 675M 1024M 5.9
5.4. Additional Applications + GM EM 675M 64M 5.8
+ GM EM 675M 512M 3.6
In this section, we present experimental results demonstrat-
ing that the Memory Bank employed by GM EM exhibits
both cross-dataset transferability and adaptability to down- Transferable memory bank across datasets.
stream text-to-image (T2I) tasks. Appendix E demonstrates the transferability and general-

9
ization of the memory bank used to guide GM EM across Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh.
different datasets. Specifically, we train GM EMIN64 and Video generation models as world simulators. OpenAI Blog,
GM EMCIFAR models on ImageNet 64 × 64 and CIFAR-10, 2024. 1
respectively, resulting in memory banks of MIN64 and [5] Nicolas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski,
MCIFAR . We then directly apply MIN64 to guide the Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito,
sampling process of GM EMCIFAR . Our results show that and Eric Wallace. Extracting training data from diffusion
models. In 32nd USENIX Security Symposium (USENIX
GM EMCIFAR is still able to generate information consistent
Security 23), pages 5253–5270, 2023. 8
with the knowledge provided by MIN64 .
[6] Hila Chefer, Yuval Alaluf, Yael Vinker, Lior Wolf, and Daniel
Cohen-Or. Attend-and-excite: Attention-based semantic guid-
Application to T2I generation. We further demonstrate ance for text-to-image diffusion models. ACM Transactions
that GM EM can be effectively applied to T2I generation tasks on Graphics (TOG), 42(4):1–10, 2023. 3
by encoding text prompts into memory snippets ŝ, which are [7] Junsong Chen, Chongjian Ge, Enze Xie, Yue Wu, Lewei
subsequently utilized by GM EM to generate corresponding Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu,
images. and Zhenguo Li. Pixart-\sigma: Weak-to-strong training of
This integration is achieved through a two-step process: diffusion transformer for 4k text-to-image generation. arXiv
(a) Pretraining the mapping function: We pretrain a map- preprint arXiv:2403.04692, 2024. 6
ping function ψ that transforms text prompt distributions [8] Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He. De-
constructing denoising diffusion models for self-supervised
p(text) into memory snippet distributions p(s) using a
learning. arXiv preprint arXiv:2401.14404, 2024. 1, 2, 4
simple contrastive loss [54].
[9] Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune
(b) Image generation: We then employ the pretrained map-
Gwon, and Sungroh Yoon. Ilvr: Conditioning method for
ping function ψ in conjunction with the trained GM EM denoising diffusion probabilistic models. arXiv preprint
model θ to generate images. arXiv:2108.02938, 2021. 1
Comprehensive descriptions of the T2I generation pro- [10] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li
cess, experimental setup, and the resulting generated images Fei-Fei. Imagenet: A large-scale hierarchical image database.
are provided in Appendix E . In 2009 IEEE conference on computer vision and pattern
recognition, pages 248–255. Ieee, 2009. 6
6. Conclusion [11] Jacob Devlin. Bert: Pre-training of deep bidirectional
transformers for language understanding. arXiv preprint
In this paper, we introduce a new paradigm for diffusion mod- arXiv:1810.04805, 2018. 5, 6
els, GM EM, which accelerates both training and sampling [12] Prafulla Dhariwal and Alexander Nichol. Diffusion models
processes by incorporating a memory bank. Our method beat gans on image synthesis. Advances in neural information
achieves SoTA performance on CIFAR-10 and ImageNet processing systems, 34:8780–8794, 2021. 2, 6, 7
64×64, improving training efficiency by more than 10× and [13] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov,
sampling efficiency by over 5× across datasets. Furthermore, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner,
we demonstrate that our memory bank can be transferred Mostafa Dehghani, Matthias Minderer, Georg Heigold, Syl-
across different datasets and can be easily adapted to T2I vain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is
tasks by generating a memory bank, further showcasing the worth 16x16 words: Transformers for image recognition at
flexibility and generalization capability of GM EM. scale. In International Conference on Learning Representa-
tions, 2021. 2
References [14] Conor Durkan and Yang Song. On maximum likelihood
training of score-based generative models. arXiv e-prints,
[1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden- pages arXiv–2101, 2021. 2
Eijnden. Stochastic interpolants: A unifying framework for [15] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim En-
flows and diffusions. arXiv preprint arXiv:2303.08797, 2023. tezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz,
8 Axel Sauer, Frederic Boesel, et al. Scaling rectified flow
[2] Florian Bordes, Randall Balestriero, and Pascal Vincent. High transformers for high-resolution image synthesis. 2024. 1, 7
fidelity visualization of what your self-supervised represen- [16] Weixi Feng, Xuehai He, Tsu-Jui Fu, Varun Jampani, Arjun
tation knows about. arXiv preprint arXiv:2112.09164, 2021. Akula, Pradyumna Narayana, Sugato Basu, Xin Eric Wang,
3 and William Yang Wang. Training-free structured diffusion
[3] Florian Bordes, Randall Balestriero, and Pascal Vincent. High guidance for compositional text-to-image synthesis. arXiv
fidelity visualization of what your self-supervised represen- preprint arXiv:2212.05032, 2022. 3
tation knows about. Transactions on Machine Learning Re- [17] Michael Fuest, Pingchuan Ma, Ming Gui, Johannes S Fis-
search, 2022. 5 cher, Vincent Tao Hu, and Bjorn Ommer. Diffusion mod-
[4] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei els and representation learning: A survey. arXiv preprint
Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric arXiv:2407.00783, 2024. 2

10
[18] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng [32] Diederik P Kingma. Auto-encoding variational bayes. arXiv
Yan. Mdtv2: Masked diffusion transformer is a strong image preprint arXiv:1312.6114, 2013. 2
synthesizer. arXiv preprint arXiv:2303.14389, 2023. 7 [33] Diederik P Kingma and Jimmy Ba. Adam: A method for
[19] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing stochastic optimization. arXiv preprint arXiv:1412.6980,
Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and 2014. 2
Yoshua Bengio. Generative adversarial nets. Advances in [34] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-
neural information processing systems, 27, 2014. 2 10 and cifar-100 datasets. URl: [Link] cs. toronto.
[20] Ligong Han, Yinxiao Li, Han Zhang, Peyman Milanfar, Dim- edu/kriz/cifar. html, 6(1):1, 2009. 6
itris Metaxas, and Feng Yang. Svdiff: Compact parame- [35] Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins,
ter space for diffusion fine-tuning. In Proceedings of the Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad
IEEE/CVF International Conference on Computer Vision, Ghavamzadeh, and Shixiang Shane Gu. Aligning text-
pages 7323–7334, 2023. 6 to-image models using human feedback. arXiv preprint
[21] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr arXiv:2302.12192, 2023. 3
Dollár, and Ross Girshick. Masked autoencoders are scalable [36] Alexander C Li, Mihir Prabhudesai, Shivam Duggal, Ellis
vision learners. In Proceedings of the IEEE/CVF conference Brown, and Deepak Pathak. Your diffusion model is secretly
on computer vision and pattern recognition, pages 16000– a zero-shot classifier. In ICCV, 2023. 2
16009, 2022. 2 [37] Tianhong Li, Dina Katabi, and Kaiming He. Self-conditioned
[22] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bern- image generation via generating representations. arXiv
hard Nessler, and Sepp Hochreiter. Gans trained by a two preprint arXiv:2312.03701, 2023. 1
time-scale update rule converge to a local nash equilibrium. [38] Tianhong Li, Dina Katabi, and Kaiming He. Return of uncon-
Advances in neural information processing systems, 30, 2017. ditional generation: A self-supervised representation genera-
6 tion method, 2024b. URL [Link] org/abs/2312.03701,
[23] Masuyuki Hitsuda. Representation of gaussian processes 2023. 2
equivalent to wiener process. 1968. 4 [39] Tianhong Li, Dina Katabi, and Kaiming He. Return of uncon-
[24] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffu- ditional generation: A self-supervised representation genera-
sion probabilistic models. Advances in neural information tion method, 2024. 3
processing systems, 33:6840–6851, 2020. 1, 2, 4, 5, 6, 7, 8 [40] Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and
[25] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple Kaiming He. Autoregressive image generation without vector
diffusion: End-to-end diffusion for high resolution images. quantization. arXiv preprint arXiv:2406.11838, 2024. 1
In International Conference on Machine Learning, pages [41] Yanqing Liu, Jianyang Gu, Kai Wang, Zheng Zhu, Wei Jiang,
13213–13232. PMLR, 2023. 6, 7 and Yang You. Dream: Efficient dataset distillation by repre-
[26] Vincent Tao Hu, David W Zhang, Yuki M Asano, Gertjan J sentative matching. arXiv preprint arXiv:2302.14416, 2023.
Burghouts, and Cees GM Snoek. Self-guided diffusion mod- 7
els. In Proceedings of the IEEE/CVF Conference on Com- [42] Cheng Lu and Yang Song. Simplifying, stabilizing and
puter Vision and Pattern Recognition, pages 18413–18422, scaling continuous-time consistency models. arXiv preprint
2023. 3 arXiv:2410.11081, 2024. 2, 5, 6
[27] Zahra Kadkhodaie, Florentin Guth, Eero P Simoncelli, and [43] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan
Stéphane Mallat. Generalization in diffusion models arises Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion
from geometry-adaptive harmonic representation. arXiv probabilistic model sampling in around 10 steps. Advances in
preprint arXiv:2310.02557, 2023. 2, 4 Neural Information Processing Systems, 35:5775–5787, 2022.
[28] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. 6, 7
Elucidating the design space of diffusion-based generative [44] Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M
models. Advances in neural information processing systems, Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring
35:26565–26577, 2022. 2, 6, 7, 8 flow and diffusion-based generative models with scalable
[29] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, interpolant transformers. arXiv preprint arXiv:2401.08740,
Timo Aila, and Samuli Laine. Analyzing and improving the 2024. 1, 2, 3, 5, 6, 7, 8
training dynamics of diffusion models. In Proceedings of [45] Sarthak Mittal, Korbinian Abstreiter, Stefan Bauer, Bernhard
the IEEE/CVF Conference on Computer Vision and Pattern Schölkopf, and Arash Mehrjou. Diffusion based representa-
Recognition, pages 24174–24184, 2024. 2, 6, 7 tion learning. pages 24963–24982. PMLR, 2023. 2
[30] Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Mu- [46] Soumik Mukhopadhyay, Matthew Gwilliam, Vatsal Agar-
rata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki wal, Namitha Padmanabhan, Archana Swaminathan, Srinidhi
Mitsufuji, and Stefano Ermon. Consistency trajectory models: Hegde, Tianyi Zhou, and Abhinav Shrivastava. Diffusion
Learning probability flow ode trajectory of diffusion. 2023. models beat GANs on image classification. 2021. 2
6, 7 [47] Alexander Quinn Nichol and Prafulla Dhariwal. Improved
[31] Diederik Kingma and Ruiqi Gao. Understanding diffusion ob- denoising diffusion probabilistic models. In International
jectives as the elbo with simple data augmentation. Advances conference on machine learning, pages 8162–8171. PMLR,
in Neural Information Processing Systems, 36, 2024. 7 2021. 1, 7

11
[48] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, ference on computer and communications security, pages
Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel 1310–1321, 2015. 8
Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: [59] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising
Learning robust visual features without supervision. arXiv diffusion implicit models. arXiv preprint arXiv:2010.02502,
preprint arXiv:2304.07193, 2023. 2 2020. 1, 4, 6
[49] William Peebles and Saining Xie. Scalable diffusion models [60] Yang Song and Stefano Ermon. Improved techniques for
with transformers. In Proceedings of the IEEE/CVF Inter- training score-based generative models. Advances in neural
national Conference on Computer Vision, pages 4195–4205, information processing systems, 33:12438–12448, 2020. 2, 4
2023. 4, 6, 7 [61] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Ab-
[50] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, hishek Kumar, Stefano Ermon, and Ben Poole. Score-based
Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. generative modeling through stochastic differential equations.
SDXL: Improving latent diffusion models for high-resolution arXiv preprint arXiv:2011.13456, 2020. 2, 3, 4, 6, 7
image synthesis. arXiv preprint arXiv:2307.01952, 2023. 1 [62] Peng Sun, Yi Jiang, and Tao Lin. Efficiency for free:
[51] Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Ideal data are transportable representations. arXiv preprint
Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Chih- arXiv:2405.14669, 2024. 5
Yao Ma, Ching-Yao Chuang, David Yan, Dhruv Choudhary, [63] Luming Tang, Menglin Jia, Qianqian Wang, Cheng Perng
Dingkang Wang, Geet Sethi, Guan Pang, Haoyu Ma, Ishan Phoo, and Bharath Hariharan. Emergent correspondence from
Misra, Ji Hou, Jialiang Wang, Kiran Jagadeesh, Kunpeng image diffusion. Advances in Neural Information Processing
Li, Luxin Zhang, Mannat Singh, Mary Williamson, Matt Le, Systems, 36:1363–1389, 2023. 2
Mitesh Kumar Singh, Peizhao Zhang, Peter Vajda, Quentin [64] Zifeng Wang, Zizhao Zhang, Chen-Yu Lee, Han Zhang,
Duval, Rohit Girdhar, Roshan Sumbaly, Sai Saketh Ramb- Ruoxi Sun, Xiaoqi Ren, Guolong Su, Vincent Perot, Jen-
hatla, Sam Tsai, Samaneh Azadi, Samyak Datta, Sanyuan nifer Dy, and Tomas Pfister. Learning to prompt for continual
Chen, Sean Bell, Sharadh Ramaswamy, Shelly Sheynin, Sid- learning. In Proceedings of the IEEE/CVF conference on com-
dharth Bhattacharya, Tao Xu, Tingbo Hou, Wei-Ning Hsu, Xi puter vision and pattern recognition, pages 139–149, 2022.
Yin, Xiaoliang Dai, Yaniv Taigman, Yaqiao Luo, Yen-Cheng 8
Liu, Yi-Chiao Wu, Yue Zhao, Yuval Kirstain, Zecheng He, [65] Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu
and Zijian He. MovieGen: A cast of media foundation models. Chen, and Mingyuan Zhou. Diffusion-gan: Training gans
Meta AI Blog Post, 2024. 1 with diffusion. arXiv preprint arXiv:2206.02262, 2022. 6, 7
[52] William H Press. Numerical recipes 3rd edition: The art of [66] Weilai Xiang, Hongyu Yang, Di Huang, and Yunhong Wang.
scientific computing. Cambridge university press, 2007. 5, 6 Denoising diffusion autoencoders are unified self-supervised
[53] Leigang Qu, Wenjie Wang, Yongqi Li, Hanwang Zhang, learners. In ICCV, 2023. 2, 4
Liqiang Nie, and Tat-Seng Chua. Discriminative probing [67] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling
and tuning for text-to-image generation. In Proceedings of the generative learning trilemma with denoising diffusion
the IEEE/CVF Conference on Computer Vision and Pattern gans. arXiv preprint arXiv:2112.07804, 2021. 2, 6, 7
Recognition, pages 7434–7444, 2024. 3 [68] Xingyi Yang and Xinchao Wang. Diffusion model as represen-
[54] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya tation learner. In Proceedings of the IEEE/CVF International
Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Conference on Computer Vision, pages 18938–18949, 2023.
Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning 1, 2, 4
transferable visual models from natural language supervi- [69] Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang,
sion. In International conference on machine learning, pages Eli Shechtman, Fredo Durand, and William T Freeman. Im-
8748–8763. PMLR, 2021. 5, 10, 2, 7 proved distribution matching distillation for fast image syn-
[55] Robin Rombach, Andreas Blattmann, Dominik Lorenz, thesis. arXiv preprint arXiv:2405.14867, 2024. 6, 7
Patrick Esser, and Björn Ommer. High-resolution image [70] Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong,
synthesis with latent diffusion models. In Proceedings of Jonathan Huang, Jinwoo Shin, and Saining Xie. Representa-
the IEEE/CVF conference on computer vision and pattern tion alignment for generation: Training diffusion transformers
recognition, pages 10684–10695, 2022. 6, 2 is easier than you think. arXiv preprint arXiv:2410.06940,
[56] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, 2024. 1, 2, 3, 4, 5, 6, 7, 8
Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael [71] Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anand-
Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Pho- kumar. Fast training of diffusion models with masked trans-
torealistic text-to-image diffusion models with deep language formers. arXiv preprint arXiv:2306.09305, 2023. 7
understanding. 2022. 1 [72] Mingyuan Zhou, Huangjie Zheng, Zhendong Wang,
[57] Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan- Mingzhang Yin, and Hai Huang. Score identity distillation:
xl: Scaling stylegan to large diverse datasets. In ACM SIG- Exponentially fast distillation of pretrained diffusion models
GRAPH 2022 conference proceedings, pages 1–10, 2022. 2, for one-step generation. In Forty-first International Confer-
7 ence on Machine Learning, 2024. 6, 7
[58] Reza Shokri and Vitaly Shmatikov. Privacy-preserving deep [73] Rui Zhu, Yingwei Pan, Yehao Li, Ting Yao, Zhenglong Sun,
learning. In Proceedings of the 22nd ACM SIGSAC con- Tao Mei, and Chang Wen Chen. Sd-dit: Unleashing the power

12
 of self-supervised discrimination in diffusion transformer.
In Proceedings of the IEEE/CVF Conference on Computer
Vision and Pattern Recognition, pages 8435–8445, 2024. 7
[74] Roland S Zimmermann, Yash Sharma, Steffen Schneider,
Matthias Bethge, and Wieland Brendel. Contrastive learning
inverts the data generating process. In International Con-
ference on Machine Learning, pages 12979–12990. PMLR,
2021. 5

13
 Generative Modeling with Explicit Memory
Supplementary Material
We include additional analysis, experiments details, and Table 5. Bank ablation study. We also conduct an ablation study
results of downstream tasks in this supplementary material. with respect to the bank size on CIFAR-10. Note that ↓ indicates
In Section A, we provide a detailed analysis and experimen- that lower values are preferred.
tal settings of the memory bank. In Section B, we present Model Bank size NI↓ FID↓
additional experiment, including the ablation of different
encoders. In Section C, we offer further implementation SiT-XL/2
details on the training settings of GM EM on the CIFAR-10 + (GM EM) 50K 22.4M 1.22
dataset, ImageNet 64×64 dataset, and ImageNet 256×256 + (GM EM) 25K 22.4M 1.50
dataset. In Section D, we provide detailed analysis on the + (GM EM) 5K 22.4M 3.90
interpolation experiments. In Section E, we demonstrate that
the bank learned from one dataset can be easily transferred output hidden dimension of the encoder. Additional, we set
to another dataset without any additional training. Besides, the memory bank size n to 50k for CIFAR-10 and 1.2M for
in Section F, we also show that GM EM can be seamlessly ap- ImageNet 64×64 and ImageNet 256×256.
plied to text-to-image generation tasks. Finally, in Section G, When constructing the memory bank, we use the Eu-
we supplement the derivation of the objective functions for clidean distance as the similarity measure between the fea-
flow-based and diffusion-based models. tures output by encoder and the memory bank and use MSE
as loss. We choose the Adam optimizer with a learning rate
A. Memory Bank of 1e-2, β1 of 0.9, and β2 of 0.999. The memory bank is
trained for 80 epochs on CIFAR-10 with the batch size of
In this section we provide a detailed analysis of the memory 64.
bank, including the theoratic analysis, building details of the
memory bank and the ablative study of the memory bank A.3. Building
size related to performance on downstream tasks. We first detailed the building process of the memory bank.
Initially, the memory bank is populated with features and
A.1. Analysis
then fine-tuned using the outputs from the encoder. Specifi-
To ensure that for every input x sampled from the data dis- cally, during fine-tuning, we compute batched features from
tribution D, the normalized feature f (x)/∥f (x)∥ corresponds both the encoder and the memory bank. We then calculate
exactly to a vector s present in the matrix M, we require the similarity between these batched features and the snip-
that:   pets within the memory bank, and compute the loss based
f (x)
min ∥s − ∥ = 0, ∀x ∼ D . (11) on the most similar snippet. Finally, we update the memory
s∈M ∥f (x)∥ bank features according to the computed similarities.
Since we enforce ∥s∥ = 1, the condition in (11) implies that: We present the results of our ablation study on memory
  bank size in Table 5 . Overall, a larger memory bank leads
f (x) to better performance, as indicated by lower FID scores.
max ∥s · ∥ = 1, ∀x ∼ D . (12)
s∈M ∥f (x)∥ Specifically, when the bank size is 50,000—equivalent to
Consequently, our objective is to satisfy: the entire dataset—the performance is optimal. Reducing
  the bank size to 25k results in a slight performance decrease
f (x) to FID = 1.50, yet it still surpasses the baseline. However,
max M = 1, ∀x ∼ D . (13)
∥f (x)∥ further reducing the bank size to 5k significantly degrades
performance, indicating the importance of maintaining a
This leads us to formulate the optimization problem as fol- relatively large memory bank for optimal results.
lows: Overall, GM EM requires less than 150MB of additional
( 2 )
 
f (x)
 storage (the size of a 50K memory bank) to achieve a 10×
arg min Ex∼D max M −1 . faster training speed and a 5× faster inference speed on
M∈Rn×m ∥f (x)∥
downstream tasks.
(14)
A.2. Setting B. Supplementary Experiments
As described in Table 6 and Table 7 , across all experi- In this seciton, we include a supplementary experiments that
ments, we set bank unit m to 768 to match the size of the apply GM EM to further validate the effectiveness of GM EM

1
on various downstream tasks. For latent space generation, we pre-compute compressed
latent vectors from raw images using the Stable Diffusion
B.1. Ablation Study on the Encoder VAE [55] and utilize these latent vectors as input. In con-
We conduct an ablation study to investigate the impact of dif- trast, for pixel space generation, we directly use the raw
ferent encoders on the performance of GM EM. Specifically, pixel data as input. Although we experimented with data
we choose the visual encoder from CLIP [54] as an alterna- augmentations such as flipping, we found that they did not
tive to show that GM EM can be applied across encoders. significantly improve performance. Therefore, we do not
apply any data augmentation in our experiments.
Experimental Setup To avoid redundant computational For projecting memory snippets into the backbone hidden
overhead, instead of training GM EM from scratch, we train dimension, we utilize a three-layer MLP with SiLU activa-
only a multilayer perceptron (MLP) with dimensions 768 × tions for diffusion transformers, following Yu et al. [70]. We
768 to map the features output by the CLIP visual encoder provide detailed hyperparameter configurations in Tables 6
to the memory bank. We utilize the Adam optimizer with a and 7.
learning rate of 1 × 10−4 and train the MLP for one epoch
on the ImageNet 256 × 256 dataset. Encoder We use Dinov2-B [48] as the encoder across
all experiments, as it has been shown to significantly en-
Results We present the results of the ablation study in hance the learning of better representations in diffusion mod-
Figure 6. els [70]. Dinov2-B offers superior performance, making
it an ideal choice for facilitating the efficient training for
C. Implementation Details constructing the memory bank.
C.1. Diffusion transformer architecture
We closely follow the architecture used in REPA [70] and Memory Bank For the memory bank size, we set it to 50K
SiT [44]. Similar to a Vision Transformer [13], In this archi- for CIFAR-10 and 1.2M for ImageNet 64×64 and ImageNet
tecture, the input image is divided into patches, reshaped into 256×256. Note it is essential to maintain a relatively large
a one-dimensional sequence of length N , and then processed memory bank for optimal performance, as demonstrated in
by the model. Unlike the original SiT, REPA includes ad- Section A. Further results on the impact of memory bank
ditional modulation layers called AdaIN-zero layers at each size on performance are provided in Section A, where we de-
attention block. These layers scale and shift each hidden tailed how varying the memory bank size affects the model’s
state based on the given timestep and additional conditions. efficacy.
For latent space generation, similar to REPA, our archi-
tecture uses a downsampled latent image z = E(x) as input, Computing Resources All models are primarily trained
where x is an RGB image and E is the encoder of the Stable on NVIDIA H800 8-GPU setups, each equipped with 80GB
Diffusion Variational Autoencoder (VAE) [55]. For pixel of memory. The training speed for GM EM-XL is approxi-
space generation, we remove the encoder and directly use mately 1.15 seconds per 1KNI.
the RGB image as input. Specifically, we modify the original
SiT by changing the number of channels from 4 to 3 and
directly feed the transformed RGB image into the model. For
D. Interpolation
ImageNet 64 × 64, as an example of pixel space generation, D.1. GM EM is Not a Memory Machine
we aslo adjust the patch size of SiT from 4 to 2 to maintain
a consistent sequence length of N = 256 patches. In this section, we provide further evidence that GM EM is
Furthermore, unlike the original REPA, instead of using not merely a memory machine. Specifically, we demonstrate
class label embeddings, we utilize the CLS token from the that GM EM can generate high-quality samples even for a
target representation output by the encoder as an additional memory snippet ŝ that never appeared in the training set.
condition for the diffusion transformer.
D.2. Interpolation Between Memory Snippets
C.2. Hyperprarameters To investigate this ability, we conduct an interpolation exper-
Additional implementation details We implement our iment. We use the ImageNet 256 × 256 dataset and employ
model based on the original REPA implementation [70]. We the model checkpoint obtained after 180M (NI), as described
use the AdamW optimizer [33] with a constant learning rate in Table 7 .
of 1 × 10−4 , β1 = 0.9, and β2 = 0.999, without weight We randomly select two memory snippets s1 and s2 from
decay. To accelerate training, we employ mixed-precision the memory bank M. We then create nine interpolated snip-
(fp16) computation along with gradient clipping. pets ŝi by linearly interpolating between s1 and s2 with

2
Figure 6. Ablation study on the encoder: Image generation using the CLIP visual encoder with GM EM applied to SiT-XL/2, trained on the
ImageNet 256 × 256 dataset.

interpolation coefficients αi ranging from 0.1 to 0.9 in incre- between a green snake and a long-faced dog results in a
ments of 0.1. The interpolated snippets are defined as: green reptilian creature that resembles both the snake and
the dog. This indicates that when the model encounters
ŝi = (1 − αi )s1 + αi s2 , αi = 0.1i, i = 1, 2, . . . , 9. unseen memory snippets, it can utilize the smooth latent
Each interpolated memory snippet ŝi is then fed into the space to generate images similar to those it has previously
transformer block to generate images. encountered.
The third and last rows showcase even more imaginative
D.3. Interpolation results interpolations. Interpolating between a monkey and barbed
The results of this interpolation experiment are presented wire results in an image of a monkey in a cage, while a
in Figure 7 and Figure 8 . We observe that the generated dog and a red hat can be interpolated into a dog with a black
images from the interpolated memory snippets ŝi are of gentleman’s hat. These outcomes suggest that the similarities
high quality and exhibit smooth transitions between the two captured by the model are not limited to visual resemblance
original memory snippets s1 and s2 . but also encompass more abstract semantic similarities in
In the first row of Figure 7 , we interpolate between an the latent space.
ape and a dog. The dog’s face gradually transforms into a We believe that this semantic similarity arises because
smoother visage, adapting to resemble the ape. This demon- our memory bank introduces additional semantic informa-
strates representation space learned by GM EM is semanti- tion, enabling the model to better understand the content
cally smooth. Surprisingly, in the second row, interpolating of images. Consequently, the model generates images that

3
Table 6. Training settings of CIFAR-10. We provide the training settings for all models and training algorithms on the CIFAR-10 dataset.

Model Size
B L XL
Model details
Batch size 128 128 128
Training iterations 200K 200k 200k
Learning rate 1e-4 1e-4 1e-4
Optimizer Adam Adam Adam
Adam β1 0.9 0.9 0.9
Adam β2 0.999 0.999 0.999
Interpolants
αt 1−t 1−t 1−t
σt t t t
ωt σt σt σt
Training Objective v-prediction v-prediction v-prediction
Sampler Euler Euler Euler
Sampling steps 50 50 50
Classifier-free Guidance × × ×
Training details of backbone
Capacity(Mparams) 130 458 675
Input dim. 32×32×3 32×32×3 32×32×3
Num. layers 12 24 28
Hidden dim. 768 1,024 1,152
Num. heads 12 12 16
Training details of GM EM
Bank size 50k 50k 50k
Bank similarity measure Euclidean Euclidean Euclidean
Bank objective MSE MSE MSE
Encoder f (x) DINOv2-B DINOv2-B DINOv2-B

Figure 7. Interpolation between memory snippets. The first and last columns show the original memory snippets s1 and s2 , respectively. The
remaining columns show the generated images from the interpolated memory snippets ŝi .

align more closely with human intuition, rather than merely ing a Memory Bank extracted from low-resolution images
memorizing the images corresponding to each snippet. to high-resolution models (e.g., Latent Diffusion Models)
may result in decreased image sharpness due to information
E. Transferability of the Memory Bank bottlenecks, it can still enhance the diversity of the generated
image.
In this section, we demonstrate the transferability of the
Memory Bank across different models. Specifically, we
E.1. Experimental Setup
show that the Memory Bank can be transferred between To investigate the transferability, we trained a Memory Bank
GM EM models trained on different datasets. While apply- MCIFAR on the CIFAR-10 dataset and directly transferred

4
Table 7. Training settings. We present the training settings for all models and training algorithms on the ImageNet 64×64 dataset (left) and
the ImageNet 256×256 dataset (right).

Model Size Model Size

B L XL B L XL
Model details Model details
Batch size 128 128 128 Batch size 256 256 256
Training iterations 400K 400k 400k Training iterations 700K 700k 700k
Learning rate 1e-4 1e-4 1e-4 Learning rate 1e-4 1e-4 1e-4
Optimizer Adam Adam Adam Optimizer Adam Adam Adam
Adam β1 0.9 0.9 0.9 Adam β1 0.9 0.9 0.9
Adam β2 0.999 0.999 0.999 Adam β2 0.999 0.999 0.999
Interpolants Interpolants
αt 1−t 1−t 1−t αt 1−t 1−t 1−t
σt t t t σt t t t
ωt σt σt σt ωt σt σt σt
Training Objective v-prediction v-prediction v-prediction Training Objective v-prediction v-prediction v-prediction
Sampler Euler Euler Euler Sampler Euler Euler Euler
Sampling steps 50 50 50 Sampling steps 50 50 50
Classifier-free Guidance × × × Classifier-free Guidance × × ×
Training details of backbone Training details of backbone
Capacity (Mparams) 130 458 675 Capacity(Mparams) 130 458 675
Input dim. 64×64×3 64×64×3 64×64×3 Input dim. 32×32×3 32×32×3 32×32×3
Num. layers 12 24 28 Num. layers 12 24 28
Hidden dim. 768 1,024 1,152 Hidden dim. 768 1,024 1,152
Num. heads 12 12 16 Num. heads 12 12 16
Training details of GM EM Training details of GM EM
Bank size 1.2M 1.2M 1.2M Bank size 50k 50k 50k
Bank similarity measure Euclidean Euclidean Euclidean Bank similarity measure Euclidean Euclidean Euclidean
Bank objective MSE MSE MSE Bank objective MSE MSE MSE
Encoder f (x) DINOv2-B DINOv2-B DINOv2-B Encoder f (x) DINOv2-B DINOv2-B DINOv2-B

Figure 8. A more elaborate interpolation experiment. The first and last columns show the original memory snippets s1 and s2 , respectively.
The remaining columns show the generated images from the interpolated memory snippets ŝi . Different row stands for different noise
applied when generating the images.

it to a model trained on ImageNet 256 × 256 GM EMIN256 E.2. Results

to guide image generation. We used the checkpoint from the
ImageNet model at 180M (NI) for generation. The detailed Figure 9 presents the generation results of our method on
experimental settings are provided in Table 6 and Table 7 . the ImageNet 256 × 256 dataset. The images demonstrate
that the transferred Memory Bank can effectively guide the
high-resolution model. Though sharpness is limited due to

5
Figure 9. Transferability of the Memory Bank. Each row corresponding to a specifical class in CIFAR-10. Specifically, the class is from top
to bottom: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, and truck.

information bottlenecks in memory snippets, it can improv- mapped to Memory Bank snippets through a trained MLP,
ing the diversity of the generated images. and then fed into the SiT backbone to generate the corre-
sponding images. This method enables straightforward and
F. Text-to-Image Generation effective generation of images from textual descriptions.

We propose that GM EM can be seamlessly applied to text- F.1. Experimental Setup

to-image generation tasks. This extension only requires In our text-to-image experiments, we utilize BERT-base [11]
training a mapping from textual CLS features to the Memory as the text encoder, using class names as textual supervision
Bank, which can be efficiently implemented using a simple signals and the CLS token as the textual feature represen-
multilayer perceptron (MLP). tation. We employ the model trained on CIFAR-10 with
During inference, we extract CLS features from the Text 22.5M (NI) as the generator.
Encoder using class names as input. These features are For mapping text features to the Memory Bank, we use

6
 In the DDPM framework introduced by Ho et al. [24], the
reverse process p(xt−1 |xt ) is parameterized as:

p(xt−1 |xt ) =
  !
1 βt
N xt−1 ; √ xt − √ εθ (xt , t) , Σθ (xt , t)
αt 1 − ᾱt
Qt
where αt = 1 − βt , ᾱt = i=1 αi , εθ (xt , t) is a neural
network parameterized by θ and Σθ (xt , t) represents the
learned variance.
The model is trained using a simple denoising autoen-
coder objective:
h i
2
Lsimple = Ex0 ,ε,t ∥ε − εθ (xt , t)∥2 ,

where ε is sampled from a standard normal distribution and

t is uniformly sampled from {1, . . . , T }.
For the variance Σθ (xt , t), Ho et al. [24] initially set it
to σt2 I with βt = σt2 . However, Nichol and Dhariwal [47]
Figure 10. Text-to-image generation results. Each row corresponds demonstrated that performance improves when Σθ (xt , t) is
to a specific class in CIFAR-10. Specifically, the classes are from learned jointly with εθ (xt , t). They propose optimizing the
top to bottom: airplane, automobile, bird, cat, deer, dog, frog, horse, variational lower bound (VLB) objective:
ship, and truck.
 
Lvlb = exp v log βt + (1 − v) log β̃t ,
a simple MLP with dimensions 768×768. We train this
mapping using the Adam optimizer with a learning rate of where v is a per-dimension component from the model out-
1 × 10−4 , optimizing the CLIP Loss [54] over 60 epochs. put and β̃t = 1− ᾱt−1
1−ᾱt βt .
By choosing a sufficiently large T and an appropri-
F.2. Experimental Results ate schedule for βt , the distribution p(xT ) approaches an
isotropic Gaussian. This allows for sample generation by
We present the generated images in Figure 10 . The results
starting from random noise and iteratively applying the
demonstrate that our method can effectively generate images
learned reverse process p(xt−1 |xt ) to obtain a data sample
that correspond closely to the provided textual descriptions.
x0 citepho2020denoising.
G. Derivation G.2. Stochastic interpolating
In this paper, we delve into two types of generative mod- In contrast to DDPM, flow-based models [15, 41] address
els that learn the target distribution by training variants of continuous time-dependent processes involving data samples
denoising autoencoders: denoising diffusion probabilistic x∗ ∼ p(x) and Gaussian noise ε ∼ N (0, I) over the interval
models (DDPM) and stochastic interpolants. t ∈ [0, 1]. The process is formulated as:

G.1. Denoising Diffusion Probabilistic Models

Diffusion models [24] aim to model a target distribution p(x) xt = αt x0 + σt ε, with α0 = σ1 = 1, α1 = σ0 = 0,
by learning a gradual denoising process that transitions from where αt decreases and σt increases as functions of t.
a Gaussian distribution N (0, I) to p(x). The core idea is to There exists a probability flow ordinary differential equation
learn the reverse process p(xt−1 |xt ) of a predefined forward (PF-ODE) characterized by a velocity field ẋt = v(xt , t),
process q(xt |x0 ), which incrementally adds Gaussian noise ensuring that the distribution at time t matches the marginal
to the data starting from x0 ∼ p(x) over T time steps. pt (x).
The forward process q(xt |xt−1 ) is defined as: The velocity v(x, t) is expressed as a combination of two
 p  conditional expectations:
q(xt |xt−1 ) = N xt ; 1 − βt x0 , βt2 I ,

where βt ∈ (0, 1) are small, predefined hyperparameters. v(x, t) = E[ẋt | xt = x] = α̇t E[x∗ | xt = x]+σ̇t E[ε | xt = x],

7
 which can be approximated by a model vθ (xt , t) through coefficient wt are implicitly defined by the forward diffusion
minimizing the training objective: process, potentially leading to a complex design space in
score-based diffusion models [28].
h i
2
Lvelocity (θ) = Ex∗ ,ε,t ∥vθ (xt , t) − α̇t x∗ − σ̇t ε∥ .

This approach aligns with the reverse stochastic differen-

tial equation (SDE):

1 √
dxt = v(xt , t) dt − wt s(xt , t) dt + wt dW̄t ,
2
where the score function s(xt , t) is similarly defined as:
1
s(xt , t) = − E[ε | xt = x].
σt
To approximate s(xt , t), one can use a model sθ (xt , t)
with the training objective:
h i
2
Lscore (θ) = Ex∗ ,ε,t ∥σt sθ (xt , t) + ε∥ .

Since s(x, t) can be directly computed from v(x, t) for

t > 0 using the relation:

1 αt v(x, t) − α̇t x
s(x, t) = · ,
σt α̇t σt − αt σ̇t
it is sufficient to estimate either the velocity v(x, t) or the
score s(x, t).
According to Albergo et al. [1], stochastic interpolants
satisfy the following conditions when αt and σt are chosen
such that: 1. αt2 + σt2 > 0 for all t ∈ [0, 1], 2. Both αt
and σt are differentiable over the interval [0, 1], 3. Boundary
conditions are met: α1 = σ0 = 0 and α0 = σ1 = 1.
These conditions ensure an unbiased interpolation be-
tween x0 and ε. Consequently, simple interpolants can be
utilized by defining αt and σt as straightforward functions
during training and inference. Examples include linear inter-
polants with αt = 1 − t and σt = t,
or variance-preserving
(VP) interpolants with αt = cos π2 t and σt = sin π2 t .

An additional advantage of stochastic interpolants is that

the diffusion coefficient wt remains independent when train-
ing either the score or velocity models. This independence
allows wt to be explicitly chosen after training during the
sampling phase using the reverse SDE.
It’s noteworthy that existing score-based diffusion mod-
els, including DDPM [24], can be interpreted within an SDE
framework. Specifically, their forward diffusion processes
can be viewed as predefined (discretized) forward SDEs that
converge to an equilibrium distribution N (0, I) as t → ∞.
Training is conducted over [0, T ] with a sufficiently large
T (e.g., T = 1000) to ensure that p(xT ) approximates an
isotropic Gaussian. Generation involves solving the corre-
sponding reverse SDE, starting from random Gaussian noise
xT ∼ N (0, I). In this context, αt , σt , and the diffusion