Efficient Personalized Text-to-image Generation by Leveraging Textual Subspace

The goal of deep generative models, such as Flow-based Models (Dinh et al., 2016; Du et al., 2022), VAE (Burgess et al., 2018), GAN (Goodfellow et al., 2020) and Diffusion Models (Dhariwal & Nichol, 2021), is to approximate an unknown data distribution by explicitly or implicitly parameterizing a model distribution using the training data.

As a class of deep generative models which has been shown to produce high-quality images, Diffusion Models (Dhariwal & Nichol, 2021; Ho et al., 2020; Nichol & Dhariwal, 2021; Song et al., 2020) synthesize data via an iterative denoising process. As suggested in (Ho et al., 2020), a reweighted variational bound can be utilized as a simplified training objective and the sampling process aims to predict the corresponding noise added in forward process. The denoising objective is finally realized as a mean squared error:

where $x_{0}$ are training data with conditions $c$, $t\sim\mathbb{U}(0,1)$, $\epsilon\sim\mathbb{N}(0,\textbf{I})$ is the gaussian noise sampled in the forward process, $\alpha_{t},\sigma_{t}$ are pre-defined scalar functions of time step $t$, $\epsilon_{\theta}$ is the parameterized reverse process with trainable parameter $\theta$.

Recently, an introduced class of Diffusion Models named Latent Diffusion Models (Rombach et al., 2022a) raises the interest of the community, which leverages a pre-trained autoencoder to map the images from pixel space to a more efficient latent space that significantly accelerates training and reduces the memory. Latent Diffusion Models consist of two core models. Firstly, a pre-trained autoencoder is extracted, which consists of an encoder $\mathcal{E}$ and a decoder $\mathcal{D}$. The encoder $\mathcal{E}$ maps the images $x_{0}\sim p(x)$ into a latent code $z_{0}$ in a low-dimensional latent space. The decoder $\mathcal{D}$ learns to map it back to the pixel space, such that $\mathcal{D}(\mathcal{E}(x))\approx x$. Secondly, a diffusion model is trained on this latent space. The denoising objective now becomes

where $z_{0}=\mathcal{E}(x_{0})$ is the latent code encoded by the pre-trained encoder $\mathcal{E}$.

For conditional synthesis, in order to improve sample quality while reducing diversity, classifier guidance (Dhariwal & Nichol, 2021) is proposed to use gradients from a pre-trained model $p(c|z_{t})$, where $z_{t}:=\alpha_{t}z_{0}+\sigma_{t}\epsilon$. Classifier-free guidance (Ho & Salimans, 2022) is an alternative approach that avoids this pre-trained model by instead jointly training a single diffusion model on conditional and unconditional objectives via randomly dropping $c$ during training with probability $1-w$, where $w$ is the guidance scale offering a tradeoff between sample quality and diversity. The modified predictive model is shown as follows: $\hat{\epsilon}_{\theta}(z_{t},c)=w\epsilon_{\theta}(z_{t},c)+(1-w)\epsilon_{\theta}(z_{t},\phi)$, where $\phi=c_{\theta}($“ ”$)$ is the embedding of a null text and $c_{\theta}$ is the pre-trained text encoder such as BERT (Devlin et al., 2018) and CLIP (Radford et al., 2021).

Recent large-scale text-to-image models such as Stable-Diffusion (Rombach et al., 2022a), GLIDE (Nichol et al., 2021) and Imagen (Saharia et al., 2022) have demonstrated unprecedented semantic generation. We implement our method based on Stable-Diffusion, which is a publicly available 1.4 billion parameters text-to-image diffusion model pre-trained on the LAION-400M dataset (Schuhmann et al., 2021). Here $c$ is the processed text condition.

Typical text encoder models include three steps to process the input text. Firstly, a textual prompt is input by the users and split by a tokenizer to transform each word or sub-word into a token, which is an index in a pre-defined language vocabulary. Secondly, each token is mapped to a unique text embedding vector, which can be retrieved through an index-based lookup (Gal et al., 2022). The embedding vector is then concatenated and transformed by the CLIP text encoder to obtain the text condition $c$.

As the demand for personalized generation continues to grow, personalized generation has become a prominent factor in the field of machine learning, such as recommendation systems (Amat et al., 2018) and language models (Cattiau, 2022). Within the vision community, adapting models to a specific object or style is gradually becoming a target of interest. Users often wish to input personalized real images to parameterize a “concept” from them and combine it with large amounts of textual prompt to create new combination of images.

A recent proposed method Textual Inversion (Gal et al., 2022) choose the text embedding space as the location of the “concept”. It intervenes in the embedding process and uses a learned embedding $\mathbf{v}$ to represent the concept, in essence “injecting” the concept into the language vocabulary. Specifically, it defines a placeholder string $S$ (such as “A photo of $*$”) as the textual prompt, where “$*$” is the pseudo-word corresponding to the target embedding $\mathbf{v}$ it wishes to learn. The embedding matrix $y\in\mathbb{R}^{N\times d}$ can be obtained by concatenating $\mathbf{v}$ with other frozen embeddings (such as the corresponding embeddings of “a”, “photo” and “of” in the example), where $N$ is the number of words in the placeholder string⁵⁵5In practice, the embedding matrix $y\in\mathbb{R}^{N_{max}\times d}$ where $N_{max}$ is the pre-defined maximum number of words in a sentence and other $N_{max}-N$ words are filled with terminator. For clarity of expression, we only consider the words with practical meaning in the main text. and $d$ is the dimension of the embedding space. The above process is defined as (Gal et al., 2022): $y\leftarrow Combine(S,\mathbf{v})$.

The optimization goal is defined as: $\mathop{\arg\min}\limits_{\mathbf{v}}\mathbb{E}_{z_{0},\epsilon,t}[\|\epsilon-\epsilon_{\theta}(\alpha_{t}z_{0}+\sigma_{t}\epsilon,c_{\theta}(y))\|_{2}^{2}]$, where $z_{0}$, $\epsilon$ and $t$ are defined in Equation (1) and (2). Please note that $y$ is a function of $\mathbf{v}$. Although Textual Inversion can extract a single concept formed by $3\text{-}5$ images and reconstruct it faithfully, it can not be combined with textual prompt flexibly since it solely considers the performance of the reconstruction task during optimization. Also, it searches the target embedding in the high-dimensional embedding space, which is time-consuming and difficult to converge.

To address the issues above, we observe that the pre-trained embeddings in the vocabulary is expressive enough to represent any introduced concept. Therefore, we explicitly define a projection matrix to efficiently optimize the target embedding in a low-dimensional tetxual subspace, which speeds up convergence and better preserves the text similarity of the learned embedding.

We first state that any vector in the embedding space $\mathbb{E}$ can be represented by the embeddings in the vocabulary $V$, as shown in the following theorem whose proof can be found in Appendix B.

As stated in Theorem 1, any vector in $\mathbb{E}$ can be represented by a linear combination of the embeddings in the vocabulary $V$. Now we define the weight vector $\mathbf{w}=(w_{1},\dots,w_{|V|})^{T}$ with each component corresponding to a embedding in $V$, and the embedding $\mathbf{v}=\sum_{i=1}^{|V|}w_{i}\mathbf{v}_{i}$. Since the users input an initial embedding $\mathbf{u}\in V$, we wish the start point in $\mathbb{E}$ to be the same as $\mathbf{u}$ in our algorithm. Thus, we initialize the weights as:

where $i_{\mathbf{u}}$ denotes the index corresponding to $\mathbf{u}$. Then the embedding $\mathbf{v}$ to be learned can now be initialized as:

Subsequently, it can be combined with the placeholder string to form the embedding matrix $y$ as stated in Section 2.3. The reconstruction task can be formulated as the following optimization problem:

where $z_{0}$, $\epsilon$, $t$, $\alpha_{t}$ and $\sigma_{t}$ are detailed in Equation (1) and (2), $c_{\theta}$ is the text encoder defined in Section 2.1, $y$ is a function of $\mathbf{v}$. To solve problem (5), we iteratively update the weight vector $\mathbf{w}$ by using gradient descent with initial point $\mathbf{w}^{0}$, so that the embedding $\mathbf{v}$ is also updated.

While it is expressive enough to represent any concept in the embedding space $\mathbb{E}$, most weights and vectors are unnecessary since the rank $r(A_{V})=d$ as stated in Theorem 1, where $A_{V}=[\mathbf{v}_{1},\dots,\mathbf{v}_{|V|}]\in\mathbb{R}^{d\times|V|}$. Thus, we only need at most $d$ vectors to optimize with $\mathbf{u}$ included, which corresponds to selecting $d$ linearly-independent vectors $\mathbf{v}_{i_{1}},\mathbf{v}_{i_{2}},\dots,\mathbf{u},\dots,\mathbf{v}_{i_{d}}$ from vocabulary $V$.

As detailed in Section 3.1, any vector in the embedding space $\mathbb{E}$ can be obtained using $d$ vectors $\mathbf{v}_{i_{1}},\mathbf{v}_{i_{2}},\dots,\mathbf{u},\dots,\mathbf{v}_{i_{d}}$. However, optimizing the embedding $\mathbf{v}$ by solving problem (5) solely target the reconstruction of input images, leading to low text similarity (Gal et al., 2022). Besides, solving problem (5) requires to search in the whole high-dimensional embedding space $\mathbb{E}$, which results in time-consuming training process and slow convergence.

It is natural to construct a textual subspace with high text similarity, in which the searched embedding is able to capture the details of the input image. To this end, vectors with high semantic relevance to $\mathbf{u}$ should be included. As suggested in (Mikolov et al., 2013a; b; Le & Mikolov, 2014; Goldberg & Levy, 2014; Rong, 2014), the vector distance (denoted by $\mathscr{F}$) between the word embeddings in $V$, such as dot product, cosine similarity and $L_{2}$ norm, can be employed as the semantic similarity of the corresponding words. Now, we are ready to give the distance vector ${dist}_{V}$ to calculate the distance between $\mathbf{u}$ and any embedding in $V$, which is given by

Next, we re-order ${dist}_{V}$ and the top $M$ vectors are selected as the basis vectors, where at most $d_{1}\leq M$ vectors are linearly-independent among them and $d_{1}$ is the dimension of the textual subspace. The choice of $d_{1}$ and $\mathscr{F}$ are further discussed in Section 5 and Appendix F. The details of selection strategy are presented in Algorithm 1.

Finally, we derive that for single-step optimization scenario, the difference of the embedding update between Textual Inversion and BaTex corresponds to a matrix transformation with rank $d_{1}$, which is the dimension of the textual subspace. The formal theorem is presented as follows, and its proof is included in Appendix C.

It can be seen from Theorem 2 that BaTex actually defines a transformation from $\mathbb{R}^{d}$ to $\mathbb{R}^{d_{1}}$ using the selection strategy stated in Algorithm 1, which intuitively benefits for optimization process since $B_{V}$ is formed by the pre-trained embeddings, showing that BaTex explicitly extracts more information from the vocabulary $V$.

In this subsection, we present the experimental settings, and more details can be found in Appendix E.

We compare the proposed BaTex with Textual Inversion (TI) (Gal et al., 2022), the original method for personalized generation which lies in the category of “Embedding Optimization”. To analyze the quality of learned embeddings, we follow the most commonly used metrics in TI and measure the performance by computing the CLIP-space scores (Hessel et al., 2021).

We first show that learning in a textual subspace significantly improves the text similarity of learned embedding while retaining the ability to reconstruct the input image. The results of text-guided synthesis are shown in Figure 2. As can be seen, for complex input textual prompt with additional text conditions, our method completely captures the input concept and naturally combines it with known concepts.

Additionally, We show the effectiveness of our method by composing two concepts together and introducing additional text conditions (shown in bolded text). The results are shown in Figure 3. It can be seen that BaTex not only allows for the lossless combination of two distinct concepts, but also faithfully generates the additional text conditions.