Unsupervised Representation Learning from Pre-trained Diffusion Probabilistic Models
Abstract
Diffusion Probabilistic Models (DPMs) have shown a powerful capacity of generating high-quality image samples.
Recently, diffusion autoencoders (Diff-AE) are proposed to explore DPMs for representation learning via autoencoding and succeed in various downstream tasks.
Their key idea is to jointly train an encoder for discovering meaningful representations from images and a conditional DPM as the decoder for reconstructing images.
Considering that training DPMs from scratch will take a long time and there have existed numerous pre-trained DPMs, we propose P re-trained D
PM
A uto E ncoding ( PDAE ), a general method to adapt existing pre-trained DPMs to the decoders for image reconstruction, with better training efficiency and performance than Diff-AE.
Specifically, we find that the reanson that pre-trained DPMs fail to reconstruct an image from its latent variables is due to the information loss of forward process, which causes a gap between their predicted posterior mean and the true one.
From this perspective, the classifierguided sampling method can be explained as computing an extra mean shift to fill the gap, reconstructing the lost class information in samples.
These imply that the gap corresponds to the lost information of the image, and we can reconstruct the image by filling the gap.
Drawing inspiration from this, we employ a trainable model to predict a mean shift according to encoded representation and train it to fill as much gap as possible, in this way, the encoder is forced to learn as much information as possible from images to help the filling.
By resuing a part of network of pre-trained DPMs and redesigning the weighting scheme of diffusion loss,
PDAE can learn meaningful representations from images efficiently.
Extensive experiments denonstrate the effectiveness, efficiency and flexibility of PDAE.
Our implementation is available at this link.
Introduction
Deep generative models such as variational autoencoders [25, 38], generative adversarial networks [13], autoregressive models [50, 48], normalizing flows [37, 23] and energy-based models [9, 45] have shown remarkable capacity to synthesize striking image samples.
Recently, another kind of generative models, Diffusion Probabilistic Models (DPMs) [43, 14] are further developed and becoming popular for their stable training process and state-of-the-art sample quality [8].
Although a large number of degrees of freedom in implementation, the DPMs discussed in this paper will refer exclusively to those trained by the denoising method proposed in DDPMs [14].
Unsupervised representation learning via generative modeling is a popular topic in computer vision.
Latent variable generative models, such as GANs [13] and VAEs [25, 38], are a natural candidate for this, since they inherently involve a latent representation of the data they generate.
Likewise,
DPMs inherently yield latent variables through the forward process.
However, these latent variables lack high-level semantic information because they are just a sequence of spatially corrupted images.
In light of this, diffusion autoencoders (Diff-AE) [35] explore DPMs for representation learning via autoencoding.
Specifically, they employ an encoder for discovering meaningful representations from images and a conditional DPM as the decoder for image reconstruction by taking the encoded representations as input conditions.
Diff-AE is competitive with the state-of-the-art model on image reconstruction and capable of various downstream tasks.
Following the paradigm of autoencoders, PDAE aims to adapt existing pre-trained DPMs to the decoders for image reconstruction and benefit from it.
Generally, pre-trained DPMs cannot accurately predict the posterior mean of x t − 1 from x t in the reverse process due to the information loss of forward process, which results in a gap between their predicted posterior mean and the true one.
This is the reason that they fail to reconstruct an image ( x 0 ) from its latent variables ( x t ).
From this perspective, the classifier-guided sampling method [8] can be explained as reconstructing the lost class information in samples by shifting the predicted posterior mean with an extra item computed by the gradient of a classifier to fill the gap.
Drawing inspiration from this method that uses the prior knowledges (class label) to fill the gap, we aim to inversely extract the knowledges from the gap, i.e., learn representations that can help to fill the gap.
In light of this, we employ a novel gradient estimator to predict the mean shift accoriding to encoded representations and train it to fill as much gap as possible, in this way, the encoder is forced to learn as much information as possible from images to help the filling.
PDAE follows this principle to build an autoencoder based on pre-trained DPMs.
Furthermore, we find that the posterior mean gap in different time stages contain different levels of information, so we redesign the weighting scheme of diffusion loss to encourage the model to learn rich representations efficiently.
We also reuse a part of network of pre-trained DPMs to accelerate the convergence of our model.
Based on pre-trained DPMs, PDAE only needs less than half of the training time that Diff-AE costs to complete the representation learning but still outperforms Diff-AE.
Moreover, PDAE also enables some other interesting features.
Background 2.1 Denoising Diffusion Probabilistic Models
DDPMs [14] employ a forward process that starts from the data distribution q ( x 0 ) and sequentially corrupts it to N ( 0 , I ) with Markov diffusion kernels q ( x t | x t − 1 ) defined by a fixed variance schedule { β t } Tt =1 .
The process can be expressed by:
[Equation]
where { x t } Tt =1 are latent variables of DDPMs.
According to the rule of the sum of normally distributed random variables, we can directly sample x t from x 0 for arbitrary t with q ( x t | x 0 ) =
N ( x t ; √ ¯ α t x 0 , (1 − ¯ α t ) I ) , where α t = 1 − β t and ¯ α t = (cid:81) t i =1 α i .
The reverse (generative) process is defined as another Markov chain parameterized by θ to describe the same but reverse process, denoising an arbitrary Gaussian noise to a clean data sample:
[Equation]
where p ( x T ) =
N ( x T ; 0 , I ) .
It employs p θ ( x t − 1 | x t ) of Gaussian form because the reversal of the diffusion process has the identical functional form as the forward process when β t is small [11, 43].
The generative distribution can be represented as p θ ( x 0 ) =
(cid:82) p θ ( x 0: T ) d x 1: T .
Training is performed to maximize the model log likelihood (cid:82) q ( x 0 ) log p θ ( x 0 )
d x 0 by minimizing the variational upper bound of the negative one.
The final objective is derived by some reparameterization and simplication [14]:
[Equation]
(3) where ϵ θ is a function approximator to predict ϵ from x t .
2.2 Denoising Diffusion Implicit Models
DDIMs [44] define a non-Markov forward process that leads to the same training objective with
DDPMs, but the corresponding reverse process can be much more flexible and faster to sample from.
Specifically, one can sample x t − 1 from x t using the ϵ θ of some pre-trained DDPMs via:
[Equation]
where ϵ t ∼ N ( 0 , I ) and σ t controls the stochasticity of forward process.
The strides greater thanare allowed for accelerated sampling.
When σ t = 0 , the generative process becomes deterministic, which is named as DDIMs.
2.3 Classifier-guided Sampling Method
Classifier-guided sampling method [43, 46, 8] shows that one can train a classifier p ϕ
( y | x t ) on noisy data and use its gradient ∇ x t log p ϕ
( y | x t ) to guide some pre-trained unconditional DDPM to sample towards specified class y .
The conditional reverse process can be approximated by a Gaussian similar to that of the unconditional one in Eq.(2), but with a shifted mean:
[Equation]
For deterministic sampling methods like DDIMs, one can use score-based conditioning trick [46, 45] to define a new function approximator for conditional sampling:
[Equation]
More generally, any similarity estimator between noisy data and conditions can be applied for guided sampling, such as noisy-CLIP guidance [32, 30].
Method 3.1 Forward Process Posterior Mean Gap
Generally, one will train unconditional and conditional DPMs by respectively learning p θ ( x t − 1 | x t ) =
N ( x t − 1 ; µ θ ( x t , t ) , Σ θ ( x t , t )) and p θ ( x t − 1 | x t , y ) =
N ( x t − 1 ; µ θ ( x t , y , t ) , Σ θ ( x t , y , t )) to approximate the same forward process posterior q ( x t − 1 | x t , x 0 )
=
N ( x t − 1 ; (cid:101) µ t ( x t , x 0 ) , 1 − ¯ α t − 1 1 − ¯ α t β t I ) .
Here y is some condition that contains some prior knowledges of corresponding x 0 , such as class label.
Assuming that both Σ θ is set to untrained time dependent constants, under the same experimental settings, the conditional DPMs will reach a lower optimized diffusion loss.
The experiment in Figure 1 can prove this fact, which means that µ θ ( x t , y , t ) is closer to (cid:101) µ t ( x t , x 0 ) than µ θ ( x t , t ) .
This implies that there exists a gap between the posterior mean predicted by the unconditional
DPMs (cid:0) µ θ ( x t , t ) (cid:1) and the true one (cid:0)(cid:101) µ t ( x t , x 0 ) (cid:1) .
Essentially, the posterior mean gap is caused by the information loss of forward process so that the reverse process cannot recover it in x t − 1 only according to x t .
If we introduce some knowledges of x 0 for DPMs, like y here, the gap will be smaller.
The more information of x 0 that y contains, the smaller the gap is.
Moreover, according to Eq.(5), the Gaussian mean of classifier-guided conditional reverse process contains an extra shift item compared with that of the unconditional one.
From the perspective of posterior mean gap, the mean shift item can partially fill the gap and help the reverse process to reconstruct the lost class information in samples.
In theory, if y in Eq.(5) contains all information of x 0 , the mean shift will fully fill the gap and guide the reverse process to reconstruct x 0 .
On the other hand, if we employ a model to predict mean shift accoriding to our encoded representations z and train it to fill as much gap as possible, the encoder will be forced to learn as much information as possible from x 0 to help the filling.
The more the gap is filled, the more accurate the mean shift is, the more perfect the reconstruction is, and the more information of x 0 that z contains.
PDAE follows this principle to build an autoencoder based on pre-trained DPMs.
3.2 Unsupervised Representation Learning by Filling the Gap
Following the paradigm of autoencoders, we employ an encoder z
= E φ ( x 0 ) for learning compact and meaningful representations from input images and adapt a pre-trained unconditional DPM p θ ( x t − 1 | x t ) =
N ( x t − 1 ; µ θ ( x t , t ) , Σ θ ( x t , t )) to the decoder for image reconstruction.
[Figure]
Specifically, we employ a gradient estimator G ψ ( x t , z , t ) to simulate ∇ x t log p ( z | x t ) , where p ( z | x t ) is some implicit classifier that we will not use explicitly, and use it to assemble a conditional DPM p θ,ψ ( x t − 1 | x t , z )
=
N ( x t − 1 ; µ θ ( x t , t )+ Σ θ ( x t , t ) · G ψ ( x t , z , t ) , Σ θ ( x t , t )) as the decoder.
Then we train it like a regular conditional DPM by optimizing following derived objective (assuming the ϵ -prediction parameterization is adopted):
[Equation]
where x t = √ ¯ α t x 0 + √ 1 − ¯ α t ϵ and λ t is a new weighting scheme that we will discuss in Section 3.4.
Note that we use pre-trained DPMs so that θ are frozen during the optimization.
Usually we set Σ θ = 1 − ¯ α t − 1 1 − ¯ α t β t I to untrained time-dependent constants.
The optimization is equivalent to minimizing (cid:13)(cid:13) Σ θ ( x t , t ) · G ψ ( x t , z , t ) − (cid:0)(cid:101) µ t ( x t , x 0 ) − µ θ ( x t , t ) (cid:1)(cid:13)(cid:13) 2 , which forces the predicted mean shift Σ θ ( x t , t ) · G ψ ( x t , z , t ) to fill the posterior mean gap (cid:101) µ t ( x t , x 0 ) − µ θ ( x t , t ) .
With trained G ψ ( x t , z , t ) , we can treat it as the score of an optimal classifier p ( z | x t ) and use the classifier-guided sampling method in Eq.(5) for DDPM sampling or use the modified function approximator ˆ ϵ θ in Eq.(6) for DDIM sampling, based on pre-trained ϵ θ ( x t , t ) .
We put detailed algorithm procedures in Appendix A.
Except the semantic latent code z , we can infer a stochastic latent code x T [35] by running the deterministic generative process of DDIMs in reverse:
[Equation]
This is optional, but helpful to near-exact reconstruction and real-image manipulation for reconstructing minor details of input images when using DDIM sampling.
We also train a latent DPM p ω ( z t − 1 | z t ) to model the learned semantic latent space, same with that in Diff-AE [35].
With a trained latent DPM, we can sample some semantic latent code z from it to help pre-trained DPMs to achieve faster and better unconditional sampling.
3.3 Network Design
Figure 2 shows the network and data flow of PDAE.
For encoder E φ , we simply use stacked convolution layers and a linear layer to map x 0 to a vector z .
For gradient estimator G ψ , we use
U-Net [39] similar to the function approximator ϵ θ of pre-trained DPM.
Considering that ϵ θ also takes x t and t as input, we can further leverage the knowledges of pre-trained DPM by reusing its trained encoder part and time embedding layer, so that we only need to employ new middle blocks, decoder part and output blocks of U-Net for G ψ .
To incorporate z into them, we follow [8] to extend
Group Normalization [53] by applying scaling & shifting twice on normalized feature maps:
[Equation]
where [ t s , t b ] and [ z s , z b ] are obtained from a linear projection of t and z , respectively.
Note that we still use skip connections from reused encoder to new decoder.
In this way, G ψ is totally determined by pre-trained DPM and can be universally applied to different U-Net architectures.
[Figure]
for different time stages.
We perform a 50 -step-gridsearch for ( t 1 , t 2 ) paris to find the shortest criticalstage that can ensure high accuracy of conditional generation.
For MNIST [27], it is (400 , 700) .
ative to
3.4 Weighting Scheme Redesign
P2-weighting [7] have shown that the weighting scheme of diffusion loss can greatly affect the performance of DPMs.
We investigate similar phenomenon for representation learning in Figure 3.
Specifically, we train an unconditional DPM and a noisy classifier on MNIST [27], and divide the diffusion forward process into three stages: early-stage between 0 and t 1 , critical-stage between tand t 2 and late-stage between t 2 and T , as shown in the top row.
Then we design a mixed sampling procedure that employs unconditional sampling but switches to classifier-guided sampling only during the specified stage.
The bottom three rows show the samples generated by three different mixed sampling procedures, where each row only employs classifier-guided sampling during the specified stage on the right.
As we can see, only the samples guided by the classifier during critical-stage match the input class labels.
We can conclude that the mean shift during critical-stage contains more crucial information to reconstruct the input class label in samples than the other two stages.
From the view of diffusion trajectories, the sampling trajectories are separated from each other during critical-stage and they need the mean shift to guide them towards specified direction, otherwise it will be determined by the stochasticity of Langevin dynamics.
Therefore, we opt to down-weight the objective function for the t in early- and late-stage to encourage the model to learn rich representations from critical-stage.
Inspired by P2-weighting [7], we redesign a weighting scheme of diffusion loss ( λ t in Eq.(7)) in terms of signal-to-noise ratio [24] (SNR ( t ) =
¯ α t 1 − ¯ α t ):
[Equation]
where the first item is for early-stage and the second one is for late-stage.
γ is a hyperparameter that balances the strength of down-weighting between two items.
Empirically we set γ = 0 .
Figureshows the normalized weighting schemes of diffusion loss for different DPMs relative to the true variational lower bound loss.
Compared with other DPMs, our weighting scheme down-weights the diffusion loss for both low and high SNR.
We find this weighting scheme is indispensable for PDAE, otherwise the loss will not converge.
Experiments
To compare PDAE with Diff-AE [35], we follow their experiments with the same settings.
Moreover, we also show that PDAE enables some added features.
For fair comparison, we use the baseline DPMs provided by official Diff-AE implementation as our pre-trained models (also as our baselines), which have the same network architectures (hyperparameters) with their Diff-AE models.
For brevity, we use the notation such as "FFHQ128-130M-z512-64M" to name our model, which means that we use a baseline DPM pre-trained with 130M images and leverage it for PDAE training with 64M images, on 128 × 128 FFHQ dataset, with the semantic latent code z of 512 - d .
We put all implementation details in Appendix B and additional samples of following experiments in Appendix C.
4.1 Training Efficiency
We demonstrate the better training efficiency of PDAE compared with Diff-AE from two aspects: training time and times.
For training time, we train both models with the same network architectures (hyperparameters) on 128 × 128 image dataset using 4 Nvidia A100-SXM4 GPUs for distributed training and set batch size to 128 ( 32 for each GPU) to calculate their training throughput (imgs/sec./A100).
PDAE achieves a throughput of 81.57 and Diff-AE achieves that of 75.41.
Owing to the reuse of the U-Net encoder part of pre-trained DPM, PDAE has less trainable parameters and achieves a higher training throughput than Diff-AE.
For training times, we find that PDAE needs about 13 ∼of the number of training batches (images) that Diff-AE needs for loss convergence.
We think this is because that modeling the posterior mean gap based on pre-trained DPMs is easier than modeling a conditional DPM from scratch.
The network reuse and the weighting scheme redesign also help.
As a result, based on pre-trained DPMs, PDAE needs less than half of the training time that Diff-AE costs to complete the representation learning.
4.2 Learned Mean Shift Fills Posterior Mean Gap
Similar to Diff-AE, we train a model of "FFHQ128-130M-z512-64M" on FFHQ [21] and evaluate it on unseen CelebA-HQ [20].
We show that our learned mean shift can fill the posterior mean gap with qualitative and quantitative results in Figure 5.
Specifically, we select some images x 0 from
CelebA-HQ, sample x t = √
¯ α t x 0 + √ 1 − ¯ α t ϵ for different t and predict ˆ x 0 from x t by denoising them for only one step (i.e., ˆ x
0 = x
t −√ 1 − ¯
α t ˆ ϵ √ ¯ α t ), using pre-trained DPM and PDAE respectively.
As we can see in the figure (left), even for large t , PDAE can predict accurate noise from x t and reconstruct plausible images, which shows that the predicted mean shift fills the posterior mean gap and the learned representation helps to recover the lost information of forward process.
Furthermore, we randomly select 1000 images from CelebA-HQ, sample x t = √ ¯ α t x 0 + √ 1 − ¯ α t ϵ and calculate their average posterior mean gap for each step using pre-trained DPM: ∥ (cid:101) µ t ( x t , x 0 ) − µ θ ( x t , t ) ∥and PDAE: ∥ (cid:101) µ t ( x t , x 0 ) − ( µ θ ( x t , t )
+ Σ θ ( x t , t ) · G ψ ( x t , E φ ( x 0 ) , t ))
∥ 2 respectively, shown in the figure (right).
As we can see, PDAE models the mean shift that significantly fills the posterior mean gap.
4.3 Autoencoding Reconstruction
We use "FFHQ128-130M-z512-64M" to compute the semantic latent code for CelebA-HQ images and run PDAE generative process of DDIM and DDPM for reconstruction respectively.
We present some autoencoding reconstruction examples in Figure 6.
As we can see, both methods generate samples with similar contents to the input.
Some stochastic variations [35] occur in minor details of hair, eye and skin when introducing stochasticity.
Due to the similar performance between DDPM and DDIM with random x T , we will always use DDIM sampling method in later experiments.
We can get a near-exact reconstruction if we use the stochastic latent code inferred from aforementioned
ODE, which further proves that the stochastic latent code controls the local details.
[Figure]
[Table]
To further evaluate the autoencoding reconstruction quality of PDAE
[Table]
we conduct the same quantitative experiments with Diff-AE.
Specifically, we use "FFHQ128-130M-z512-64M" to encode-andreconstruct all 30k images of CelebA-HQ and evaluate the reconstruction quality with their average
SSIM [52], LPIPS [55] and MSE.
We use the same baselines described in [35], and the results are shown in Table 1.
We can see that PDAE is competitive with the state-of-the-art NVAE even with much less latent dimensionality and also outperforms Diff-AE in all metrics except the LPIPS for random x T .
Moreover, PDAE only needs about half of the training times that Diff-AE needs for representation learning, which shows that PDAE can learn richer representations from images more efficiently based on pre-trained DPM.
4.4 Interpolation of Semantic Latent Codes and Trajectories
Given two samples x 10 and x 20 from CelebA-HQ, we first use "FFHQ128-130M-z512-64M" to encode them into ( z 1 , x 1 T ) and ( z 2 , x 2 T ) and then run PDAE generative process of DDIM starting from Slerp ( x 1 T , x 2 T ; λ ) under the guidance of G ψ (cid:0) x t , Lerp ( z 1 , z 2 ; λ ) , t (cid:1) with 100 steps, expecting smooth transitions along λ .
Moreover, from the view of the diffusion trajectories, PDAE generates desired samples by shifting the unconditional sampling trajectories towards the spatial direction predicted by G ψ ( x t , z , t ) .
This enables PDAE to directly interpolate between two different sampling trajectories.
Intuitively, the spatial direction predicted by the linear interpolation of two semantic latent codes, G ψ (cid:0) x t , Lerp ( z 1 , z 2 ; λ ) , t (cid:1) , should be equivalent to the linear interpolation of two spatial directions predicted by respective semantic latent code, Lerp (cid:0) G ψ ( x t , z 1 , t ) , G ψ ( x t , z 2 , t ); λ (cid:1) .
We present some examples of these two kinds of interpolation methods in Figure 7.
As we can see, both methods generate similar samples that smoothly transition from one endpoint to the other, which means that G ψ (cid:0) x t , Lerp ( z 1 , z 2 ; λ ) , t (cid:1) ≈ Lerp (cid:0) G ψ ( x t , z 1 , t ) , G ψ ( x t , z 2 , t ); λ (cid:1) , so that
G ψ ( x t , z , t ) can be seen as a function of z analogous to a linear map.
The linearity guarantees a meaningful semantic latent space that represents the semantic spatial change of image by a linear change of latent code.
[Figure]
[Figure]
4.5 Attribute Manipulation
We can further explore the learned semantic latent space in a supervised way.
To illustrate this, we conduct attribute manipulation experiments based on "FFHQ128-130M-z512-64M".
Specifically, we first encode an image to its semantic latent code, then move it along the learned direction and finally decode it to the manipulated image.
Similar to Diff-AE [35], we train a linear classifier to separate the semantic latent codes of the images with different attribute labels and use the normal vector of separating hyperplane (i e. the weight of linear classifier) as the direction vector.
We use the images and attribute labels from CelebA-HQ to train the linera classifier.
We present some attribute manipulation examples in Figure 11.
As we can see, PDAE succeeds in manipulating images by moving their semantic latent codes along the direction of desired attribute with different scales.
Like
Diff-AE, PDAE can change attribute-relevant features while keeping other irrelevant details almost stationary if using the inferred x T of input image.
4.6 Truncation-like Effect
According to [8, 15], we can obtain a truncation-like effect in DPMs by scaling the strength of classifier guidance.
We have assumed that G ψ ( x t , z , t ) trained by filling the posterior mean gap simulates the gradient of some implicit classifier, and it can actually work as desired.
In theory, it can also be applied in truncation-like effect.
To illustrate this, we directly incorporate the class label into
G ψ ( x t , y , t ) and train it to fill the gap.
Specifically, we train a model of "ImageNet64-77M-y-38M" and use DDIM sampling method with 100 steps to generate 50k samples, guided by the predicted mean shift with different scales for a truncation-like effect.
Figure 9 shows the sample quality
[Table]
effects of sweeping over the scale.
As we can see it achieves the truncation-like effect similar to that of classifier-guided sampling method, which helps us to build connections between filling the posterior mean gap and classifier-guided sampling method.
The gradient estimator trained by filling the posterior mean gap is an alternative to the noisy classifier.
4.7 Few-shot Conditional Generation
Followig D2C [42], we train a model of "CelebA64-72M-z512-38M" on CelebA [20] and aim to achieve conditional sampling given a small number of labeled images.
To achieve this, we train a latent DPM p ω ( z t − 1 | z t ) on semantic latent space and a latent classifier p η ( y | z ) using given labeled images.
For binary scenario, the images are labeled by a binary class (100 samples, 50 for each class).
[Table]
For PU scenario, the images are either labeled positive or unlabeled ( positively labeled and 10k unlabeled samples).
Then we sample z from p ω ( z t − 1 | z t ) and accept it with the probability of p η ( y | z ) .
We use the accepted z to generate 5k samples for every class and compute the FID score between these samples and all images belonging to corresponding class in dataset.
We compare PDAE with DiffAE and D2C. We also use the naive approach that computes the FID score between the training images and the corresponding subset of images in dataset.
Table 3 shows that
PDAE achieves better FID scores than Diff-AE and D2C.
4.8 Improved Unconditional Sampling
As shown in Section 4.2, with the help of z , for the same x t , PDAE can generate plausible images in only one step.
If we can get z in advance, PDAE can generate better image samples using fewer steps than pre-trained DPMs.
Therefore, we train a latent DPM on semantic latent space and sample z from it.
Unlike Diff-AE that must take z as input for sampling, PDAE uses an independent gradient estimator as a corrector of the pre-trained DPM for sampling.
We find that only using pre-trained DPMs in the last few sampling steps can achieve better sample quality, which may be because that the gradient estimator is sensitive to z during early-stage and the stochasticity of sampled z will lead to out-of-domain samples.
Empirically, we carry out this strategy in the last 20% sampling steps.
We evaluate unconditional sampling result on "FFHQ128-130M-z512-64M", "Horse128-130M-z51264M", "Bedroom128-120M-z512-70M" and "CelebA64-72M-z512-38M" using DDIM sampling method with different steps.
For each dataset, we calculate the FID scores between 50k generated samples and 50k real images randomly selected from dataset.
Table 2 shows that PDAE significantly improves the sample quality of pre-trained DPMs and outperforms Diff-AE.
Note that PDAE can be applied for any pre-trained DPMs as an auxiliary booster to improve their sample quality.
Related Work
Our work is based on an emerging latent variable generative model known as Diffusion Probabilistic
Models (DPMs) [43, 14], which are now popular for their stable training process and competitive sample quality.
Numerous studies [33, 24, 8, 15, 44, 19, 46, 29] and applications [5, 26, 18, 31, 56, 6, 28, 41, 3, 16, 17] have further significantly improved and expanded DPMs.
Unsupervised representation learning via generative modeling is a popular topic in computer vision.
Latent variable generative models, such as GANs [13], VAEs [25, 38], and DPMs, are a natural candidate for this, since they inherently involve a latent representation of the data they generate.
For GANs, due to its lack of inference functionality, one have to extract the representations for any given real samples by an extra technique called GAN Inversion [54], which invert samples back into the latent space of trained GANs.
Existing inversion methods [57, 34, 4, 1, 2, 51] either have limited reconstruction quality or need significantly higher computational cost.
VAEs explicitly learn representations for samples, but still face representation-generation trade-off challenges [49, 42].
VQ-VAE [49, 36] and D2C [42] overcome these problems by modeling latent variables post-hoc in different ways.
DPMs also yield latent variables through the forward process.
However, these latent variables lack high-level semantic information because they are just a sequence of spatially corrupted images.
In light of this, diffusion autoencoders (Diff-AE) [35] explore DPMs for representation learning via autoencoding.
Specifically, they jointly train an encoder for discovering meaningful representations from images and a conditional DPM as the decoder for image reconstruction by treating the representations as input conditions.
Diff-AE is competitive with the state-of-the-art model on image reconstruction and capable of various downstream tasks.
Compared with Diff-AE, PDAE leverages existing pre-trained DPMs for representation learning also via autoencoding, but with better training efficiency and performance.
A concurrent work with the similar idea is the textual inversion of pre-trained text-to-image DPMs [12, 40].
Specifically, given only 3-5 images of a user-provided concept, like an object or a style, they learn to represent it through new "words" in the embedding space of the frozen text-to-image DPMs.
These learned "words" can be further composed into natural language sentences, guiding personalized creation in an intuitive way.
From the perspective of posterior mean gap, for the given new concept, textual inversion optimizes its corresponding new "words" embedding vector to find a best textual condition (cid:0) c (cid:1) , so that which can be fed into pre-trained text-to-image DPMs (cid:0) ϵ θ ( x t , c , t ) (cid:1) to fill as much gap (cid:0) ϵ − ϵ θ ( x t , ∅ , t ) (cid:1) as possible.
Conclusion
In conclusion, we present a general method called PDAE that leverages pre-trained DPMs for representation learning via autoencoding and achieves better training efficiency and performance than
Diff-AE [35].
Our key idea is based on the concept of posterior mean gap and its connections with classifier-guided sampling method [43, 46, 8].
A concurrent work, textual inversion of pre-trained text-to-image DPMs [12, 40], can also be explained from this perspective.
We think the idea can be further explored to extract knowledges from pre-trained DPMs, such as interpretable direction discovery [51], and we leave it as future work.
Acknowledgments and Disclosure of Funding
This work was supported in part by the National Natural Science Foundation of China (Grant No.
62072397 and No.61836002), Zhejiang Natural Science Foundation (LR19F020006) and Yiwise.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b)
Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[No] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Algorithm
Algorithm 1 shows the training procedure and Algorithm 2 3 show the DDPM and DDIM sampling procedure of PDAE, respectively.
Algorithm
Prepare: dataset distribution p pre θ
Initialize: encoder E φ , gradient-estimator G ψ .
Run: repeat x 0 ∼
[Equation]
x 0 )
[Equation]
[Equation] [Equation]
[Equation] [Equation]
I
[Equation]
Update φ and ψ by taking gradient descent step on
[Equation]
until converged ;
Usually we set Σ θ = σ 2 t I = 1 − ¯ α t − 1 1 − ¯ α t β t I to untrained time-dependent constants.
Algorithm 2: DDPM Sampling(Autoencoding)
Prepare: pre-trained DPM ϵ θ , trained encoder E φ , trained gradient estimator G ψ ,
Input: sample x
Run: z = E φ ( x ) x T ∼ N ( 0 , I ) for t =
T to 1 do ϵ ∼ N ( 0 , I ) if t ≥ 2 , else ϵ =
[Equation]
−
[Equation]
return x
Algorithm 3:
Sampling(Autoencoding)
Prepare: pre-trained DPM ϵ θ , trained encoder E φ , trained gradient estimator G ψ ,
Input: sample x , sampling sequence { t i } Si =1 where t 1 = 0 and t S =
T
Run: z = E φ ( x ) x T ∼ N ( 0 , I ) or use inferred x T for i =
S to 2 do ˆ ϵ θ
( x t i , t i ) =
ϵ θ ( x t i , t i ) − √ 1 − ¯ α t i · G ψ ( x t i , z , t i ) x
t i − 1 = √
¯ α t i −(cid:18)
x t − √ 1 − ¯ α ti ˆ ϵ θ ( x ti ,t i )
√ ¯ α ti (cid:19)
+ (cid:112) 1 − ¯ α t
i − 1 · ˆ ϵ θ ( x t i , t i ) return x
B Implementation Details B.1 Network Architecture
Table 4 shows the network architecture of pre-trained DPMs we use.
G ψ is completely determined by pre-trained DPMs.
For E φ , we use stacked GroupNorm-SiLU-Conv layers to convert input images into 256 × 4 × 4 feature maps and a linear layer to project it into z .
We also use an attention block at 16 × 16 resolution.
For latent DPMs, we use deep MLPs as denoising network same with that of Diff-AE [35].
Tableshows the network architecture.
Specifically, we calculate z = E φ ( x 0 ) for all x 0 from dataset and
[Table] [Table]
normalize them.
Then we learn the latent DPMs p ω ( z t − 1 | z t ) by optimizing:
[Equation]
(11) where ϵ t ∼ N ( 0 , I ) and z t = √
¯ α t z + √ 1 − ¯ α t ϵ t .
Note that we use summing denoising loss of all
T steps for one batch.
The sampled z will be denormalized.
During the training of PDAE
[Table]
we set batch size as 128 for all datasets.
We always set learning rate as 1 e − 4 and use 512 - d z .
We use four Nvidia A100-SXM4 GPUs for all experiments.
B.2 Experimental Details
For attribute manipulation, we train a linear classifier to separate the semantic latent codes of the images with different attribute labels.
During manipulation, we first normalize z = E φ ( x 0 ) to zero mean and unit variance, then move it towards the normal vector of separating hyperplane (i.e. the weight of linear classifier) with different scales, finally denormalize it for sampling.
For few-shot conditional generation, we follow [42] to train PU classifier by oversampling positively labeled samples to balance the batch samples.
During conditional generation of class y , for a sampled z , we reject it when p η ( y | z ) < 0 .
5 and accept it with the probability of p η ( y | z ) when p η ( y | z ) ≥ 0 .
C Additional Samples C.1 Learned Mean Shift Fills Posterior Mean Gap
Figure 1 2 3 show the predicted ˆ x 0 by denoising x t = √ ¯ α t x 0 + √ 1 − ¯ α t ϵ for only one step using different models.
Figure 4 shows the calculated average posterior mean gap for ∥ (cid:101) µ t ( x t , x 0 ) − µ θ ( x t , t ) ∥ 2 and ∥ (cid:101) µ t ( x t , x 0 ) − ( µ θ ( x t , t ) + Σ θ ( x t , t ) · G ψ ( x t , E φ ( x 0 ) , t ))
∥ 2 .
As we can see,
PDAE can predict the mean shift that indeed fills the posterior mean gap.
C.2 Autoencoding Reconstruction
Figure 5 6 7 show some autoencoding reconstruction examples using different models.
As we can see, the deterministic method can almost reconstruct the input images even with only 100 steps and both stochastic methods can generate samples with similar contents to the input except some minor details, such as sheet pattern and wrinkle for LSUN-Bedroom; horse eye, spot and mane for LSUN-Horse.
C.3 Interpolation of Semantic Latent Codes and Trajectories
Figure 8 9 10 show some examples of two kinds of interpolation methods using different models.
Due to complex scenes for images of LSUN-Bedroom and LSUN-Horse, we manually select some spatially-similar image pairs for interpolation.
As we can see, both methods generate similar samples that smoothly transition from one endpoint to the other.
C.4 Attribute Manipulation
Figure 11 shows some attribute manipulation examples.
As we can see, PDAE succeeds in manipulating images by moving their semantic latent codes along the direction of desired attribute with different scales.
Like Diff-AE, PDAE can change attribute-relevant features while keeping other irrelevant details almost stationary if using the inferred x T of input image.
C.5 Few-shot Conditional Generation
We present some samples for 4 PU scenarios of few-shot conditional generation in Figure 12.
As we can see, PDAE can generate samples belonging to specified class for different few-shot scenarios, which shows that our semantic latent codes are easy to classify even with a very small number of labeled samples.
C.6 Visualization of Mean Shift
We visualize some examples of G ψ ( x t , E φ ( x 0 ) , t ) in Figure 13.
As we can see, the gradient estimator learns a mean shift direction towards x 0 for each x t .
D Limitations and Potential Negative Societal Impacts
Although better training efficiency, PDAE has a slower inference speed than Diff-AE due to an extra gradient estimator, which also needs more memory and storage space.
Slow generation speed is a common problem for DPM-based works.
Although many studies have been able to achieve decent performance with few reverse steps, they still lag behind VAEs and
GANs, which only need a single network pass.
Furthermore, almost perfect PDAE reconstruction needs hundreds of extra forward steps to infer the stochastic latent code.
Moreover, we have found that the weighting scheme of diffusion loss is indispensable to PDAE, but we haven’t explored its mechanism, which may help to further improve the efficiency and performance of PDAE.
We leave empirical and theoretical investigations of this aspect as future work.
Potential negative impacts of our work mainly involve deepfakes, which leverage powerful generative techniques from machine learning and artificial intelligence to create synthetic media, which may be used for hoaxes, fraud, bullying or revenge.
Although some synthetic samples are hard to distinguish, researchers have developed algorithms similar to the ones used to build the deepfake to detect them with high accuracy.
Some other techniques such as blockchain and digitally signing can help platforms to verify the source of the media.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]