Exploring Compositional Visual Generation with Latent Classifier Guidance

A diffusion model is a type of deep latent variable model that approximates an unknown data distribution $p(x)$ through smooth, iterative denoising steps. It maps a pre-defined noise distribution to the data distribution using the following formula: $p_{\theta}(x_{0})=\int p_{\theta}(x_{0:T})dx_{1:T}$, where $x_{1:T}$ are the latent variables with the same dimensionality as the data $x_{0}$. The forward diffusion process, resembling a parameter-free encoder, is a Markov chain $q(x_{1:T}|x_{0})=\prod_{t=1}^{T}q(x_{t}|x_{t-1})$, where each $q(x_{t}|x_{t-1})$ is typically a Gaussian distribution. The forward process perturbs inputs according to a pre-defined schedule, and the transformed data distribution $q(x_{t}|x_{0})$ will gradually converge to a standard Gaussian $\mathcal{N}(x_{T};\mathbf{0},\mathbf{I})$. The reverse sampling process, resembling a hierarchical decoder, is composed of a sequence of de-noising steps $p_{\theta}(x_{t-1}|x_{t})$, which is parameterized by a deep neural network with parameter $\theta$.

During training, the input images are corrupted by the forward process, and the diffusion model is trained to reconstruct the original images from the corrupted inputs. Specifically, for de-noising diffusion probabilistic models (DDPM) [14], the training objective is formulated as a re-weighted variational bound by treating DDPMs as VAEs, while for scored-based generative models [41], the objective is derived using score matching. Once trained, to generate samples from the learned distribution, one first samples $x_{T}$ from a standard Gaussian and then uses the reverse process to transform it into the image space.

Here, we focus on leveraging the latent space of a pre-trained generative model. Specifically, we train a diffusion model to approximate the latent distribution $p(z)$ of a pre-trained generator $G$ that maps a latent space $\mathcal{Z}$ to the image space $\mathcal{X}$. Modeling the latent space has several advantages over modeling the image space. For instance, the latent space enjoys properties such as disentanglement [6], which can facilitate more controllable manipulations of the generated images. Additionally, using various guidance techniques in the latent space is often more feasible since training latent guidance terms is generally easier than training other manipulation methods in image space [36].

Conditional generation with diffusion models relies on perturbing unconditional generation with user-specified guidance terms, namely classifier guidance [39, 42, 9] and classifier-free guidance [15]. Although classifier-free guidance performs competitively in image space and is sometimes more favorable than classifier guidance [15, 23], we argue that using classifier guidance in latent diffusion models has its unique advantages. Regarding the under-performance of classifier guidance in the image space, one popular suspicion is that image classifiers tend to learn shortcuts from suspicious correlations. For example, a deep neural network classifier on the attribute “old” can be misguided by “white hair” and ignore its holistic features. This problem is alleviated in a compact, even disentangled latent space, if the semantic directions of ‘old’ and ‘white hair’ are orthogonal. Also, deep image classifiers are typically vulnerable to adversarial attacks, while latent classifiers with much few parameters suffer less from this problem. Another benefit is that classifiers are usually easier to train than diffusion models used in classifier-free guidance. Finally, when the classifiers are linear, classifier guidance resembles linear arithmetic methods, as we will show in Section 3.4.

The goal of conditional generation is to model the conditional distribution $p(z|y)$ where $y$ is the conditions or attributes. By Bayes rules $p(z_{t}|y)=\nicefrac{{p(z_{t})p(y|z_{t})}}{{p(y)}}$, the score of the conditional probability $\nabla_{z_{t}}\log p(z_{t}|y)$ can be factorized as the unconditional score $\nabla_{z_{t}}\log p(z_{t})$ and the gradient flow $\nabla_{z_{t}}\log p(y|z_{t})$. Therefore, one simply needs an unconditional latent diffusion model and a latent classifier to model the conditional score, known as classifier guidance. In practice, the classifier guidance term is usually scaled by a factor $\alpha$, such that $\nabla_{z_{t}}\log p(z_{t}|y)=\nabla_{z_{t}}\log p(z_{t})+\alpha\nabla_{z_{t}}\log p(y|z_{t})$. The factor $\alpha$ serves as a temperature parameter which adds another layer of controllability to the sharpness of the posterior distribution $p(y|z_{t})$.

Compositional generation can be considered as conditional generation with multiple conditions and the relations among them. In this paper, we consider two relations, conjunction “AND” and negation “NOT”. For the conjunction of attributes $y^{1}\wedge y^{2}\wedge...\wedge y^{n}$, assuming the conditions to be independent of each other, we can simply factorize the compositional log probability as

And with attribute negations $y^{1}\wedge...\wedge y^{m-1}\wedge\overline{y^{m}}\wedge...\wedge\overline{y^{n}}$, without loss of generality, we can factorize the log probability similarly

While classifier guidance is useful for compositional generation, there is no guarantee that those results will be anything similar to the original image when doing manipulations. This is because the generation is not conditioned on the original image. As there is no constraint on the specific form of the posterior [39], conditioning on the original image amounts to adding a new guidance term $\gamma_{t}\nabla_{z}\log p(\hat{z}|z)$, where $\hat{z}$ is the latent of the image to be manipulated. For conjunction relations as in Eq. (1), the overall score function for manipulation then becomes

similarly for Eq. (2) with the presence of negation. When $p(\hat{z}|z_{t})$ is modeled by an isotropic Gaussian distribution, the new guidance term $\gamma_{t}\nabla_{z}\log p(\hat{z}|z)$ behaves as a regularization term $\nabla_{z_{t}}{\|z_{t}-\hat{z}\|}_{2}^{2}$.

A true compositional model should be able to easily encompass new attributes without re-training the whole model. Indeed, the training of the unconditional diffusion models and the latent classifiers can be decoupled, and such training amounts to maximizing the evidence lower bound (ELBO) of the conditional log-likelihood. This means that encompassing new attributes simply requires training classifiers on them, and the latent diffusion model as well as used classifiers can be recycled.

We take DDPM as our example and begin with unconditional generation.

See [14] for the detailed proof.

Lemma. 2 states that training unconditional diffusion models and their latent classifiers is equivalent to maximizing the ELBO of joint log-likelihood of $z$ and $y$ up to a constant.

The regularized guidance manipulates a given latent $\hat{z}$ in a non-linear fashion, but it degrades to linear manipulation with additional assumptions. We take the case where there are only conjunction relations as an example and consider Eq. (3).

For attribute negation, the solution perturbs $\hat{z}$ towards the negative direction of the classifiers. This is a natural multi-attributes generalization of the vector arithmetic method, and we refer it as the linear version of latent classifier guidance in later comparisons.

We first evaluate the ability of latent classifier guidance to generate images with multiple desired attributes. For high-resolution images (1024$\times$1024), we select StyleGAN2 as the pre-trained generator; for low-resolution images (256$\times$256), we use Diffusion Autoencoder.

We compare our proposed method, which we refer as LCG (Latent Classifier Guidance) from below, with StyleFlow [3] and LACE [27]. Results of StyleFlow are directly taken from [27] for the conjunction of “gender”, “smile” and “age”, while it cannot handle the other compositional task where negation relations are involved. To compare with LACE, we use their official implementation. Note that LACE is not applicable for Diffusion Autoencoder as its semantic latent space is not endowed with a parameterized distribution as the $\mathcal{Z}_{\text{s}}$ or $\mathcal{W}_{\text{s}}$ space of StyleGAN2. However, we can still apply latent classifier guidance because it can be easily fitted by diffusion models.

The quantitative comparison is shown in Table 2 and the qualitative comparison shown in Fig. 1 and Fig. 2. While for the quantitative results target conditions are randomly sampled for each attribute, for the qualitative results, we use fixed targets for the sake of visualization. As we can see, with latent classifier guidance, using simple linear arithmetic (LCG-Linear) and latent diffusion models (LCG-Diffusion) both perform competitively against previous non-linear methods.

We evaluate latent classifier guidance on manipulating both synthetic and real images.

One main challenge of manipulating multiple attributes is maintaining the non-targetted attributes. For compositional generation, these are the out-of-scope attributes; for sequential editing, these also include the previously manipulated attributes. To tackle the challenge, previous methods either design a protection with linear manipulation such as InterfaceGAN [36], or use non-linear manipulation such as StyleFlow and LACE. As we have shown in previous sections, our LCG-linear can be very competitive against other non-linear methods despite its simplicity.

To understand its power, we examine the linear classifiers that are learned for manipulation. Here we use the heatmap to visualize the correlation between each pair of linear classifiers learned on the semantic latent space of Diffusion Autoencoder, as shown in Fig. 4. As we can see, the semantic latent space of Diffusion Autoencoder is favorably disentangled and thus the linear classifiers show strong orthogonality frequently. This means that even without a specific protection mechanism as in [36], LCG-linear is still capable of preserving the identity in most cases. In Table 4, we list the changes of conditional accuracy of other attributes when a single attribute is linearly manipulated. The four edits correspond to the four attributes “yaw”, “smile”, “age”, and “glasses” respectively. The small changes indicate that the attributes are well disentangled in such generative models, and further explain the efficacy of LCG-linear.

Despite the complicity of the non-linear diffusion based method, it does not always perform favorably against the linear version. The main motivation for non-linear methods, as argued in [3], is that linear manipulation often moves a latent code outside the latent distribution which leads to low-quality generation. Indeed, in linear manipulation, we assume a non-informative $p(z_{t})$ which favors different regions of sample space equally, regardless of the actual latent distribution. This indicates that non-linear control is likely to profit when the generation needs to traverse low density region or is simply out-of-distribution.

An example is sequential editing. Sequential editing is more prone to low density region, as the new edit is conditioned on previous edits that possibly have already guided the latent to low density regions. To see this, imagine a 3-D Guassian distribution where each axis represents an attribute and we want to guide a sample point from $[-1,-1,-1]$ to $[1,1,1]$. Compositional generation is analogous to a direct path $[-1,-1,-1]\rightarrow[1,1,1]$, while sequential editing is analogous to a path $[-1,-1,-1]\rightarrow[-1,-1,1]\rightarrow\rightarrow[-1,1,1]\rightarrow[1,1,1]$ that traverses more low density regions. As we can in Table 3, LCG-diffusion outperforms LCG-linear on FID, generating more realistic images. Characterizing the latent distribution with diffusion model in this case is favorable than a non-informative one, as the diffusion model always pulls the sample toward high density region thus preventing it going out-of-distribution. A downside of this is that, as we can see from the ID score, keeping images realistic is at the cost of losing identity preservation. Improving the identity preservation should serve as an interesting topic for future exploration.