Divide and Compose with Score Based Generative Models

Sohl-Dickstein et al. [25] and Ho et al. [11] proposed to design a generative model, called Denoising diffusion probabilistic model (DDPM) from a Bayesian perspective. Imagine we sample an image from the data distribution, ${x}_{0}\sim p_{0}$. Consider the data corruption sequence where we incrementally add Gaussian noise to the image until it turns into complete noise. This forms a Markov chain and the joint distribution of the forward series is given by $p_{0}(x_{0})\prod_{t=1}^{t=T}p_{t-1,t}(x_{t})|p_{t-1,t}(x_{t-1})$. Then, a reverse Markov chain is considered were $p_{\theta}(x_{t-1}|x_{t})$ is conditionally Gaussian such that the reverse process joint distribution is given by $p_{\theta}(x_{T})\prod_{t=T}^{t=1}p_{\theta}(x_{t-1}|x_{t})$. The DDPM algorithm optimizes the evidence lower bound of the data likelihood such that when optimization is complete, the reverse joint distribution coincides the forward joint distribution. Since the reverse conditional distributions are parameterized by the error functions $\epsilon_{\theta}$, the algorithm, essentially, boils down to optimizing the $\epsilon_{\theta}$. One key trick proposed in DDPM is that the complex loss obtained from the ELBO can be neatly expressed as an extremely simple loss function as follows:

$\displaystyle\mathcal{L}_{\theta}^{ddpm}=\mathbb{E}_{t,x_{0},\epsilon}\big{\{}\lambda_{t}||\epsilon_{\theta}(x_{t},t)-\epsilon||^{2}\big{\}}$

(2)

where, $x_{t}=\sqrt{\alpha_{t}}x_{0}+\sqrt{1-\alpha_{t}}\epsilon$ is a sample from distribution $p_{t}(x|x_{0})=\mathcal{N}(x;\sqrt{\alpha_{t}}x_{0},1-\alpha_{t})$, the marginal distribution at time $t$ and $\epsilon$ is a random vector from an isotropic Gaussian distribution. Note that as we incrementally add noise to $x_{0}$, the marginal distribution $p_{t}$ at time $t$ can be expressed as a Gaussian distribution conditioned on $x_{0}$. The mean at time $t$ has diluted by a factor of $\sqrt{\alpha_{t}}$ and variance is $1-\alpha_{t}$. Also, the $\lambda_{t}$ in eq.(2) is a function of $\alpha_{t}$s and noise added at each time instant. Once we have the optimized network, $\epsilon_{\theta}$, the image generation process is nothing but following the reverse conditional distribution $p_{\theta}(x_{t-1}|x_{t})$ starting from the isotropic Gaussian noise $p_{\theta}(x_{T})$.

Song et al. [27] proposed a different generative model by estimating the gradient of data distribution, called the score function, whose sampling looked similar to that of DDPM. Even though the score based generative model seemed similar to diffusion model, DDPM, the connection was unclear until another paper due to Song et al. [28], which showed that there is a deeper connection between the two networks: the error network in DDPM, $\epsilon_{\theta}$ is same as the score function $s_{\theta}$ in score generative model. They establish that DDPM essentially performs score matching [15, 29]. They further develop this line of argument showing that the continuous version of score based generative model can be obtained by reversing a stochastic differential equation (SDE). Specifically, they generalize the forward process of adding noise to a continuous setting where the forward process takes the form of a stochastic differential equation (SDE). Similarly, the reverse process of starting from the isotropic Gaussian and incrementally removing noise can also be shown to be a stochastic different equation continuous in time:

	$\displaystyle\text{FOR}:dx$	$\displaystyle=f(x,t)dt+g_{t}dw$		(3)
	$\displaystyle\text{REV}:dx$	$\displaystyle=[f(x,t)-g_{t}^{2}\nabla_{x}\log p_{t}(x)]dt+g_{t}d\bar{w}$		(4)

Luckily, the reverse stochastic differential equation only needs the score function at each timestamp $t$ in addition to other functions $f,g$ from the forward equation. Using the same score matching idea, the score function can be first trained with the following loss function

$\displaystyle\mathcal{L}_{\theta}=\mathbb{E}_{t}\big{\{}\lambda_{t}\mathbb{E}_{x_{0}}\mathbb{E}_{x_{t}|x_{0}}[||s_{\theta}(x_{t},t)-\nabla_{x_{t}}\log p_{0,t}(x_{t}|x_{0})||^{2}]\big{\}}$

(5)

Note that this loss is same as eq.(2) once we use the fact that $p_{t}$ is conditionally Gaussian. Once the score function is learnt, the generative model is given by replacing $s=\nabla_{x}\log p$ in the reverse time stochastic differential equation:

$\displaystyle dx=[f(x,t)-g_{t}^{2}s_{\theta}(x,t)]dt+g_{t}d\bar{w}$

(6)

In practice, we need to discretize eq. (6) to obtain a generative model. We achieve this through Euler-Maruyama discretization as suggested in [28]:

$\displaystyle x(t-\delta t)=[f(x,t)-g_{t}^{2}s_{\theta}(x,t)](-\delta t)+g_{t}\sqrt{\delta t}z$

(7)

where $z\sim\mathcal{N}(0,I)$ is a random sample from standard Gaussian distribution.

Following the unconditional generative models, a few conditional generative models have been developed. Most of these models, however, require some form of supervision regarding the group on which to condition the generation. For example, Dhariwal et al. [8] developed conditional generative models based on class label by modifying score functions with extra term representing the gradient of log of classifier likelihood. These methods are known as classifier guided conditional generative models. Later, classifier-free models [12] were also proposed who eschewed the idea of training a classifier altogether. Nevertheless, they also need class label to train.

Towards the direction of unsupervised learning, Diffusion autoencoder (DiffAE) [22] learns the conditional distribution based on latent vector obtained from the encoder, which is in the line of our work. Our work differs from theirs in the fundamental notion of what represents the component. In our model, we encode different score components from an image, where each component sort of captures one concept (energy function). Later we generate images by recombining these components.

We can derive a likelihood formulation to train the score based generative model based on Feynman-Kac Theorem [16], as shown in [13, 5, 26]. We start from the likelihood formulation of training score based generative model.

		$\displaystyle\log p(x_{0})\geq\mathcal{L}_{VLB}(x_{0},\theta)=E_{p_{T}}[\log p_{\theta}(x_{T})]$
		$\displaystyle-\int_{0}^{T}\mathbb{E}_{p_{t}}\bigg{[}\frac{1}{2}g_{t}^{2}||s_{\theta}(x_{t},t)||^{2}+\nabla_{x}(g_{t}^{2}s_{\theta}(x_{t},t)-f)\bigg{]}dt$		(8)

Taking expectation, we arrive at,

$\displaystyle\mathcal{L}_{EVLB}(\theta)=\mathbb{E}_{{x_{0}}\sim p_{0}}[\mathcal{L}_{VLB}(x_{0},\theta)]$

(9)

From eq.(9), we can derive the same loss function as in eq.(5) by using a rough approximation of the integral with the discrete summation and equivalence between different score matching [29, 15](as shown in [13]). Therefore, unconditional score estimation can be thought as a crude approximation of the expected likelihood maximization. Nevertheless, expressing the likelihood as the integral in eq.(4.1) is much more illuminating and powerful. Observe that eq.(9) is obtained by taking expectation of the likelihood of individual data in eq.(4.1). It is this expectation which leads to learning a common unconditional score function. We can forego expectation to train score functions for each data point. For any data $x_{\zeta}$, we can directly optimize its likelihood by training a score function $s_{\theta,\zeta}$ to optimize the lower bound on the right hand side of eq.(4.1). Specifically, we train an encoder, Enc and score function, $s$ as follows:

		$\displaystyle\zeta=\texttt{Enc}_{\theta}(x_{\zeta})$		(10)
		$\displaystyle\log p(x_{\zeta})\geq=E_{p_{T}}[\log p_{\theta}(x_{T})|x_{0}=x_{\zeta}]$
		$\displaystyle-\int_{0}^{T}\mathbb{E}_{p_{t}}\bigg{[}\frac{1}{2}g_{t}^{2}||s_{\theta}||^{2}+\nabla_{x}(g_{t}^{2}s_{\theta,\zeta}-f)|x_{0}=x_{\zeta}\bigg{]}dt$		(11)

Eq.(10) and eq.(4.1) completes our autoencoder model that maximizes the lower bound on log likelihood of individual image, $x_{\zeta}$.

As motivated in the introduction (eq.(1)), we want to decompose each score function into multiple components, which intuitively represent different energy components (negative of gradient of energy to be precise). We decompose each score function into $K$ components such that $s_{\theta,\zeta}=(s_{\theta,\zeta}^{(1)}+s_{\theta,\zeta}^{(2)}+...+s_{\theta,\zeta}^{(K)})/K$. This decomposition requires $K$ different score functions. A computationally efficient way would be to share the weight of score functions but use different latent vectors for different components, which also gives a structure to the latent space. Therefore, we decompose the score function as:

$\displaystyle s_{\theta,\zeta}=(s_{\theta,\zeta^{(1)}}+s_{\theta,\zeta^{(2)}}+...+s_{\theta,\zeta^{(K)}})/K$

(12)

Note that now the burden of learning different energy components has been shifted to the latent vectors together with a shared conditional score function $s_{\theta,\zeta^{(k)}}$.

The encoder encodes each image $x_{\zeta}$ into $K$ latent vectors $\zeta^{(K)}$. The summed score function given by eq.(12) is used to maximize the lower bound on the log likelihood of each $x_{\zeta}$. We also approximate this integral with discrete summation and invoke the equivalence of implicit score matching [15] and denoising score matching [29, 28] to obtain an approximation [13] as the following loss function:

$\displaystyle\mathcal{L}_{\theta,\zeta}=\mathbb{E}_{t}\big{\{}\lambda_{t}\mathbb{E}_{x_{t}|x_{\zeta}}[||s_{\theta,\zeta}(x_{t},t)-\nabla_{x_{t}}\log p_{0,t}(x_{t}|x_{\zeta})||^{2}]\big{\}}$

(13)

We jointly train $K$ encoders and the score loss in eq.(13), where the score is given by eq.(12). To design a score function as a function of the the latent vector $\zeta$, we use the adaptive Group Norm (AdaGN) strategy as described in [8]. Once the model is trained, we generate image by first sampling noise $z$ from the standard Gaussian distribution and iteratively applying Euler-Maruyama discretization of the SDE as described by eq.(7).

We also experiment with generating samples by combining score components with unconditional score function similar to [12]. For this, we can separately train conditional and unconditional score functions. However, to make it computationally cheap, we pass a vector of ones 1 to realize the unconditional score function $s_{u}=s_{\theta,\textbf{1}}$, similar to the trick used in [12]. We combine conditional and unconditional score functions with linear combination of coefficients as described in experiments.

The score network is the modified version of UNet architecture as described in Dhariwal et al.[8]. To condition on the latent vector obtained from encoder, we use the AdaGN as described in [8], which is inspired from adaptive instance normalization (AdaIN) [14], but uses group normalization [30] instead of instance norm. Our architecture is similar to DiffAE’s adaptation of AdaGN, i.e.

$\displaystyle\texttt{AdaGN}(h,t,\zeta)=\zeta_{s}(t_{s}.\texttt{GroupNorm}(h)+t_{b})$

(14)

where $h$ is normalized feature map at different layers, $t_{s},t_{b}$ are obtained from the time embedding by applying MLP and $\zeta_{s}$ is affine transformation of the latent vector $\zeta$.

We burrow encoder from DiffAE [22], i.e. the first part of UNet architecture used in score function. We experiment with two values of K: $K=\{3,5\}$. The encoder encodes K vectors of length $128$, which condition the score functions.

Due to the computing constraints, we train for 150K iterations using Adam optimizer and batch size of 32. We experiment on four datasets in computer vision: 1) Celeb-A, 2) LSUN-outdoor_church, 3) Cifar-10 and 4) SVHN. We use image size of $32\times 32$ for Cifar-10 and SVHN and $48\times 48$ for LSUN and Celeb-A, again due to resource constraints.

In Fig.(1), we show auto encoding of samples from four datasets. We input a batch of 16 images, all of which are the same. Figure shows that the reconstruction is very close to the ground truth, but, at the same time, there are natural variations.

Visualization of different score components can be insightful about what each component is capturing from the image. It is unclear what is the best way to visualize score components which are essentially represented as matrices. We demonstrate that generating image samples using each component can be the best visualization strategy of each components. To be precise, if $s_{\theta,\zeta}=(s_{\theta,\zeta^{(1)}}+s_{\theta,\zeta^{(2)}}+s_{\theta,\zeta^{(3)}})/3$, then we take each component, say $s_{\theta,\zeta^{(k)}}$ and generate images using the Euler Maruyama discretization in eq. (7). The samples generated using three components are plotted in Fig.2 where each row is a component. The input image of these components are the same images as in Fig.1. Note that even though input image and the latent vectors are the same, there are natural variations in samples form each component.

Several observations are in order. First, this method can result in human interpretable components in some cases, but not always. For example, the third component in the svhn image is clearly capturing digits, however, first and second components are capturing some abstract texture or lighting related information. Similarly, second component of the bird is clearly focusing on the color, but first one is capturing some abstract shape information.

This score decomposition forces us to rethink the definition of components of an image. In classical works, disentanglement has been seen as some kind of statistical property of the distribution and methods were proposed to enforce such statistical pattern, for example, independence. Here, however, we see that components could capture interesting pattern that may or may not be independent or interpretable. We can certainly say that the three components are different, but at the same time capture something about the input image. Therefore, they are different disentangled components. Yet, they may or may not possess statistical independence property or complete human interpretability.

We manipulate images using these learned components, with the help of an unconditional score function. In Fig.3, we retain two components untouched while we linearly interpolate one component with the unconditional score function to generate images. That is, if $s^{(1)},s^{(2)},s^{(3)}$ are three components, we generate image with the following score function:

$\displaystyle s_{\alpha}^{interp}=(\alpha*s^{(1)}+s^{(2)}+s^{(3)})/3+(1-\alpha)s_{u}/3$

(15)

In Fig.3, we observe that the first component is associated with the background information and less in the church building. When this component is interpolated, we see that the information like yellow color (related to second component) and building architectures (related to third components) are preserved. As we go from $\alpha=1$ to $\alpha=0.1$, we see that the image changes while preserving the yellow color of the church and with different building architectures, but it changes the tree and the background making it more diverse. At $\alpha=0.1$, we see that a lot of tree has been removed and replaced with something else to create diverse background setting.

In Fig.4, we observe similar effect. The first component is changed in second row while retaining second and third component. From the first row, it is clear that the first component is associated with smile. Hence, as we go from $\alpha=1$ to $\alpha=0.1$, the sample images are more diverse in terms of smile. At $\alpha=0.1$, we see a few instances where the mouth is even closed.

Similarly, in Fig.5, we change second and third component while retaining the first one. Here also, the first component is associated with the smile while second and third are associated with shape features and hair. We fix the first component and interpolate second and third component as follows:

$\displaystyle s_{\alpha}^{inter}=\frac{(\alpha*(s^{(1)}+*s^{(2)})+s^{(3)})}{3}+\frac{2(1-\alpha)s_{u}}{3}$

(16)

In doing so, when we go from $\alpha=1$ to $\alpha=0.1$, smile is preserved while other features change. Hence, we see image with a lot of diversity in terms of shape, color and hair while preserving smile as we go towards $\alpha=0.1$.

In Fig.6, we tune the weight of two score components $s^{(1)}$ and $s^{(2)}$. When the $s^{(2)}$ is weighted heavily, the greenish hue and the corresponding shape of the church are dominating. As we moved more heavily towards $s^{(1)}$, we started seeing more of a mix between $s^{(1)}$ and $s^{3}$.

Fig.7 compares five components against three components learned in Cifar10 dataset. Components 4, 5 and 1 in five components are similar to the three in first row. However, Component 3 in second row is different and seems to capture the colored pattern with shape in the middle.