Learning to Customize Text-to-Image Diffusion In Diverse Context

We apply our approach to various text-to-image baseline models (Gal et al., 2022; Ruiz et al., 2023; Voynov et al., 2023; Kumari et al., 2023) that are based on Stable Diffusion (SD) (Rombach et al., 2022). SD consists of a CLIP text encoder $\Gamma$ that encodes an input text t into a sequence of input token embeddings, denoted as Tokenize, $\mathbf{P}=\text{Tokenize}(\textbf{t})$, then outputs corresponding text embedding $\textbf{C}=\Gamma(\mathbf{P})$ using self-attention layers. A Variational Auto Encoder (VAE) of SD $\mathcal{E}$ encodes an image $x$ to a lower dimensional latent $\mathbf{z}=\mathcal{E}(\mathbf{x})$.

During training, given a timestep $t\sim\text{Uniform}[0,\text{T}-1]$, a random noise map $\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ is added to the latent map to get a noised latent map $\mathbf{z}_{t}=\alpha_{t}\mathbf{z}+\sigma_{t}\epsilon$, where $\alpha_{t}$ and $\sigma_{t}$ denote the noise scheduling coefficients. Then, the diffusion U-Net $\epsilon_{\theta}$ is trained to minimize the following objective for denoising,

Utilizing a text prompt $\widetilde{\textbf{t}}$ that incorporates a concept token [v] (e.g., “a picture of a [v] dog”), the tokenized input embedding is encoded with CLIP text encoder $\widetilde{\textbf{C}}=\Gamma(\widetilde{\textbf{P}})$. Following the text encoding, the denoising objective is computed as below,

where $\mathcal{L}_{\text{Custom}}$ denotes the denoising loss utilized for model customization, with $\textbf{z}=\mathcal{E}(\textbf{x})$ encoding an image x sampled from a small set of personal concept images. Depending on the baseline method, a different set of parameters are optimized. For Textual Inversion (TI) (Gal et al., 2022), only the concept token embedding is optimized with respect to Eqn. 2. We propose a method that does not require any architectural modification, hence it is highly compatible with baselines. We demonstrate this applicability by applying our method to different baselines. Unless otherwise specified, we illustrate our method using TI.

In order to enhance the concept token embedding with context-rich text prompts, we adopt Masked Language Modeling (MLM) during the model customization. The overall process is illustrated in Figure 3. We elaborate on the corresponding details in the following steps,

(i) A text prompt t, drawn from a contextually diverse prompt set that includes the concept token, e.g., “a picture of [v] dog at Eiffel Tower”, is tokenized and mapped to a prompt embedding,

where $\textbf{P}\in\mathbb{R}^{L\times d}$, $L$ is the number of tokens, $d$ is the feature dimension of prompt embeddings.

(ii) A subset of prompt embedding P is randomly selected and those selected tokens are replaced by a mask token embedding $p_{\text{mask}}$ with the probability $\rho_{\text{mask}}$, yielding $\textbf{P}_{\text{masked}}=\text{RandomMask}(\textbf{P},\rho_{\text{mask}})$. Then, text embedding is obtained from CLIP text encoder $\Gamma$,

where $\textbf{C}_{\text{masked}}=\{c_{i}\}$, with $c_{i}\in\mathbb{R}^{d}$ denoting one element in token embedding.

(iii) Finally, we predict the label of the masked token $\hat{\mathbf{y}}=\psi(\textbf{C}_{\text{masked}})$ and calculate the MLM loss,

where $\psi$ denotes the classification network with self-attention layers.

Notably, the attention layer computes the output token embedding O as the linear combination of input embeddings, with the weights determined by the self-attention map $\textbf{A}^{\text{self}}\in\mathbb{R}^{L\times L}$, implying that the output is the contextualization of the input. For the $i_{m}$-th output token, i.e., the masked token, it can be formulated as shown below,

where $j_{*}$ is the index of the concept token, and V is the value matrix.

By optimizing the concept token embedding to minimize $\mathcal{L}_{\text{MLM}}$, the concept token is guided to learn the diverse context, as the MLM encourages the utilization of surrounding contexts for mask prediction. The overall process is illustrated in Figure 3. We later show that customization with this additional objective results in regularizing the text embeddings from overfitting to the concept token, which eventually leads to semantically enhanced image generation (Section 4.4).

Prompt Set Construction. To achieve customization with diverse contexts, we construct a set of context-rich prompts that incorporate the special concept token. Inspired by recent work (Brooks et al., 2023), we leverage a pretrained large language model (OpenAI, 2023; Brown et al., 2020) to minimize the effort in manually crafting a large set of prompts. For this, we query the LLM to generate a list of contexts of different types, e.g., background or subjection variation. For detailed descriptions of the prompt set construction process, refer to the Appendix Section A.2.

Pretraining. Although the CLIP text encoder of SD has the linguistic capability of comprehending text prompts, it is solely trained with contrastive learning objectives (Radford et al., 2021), and does not support MLM. Therefore, before proceeding with fine-tuning for customization, we first pretrain a network, namely, a contextualizer $\psi$ to incorporate the MLM capability. We provide the pretraining procedure of the contextualizer $\psi$ in Algorithm 1. During the pretraining of $\psi$, the concept token is not involved. We only train the mask embedding and the layers of the contextualizer. The CLIP text encoder and diffusion U-Net remain fixed.

Finetuning. Utilizing the contextually diverse prompt set, we proceed with the model customization (Figure 3). The model is optimized by minimizing the denoising objective $\mathcal{L}_{\text{Diff}}$ (Eqn. 2) and the MLM loss $\mathcal{L}_{\text{MLM}}$ (Eqn. 5). Two different types of prompts are utilized for each objective. Text embeddings $\widetilde{\textbf{C}}$ encoded from a context-simple text prompt $\widetilde{\textbf{t}}$ (e.g., “a picture of [v] dog”) for $\mathcal{L}_{\text{Diff}}$, and text embedings $\textbf{C}_{\text{masked}}$ encoded from a text context-rich text prompt t (e.g., “a [v] dog at Eiffel Tower”) for $\mathcal{L}_{\text{MLM}}$. The overall learning objective can be formulated as follows,

where $\mathbf{z}$ denotes the noised latent of personal concept image, $t$ denotes the timestep, and $\lambda$ denotes the weight for the MLM loss. Note that, as the MLM objective does not require any images, the concept token embedding learns contextual semantics without constructing corresponding images. Additionally, our approach does not require any architectural modification of SD, it is highly compatible with existing text-to-image approaches. Hence, we combine our approach with different baseline methods and demonstrate its generalizability. The overall training procedure is described in Algorithm 2. We refer to Appendix Section A.1 for additional details of training.

In this section, we illustrate how adopting the MLM with diverse contexts during the model customization leads to semantically enhanced textual representation, which ultimately translates to improved image generation. We first validate the following to explain how our method leads to semantically enhanced textual representation,

• •

  The model overfits to the personal concept, when the semantics of the context tokens (i.e., non-personal) become similar to the concept token.

• •

  The semantics of the context tokens get distinct from the concept token as the diverse contextual semantics are learned with MLM.

Most customization methods that fine-tune the model with limited context (Gal et al., 2022; Ruiz et al., 2023; Voynov et al., 2023; Kumari et al., 2023) often suffer from concept overfitting (Zeng et al., 2024), resulting in generated images that primarily contain the personal concept without adhering to the prompt. We next analyze that concept overfitting leads to a high similarity between the text embeddings of context tokens and the concept token, ultimately causing a loss of contextual semantics.

In contrast, the MLM regularizes the loss of contextual information and guides the learning of diverse contexts (Figure 4, right). Specifically, we focus on two types of text embeddings for the concept token: (i) the context token embedding derived from the prompt with the concept token (e.g., “a picture of [v] dog at Eiffel Tower”), referred to as $c_{b}$; and (ii) the context embedding from the prompt without the concept token (e.g., “a picture of Eiffel Tower”), denoted as $\hat{c}_{b}$.

We provide the proof of Proposition 1 and 2 in Appendix A.4.

To empirically validate this, we provide a cosine similarity analysis using a set of 200 text prompts for 7 different concepts in varying contexts (Table 1). We analyze the cosine similarity, $sim_{1}=cos(c_{i_{*}},c_{b})$, and $sim_{2}=cos(c_{b},\hat{c}_{b})$, respectively. The result shows that the baseline approach leads to high similarity between the concept-context tokens ($sim_{1}$) and low similarity in the context token from the two prompts ($sim_{2}$). The baseline result implies that the contextual semantics are not only getting similar to the concept (high $sim_{1}$) but also implies that the context token loses its semantics as it becomes dissimilar to the semantics that is preserved in the context token (low $sim_{2}$). In contrast, we observe the opposite trend from our results, which implies that the MLM encourages the contextual semantics to be distinct (low $sim_{1}$) and preserves the semantics (high $sim_{2}$).

The cross-attention map plays a key role in controlling the overall image generation (Hertz et al., 2022; Chefer et al., 2023), where the attended region corresponds to the area most influenced by the token. Next, we provide an analysis of the cross-attention map to illustrate how the aforementioned semantic enhancement in textual space can be transferred to image space, thereby improving the prompt fidelity of image generation.

Let $\textbf{Q}_{\mathcal{I}}$, $\textbf{K}_{\mathcal{T}}$ and $\textbf{V}_{\mathcal{T}}$ denote the Query, Key and Value of the cross-attention layer, projected from image $\mathcal{I}$, and text $\mathcal{T}$. During the denoising process, the cross-attention map $\textbf{A}^{\text{cross}}=\text{Softmax}(\frac{\textbf{Q}_{\mathcal{I}}\textbf{K}_{\mathcal{T}}^{\top}}{\sqrt{d}})$ is computed between $\textbf{Q}_{\mathcal{I}}\in\mathbb{R}^{|queries|\times d}$ and $\textbf{K}_{\mathcal{T}}\in\mathbb{R}^{L\times d}$, where $|queries|$ denotes the number of image tokens, $L$ denotes the number of text tokens, and $d$ denotes the dimension of each image/text token embedding.

We provide the proof of Proposition 4 in Appendix A.5.

We visualize the cross-attention maps of the tokens and compare the results of the baseline (Gal et al., 2022) and ours to further validate our claims (Figure 5). The result indicates that concept overfitting of the text embedding in the baseline approach leads to cross-attention maps of the concept token and the context token being more closely distributed, leading to semantically degraded image generation. Using the same prompt set used for cosine similarity analysis that contains both context token and concept token, we measure the symmetric KL divergence between the cross-attention maps of the concept tokens $c_{i_{*}}$ and the context tokens $c_{b}$ within the same prompt: $SKL=\frac{1}{2}D_{\text{KL}}(\textbf{A}^{\text{cross}}[:,i_{*}]||\textbf{A}^{\text{cross}}[:,b])+\frac{1}{2}D_{\text{KL}}(\textbf{A}^{\text{cross}}[:,b]||\textbf{A}^{\text{cross}}[:,i_{*}])$, where $D_{\text{KL}}$ is the Kullback-Leibler (KL) divergence, $\textbf{A}^{\text{cross}}[:,k]\in\mathbb{R}^{|queries|}$ denotes the cross-attention map of the $k$-th token. The results show that the baseline approach produces a more similar distribution of the cross-attention maps between the concept and context tokens (Table 1, $SKL$). This finding, along with our visual evidence (Figure 4), indicates that semantic enhancement in textual space leads to enhancement in image space, resulting in improved prompt fidelity of the image generation.

In this section, we conduct ablation studies to provide deeper insights into our method and validate its effectiveness. To solely compare the effect of adjusting the textual space, we have UNet and CLIP fixed and only update the concept token embedding.

Impact of $\lambda$. We first study the influence of MLM loss weight (Table 4). We adopt TI as our baseline and compare the prompt and subject fidelity with the model trained with different $\lambda$. We note the general trend of improved CLIP-T score as the model is trained with increased $\lambda$. Additionally, the result indicates as the $\lambda$ increases, subject fidelity can be slightly decreased. We hypothesize that this trade-off arises due to the model prioritizing textual context alignment over subject preservation at higher $\lambda$ values, as the MLM objective encourages the model to focus more on capturing the semantic relationships of the contexts.

We present the additional qualitative comparison results of the proposed method in (Figure 7). In general, the baseline method that integrates our approach leads to improvement in prompt fidelity. In this visualization results, the baseline approach shows a higher tendency to neglect the context semantics. We analyze that the baseline approach loses the semantics of the context in textual space, which leads to the loss of semantics of the generated images. In contrast, the adoption of MLM leads to the preservation of the contextual semantics in text embedding, resulting in images with enhanced semantics with higher prompt fidelity. We provide additional qualitative results of the proposed method in Appendix Section A.6.