Instructing Text-to-Image Diffusion Models via Classifier-Guided Semantic Optimization

1 Introduction

Faithful image reconstruction is a fundamental prerequisite for image editing tasks. Recently, diffusion-based [Ho et al., 2020; Song et al., 2020; Rombach et al., 2022; Song et al., 2021] generative models have demonstrated significant advantages over GANs [Goodfellow et al., 2014] in the field of image generation [Dhariwal and Nichol, 2021]. These approaches have emerged as powerful tools for image editing due to their generative modeling capabilities. However, exploring the decoupled semantic subspace of diffusion models remains a critical yet challenging area of research.

Most of the initial works on diffusion-based image editing require training or fine-tuning the model and focus on Trainable Frozen SD Component Classifier BackPropagation a single domain of data [Preechakul et al., 2022; Lu et al., 2024]. The rapid rise in popularity of text-to-image diffusion models, such as Stable Diffusion (SD) [Rombach et al., 2022], DeepFloyd IF [Saharia et al., 2022], and Latent Consistency Models [Luo et al., 2023], has further inspired researchers to achieve decoupled image editing based on text or other conditions (e.g., semantic segmentation) without requiring additional training or fine-tuning [Brack et al., 2023; Lee et al., 2024]. Unlike small, single-domain diffusion models, these advanced models enable cross-domain generation, broadening their applicability across diverse image editing scenarios, have become a hot spot of current research. However, existing solutions are mainly based on text prompts [Brack et al., 2023; Wu and De la Torre, 2023; Lee et al., 2024], which present several challenges. Creating appropriate text prompts can be inherently difficult, and prompts with identical semantics often yield significantly divergent outcomes. Furthermore, achieving disentangled editing through text prompts alone is highly challenging, as they frequently affect unrelated regions of the image. Recent works that eschew text prompts proposed unsupervised ways to discover editing directions in the latent space of diffusion models [Chen et al., 2024; Dalva and Yanardag, 2024] also exhibit limitations, particularly in terms of generalization. For instance, while these methods explore editable directions using real face data, their applicability often diminishes when extended to other types of images, such as artistic or animestyle faces. Moreover, they can only obtain a limited number of subtle editing directions, often lack the ability to specify attributes, and rely on the clear structure or inherent variance within the dataset itself.

∗Completed during the internship at A*STAR and SGIT AI Lab

†Corresponding author

Figure 1: ClAssifier-guide Semantic Optimization (CASO). The trainable continuous semantic embedding for the target attribute a, guides Stable Diffusion for desired edits.

In this paper, we propose optimizing continuous semantic embeddings through attribute classifier gradients to guide text-to-image models for targeted editing. Our framework is shown in Fig. 1. It can identify disentangled and generalizable editing directions. Our contributions are as follows.

• We propose ClAssifier-Guided Semantic Optimization (CASO), a lightweight and efficient framework for image editing, which simply tunes trainable semantic embeddings and does not require training or fine-tuning of diffusion models.
• We theoretically and experimentally demonstrate that CASO learns semantic embeddings containing attribute class mean information, thereby enabling precise attribute editing.
• Experimental results show that our editing directions are highly disentangled and exhibit strong generalization, outperforming various existing approaches.

2 RelatedWork

Deep generative models have achieved great success in image editing. Various types of generative models, such as VAEs [Kingma, 2013], GANs [Goodfellow et al., 2014], flow-based models [Dinh et al., 2014], and diffusion models [Ho et al., 2020; Song et al., 2020], have been adopted as editing frameworks. Among these, the recent state-of-the-art editing models are based on diffusion models. In the following, we will focus on reviewing work on diffusion based image editing.

Diffusion models perform iterative denoising based on a sequence of latent variables, which lack explicit semantic representations like those in GANs. So identifying the disentangled semantic space in diffusion models remains a significant challenge. Some image editing methods based on diffusion models require retraining or fine-tuning a diffusion model

[Kim et al., 2022; Wang et al., 2023; Preechakul et al., 2022; Valevski et al., 2022]. For example, DiffAE [Preechakul et al., 2022] and HDAE [Lu et al., 2024] use diffusion model as the representation encoders and enable image editing by modifying representations. These methods are generally only suitable for single-domain (such as human face) diffusion models. DiffusionCLIP [Kim et al., 2022] fine-tunes the diffusion model with the help of the CLIP [Radford et al., 2021]. UniTune [Valevski et al., 2022] fine-tunes the diffusion model on a single base image, encouraging the model to generate images similar to the base image. However, training or finetuning the diffusion model can be somewhat expensive.

Recent advances [Yang et al., 2023; Brack et al., 2023; Tumanyan et al., 2023; Wu and De la Torre, 2023; Hertz et al., 2022a; Mokady et al., 2022] avoid training or fine-tuning diffusion models. Some methods (SEGA [Brack et al., 2023] and OIG [Lee et al., 2024]) use text prompts to guide editing. OIG uses CLIP text encoder [Radford et al., 2021] to ensure semantic correlation and extracts the feature map of UNet to ensure structural similarity. SEGA generates editing noise from text prompts and then extracts a small portion of this noise (by default, the top 5%, based on empirical observations) for editing. Unfortunately, the 5% threshold can only be set empirically, which is an unstable way to edit and can cause image quality loss. Additionally, some work optimize text-based embeddings to faithfully reconstruct the input image, so as to achieve disentangled edits (the editing does not change the irrelevant features) with classifier-free guidance. Null-Text Inverse [Mokady et al., 2022] fine-tunes the nulltext embedding to align the sampling and inversion trajectories, enabling accurate image reconstruction, but it relies on a user-provided source prompt. Prompt Tuning Inversion [Dong et al., 2023] introduces a faster approach by encoding image information into a learnable text embedding and performing editing through linear interpolation with the target embedding. However, these methods prioritize reconstruction quality, whereas our goal is to obtain an optimal semantic embedding for effective image editing.

NoiseCLR [Dalva and Yanardag, 2024] and LOCO [Chen et al., 2024] propose unsupervised methods to discover interpretable directions in pre-trained diffusion models, but they suffer from several shortcomings. The unsupervised learned direction is uncontrollable and may not converge to userdesired direction. And no more than 10% of the directions learned by NoiseCLR are semantic meaningful. LOCO is only suitable for local editing, and some major structural changes such as cat→dog are not feasible.

3 Method

In this section, we describe our proposed method. First, we briefly discuss the background of latent diffusion models, and then introduce our method.

3.1 Latent Diffusion and Classifier-free Guidance

Diffusion models are generative models that produce data samples by fitting the data distribution through an iterative denoising process [Ho et al., 2020]. Stable Diffusion (SD) [Rombach et al., 2022] is a type of Latent Diffusion Model (LDM) that operates in the latent space of image data, where the conversion between latent code z and raw image data x is performed by the VAE encoder E and decoder D. The training loss is defined as:

where t is uniformly sampled from {1, . . . , T} and zt is a noisy version of z0 = E(x).

Classifier-free guidance [Ho and Salimans, 2022] trains a conditional diffusion model. Since the condition is left blank with a certain probability during training, it also supports unconditional generation. The final predicted noise with classifier-free guidance in SD is defined as:

where c is the condition, ϕ is null text and λ is guidance scale. For simplicity, we use ϵθ(zt) instead of ϵθ(zt, t) to represent the predicted noise for timestep t, as t is implicitly denoted with variable zt.

3.2 Classifier-guided Semantic Optimization

We can set c as a suitable text to guide the editing. However, as shown in Appendix.B, even when sampling with the same random seed, images generated from different text prompts with the same semantic meaning exhibit significant differences (the first four images). Furthermore, adding additional descriptions to the text also leads to substantial changes in the overall appearance of the generated images (the first and last images). Text guidance struggles to achieve disentangled edits, namely, preserving the attribute excluding details. To achieve precise image editing, we propose a lightweight approach to obtain semantic embeddings as conditional guidance based on a ready-to-use attribute classifier.

Suppose the target editing attribute has K classes, i.e., a ∈ {1, . . . ,K}. Define the ready-to-use attribute classifier as F and the images generated by the Stable Diffusion editing process as ˆx = G(x, ea), where x is the input image and ea is the embedding for each class a of the attribute. We define an editing loss to optimize {ea}Ka=1 as follows:

where ℓc is the classification loss that enforces generated images conditioned on ea to be classified as the corresponding attribute class a.

The attribute embedding {ea}Ka=1 obtained above contains the attribute class information at the dataset level, enabling image editing with precise semantics. We provide a theoretical justification in the following.

Denoting all layers of the classifier except the last as F-1, and the last layer with weight wa and bias ba for each attribute class a, the classifier for predicting a class label can be expressed as arg maxa′ (⟨wa′ , F-1(x)⟩ + ba′ ).

Definition 1 (Globally-centered attribute class mean µa). We define the globally-centered attribute class mean µa as:

where a = 1, 2, . . . , K and hi,a = F-1(xi,a) is the last-layer feature of the i-th sample in attribute class a for training the classifier.

Theorem 1. [Han et al., 2021] For a sufficiently large classifier network, the last layer of the classifier wa will converge to the globally-centered class mean µa, namely,

Theorem 1 describes the neural collapse phenomenon of training a deep classifier network [Han et al., 2021; Papyan et al., 2020; Yang et al., 2022]. We apply it to demonstrate that our classifier-guided semantic embedding encodes the attribute class mean.

Proposition 2. Give a sufficiently large classifier network F which is well-trained on the training data. When optimizing {ea}Ka=1 using Eq. (3) until converge while fixing F, the neural collapse phenomenon still happens, namely,

where a = 1, 2, . . . ,K. In particular,

which denotes the globally-centered attribute class mean of generated images.

The proof is similar to Theorem 1 given the same classifier. We give a detailed proof in the Appendix.

With Eq. (4) and Eq. (5), we can easily obtain µ′a (ea) ∥µ′a (ea)∥2 → µa ∥µa∥2 . This means the optimal attribute embedding ea in Eq. (3) is determined by the attribute class mean µa. We also verify this in our empirical study.

Algorithm 1 Training algorithm

Input: Training data X = {xi }Ni=1 , batch size B, denoising UNet ϵθ, VAE encoder E, VAE decoder D, classifier F.

Output: Semantic embedding ea

• 1: repeat
• 2: Sample a batch of images {xi }Bi=1 .
• 3: Get latent codes {zi0 = E(xi )}Bi=1 .
• 4: Following Eq. (8), generate noise latent codes {ziL}Bi=1 .
• 5: Following Eq. (10), generate edited images {ˆxi }Bi=1 .
• 6: Train the condition embedding ea according to Eq. (13).
• 7: until converged

Image Editing by Diffusion Models

The typical image editing pipeline using LDMs incorporates conditioning during the denoising process, following these steps [Mokady et al., 2022]:

• (i) The input image x is encoded into its latent code z0 = E(x) using the VAE encoder E.

(ii) The clean latent code z0 is converted to its noisy counterpart zL through DDIM inversion (Eq. (6)), a deterministic mapping process.

where αt is the time dependent scale.

(iii) Classifier-free guidance (Eq. (2)) is applied during the denoising process using the condition ea:

where t = L, L-1, . . . , 1 and ˜ ϵθ(zt, ea) is calculated according to Eq. (2). We initialize ˆ zL = zL. (iv) The edited image is obtained by decoding the denoised latent code D(ˆ z0).

For training, the edited image D(ˆ z0) is used as G(x, ea) in Eq. (3) to update the embedding. However, steps (ii) and (iii) require multiple iterations, which impose significant time and memory overhead for the training phase. To improve training efficiency, we introduce the following approximations:

1⃝ For step (ii), we just add random noise to it directly via diffusion forward process by Eq. (8).

where L is a hyperparameter in this paper with 0 < L ≤ T and A is the operator that transfers z0 to its noisy version.

2⃝ For steps (iii) and (iv) to obtain a clean edited image for calculating the classification loss (Eq. (3)), we propose computing the loss on the decoded image from ˆ zL rather than ˆ z0. To prevent the classifier from handling noisy data p(a | D(ˆ zL)), we make the following approximation:

Eyebrows

No Wavy-Hair Wavy-Hair

Input

Figure 4: CASO Interpolation Results. Our method allow users to implement fine-grained editing and bidirectional editing by simply changing the classifier free guidance scale.

where Eˆ z0∼p(ˆ z0|ˆ zL) [ˆ z0] = ˆ ϵθ(ˆ zL,ea) √αL .

Now the decoded image ˆx in Eq. (3) can be:

Notably, Eq. (10) still allows us to incorporate the conditional guidance ea.

Remark 1. The approximation error of Eq. (9) can be quantified with the Jensen gap [Chung et al., 2022]. Its upper bound is:

where d is environmental dimension (dimension of z), σ2 is the ambient noise when the classifier predicts the label. And when σ → 0, d √2πσ2 e-1/2σ2 → 0. ∥∇zF(D(z))∥ := maxz′ ∥∇z′F(D(z′ ))∥ and Q := ∫ ∥z0 - ˆ z0∥p(z0|zL) dz0. ∥∇zF(D(z))∥ is generally finite [Chung et al., 2022] and when we choose the value of L carefully, the upper bound on J can be sufficiently small. When L is too large, ∥z0 - ˆ z0∥ in Q becomes excessively large, reducing classifier accuracy. Conversely, when L is too small, zL is close to clean data, making it difficult to edit the image at this time.

Reconstruction loss To preserve the attribute excluding details, we add a latent reconstruction loss:

where ℓr is the reconstruction loss. The final training objective is as follows:

We show the impact of this part of the loss function on the results in Appendix. Our training algorithm is Algorithm.1.

4 Experiment

To demonstrate the advancement and generalization of our method, we perform the following experiments.

4.1 Experimental Setup

We use Stable Diffusion-v1.5 1 for all experiments following [Dalva and Yanardag, 2024]. The datasets used include: FFHQ [Karras et al., 2019], AFHQ [Choi et al., 2020], CelebAHQ [Karras, 2017] and Stanford Cars datasets [Krause et al., 2013]. For more details, please refer to Appendix.

Timesteps are also critical for editing quality and disentanglement [Wu et al., 2023; Hertz et al., 2022b]. During training process, L is set to 0.3T for human face and 0.4T for others. During editing, for subtle features like eyebrows, we start to apply our direction from t ∈ [0.1T, 0.3T], while for some coarse-grained changes like species, editing at earlier timesteps is required (t ∈ [0.8T, 0.9T ]). The results we show in the main text are all done with timesteps T = 50. Our classifiers are all obtained by simply fine-tuning a VGG16 model [Simonyan and Zisserman, 2014]. With a well-trained classifier, only 100-200 images are needed to train the embedding.

We compare our method with: DiffAE [Preechakul et al., 2022], Cycle-Diff [Wu and De la Torre, 2023], SEGA [Brack et al., 2023], Comp-Diff [Liu et al., 2022], NoiseCLR [Dalva and Yanardag, 2024] and T-LOCO 2 [Chen et al., 2024].

4.2 Qualitative and Quantitative Results

Our qualitative results are shown in Fig. 2. Although our approach is trained on real world images, it exhibits remarkable generalization on images of other styles, demonstrates that our method generalizes beyond previous works. Notably, our edits are not confined to the data seen by the classifier. For example, if we use a well-trained dog and cat classifier learn semantic embedding for "cat", it enables edits not only for cat and dog, but also all of others. The results demonstrate that our embeddings can capture precise semantics for the target attribute.

We compare the generalization performance with other methods for attribute "Mustache". This is shown in Fig. 3. Regardless of the type of image, our model consistently achieves superior performance.

1 https://huggingface.co/stable-diffusion-v1-5/stable-diffusionv1-5

2 LOCO is purely unsupervised editing method. We found it is hard to obtain specified attribute editing directions for comparison. So we adopt its text extension version (T-LOCO) for the experiments.

Figure 5: CASOMuti-attribute Edit. In complex and challenging scenarios, our method can still achieve perfect editing.

Table 1 and Table 2 shows quantitative results. We calculate the LPIPS metric (lower is better) [Zhang et al., 2018] for editing different attributes on human face data and species conversion between animal data 3 . Our method achieve lower LPIPS than the other methods, indicating greater coherence while performing the edits.

Interpolating Edit Our approach can achieve fine-grained attribute editing and reverse editing by simply changing the classifier free guidance scale λ. For instance, using the "Bushy Eyebrows" embedding, users can create the effect of "Sparse Eyebrows" by setting the λ to a negative value.

3 For pecies conversion between animal data, we do not compare our method with NoiseCLR, DiffAE and T-LOCO. This is because NoiseCLR and DiffAE do not support such editing. And T-LOCO performs poorly on this structure changes greatly task and the statistical results are not meaningful. We show examples in the Appendix.

Figure 6: Generation Result. We use different direction: learned from NoiseCLR, learned from ours, and two kinds of text prompt (see Appendix) to guide the generation of images from Gaussian noise with corresponding attributes and calculate their IS scores.

Some examples are shown in Fig. 4. A similar effect can be achieved by changing the time at which the condition embedding is added, but this does not allow for reverse editing. We show it in the Appendix.

Multi-attribute Edit Our method allows for multi-attribute editing by simply concatenating multiple embeddings. Embedding available on different datasets can be combined (e.g. human face and animal). Experiments in Fig. 5a show that the position order of embedding stitching does not affect the result of image editing, which also proves the semantic decoupling of the learned orientation. We show editing results in difficult scenarios: In Fig. 5a, the target region to be edited has occlusion. In Fig. 5b middle, an artwork image, the face region to be edited is only a small part of the whole image. These are complex scenarios where our model still performs the editing task perfectly.

Single-Step Edit To further verify the effectiveness and consistency of our method, we conduct a single-step editing experiment, as shown in Fig. 5a. In this experiment, guidance information is applied in only one step of the denoising process, while all other steps before and after rely on unconditional embeddings. In contrast, most existing methods require conditional guidance across multiple denoising steps. Experimental results in Fig. 5a show that our outcomes are highly consistent, regardless of whether single-step editing is used. This demonstrates that our editing method is powerful enough that images generated by a single-step editing already exhibit the target properties.

4.3 Accurate Semantics

(a) Ours

(b) SEGA

Figure 7: T-SNE of cat→dog image edit. Red dot: real cat images; blue dots: real dog images; black dots: real cat images randomly sampled for editing; green dots: the generated images (to be edited as dog from the sampled cat images).

To demonstrate that our model can learn the most precise semantic direction, we directly sample from Gaussian noise with different guidance (ours, NoiseCLR and text). The results are shown in Fig. 6a. The generated images demonstrate that our learned embeddings more accurately capture semantic features. For instance, comparing the generated images guided by "Bushy Eyebrows" learned using our method and NoiseCLR separately, the faces generated by our method consistently exhibit the "Bushy Eyebrows" property, while the images generated by NoiseCLR do not have the desired semantics. We attribute the ability of NoiseCLR to perform attribute editing despite learning incorrect semantics to the coincidence of feature subspaces in the data space, which is also mentioned in SEGA [Brack et al., 2023]. Additionally, we evaluate the IS (Inception Score) [Salimans et al., 2016] of different guided sampling images (higher is better), which is shown in Fig. 6b. Some IS scores are not reported, because it is not an ideal edit direction (Fig. 6a) or the authors do not open source its weight 4 . Experiments show that our guidance can effectively improve the quality of image generation.

We conduct another experiment in which we randomly sample from the cat dataset, edit the images into dogs using different methods, including our approach and SEGA [Brack et al., 2023] (the leading text-based image editing work at present). We visualize the features obtained by the classifier (i.e., F-1(x)) using T-SNE, as shown in Fig. 7. Notably, the features of images edited by our CASO locate at the dog class mean, consistent with our theoretical result. This demonstrate that our embedding indeed captures accurate attribute semantics at the dataset level, enabling superior editing performance.

4 We tried to reproduce it. But because the direction of convergence obtained by such unsupervised methods is not controllable, we did not get the direction of interest based on their method.

4.4 Not Only Edit

Many editing works based on pre-trained diffusion models struggle to achieve satisfactory results due to the model's limited ability to accurately reconstruct certain images. In our experiments, we demonstrate that our method not only enables effective image editing but also enhances reconstruction quality, as shown in Fig. 8. We argue that the embeddings learned by our approach effectively capture both the overall structural information of the edited object and its target attributes.

When editing an image using attributes identical to those in the original (e.g., guiding the denoising of a cat image with the embedding of "cat"), the task effectively becomes reconstruction. This leads to higher-quality results and faster convergence. As shown in Fig. 8, our method enhances image quality in tasks such as car color editing and cat image reconstruction. This also explains the high quality of the generated images in 4.3. However, we acknowledge that an excessively large λ may introduce distortions in the reconstructed image. More examples can be found in the Appendix.

We calculated the FID metrics [Heusel et al., 2017] for AFHQ Cat and Dog datasets under unconditional reconstruction and guided reconstruction with "cat" and "dog" embeddings, respectively. The results are presented in Table 3 and Fig. 8. Experiments show that our embeddings significantly enhance the quality of reconstruction for images within the same class.

5 Conclusion

We propose an image editing framework that optimizes semantic embeddings with a few image inputs, guided by a classifier, to deliver high-quality, disentangled edits and superior generalization compared to existing methods. However, since it builds upon the pre-trained Stable Diffusion model, its effectiveness is limited by Stable Diffusion. Additionally, the potential for malicious misuse, such as deepfake creation, is a concern [Korshunov and Marcel, 2018]. We suggest implementing strict access controls, maintaining usage logs, and promoting responsible use to ensure ethical deployment.

Acknowledgments

The work was supported by the National Key R&D Program of China (2023YFC3306100), National Natural Science Foundation of China (62172324, 62272379, 62403429), National Research Foundation, Key R&D in Shaanxi Province (2023-YBGY-269, 2022-QCY-LL33HZ), Xixian New Area Science and Technology Plan Project (RGZN-2023-002, 2022 ZDJS-001), Zhejiang Provincial Natural Science Foundation (LQN25F030008). This paper was also supported by the National Research Foundation, Singapore under its National Large Language Models Funding Initiative (AISG Award No: AISG-NMLP-2024-004). Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore.