Unsupervised Controllable Generation with Score-based Diffusion Models: Disentangled Latent Code Guidance
Abstract
From the impressive empirical success of Score-based diffusion models, it is recently spotlighted in generative models.
In real-world applications, the controllable generation enriches the impact of diffusion models.
This paper aims to solve the challenge by presenting the method of control in an unsupervised manner.
We propose the Latent Code Guidance Diffusion Model (LCG-DM), which is the first approach to apply disentanglement on Score-based diffusion models.
Disentangled latent code can be considered as a pseudo-label, since it separately expresses semantic information in each dimension.
LCG-DM is a Score-based diffusion model that reflects disentangled latent code as the condition.
LCG-DM shows the best performance among baselines in terms of both sample quality and disentanglement on dSprites dataset.
LCG-DM can manipulate images on CelebA dataset, with comparable FID performance compared to non-disentangling Score-based diffusion models.
Furthermore, we provide experimental results of scaling method that reflects more on pseudo-label with MNIST dataset.
Introduction
Deep generative models (DGMs) estimate the underlying data distribution from a finite number of observed data samples [23].
This estimation comes from the iterative minimization of the statistical divergence between the model distribution and the data distribution [16].
One branch of DGM starts the divergence minimization by explicitly specifying the model distribution and its transformations, i e. autoregressive models [18], Variational Autoencoder (VAE) [12, 4, 28], Flow-models [21], Scorebased diffusion models [27, 24, 10], etc, and this paper refers to these DGM variants as likelihoodbased models.
Among them, the Score-based diffusion models have been recently developed with the high quality of sampled data instances, which Generative Adversarial Network(GAN), a likelihoodfree model, had been considered the best approach.
Given the better sample quality with an explicit density model by the Score-based diffusion models, researchers are expanding research questions on diffusion models to control [19, 22], guide [5, 8], and disentangle the model distribution while maintaining the sample quality.
The control and the disentanglement on the model distribution are enabling factors to utilize DGM in many industrial applications [20, 6].
Meanwhile, the controllable generation [1, 22] requires complete joint annotations on data instances, so the joint condition will specify the correlation between individual condition random variables.
Because of the large joint condition space and limited data samples to sufficiently shape the density on the condition space, the disentangled representation learning (DRL) on the condition variables makes controllable generation feasible with limited data instances [7, 11, 2, 3, 13].
Moreover, DRL enables synthesizing a data instance on a certain condition dimension with minimal impact on the other conditions.
This disentanglement statistically results in independence across the condition random variables, and the learning objectives of DRL require mechanisms to ensure such independence.
There are several approaches of DRL, mostly originating from VAE [7, 11, 2] and GAN [3, 13] frameworks.
In VAE, latent variables are regularized to be dimension-wise independent when they are inferred.
However, the regularization for the disentanglement significantly reduces sample quality.
In GAN, the regularization by mutual information makes the latent code disentangled, yet this regularization on GAN still degrades the sample quality with further learning instability.
Given the advanced sample quality from the Score-based diffusion model, it is natural progress to formulate a DRL variation of diffusion models.
This paper proposes Latent Code Guidance Diffusion
Model (LCG-DM), enabling unsupervised controllable sample generation, which is the first DRL model in the community of Score-based diffusion models.
Preliminary 2.1 Score-based diffusion models
Score-based diffusion models consist of forward and reverse diffusion processes.
The forward diffusion process refers to adding noise to data, which is defined as Eq.
(1) with a drift function, f ( x t , t ) and a volatility function, g ( t ) .
x 0 denotes a data random variable and x t denotes a perturbed random variable at t .
w t refers to the standard Wiener process.
[Equation]
[Equation]
While the forward process is a simple noise additive process, diffusion models define the reverse diffusion to denoise an image from a noisy input.
Given that reverse diffusion is the reverse process of the forward diffusion, Eq.
(2) derives the reverse diffusion.
p t ( x t ) is a probability distribution of x t .
[Equation]
The score function ∇ log p t ( x t ) has every information to reconstruct the data distribution after the reverse process.
The estimation of the score function s θ ( x t , t )
≈ ∇ log p t ( x t ) is the goal of training.
[Equation]
Since it is intractable to obtain the score ∇ log p t ( x t ) from x t , some approaches approximate the score function with L dsm [25, 26], where λ ( t ) is a weighting function.
2.2 Disentangling Variational Autoencoders
A typical approach to disentangle latent variables in VAE is regularizing variational distribution q ϕ ( z ) , so z could be dimension-wise independent.
The simplest approach is β -VAE to place a higher coefficient on D KL ( q ϕ ( z | x 0 )
|| p ( z )) in the evidence lower-bound (ELBO), so the Kullback–Leibler divergence enforces q ϕ ( z | x 0 ) to be closer to N ( 0 , I ) , which is given as the prior distribution of p ( z ) .
Later, Factor-VAE [11] and β -TCVAE [2] proposed Eq.
(4) with γ > 1 showing better trade-off curve of the reconstruction sample quality and the disentanglement strength, where q ϕ ( z ) =
(cid:81)
dj
=
1 q ϕ ( z j ) .
[Equation]
While we inherit the idea of disentangling regularization from Eq.
(4) with hyper-parameter γ , we do not impose them directly on the loss function of Eq.
(4) for a generation.
Instead, LCG-DM creates a structure that uses inferred z as a condition in Score-based diffusion models.
Method
Figure 1 shows the overall structure of LCG-DM.
q ϕ ( z | x 0 ) maps from x 0 to disentangled latent code z .
The disentangled latent code vector z expresses semantic information with the independent dimensions of z .
From an alternative perspective, z becomes a pseudo-label of x 0 in the ideal setting.
Hence, we design LCG-DM to be a z -conditional Score-based diffusion model with the disentangling mechanism.
In Figure 1, s θ ( x t , t, z ) is the function to estimate the score given z , so s θ becomes the score estimation for learning.
Meanwhile, the disentangled latent code of z is not provided from the dataset, so we create a VAE structure to extract the disentangled z from q ϕ and p ψ .
This VAE structure follows Factor-VAE, and we utilize Eq.
(4) in LCG-DM, which means that we have γ -controlled disentanglement term.
Since we leverage the disentangled latent condition z in training the diffusion model of s θ , it is necessary to have stable signals from z in the diffusion model learning.
Therefore, it is natural to have a step-wise approach in training.
[Equation]
In detail, L overall in Eq.
(5) denotes the loss function of LCG-DM.
We conducted the training using two stages.
In Stage 1, the model optimizes L fvae , so that q ϕ extracts disentangled latent code z .
Once we have convergence in obtaining z , we initiate Stage 2, which optimizes L csdm , so that s θ estimates the score function of conditional distribution ∇ log p t ( x t | z ) .
We derive that the minimization of conditional denoising score matching L cdsm equals L2 minimization of conditional score matching in Appendix A.5.
[Equation]
[Figure]
Results
[Figure]
We experimented with two benchmark datasets: dSprites [15] and CelebA [14].
dSprites is chosen because of its synthetic nature, so we can comprehend and measure the disentangled conditional latent variable z and its quality.
CelebA provides real-world images for generation tasks by manipulating the disentangled latent variables
4.1 dSprites dataset
[Table]
dSprites is a synthetic dataset that has known six latent factors, i [Table] e. shape, scale, position, etc, in generating a single image [15].
Therefore, we can explicitly measure the disentanglement quality by FVM [11] and MIG [2].
We compared the performance of LCG-DM against the disentangling structures of 1) VAE variations ( βVAE [7] and FactorVAE [11]), and 2) GAN variations (InfoGAN [3] and ID-GAN [13]).
Additionally, we included DDPM++ +ST [10], which is the backbone diffusion model we use.
DDPM++ +ST is not designed for disentanglement while the proposed model, LCGDM, is the first disentanglement diffusion model to our knowledge.
[Table]
Table 1 shows the evaluation result from two perspectives: image generation and disentanglement.
LCG-DM shows the best image generation in FID, and LCG-DM also provides the most disentangled latent variable z in FVM and MIG.
Particularly, we experimented in two versions of LCG-DM: the prior version draws a sample from z , x T ∼ N ( 0 , I ) , and the encoded version draws a sample from z ∼ q ϕ and x T ∼ N ( 0 , I ) .
The encoded version is slightly better because it takes advantage in drawing the condition with a sample image x 0 via encoder q ϕ .
The outperformed FID is achieved by diffusion model structure, and the proposed conditioning structure makes higher γ minimally damages the generation.
Table 2 shows the impact of γ , which is a strength hyperparameter for disentanglement.
As γ increases, FVM and MIG become better as expected.
However, FID seems to be uncorrelated with γ , and the reason could be the separated training stages between the disentanglement by VAE and the generation by a diffusion model.
4.2 CelebA dataset
[Table]
Table 3 shows the FID evaluation across disentangling models.
We observe LCG-DM beats all baselines in its generation at the cost of encoded sampling.
However, the gap between prior and q ϕ yields a large discrepancy, and we attribute this to the inflexibility of the VAE structure.
The large variational gap is observed not in dSprites dataset, but in CelebA datasets.
This is because CelebA contains more complex information, so latent space also be more complex.
Table 4 shows the FID evaluation across diffusion models without disentanglement, except LCG-DM.
While LCG-DM shows the worst FID from the prior version, it should be noted that 1) the NLL is comparable; and 2) LCG-DM has the disentangled latent variables unlike NCSN++ [27] and DDPM++ +ST [10].
[Table]
[Table]
Finally, Figure 2 shows the manipulation of a disentangled latent in image generations.
Figure 2-(b) shows that the independent factor is smoothly changed for each z dimension.
Figure 2-(c) shows that editing of image by adjusting each dimension.
The same arrow means the change of the same dimension.
Conclusion
This paper presents the first disentanglement diffusion model, LCG-DM.
The proposed model utilizes a conditioning structure, so the impact of the disentanglement constraint minimally damages the image generation performance.
This disentanglement diffusion model is a gateway to causal inference and counterfactual generations for various machine learning tasks, i.e. fairness research.
Acknowledgements
This research was supported by Research and Development on Key Supply Network Identification,
Threat Analyses, Alternative Recommendation from Natural Language Data(NRF) funded by the
Ministry of Education(2021R1A2C200981612)
A Appendix A.1 Scaling method
It is also an important problem to reflect the condition well in the generation, in conditional diffusion models.
The estimation target of the conditional diffusion model is the score function of conditional probability ∇ log p t ( x t | z ) .
Since ∇ log p t ( x t | z ) = ∇ log p t ( x t ) + ∇ log p t ( z | x t ) , some paper propose separately estimating 1) an unconditional score function and 2) an additional network which would become a classifier when z is regarded as a class label [27].
On top of this structure separation between ∇ log p t ( x t ) and ∇ log p t ( z | x t ) , recent works propose a scaling method that uses ∇ log p t ( x t ) + λ · ∇ log p t ( z | x t ) as a conditional score function to strengthen the conditional input in the generation process [5].
Instead of separating and utilizing an adhoc network of ∇ log p t ( z | x t ) , we adopt the classifier-free guidance method without the adhoc network.
Subsequently, we estimate ∇ log p t ( x t )
≈ s θ
( x t , t, ∅ ) and ∇ log p t ( x t | z )
≈ s θ
( x t , t, z ) , which would turn our conditional score function into Eq.
(7) by following the prior work of [8].
[Equation]
Figure 3 shows the effects of scaling method.
We observed that higher λ makes more reflection of pseudo-label, but has less diversity.
A.2 Additional latent traversals on CelebA dataset
A.3 Experiment details LCG-DM
The disentangling structure in LCG-DM uses Factor-VAE with γ =20, and 10 latent dimensions.
The score network is pretrained in just the same setting as DDPM++ +ST [10], and fine-tuned with conditional embedding.
We take 150,000 iterations with 64 batch sizes for fine-tuning.
We use a probability flow ode sampler for a consistent sample from fixed x T [27].
Disentangling VAES
We measured FID, FVM, and MIG by experimenting 3 times for VAE, β -VAE, Factor-VAE, and β -TCVAE.
We use the same neural network architecture for encoder and decoder structures, which consists of stacked CNN and MLP layers.
The latent dimension was set to 10, and the disentangling coefficient was set to 10.
Results from literature
We denote † on the table, which means the value from other literature.
We refer to the results of GAN variants from [13, 9], and Score-based diffusion model variants from [10, 27].
A.4 Metrics FID
Fréchet Inception Distance (FID) indicates the distance between training data and generated data, which is the most famous metric of sample quality.
We use the clean-FID calculation module [17] with the Inception V3 network.
We measure on 50k random samples in all settings.
NLL
We measure the Negative Log Likelihood (NLL) of LCG-DM on CelebA test dataset, and we compare it with non-disentangled diffusion models in Table 3.
To measure it in LCG-DM, we calculate the upper bound of NLL as a state in Eq.
We compute log p θ ( x | z ) using a probability flow ode [10].
[Equation]
FVM
Factor VAE Metric (FVM) is a disentanglement metric that has a scale of 0 to 1, and the larger the better [11].
Let ( x (1) , ..., x ( L ) )
the random subset of training data with one ground-truth factor v k is fixed.
q ϕ maps from ( x (1) , ..., x ( L ) ) to ( z (1) , ..., z ( L ) ), then let d is the dimension of z that has minimum variance.
FVM indicates the accuracy of a majority-vote classifier that trains ( d , k ).
MIG
Mutual Information Gap (MIG) is a disentanglement metric that has a scale of 0 to 1, and the larger the better [2].
For the given ground-truth factor v k , let z i ( k ) and z j ( k ) be the top two latent variables that have the highest mutual information with v k .
MIG is the average of mutual information gap, where K indicates the number of ground-truth factors.
[Equation]
[Equation]
A.5 Derivation of conditional denoising score matching
The score network s θ ( x t , t, z ) : R n +1+ d → R n estimates score of conditional density function.
Eq.
(10) shows L2 loss of conditional score matching, but it is hard to train since ∇ log p t ( x t | z ) is intractable.
This proof shows that optimize L cdsm equals to optimize Eq.
( 10 ), since Eq.
(12) shows the same form of L cdsm .
Here, Eq.
(11) is derived from Eq.
[Equation]
z | x 0 )d x 0 d x t d z = E x 0 , x t , z [ s θ ( x t , t, z ) · ∇ log p t ( x t , z | x 0 )]
=
E x 0 , x t , z (cid:2) s θ ( x t , t, z ) · [ ∇ log p 0 t ( x t | x 0 ) + ∇ log p ( z | x 0 )] (cid:3) =
E x 0 , x t , z [ s θ ( x t , t, z ) · ∇ log p 0 t ( x t | x 0 )]