SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Timur Garipov $^{1*}$  Sebastiaan De Peuter $^{2}$  Ge Yang $^{1,4}$  Vikas Garg $^{2,5}$  Samuel Kaski $^{2,3}$  Tommi Jaakkola $^{1}$
$^{1}$ MIT CSAIL  $^{2}$ Aalto University  $^{3}$ University of Manchester  $^{4}$ Institute for Artificial Intelligence and Fundamental Interactions  $^{5}$ YaiYai Ltd
High training costs of generative models and the need to fine-tune them for specific tasks have created a strong interest in model reuse and composition. A key challenge in composing iterative generative processes, such as GFlowNets and diffusion models, is that to realize the desired target distribution, all steps of the generative process need to be coordinated, and satisfy delicate balance conditions. In this work, we propose Compositional Sculpting: a general approach for defining compositions of iterative generative processes. We then introduce a method for sampling from these compositions built on classifier guidance. We showcase ways to accomplish compositional sculpting in both GFlowNets and diffusion models. We highlight two binary operations — the harmonic mean  $(p_1 \otimes p_2)$  and the contrast  $(p_1 \bullet p_2)$  between pairs, and the generalization of these operations to multiple component distributions. We offer empirical results on image and molecular generation tasks. Project codebase: https://github.com/timgaripov/compositional-sculpting.
Large-scale general-purpose pre-training of machine learning models has produced impressive results in computer vision [1-3], image generation [4-6], natural language processing [7-11], robotics [12-14], and basic sciences [15]. By distilling vast amounts of data, such models can produce powerful inferences that lead to emergent capabilities beyond the specified training objective [16]. However, generic pre-trained models are often insufficient for specialized tasks in engineering and basic sciences. Field-adaptation via techniques such as explicit fine-tuning on bespoke datasets [17], human feedback [18], or cleverly designed prompts [19, 20] is therefore often required. Alternatively, capabilities of pre-trained models can be utilized and extended via model composition.
Compositional generation [21-28] views a complex target distribution in terms of simpler pre-trained building blocks which can be mixed and matched into a tailored solution to a specialized task. Given a set of base models capturing different properties of the data, composition provides a way to fuse these models into a single composite model with capacity greater than any individual base model. In this way it allows one to specify distributions over examples that exhibit multiple desired properties simultaneously [22]. The need to construct complex distributions adhering to multiple constraints arises in numerous practical multi-objective design problems such as molecule generation [29-31]. In this context, compositional modeling provides mechanisms for control of the resulting distribution and exploration of different trade-offs between the objectives and constraints.
Prior work on generative model composition [21, 22, 24] has developed operations for piecing together Energy-Based Models (EBMs) via algebraic manipulations of their energy functions. For example, consider two distributions  $p_1(x) \propto \exp(-E_1(x))$  and  $p_2(x) \propto \exp(-E_2(x))$  induced by energy
*Correspondence to Timur Garipov (timur@csail.mit.edu).
(a)  $p_1$
(b)  $p_2$
(c)  $p_1\otimes p_2$
(d)  $p_1\bullet p_2$
(e)  $p_1\mathbf{0}p_2$
Figure 1: Composition operators. (a,b) base distributions  $p_1$  and  $p_2$ . (c) harmonic mean of  $p_1$  and  $p_2$ . (d) contrast of  $p_1$  with  $p_2$  (e) reverse contrast  $p_1 \odot p_2$ . Lines show the contours of the PDF level sets.
functions  $E_{1}$  and  $E_{2}$ . Their product  $p_{\mathrm{prod}}(x)\propto p_1(x)p_2(x)\propto \exp (-(E_1(x) + E_2(x)))$  and negation  $p_{\mathrm{neg}}(x)\propto p_1(x) / (p_2(x))^{\gamma}\propto \exp (-(E_1(x) - \gamma E_2(x))$ ,  $\gamma >0$  correspond to operations on the underlying energy functions. The product assigns high likelihood to points  $x$  that have high likelihood under both base distributions but assigns low likelihood to points that have close-to-zero likelihood under one (or both). The negation distribution assigns high likelihood to points that are likely under  $p_1$  but unlikely under  $p_2$  and assigns low likelihood to points that are likely under  $p_2$  but unlikely under  $p_1$ .
Iterative generative processes including diffusion models [5, 32-34] and GFlowNets [35, 36] progressively refine coarse objects into cleaner ones over multiple steps. Realizing effective compositions of these models is complicated by the fact that simple alterations in their generation processes result in non-trivial changes in the distributions of the final objects. Unlike for EBMs, products and negations of diffusion models cannot be realized through simple algebraic operations on their score functions. Du et al. [28] show that the result of the addition of score functions is not equal to the score of the diffused product distribution and develop a method that corrects the sum-of-scores sampling via additional MCMC steps nested under each step of the diffusion time loop.
Jain et al. [31] develop Multi-Objective GFlowNets (MOGFNs), an extension of GFlowNets for multi-objective optimization tasks. While a vanilla GFlowNet captures a distribution induced by a single reward (objective) function  $p(x) \propto R(x)$  (see Section 2 for details), an MOGFN aims to learn a single conditional model that can realize distributions corresponding to various combinations (e.g. a convex combination) of multiple reward functions. Though a single MOGFN realizes a spectrum of compositions of base reward functions, the approach assumes access to the base rewards at training time. Moreover, MOFGNs require the set of possible composition operations to be specified at training time. In this work, we address post hoc composition of pre-trained GFlowNets (or diffusion models) and provide a way to create compositions that need not be specified in advance.
In this work, we introduce Compositional Sculpting, a general approach for the composition of pretrained models. We highlight two special examples of binary operations — harmonic mean:  $(p_1 \otimes p_2)$  and contrast:  $(p_1 \oplus p_2)$ . More general compositions are obtained as conditional distributions in a probabilistic model constructed on top of pre-trained base models. We show that these operations can be realized via classifier guidance. We provide results of empirical verification of our method on molecular generation (with GFlowNets) and image generation (with diffusion models).
Generative flow networks (GFlowNets). GFlowNets [35, 36] are an approach for generating structured objects (e.g. graphs) from a discrete space  $\mathcal{X}$ . Given a "reward function"  $R(x) \geq 0$ , a GFlowNet seeks to sample from  $p(x) = R(x)/Z$ , where  $Z = \sum_{x} R(x)$ , i.e. the model assigns larger probabilities to high-reward objects.
Starting at a fixed initial state  $s_0$ , objects  $x$  are generated through a sequence of changes corresponding to a trajectory of incomplete states  $\tau = (s_0 \to s_1 \to \ldots \to s_{n-1} \to x)$ . The structure of possible trajectories corresponds to by a DAG  $(S, \mathcal{A})$  where  $S$  is a set of states (both complete and incomplete) and  $\mathcal{A}$  is the set of directed edges (actions)  $s \to s'$ . The set of complete objects (terminal states)  $\mathcal{X}$  is a subset of  $S$ . The generation process starts at  $s_0$  and follows a parameterized stochastic "forward policy"  $P_F(s' | s; \theta)$  which for each state  $s \in S \setminus \mathcal{X}$  specifies a probability distribution over all possible successor states  $s': (s \to s') \in \mathcal{A}$ . The process terminates once a terminal state is reached.
Diffusion models. Diffusion models [5, 32-34, 37] are a family of generative models developed for continuous domains. Given an empirical (data) distribution  $\hat{p}(x) = \frac{1}{n} \sum_{i} \delta_{\hat{x}_i}(x)$  in  $\mathcal{X} = \mathbb{R}^d$ , diffusion models seek to approximate  $\hat{p}(x)$  via a generative process  $p(x)$ .
A diffusion process is a noising process that gradually destroys the original "clean" data  $x$ . Viewed as a stochastic differential equation (SDE) [34], it is a time-indexed collection of random variables  $\{x_{t}\}_{t=0}^{T}$  in  $\mathcal{X} = \mathbb{R}^{d}$  which interpolates between the data distribution  $p_0(x) = \hat{p}(x)$  at  $t = 0$  and a prior distribution  $p_T(x)$  at  $t = T$ . The evolution of  $x_{t}$  is described by the "forward SDE"  $dx_{t} = f_{t}(x_{t})dt + g_{t}dw_{t}$  with drift coefficient  $f_{t}: \mathbb{R}^{d} \to \mathbb{R}^{d}$  and diffusion coefficient  $g_{t} \in \mathbb{R}$ . Here,  $w_{t}$  is the standard Wiener process. Crucially, the coefficients  $f_{t}, g_{t}$  are generally chosen such that the prior  $p_T$  and the transition probabilities  $p_{st}(x_t|x_s), 0 \leq s < t \leq T$  have a closed form (see [34]).
Song et al. [34] invoke a result from the theory of stochastic processes [38] which gives the expression for the reverse-time process or "backward SDE":  $dx_{t} = \left[f_{t}(x_{t}) - g_{t}^{2}\nabla_{x}\log p_{t}(x_{t})\right]dt + g_{t}d\overline{w}_{t}$ , where  $\overline{w}_t$  is the standard Wiener process in reversed time. This SDE involves the known coefficients  $f_{t}, g_{t}$  and the unknown score function  $\nabla_{x}\log p_{t}(\cdot)$  of the marginal distribution  $p_t(\cdot)$  at time  $t$ . A "score network"  $s_t(x;\theta)$  (a deep neural network with parameters  $\theta$ ) is trained to approximate  $\nabla_{x}\log p_{t}(x)$ . Once  $s_t(\cdot ;\theta)$  is trained, sampling reduces to numerical integration of the backward SDE.
Classifier guidance in diffusion models. Classifier guidance [32, 39] is a technique for controllable generation in diffusion models. Suppose that each example  $x_0$  is accompanied by a discrete class label  $y$ . The goal is to sample from the conditional distribution  $p_0(x_0|y)$ . The Bayes rule  $p_t(x_t|y) \propto p_t(x_t)p_t(y|x_t)$  implies the score-function decomposition  $\nabla_{x_t}\log p_t(x_t|y) = \nabla_{x_t}\log p_t(x_t) + \nabla_{x_t}\log p_t(y|x_t)$ , where the first term is already approximated by a pre-trained unconditional diffusion model and the second term can be derived from a time-dependent classifier  $p_t(y|x_t)$ . Therefore, the stated goal can be achieved by first training the classifier  $p_t(y|x_t)$  using noisy samples  $x_t$  from the intermediate steps of the process, and then plugging in the expression for the conditional score into the backward SDE sampling process [34].
Generative model composition. In Section 1 we reviewed prior work on energy-based composition operations and Multi-Objective GFlowNets. Learning mixtures of Generative Adversarial Networks has been addressed in [40], where the mixture components are learned simultaneously, and in [41], where the components are learned one by one in an adaptive boosting fashion. Algorithms for additive and multiplicative boosting of generative models have been developed in [42].
This work focuses on the composition of pre-trained models. Assuming that each pre-trained model represents the distribution of examples demonstrating certain concepts (e.g. molecular properties), the composition of models is equivalent to concept composition (e.g. property "A" AND property "B"). The inverse problem is known as "unsupervised concept discovery", where the goal is to automatically discover composable concepts from data. Unsupervised concept discovery and concept composition methods have been proposed for energy-based models [25] and text-to-image diffusion models [43].
Controllable generation. Generative model composition is a form of post-training control of the generation process - an established area of research. A simple approach to control is training a conditional generative model  $p(x|c)$  on pairs  $(x, c)$  of objects  $x$  and conditioning information  $c$ . Annotations  $c$  can be class labels [39], text prompts [4, 6, 44], semantic maps, and images [4]. Different from out work, this assumes that generation control operations are specified at training time. Dhariwal and Nichol [39] apply classifier guidance [32] on top of (conditional) diffusion models to improve their fidelity. Ho and Salimans [45] develop classifier-free guidance by combining conditional and unconditional score functions. In ControlNet [17], an additional network is trained to enable a pre-trained diffusion model to incorporate previously unavailable conditioning information. Meng et al. [46] and Couairon et al. [47] develop semantic image editing methods which first partially noise and then denoise an image to generate an edited version, possibly conditioned on a segmentation mask [47]. Similar to conditional diffusion models, conditional GFlowNets have been used to condition generation on reward exponents [36] or combinations of multiple predefined reward functions [31]. Note that the methods developed in this work can be combined with conditional diffusion models and GFlowNets:  $p(x|c_1), \ldots, p(x|c_m)$  can act as base generative models to be composed.
Compositional generalization. The notion of compositionality has a broad spectrum of interpretations across a variety of disciplines including linguistics, cognitive science, and philosophy. Hupkes et al. [48] collect a list of aspects of compositionality from linguistic and philosophical theories and design practical tests for neural language models. Conwell and Ullman [49] empirically examine the relational understanding of DALL-E 2 [50], a text-guided image generation model, and point out limitations in the model's ability to capture relations such as "in", "on", "hanging over", etc. In this work, we focus on a narrow but well-defined type of composition where we seek to algebraically compose probability densities in a controllable fashion, such that we can emphasize or de-emphasize regions in the data space where specific base distributions have high density.
Connections between GFlowNets and diffusion models. Our method is applicable to compositions of both GFlowNets and diffusion models. This is due to deep connections between these two model families. GFlowNets were initially developed for generating discrete (structured) data [36] and diffusion models were initially developed for continuous data [5, 32]. Lahlou et al. [51] develop an extension of GFlowNets for DAGs with continuous state-action spaces. Zhang et al. [52] point out unifying connections between GFlowNets and other generative model families, including diffusion models. In this work, we articulate another aspect of the relation between GFlowNets and diffusion models: in Section 5.2 we derive the expressions for mixture GFlowNet policies and classifier-guided GFlowNet policies analogous to those derived for diffusion models in [32, 39, 53, 54].
4 Compositional Sculpting of Generative Models
Suppose we can access a number of pre-trained generative models  $\{p_i(x)\}_{i=1}^m$  over a common domain  $\mathcal{X}$ . We may wish to compose these distributions such that we can, say, draw samples that are likely to arise from  $p_1(x)$  and  $p_2(x)$ , or that are likely to arise from  $p_1(x)$  but not from  $p_2(x)$ . In other words, we wish to specify a distribution that we can shape to emphasize and de-emphasize specific base models.
4.1 Binary Composition Operations
Let us first focus on composing two base models. We could specify the composition as a weighted sum  $\widetilde{p}(x) = \sum_{i=1}^{2} \omega_i p_i(x)$  with weights  $\omega_1, \omega_2 \geq 0$  summing to one. The weights determine the prevalence of each base model in the composition, but beyond that our control is limited. We cannot emphasize regions where  $p_1$  and  $p_2$  both have high density, or de-emphasize regions where  $p_2$  has high density.
An alternative is to use conditioning to shape a prior  $\widetilde{p}(x)$  based on the base models. When we condition  $x$  on some observation  $y_1$ , the resulting posterior takes the form  $\widetilde{p}(x|y_1) \propto \widetilde{p}(y_1|x)\widetilde{p}(x)$ . Points  $x$  that match  $y_1$  according to  $\widetilde{p}(y_1|x)$  will have increased density, and the density of points that do not match it decreases. Intuitively, by defining  $y_1 \in \{1,2\}$  as the event that  $x$  was generated by a specific base model, we can shape a prior  $\widetilde{p}(x)$  based on the densities of the base models. To this end we define a uniform prior over  $y_k$  and define the conditional density  $p(x|y_1 = i)$  to represent the fact that  $x$  was generated from  $p_i(x)$ . This gives us the following model:
$$
\widetilde {p} (x \mid y _ {1} = i) = p _ {i} (x), \quad \widetilde {p} (y _ {1} = i) = 1 / 2, \quad \widetilde {p} (x) = \widetilde {p} (x \mid y _ {1} = 1) \widetilde {p} (y _ {1} = 1) + \widetilde {p} (x \mid y _ {1} = 2) \widetilde {p} (y _ {1} = 2). \tag {1}
$$
Under this model, the prior  $\widetilde{p}(x)$  is a uniform mixture of the base models. The likelihood of  $y_{1}$
$$
\widetilde {p} \left(y _ {1} = 1 | x\right) = 1 - \widetilde {p} \left(y _ {1} = 2 | x\right) = p _ {1} (x) / \left(p _ {1} (x) + p _ {2} (x)\right), \tag {2}
$$
implied by this model tells us how likely it is that  $x$  was generated by  $p_1(x)$  rather than  $p_2(x)$ . In fact, it corresponds to the output of an optimal classifier trained to tell  $p_1(x)$  and  $p_2(x)$  apart.
Our goal is to realize compositions which generate samples likely to arise from both  $p_1(x)$  and  $p_2(x)$  or from  $p_1(x)$  but not  $p_2(x)$ . Thus we introduce a second observation  $y_2 \in \{1, 2\}$  such that  $y_1$  and  $y_2$  are independent and identically distributed given  $x$ . The resulting model and inferred posterior are:
$$
\widetilde {p} (x, y _ {1}, y _ {2}) = \widetilde {p} (x) \prod_ {k = 1} ^ {2} \widetilde {p} (y _ {k} | x), \quad \widetilde {p} (x) = \frac {1}{2} p _ {1} (x) + \frac {1}{2} p _ {2} (x), \quad \widetilde {p} (y _ {k} = i | x) = p _ {i} (x) / (p _ {1} (x) + p _ {2} (x)), \tag {3}
$$
$$
\widetilde {p} (x \mid y _ {1} = i, y _ {2} = j) \propto \widetilde {p} (x) \widetilde {p} (y _ {1} = i \mid x) \widetilde {p} (y _ {2} = j \mid x) \propto p _ {i} (x) p _ {j} (x) / \left(p _ {1} (x) + p _ {2} (x)\right). \tag {4}
$$
The above posterior shows clearly how conditioning on observations  $y_{1} = i$ ,  $y_{2} = j$  has shaped the prior mixture to accentuate regions in the posterior where the observed base models  $i, j$  have high density.
Conditioning on observations  $y_{1} = 1$  and  $y_{2} = 2$ , or equivalently  $y_{1} = 2$ ,  $y_{2} = 1$ , results in the posterior
$$
(p _ {1} \otimes p _ {2}) (x) := p \left(x \mid y _ {1} = 1, y _ {2} = 2\right) \propto p _ {1} (x) p _ {2} (x) / \left(p _ {1} (x) + p _ {2} (x)\right). \tag {5}
$$
We refer to this posterior as the "harmonic mean of  $p_1$  and  $p_2$ ", and denote it as a binary operation  $p_1 \otimes p_2$ . Its value is high only at points that have high likelihood under both  $p_1(x)$  and  $p_2(x)$  at the same time (Figure 1(c)). Thus, the harmonic mean is an alternative to the product operation for EBMs. The harmonic mean is commutative ( $p_1 \otimes p_2 = p_2 \otimes p_1$ ) and is undefined when  $p_1$  and  $p_2$  have disjoint supports, since then the RHS of (5) is zero everywhere.
Conditioning on observations  $y_{1} = 1$  and  $y_{2} = 1$  results in the posterior
$$
(p _ {1} \oplus p _ {2}) (x) := \widetilde {p} (x \mid y _ {1} = 1, y _ {2} = 1) \propto (p _ {1} (x)) ^ {2} / (p _ {1} (x) + p _ {2} (x)). \tag {6}
$$
We refer to this binary operation, providing an alternative to the negation operation in EBMs, as the "contrast of  $p_1$  and  $p_2$ ", and denote it as  $p_1 \mathbf{0} p_2$ . The ratio (6) is high when  $p_1(x)$  is high and  $p_2(x)$  is low (Figure 1(d)). The contrast is not commutative ( $p_1 \mathbf{0} p_2 \neq p_2 \mathbf{0} p_1$ , unless  $p_1 = p_2$ ). We denote the reverse contrast as  $p_1 \mathbf{0} p_2 = p_2 \mathbf{0} p_1$ . Appendix C provides a detailed comparison between the contrast and negation operations, and between the harmonic mean and product operations.
Controlling the individual contributions of  $p_1$  and  $p_2$  to the composition. In order to provide more control over the extent of individual contributions of  $p_1$  and  $p_2$  to the composition, we modify model (3). Specifically, we introduce an interpolation parameter  $\alpha$  and change the likelihood of observation  $y_2$  in (3):  $\widetilde{p}(y_2 = i|x;\alpha) = (\alpha p_1(x))[i=1] \cdot ((1 - \alpha)p_2(x))[i=2] / (\alpha p_1(x) + (1 - \alpha)p_2(x))$ , where  $\alpha \in (0,1)$  and  $[\cdot]$  denotes the indicator function. Conditional distributions in this model give harmonic interpolation $^2$  and parameterized contrast:
$$
(p _ {1} \otimes_ {(1 - \alpha)} p _ {2}) (x) \propto \frac {p _ {1} (x) p _ {2} (x)}{\alpha p _ {1} (x) + (1 - \alpha) p _ {2} (x)}, \quad (p _ {1} \mathbf {0} _ {(1 - \alpha)} p _ {2}) (x) \propto \frac {(p _ {1} (x)) ^ {2}}{\alpha p _ {1} (x) + (1 - \alpha) p _ {2} (x)}. \tag {7}
$$
Operation chaining. As the operations we have introduced result in proper distributions, we can create new  $N$ -ary operations by chaining binary (and  $N$ -ary) operations together. For instance, chaining binary harmonic means gives the harmonic mean of three distributions
$$
((p _ {1} \otimes p _ {2}) \otimes p _ {3}) (x) = (p _ {1} \otimes (p _ {2} \otimes p _ {3})) (x) \propto \frac {p _ {1} (x) p _ {2} (x) p _ {3} (x)}{p _ {1} (x) p _ {2} (x) + p _ {1} (x) p _ {3} (x) + p _ {2} (x) p _ {3} (x)}. \qquad (8)
$$
4.2 Compositional Sculpting: General Approach
The above approach for realizing compositions of two base models can be generalized to compositions of  $m$  base models  $p_1(x), \ldots, p_m(x)$  controlled by  $n$  observations. Though operator chaining can also realize compositions of  $m$  base models, our generalized method allows us to specify compositions more flexibly and results in different compositions. We introduce an augmented probabilistic model  $\widetilde{p}(x, y_1, \ldots, y_n)$  as a joint distribution over the original objects  $x \in \mathcal{X}$  and  $n$  observation variables  $y_1 \in \mathcal{Y}, \ldots, y_n \in \mathcal{Y}$  where  $\mathcal{Y} = \{1, \ldots, m\}$ . By defining appropriate conditionals  $p(y_k|x)$  we can controllably shape a prior  $\widetilde{p}(x)$  into a posterior  $\widetilde{p}(x|y_1, \ldots, y_n)$ .
As in the binary case, we propose to use a uniformly-weighted mixture of the base models  $\widetilde{p}(x) = (1/m) \cdot \sum_{i=1}^{m} p_i(x)$ . The support of this mixture is the union of the supports of the base models:  $\bigcup_{i=1}^{m} \operatorname{supp}\{p_i(x)\} = \operatorname{supp}\{\widetilde{p}(x)\}$ . This is essential as the prior can only be shaped in places where it has non-zero density. As before we define the conditionals  $p(y_k = i|x)$  to correspond to the observation that  $x$  was generated by base model  $i$ . The resulting full model is
$$
\widetilde {p} (x, y _ {1}, \dots , y _ {n}) = \widetilde {p} (x) \prod_ {k = 1} ^ {n} \widetilde {p} (y _ {k} | x), \widetilde {p} (x) = \frac {1}{m} \sum_ {i = 1} ^ {m} p _ {i} (x), \widetilde {p} (y _ {k} = i) = \frac {1}{m}, \widetilde {p} (y _ {k} = i | x) = \frac {p _ {i} (x)}{\sum_ {j = 1} ^ {m} p _ {j} (x)}, \tag {9}
$$
Note that under this model the mixture can be represented as  $\widetilde{p}(x) = \sum_{y_k=1}^m \widetilde{p}(x | y_k) \widetilde{p}(y_k)$  for any  $k$ .
The inferred posterior over  $x$  for this model is
$$
\widetilde {p} (x \mid y _ {1} = i _ {1}, \dots , y _ {n} = i _ {n}) \propto \widetilde {p} (x) \widetilde {p} (y _ {1} = i _ {1}, \dots , y _ {n} = i _ {n} \mid x) \tag {10}
$$
$$
\propto \widetilde {p} (x) \prod_ {k = 1} ^ {n} \widetilde {p} \left(y _ {k} = i _ {k} \mid x\right) \propto \left(\prod_ {k = 1} ^ {n} p _ {i _ {k}} (x)\right) / \left(\sum_ {j = 1} ^ {m} p _ {j} (x)\right) ^ {n - 1}. \tag {11}
$$
the harmonic interpolation approaches  $p_1$  when  $\alpha \to 0$  and  $p_2$  when  $\alpha \to 1$ .