SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zihao Wang $^{1}$ , Lin Gui $^{1}$ , Jeffrey Negrea $^{2}$ , and Victor Veitch $^{1,2,3}$
$^{1}$ Department of Statistics, University of Chicago
$^{2}$ Data Science Institute, University of Chicago
3Google Research
This paper concerns the structure of learned representations in text-guided generative models, focusing on score-based models. A key property of such models is that they can compose disparate concepts in a 'disentangled' manner. This suggests these models have internal representations that encode concepts in a 'disentangled' manner. Here, we focus on the idea that concepts are encoded as subspaces of some representation space. We formalize what this means, show there's a natural choice for the representation, and develop a simple method for identifying the part of the representation corresponding to a given concept. In particular, this allows us to manipulate the concepts expressed by the model through algebraic manipulation of the representation. We demonstrate the idea with examples using Stable Diffusion.
Large-scale text-controlled generative models are now dominant in many parts of modern machine learning and artificial intelligence [e.g., Bro+20; Rad+21; Bom+21; Koj+22]. In these models, the user provides a prompt in natural language and the model generates samples based on this prompt—e.g., in large language models the sample is a natural language response, and in text-to-image models the sample is an image. These models have a remarkable ability to compose disparate concepts to generate coherent samples that were not seen during training. This suggests that these models have some internal representation of high-level concepts that can be manipulated in a 'disentangled' manner. Broadly, the goal of this paper is to shed light on how this concept representation works, and how it can be manipulated. We focus on text-to-image diffusion models, though many of the ideas are generally applicable.
Our starting point is the following commonly observed structure of representations:
1. Each data point  $x$  is mapped to some representation vector  $\mathrm{Rep}(x) \in \mathbb{R}^p$ .
2. High-level concepts correspond to subspaces (directions) of the representation space.
Perhaps the best known example of this structure is in word embeddings, where semantic relationships such as Rep("king")—Rep("queen")  $\approx$  Rep("man")—Rep("woman") suggest that high-level concepts (here, sex) are encoded as directions in the representation space [Mik+13]. This kind of encoding of concepts has been argued to occur in many contexts, including in the latent space of variational autoencoders [ZW20; Khe+20; MFM21] and in the latent space of language models [Bol+16; GG19; Rad+21; Elh+22]. We'll call representations of this kind arithmetically composable, because composition corresponds to arithmetic operations on the representation vectors. The goal of this paper is to develop arithmetically composable representations of text for score-based text to image models.
There are two main motivations. First, understanding the structure of the representation space is important for foundational progress on understanding the emergent behavior of text-controlled generative models. It is particularly interesting to study this question in the text-to-image setting 37th Conference on Neural Information Processing Systems (NeurIPS 2023).
(a) Rep["a portrait of mathematician"]
(b)  $(\mathbb{I} - \mathrm{proj}_{\mathrm{sex}})\mathrm{Rep}[\text{"a portrait of a mathematician"}] + \mathrm{proj}_{\mathrm{sex}}\mathrm{Rep}[\text{"a person"}]$
(c)  $(\mathbb{I} - \mathrm{proj}_{\mathrm{style}})\mathrm{Rep}[\text{"a portrait of a mathematician"}] + \mathrm{proj}_{\mathrm{style}}\mathrm{Rep}[\text{"in Fauvism style"}]$
Figure 1: We show that high-level concepts such as sex and artistic_style are encoded as subspaces of a suitably chosen representation space. This allows us to manipulate the concepts expressed by a prompt through algebraic operations on the representation of that prompt. Namely, we edit the representation projected on to the subspace corresponding to a concept. Note images are paired by random seed.
because the multi-modality of the data makes it straightforward to distinguish concepts from inputs, and because it is not clear a priori that the models themselves build in any inductive bias towards arithmetic structure. The secondary motivation is that having such a representation would allow us to manipulate the concepts expressed by the model through linear-algebraic manipulations of representation of the input text; fig. 1 illustrates this idea.
The development of the paper is as follows:
1. We develop a mathematical formalism for describing the connection between representation structures and concepts for text-controlled generative models.
2. Using this formalism, we show that the Stein score of the text-conditional distribution is an arithmetically composable representation of the input text.
3. Then, we develop concept algebra as a method for manipulating the concepts expressed by the model through algebraic manipulation of this representation. We illustrate the approach with examples manipulating a variety of concepts.
2 A Mathematical Framework for Concepts as Subspaces
The first task is to develop a precise mathematical formalism for connecting the structure of representations and high-level concepts. This is necessary for understanding when such representations are possible, how to construct them, and when they may fail. We must make precise what is a concept, how concepts relate to inputs  $x$ , and what it means to represent concepts.
Concepts The real-world process that generated the training data has the following structure. First, images  $Y$  are generated according to some real-world, physical process. Then, some human looks at each image and writes a caption describing it. Inverting this process, each text  $x$  induces some probability density  $p(y \mid x)$  over images  $Y$  based on how compatible they are with  $x$  as a caption. The (implicit) goal of the generative model is to learn this distribution.
To write the caption, the human first maps the image to a set of high-level variables summarizing the image's content, then uses these latent variables to generate the text  $X$ . Let  $C$  be the latent variable that captures all the information about the image that is relevant for a human writing a caption. So,
$$
p (y \mid X = x) = \int p (y \mid C = c) p (C = c \mid X = x) d c. \tag {1}
$$
The random variable  $C$  captures the information that is jointly relevant for both the image and caption. Variables in  $C$  include attributes such as has(mathematician or is_man, but not pixel_14_is_red. We define concepts in terms of the latent  $C$ .
Definition 2.1. A concept variable  $Z$  is a  $C$ -measurable random variable. The concept  $\mathcal{X}$  associated to  $Z$  is the sample space of  $Z$ .
Now, the full set of all possible concepts is unwieldy. Generally, we are concerned only with the concepts elicited by a particular prompt  $x$ .
Definition 2.2. A set of concepts  $\mathcal{X}_1, \ldots, \mathcal{X}_k$  is sufficient for  $x$  if  $p(c \mid X = x, Z_{1:k} = z_{1:k}) = p(c \mid Z_{1:k} = z_{1:k})$  for all  $z_1, \ldots, z_k \in \mathcal{X}_1 \times \dots \times \mathcal{X}_k$ .
For example, the concept profession would be sufficient for the prompt "A nurse". This prompt induces a distribution on many concepts (e.g., background is likely to be a hospital) but these other concepts are independent of the caption given profession = nurse. Then,
$$
p (y \mid x) = \sum_ {z _ {1: k}} p (y \mid z _ {1: k}) p \left(z _ {1: k} \mid x\right). \tag {2}
$$
Concept Distributions Following eq. (2), we can view each text  $x$  as specifying a distribution  $p(z_{1:k} \mid x)$  over latent concepts  $\mathcal{Z}_1, \ldots, \mathcal{Z}_k$ . This observation lets us make the relationship between text and concepts precise.
Definition 2.3. A concept distribution  $Q$  is a distribution over concepts. Each text  $x$  specifies a concept distribution as  $Q_{x} = p(z_{1:k} \mid x)$ .
That is, we move from viewing text as expressing specific concept values (is(mathematician = 1) to expressing probability distributions over concepts ( $Q_{x}$  (is mathematician = 1) = 0.99). The probabilistic view is more general—deterministically expressed concepts can be represented as degenerate distributions. This extra generality is necessary: for example, the prompt "a person" induces a non-degenerate distribution over the sex concept.
Concept Representations A text-controlled generative model takes in prompt text  $x$  and produces a random output  $Y$ . Implicitly, such models are maps from text strings  $x$  to the space of probability densities over  $Y$ . We'll define a representation  $\operatorname{Rep}(x) \in \mathcal{R}$  of  $x$  as any function of  $x$  that suffices to specify the output distribution. We define  $f_r(\cdot)$  as the density defined by  $r \in \mathcal{R}$ , and assume that the model learn's the true data distribution of  $Y|X = x$ :
$$
f _ {\operatorname {R e p} (x)} (y) = p (y \mid x). \tag {3}
$$
The key idea for connecting representations and concepts is to move from considering representations of prompts to representations of concept distributions.
Definition 2.4. A concept representation  $\operatorname{Rep}$  is a function that maps a concept distribution  $Q$  to a representation  $\operatorname{Rep}(Q) \in \mathcal{R}$ , where  $\mathcal{R}$  is a vector space. The representation of a prompt  $x$  is the representation of the associated concept distribution,  $\operatorname{Rep}(x) \coloneqq \operatorname{Rep}(Q_x)$ .
There are two reasons why this view is desirable. First, defining the representation in terms of the concept distribution makes the role of concepts explicit—this will allow us to explain how representation structure relates to concept structure.
The second reason is that it allows us to reason about representations that do not correspond to any prompt. Every prompt defines a concept distribution, but not every concept distribution can be defined by a prompt. This matters because we ultimately want to reason about the conceptual meaning of representation vectors created by algebraic operations on representations of prompts. Such vectors need not correspond to any prompt.
Arithmetic Compositionality We now have the tools to define what it means for a representation to be arithmetically composable. We define composability for a pair of concepts  $\mathcal{Z}$  and  $\mathcal{W}$ . In the subsequent development, our aim will be to manipulate  $\mathcal{Z}$  while leaving  $\mathcal{W}$  fixed.
Definition 2.5. A representation Rep is arithmetically composable with respect to concepts  $\mathcal{Z},\mathcal{W}$  if there are vector spaces  $\mathcal{R}_Z$  and  $\mathcal{R}_W$  such that for all concept distributions of the form  $Q(z,w) = Q_Z(z)Q_W(w)$ ,
$$
\operatorname {R e p} \left(Q _ {Z} Q _ {W}\right) = \operatorname {R e p} _ {Z} \left(Q _ {Z}\right) + \operatorname {R e p} _ {W} \left(Q _ {W}\right), \tag {4}
$$
where  $\mathrm{Rep}_Z(Q_Z) \in \mathcal{R}_Z$  and  $\mathrm{Rep}_W(Q_W) \in \mathcal{R}_W$ .
In words: we restrict to product distributions to capture the requirement that the concepts  $\mathcal{Z}$  and  $\mathcal{W}$  can be manipulated freely of each other (the typical case is that one or both of  $Q_{Z}$  and  $Q_{W}$  are degenerate, putting all their mass on a single point). Then, the definition requires that there are fixed subspaces corresponding to each concept in the sense that, e.g., changing only  $Q_{Z}$  induces a change only in  $\mathcal{R}_{Z}$ .
3 The Score Representation
We now have an abstract definition of arithmetically composable representation. The next step is to find a specific representation function that satisfies the definition.
We will study the following choice.
Definition 3.1. The score representation  $s[Q]$  of a concept distribution  $Q$  is defined by:
$$
s [ Q ] (y) := \nabla_ {y} \log \int p (y \mid z, w) Q (z, w) d z d w.
$$
The centered score representation  $\bar{s}[Q]$  is defined by  $\bar{s}[Q] := s[Q] - s[Q_0]$ .
Here,  $s[Q]$  is itself a function of  $y$  and the representation space  $\mathcal{R}$  is a vector space of functions. This is a departure from the typical view of representations as elements of  $\mathbb{R}^p$ . The score representation can be thought of as a kind of non-parametric representation vector. The centered score representation just subtracts off the representation of some baseline distribution  $Q_0$ .<sup>1</sup>
The main motivation for studying the score representation is that
$$
s [ x ] (y) := s \left[ Q _ {x} \right] (y) = \nabla_ {y} \log p (y \mid x).
$$
The importance of this observation is that  $\nabla_y\log p(y\mid x)$  is learnable from data. In fact, this score function is ultimately the basis of many controlled generation models [e.g., HJA20; Ram+22; Sah+22], because it characterizes the conditional while avoiding the need to compute the normalizing constant [HD05; SE19]. Accordingly, we can readily compute the score representation of prompts in many generative models, without any extra model training.
Causal Separability The score representation does not have arithmetically composable structure with respect to every pair of concepts. The crux of the issue is that concepts are reflected in the representation based on their effect on  $Y$ . If the way they affect  $Y$  depends fundamentally on some interaction between two concepts, the representation cannot hope to disentangle them. Thus, we must rule out this case.
Definition 3.2. We say that  $Y$  is causally separable with respect to  $\mathcal{Z}, \mathcal{W}$  if there exist unique  $Y$ -measurable variables  $Y_{\mathcal{Z}}, Y_{\mathcal{W}}$ , and  $\xi$  such that
1.  $Y = g(Y_{\mathcal{X}},Y_{\mathcal{W}},\xi)$  for some invertible and differentiable function  $g$ , and
2.  $p(y_{\mathcal{X}},y_{\mathcal{W}},\xi \mid z,w) = p(\xi)p(y_{\mathcal{X}}\mid z)p(y_{\mathcal{W}}\mid w)$
Informally, the requirement is that we can separately generate  $Y_{\mathcal{Z}}$  and  $Y_{\mathcal{W}}$  as the part of the output affected by  $\mathcal{Z}$  and  $\mathcal{W}$  (and  $\xi$  as the part of the image unrelated to  $Z$  and  $W$ ), then combine these parts to form the final image. That is, generating the visual features associated to a concept  $\mathcal{W}$  can't require us to know the value of another concept  $\mathcal{Z}$ . As an example where causal separability fails, consider the concepts of species  $\mathcal{W} = \{\text{deer, human}\}$  and sex  $\mathcal{Z} = \{\text{male, female}\}$ . It seems reasonable that there is a  $Y$ -measurable  $Y_{\mathcal{W}}$  that is the species part of the image—e.g., the presence of fur vs skin, snouts vs noses, and so forth. However, there is no part of  $Y$  that corresponds to a sex concept in a manner that's free of species. The reason is that the visual characteristics of sex are fundamentally different across species—e.g., male deer have antlers, but humans usually do not. In fig. 8 we test this example, finding that concept algebra fails in the absence of causal separability.
It turns out it suffices to rule out this case (all proofs in appendix):
<sup>1</sup>The representation space  $\mathcal{R}$  is the same for all  $Q_0$ ; the choice is arbitrary. We define  $\bar{s}$  to ensure 0 is an element of  $\mathcal{R}$ . This is for theoretical convenience; we will see that only  $s$  is required in practice.
Proposition 3.1. If  $Y$  is causally separable with respect to  $\mathcal{W}$  and  $\mathcal{Z}$ , then the centered score representation is arithmetically composable with respect to  $\mathcal{W}$  and  $\mathcal{Z}$ .
That is: the (centered) score representation is structured such that concepts correspond to subspaces of the representation space.
4 Concept Algebra
We have established that concepts correspond to subspaces of the representation space. We now consider how to manipulate concepts through algebraic operations on representations.
To modify a particular concept  $\mathcal{Z}$  we want to modify the representation only on the subspace  $\mathcal{R}_{\mathbb{Z}}$  corresponding to  $\mathcal{Z}$ . For example, consider changing the style concept to Fauvism. Intuitively, we want an operation of the form:
$$
s _ {\text {e d i t}} \leftarrow (\mathbb {I} - \operatorname {p r o j} _ {\text {s t y l e}}) s [ ^ {\prime \prime} a \text {p o r t r a i t o f m a t h e m a t h e m a t i c i a n}" ] + \operatorname {p r o j} _ {\text {s t y l e}} s [ ^ {\prime \prime} \text {F a u v i s m s t y l e}" ], \tag {5}
$$
where  $\mathrm{proj}_{\mathrm{style}}$  is the projection onto the subspace corresponding to the style concept. The idea is that the representation of the original prompt  $x_{\mathrm{orig}}$  is unchanged except on the style subspace. On the style subspace, the representation takes on the value elicited by the new prompt  $x_{\mathrm{new}} = \text{"Fauvism style"}$ .
There are two main challenges for putting this intuition into practice. First, because we are working with an infinite dimensional representation, it is unclear how to do the projection. Second, although we know that some  $\mathcal{R}_Z$  exists, we still need a way to determine it explicitly.
4.1 Concept Manipulation through Projection
Following proposition 3.1, we have that
$$
\bar {s} \left[ Q _ {Z} \times Q _ {W} \right] = \bar {s} _ {Z} \left[ Q _ {Z} \right] + \bar {s} _ {W} \left[ Q _ {W} \right], \tag {6}
$$
for some representation functions  $\bar{s}_Z$  and  $\bar{s}_W$  with range in  $\mathcal{R}_Z$  and  $\mathcal{R}_W$  respectively. We have that the  $Z$ -representation space is
$$
\mathcal {R} _ {Z} = \operatorname {s p a n} \left(\{\bar {s} _ {Z} [ Q _ {Z} ]: Q _ {Z} \text {a c o n c e p t d i s t r i b u t i o n} \}\right). \tag {7}
$$
Our goal is to find a projection onto  $\mathcal{R}_Z$ .
The first obstacle is that  $\mathcal{R}_Z$  is a function space, making algebraic operations difficult to define. The resolution is straightforward. In practice, score-based models generate samples by running a discretized (stochastic) differential equation forward in time. These algorithms only require the score function evaluated at the finite set of points. At each  $y$ , we have that  $\bar{s}(y) \in \mathbb{R}^m$ . Accordingly, by restricting attention to a single value of  $y$  at a time, we can use ordinary linear algebra to define the manipulations:
Definition 4.1. The  $Z$  subspace at  $y$  is
$$
\mathcal {R} _ {Z} (y) := \operatorname {s p a n} \left(\{\bar {s} _ {Z} [ Q _ {Z} ] (y): Q _ {Z} \text {a c o n c e p t d i s t r i b u t i o n} \}\right) \tag {8}
$$
and the  $Z$ -projection at  $y$ , denoted  $\operatorname{proj}_Z(y)$  is the projection onto this subspace.
If we can compute  $\mathrm{proj}_Z(y)$  then we can just edit the representation at each point  $y$ . That is, we transform the score function at each point:
$$
\bar {s} _ {\text {e d i t}} (y) \leftarrow (\mathbb {I} - \operatorname {p r o j} _ {Z} (y)) \bar {s} [ x _ {\text {o r i g}} ] (y) + \operatorname {p r o j} _ {Z} (y) \bar {s} [ x _ {\text {n e w}} ] (y). \tag {9}
$$
We then draw samples from the stochastic differential equation defined by  $\bar{s}_{\mathrm{edit}}$
4.2 Identifying the Concept Subspace
The remaining obstacle is that we need to explicitly identify  $\mathcal{R}_Z(y)$ , so that we can compute  $\mathrm{proj}_Z(y)$ . The problem is that the function  $\bar{s}_Z$  in eq. (6) is unknown, so we cannot compute  $\mathcal{R}_Z(y)$  directly. Our strategy for estimating the space is based on the following proposition.
Proposition 4.1. Let  $Q_W$  be any fixed distribution over the  $W$  concept and  $Q_Z^0$  be any reference distribution over  $Z$ . Then, assuming causal separability for  $\mathcal{Z}, \mathcal{W}$ ,
$$
\mathcal {R} _ {Z} (y) = \operatorname {s p a n} \left(\left\{s \left[ Q _ {Z} Q _ {W} \right] (y) - s \left[ Q _ {Z} ^ {0} Q _ {W} \right] (y): Q _ {Z} \text {a c o n c e p t d i s t r i b u t i o n} \right\}\right). \tag {10}
$$
The importance of this expression is that it does not require the unknown  $s_Z$ .
We'll use eq. (10) to identify  $\mathcal{R}_Z(y)$ . The idea is to find a basis for the subspace using prompts  $x_0, \ldots, x_k$  that elicit distributions of the form  $Q_x = Q_Z Q_W$ . For example, to identify the sex concept we use the prompts  $x_0 = \text{"a man"}$  and  $x_1 = \text{"a woman"}$ , with the idea that
$$
Q _ {x _ {0}} = \delta_ {\text {m a l e}} \times Q _ {W} \quad Q _ {x _ {1}} = \delta_ {\text {f e m a l e}} \times Q _ {W}, \tag {11}
$$
with the same marginal distribution  $Q_W$ . We then use the prompts to define the estimated subspace as
$$
\hat {\mathcal {R}} _ {Z} (y) := \operatorname {s p a n} \left(\left\{s \left[ x _ {i} \right] (y) - s \left[ x _ {0} \right] (y): i = 1, \dots , k \right\}\right). \tag {12}
$$
That is, we write  $k + 1$  prompts  $x_0, \ldots, x_k$  designed so that each elicits a different distribution over  $Z$ , but the same distribution on  $W$ . Then, the estimated subspace is given by eq. (12).
4.3 Concept Algebra
Summarizing, our approach to algebraically manipulating concepts is:
1. Find prompts  $x_0, \ldots, x_k$  such that each elicits a different distribution on  $Z$ , but the same distribution on  $W$ . That is,  $Q_{x_j} = Q_Z^j Q_W$  for each  $j$ .
2. Construct the estimated representation space  $\hat{\mathcal{R}}_Z(y)$  following eq. (12), and define  $\mathrm{proj}_Z(y)$  as the projection onto this space.
3. Sample from the discretized SDE defined by the manipulated score representation<sup>2</sup>
$$
s _ {\mathrm {e d i t}} (y) \leftarrow (\mathbb {I} - \operatorname {p r o j} _ {Z} (y)) s \left[ x _ {\text {o r i g}} \right] (y) + \operatorname {p r o j} _ {Z} (y) s \left[ x _ {\text {n e w}} \right] (y). \tag {13}
$$
Implementation of eq. (13) with the diffusion model is described in appendix A.
5 Validity of Concept Subspace Identification
The procedure described in the previous section relies on finding spanning prompts  $x_0, \ldots, x_k$  for the target concept subspace. These prompts must satisfy  $Q_{x_j} = Q_Z^j Q_W$  for some common  $Q_W$ , and we must have sufficient prompts to span the subspace. The first condition is a question of prompt design, and is often not too hard in practice. However, it is natural to wonder when it's possible to actually recover  $\mathcal{R}_Z$  using only a practical number of prompts. We give some results showing that the dimension of  $\mathcal{R}_Z(y)$  is often small, and thus can be spanned with a small number of prompts. Note that these results rely on the special structure of the score representation, and may not hold for other representations.
First, the case where  $\mathcal{X}$  is categorical with few categories:
Proposition 5.1. Assuming causal separability holds for  $\mathcal{Z},\mathcal{W}$ . If  $\mathcal{Z}$  is categorical with  $L$  possible values ( $L\geq 2$ ), then  $\dim (\mathcal{R}_Z(y))\leq L - 1$ .
This result covers concepts such as sex, which can be spanned with only two prompts.
The next result extends this to certain categorical concepts with large cardinality, such as style. The idea is that if a concept is composed of finer grained categorical concepts, each with small cardinality, then the representation space of the concept is also low-dimensional. For example, style may be composed of lower-level concepts such as color, stroke, textures, etc.
Proposition 5.2. Suppose  $\mathcal{Z}$  is composed of categorical concepts  $\{\mathcal{Z}_k\}_{k=1}^K$  each with the number of categories  $L_k$ , in the sense that  $\mathcal{Z} = \mathcal{Z}_1 \times \ldots \times \mathcal{Z}_k$ . Assume  $Y$  satisfies causal separability with respect to  $\mathcal{Z}, \mathcal{W}$ , with  $Y_{\mathcal{Z}}$  the corresponding  $Y$ -measurable variable for  $\mathcal{Z}$ . Further assume that there exists  $Y_{\mathcal{Z}}$ -measurable variables  $Y_{\mathcal{Z}_k}$  such that  $p(y_{\mathcal{Z}} \mid z) = \prod_{k=1}^{K} p(y_{\mathcal{Z}_k} \mid z_k)$ . Then
$$
\dim \left(\mathcal {R} _ {Z} (y)\right) \leq \sum_ {k = 1} ^ {K} \left(L _ {k} - 1\right) \tag {14}
$$