SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The posterior  $\widetilde{p}(x|y_1 = i_1, \ldots, y_n = i_n)$  is a composition of distributions  $\{p_i(x)\}_{i=1}^m$  that can be adjusted by choosing values for  $y_1, \ldots, y_n$ . By adding or omitting an observation  $y_k = i$  we can sculpt the posterior to our liking, emphasizing or de-emphasizing regions of  $\mathcal{X}$  where  $p_i$  has high density. The observations can be introduced with multiplicities (e.g.,  $y_1 = 1$ ,  $y_2 = 1$ ,  $y_3 = 2$ ) to further strengthen the effect. Moreover, one can choose to introduce all observations simultaneously as in (10) or sequentially as in (11). As we show below (Section 5.1 for GFlowNets; Appendix A.2 for diffusion models), the composition (10) can be realized by a sampling policy that can be expressed as a function of the pre-trained (base) sampling policies.
Special instances and general formulation. The general approach outlined in this section is not limited to choices we made to construct the model in equation (9), i.e.  $\widetilde{p}(x)$  does not have to be a uniformly weighted mixture of the base distributions,  $y_1, \ldots, y_n$  do not have to be independent and identically distributed given  $x$ , and different choices of the likelihood  $\widetilde{p}(y = i|x)$  are possible. For instance, in the model for parameterized operations (7) the likelihoods of observations  $\widetilde{p}(y_1|x)$ ,  $\widetilde{p}(y_2|x)$  differ.
5 Compositional Sculpting of Iterative Generative Processes
In this Section, we show how to apply the model above to compose GFlowNets, and how one can use classifier guidance to sample from the composition. The similar method for diffusion model composition is described in Appendix A.2.
5.1 Composition of GFlowNets
Besides a sample  $x$  from  $p_i(x)$ , a GFlowNet also generates a trajectory  $\tau$  which ends in the state  $x$ . Thus, we extend the model  $\widetilde{p}(x, y_1, \ldots, y_n)$ , described above, and introduce  $\tau$  as a variable with conditional distribution  $\widetilde{p}(\tau | y_k = i) = \prod_{t=0}^{|\tau|-1} p_{i,F}(s_{t+1} | s_t)$ , where  $p_{i,F}$  is the forward policy of the GFlowNet that samples from  $p_i$ .
Our approach for sampling from the composition is conceptually simple. Given  $m$  base GFlowNets that sample from  $p_1, \ldots, p_m$  respectively, we start by defining the prior  $\widetilde{p}(x)$  as the uniform mixture of these GFlowNets. Proposition 5.1 shows that this mixture can be realized by a policy constructed from the forward policies of the base GFlowNets. We then apply classifier guidance to this mixture to sample from the composition. Proposition 5.2 shows that classifier guidance results in a new policy which can be constructed directly from the GFlowNet being guided.
Proposition 5.1 (GFlowNet mixture policy). Suppose distributions  $p_1(x), \ldots, p_m(x)$  are realized by GFlowNets with forward policies  $p_{1,F}(\cdot|\cdot), \ldots, p_{m,F}(\cdot|\cdot)$ . Then, the mixture distribution  $p_M(x) = \sum_{i=1}^{m} \omega_i p_i(x)$  with  $\omega_1, \ldots, \omega_m \geq 0$  and  $\sum_{i=1}^{m} \omega_i = 1$  is realized by the GFlowNet forward policy
$$
p _ {M, F} \left(s ^ {\prime} \mid s\right) = \sum_ {i = 1} ^ {m} p (y = i \mid s) p _ {i, F} \left(s ^ {\prime} \mid s\right), \tag {12}
$$
where  $y$  is a random variable such that the joint distribution of a GFlowNet trajectory  $\tau$  and  $y$  is given by  $p(\tau, y = i) = \omega_i p_i(\tau)$  for  $i \in \{1, \dots, m\}$ .
Proposition 5.2 (GFlowNet classifier guidance). Consider a joint distribution  $p(x, y)$  over a discrete space  $\mathcal{X} \times \mathcal{Y}$  such that the marginal  $p(x)$  is realized by a GFlowNet with forward policy  $p_F(\cdot | \cdot)$ . Assume that the joint distribution of  $x, y$ , and GFlowNet trajectories  $\tau = (s_0 \dots s_n = x)$  decomposes as  $p(\tau, x, y) = p(\tau, x)p(y|x)$ , i.e.  $y$  is independent of the intermediate states  $\{s_i\}_{i=0}^{n-1}$  in  $\tau$  given  $x$ . Then,
1. For all non-terminal nodes  $s \in S \setminus \mathcal{X}$ , the probabilities  $p(y|s)$  satisfy
$$
p (y | s) = \sum_ {s ^ {\prime}: (s \rightarrow s ^ {\prime}) \in \mathcal {A}} p _ {F} \left(s ^ {\prime} \mid s\right) p (y \mid s ^ {\prime}). \tag {13}
$$
2. The conditional distribution  $p(x|y)$  is realized by the classifier-guided policy
$$
p _ {F} \left(s ^ {\prime} \mid s, y\right) = p _ {F} \left(s ^ {\prime} \mid s\right) \cdot p \left(y \mid s ^ {\prime}\right) / p (y \mid s). \tag {14}
$$
Note that (13) ensures that  $\sum_{s' : (s \to s') \in \mathcal{A}} p_F(s'|s, y) = 1$ .
Proposition 5.1 is analogous to results on mixtures of diffusion models (Theorem 1 of Peluchetti [53], Theorem 1 of Lipman et al. [54]). Proposition 5.2 is analogous to classifier guidance for diffusion models [32, 39]. To the best of our knowledge, our work is the first to derive both results for GFlowNets.
Theorem 5.3 summarizes our approach. The propositions and the theorem are proved in Appendix D. Theorem 5.3. Suppose distributions  $p_1(x), \ldots, p_m(x)$  are realized by GFlowNets with forward policies  $p_{1,F}(\cdot|\cdot), \ldots, p_{m,F}(\cdot|\cdot)$  respectively. Let  $y_1, \ldots, y_n$  be random variables defined by (9). Then, the conditional  $\widetilde{p}(x|y_1, \ldots, y_n)$  is realized by the forward policy
$$
p _ {F} \left(s ^ {\prime} \mid s, y _ {1}, \dots , y _ {n}\right) = \frac {\widetilde {p} \left(y _ {1} , \dots , y _ {n} \mid s ^ {\prime}\right)}{\widetilde {p} \left(y _ {1} , \dots , y _ {n} \mid s\right)} \sum_ {i = 1} ^ {m} p _ {i, F} \left(s ^ {\prime} \mid s\right) \widetilde {p} (y = i \mid s) \tag {15}
$$
Note that the result of conditioning on observations  $y_{1},\ldots ,y_{n}$  is just another GFlowNet policy. Therefore, to condition on more observations, one can apply classifier guidance repeatedly.
5.2 Classifier Training (GFlowNets)
The evaluation of policy (15) requires knowledge of the probabilities  $\widetilde{p}(y_1, \ldots, y_n|s)$ ,  $\widetilde{p}(y|s)$ . These probabilities can be estimated by a classifier fitted to trajectories sampled from the base GFlowNets. Let  $\widetilde{Q}_{\phi}(y_1, \ldots, y_n|s)$  be a classifier with parameters  $\phi$  that we wish to train to approximate the ground-truth conditional  $\widetilde{p}(y_1, \ldots, y_n|s)$ . Under model (9),  $y_1, \ldots, y_n$  are dependent given a state  $s \in S \setminus \mathcal{X}$  but, are independent given a terminal state  $x \in \mathcal{X}$ . This motivates separate treatment of terminal and non-terminal states.
Learning the terminal state classifier. For a terminal state  $x$ , the variables  $y_{1},\ldots ,y_{n}$  are independent, hence we can use the factorization  $\widetilde{Q}_{\phi}(y_1,\dots,y_n|x) = \prod_{k = 1}^{n}\widetilde{Q}_{\phi}(y_k|x)$ . Moreover, all distributions on the r.h.s. must be the same, i.e. it is enough to learn just  $\widetilde{Q}_{\phi}(y_1|x)$ . This marginal classifier can be learned by minimizing the cross-entropy loss
$$
\mathcal {L} _ {\mathrm {T}} (\phi) = \underset {(\hat {x}, \hat {y} _ {1}) \sim \widetilde {p} (y _ {1}) \widetilde {p} (x | y _ {1})} {\mathbb {E}} \left[ - \log \widetilde {Q} _ {\phi} \left(y _ {1} = \hat {y} _ {1} \mid x = \hat {x}\right) \right]. \tag {16}
$$
Here  $\widehat{y}_1$  is sampled from  $\widetilde{p}(y_1)$ , which is uniform under our choice of  $\widetilde{p}(x)$ . Then,  $\widehat{x} | (y_1 = \widehat{y}_1)$  is generated from the base GFlowNet  $p_{\widehat{y}_1}$ , since (9) implies that  $\widetilde{p}(x | y = \widehat{y}_1) = p_{\widehat{y}_1}(x)$ .
Learning the non-terminal state classifier. Given a non-terminal state  $s \in S \setminus \mathcal{X}$ , we need to model  $y_1, \ldots, y_n$  jointly, and training requires sampling tuples  $(\widehat{s}, \widehat{y}_1, \ldots, \widehat{y}_n)$ . Non-terminal states  $s$  can be generated as intermediate states in trajectories  $\tau = (s_0 \to s_1 \to \ldots \to x)$ . Given a sampled trajectory  $\widehat{\tau}$  and a set of labels  $\widehat{y}_1, \ldots, \widehat{y}_n$  we denote the trajectory cross-entropy by
$$
\ell (\widehat {\tau}, \widehat {y} _ {1}, \dots , \widehat {y} _ {n}; \phi) = \sum_ {t = 0} ^ {| \tau | - 1} \left[ - \log \widetilde {Q} _ {\phi} \left(y _ {1} = \widehat {y} _ {1}, \dots , y _ {n} = \widehat {y} _ {n} \mid s = \widehat {s} _ {t}\right) \right]. \tag {17}
$$
Pairs  $(\widehat{\tau}, \widehat{y}_1)$  can be generated in the same way as in the terminal classifier training above:  $1) \widehat{y}_1 \sim \widetilde{p}(y_1)$ ;  $2) \widehat{\tau} \sim p_{\widehat{y}_1}(\tau)$ . Sampling  $\widehat{y}_2, \ldots, \widehat{y}_n$  given  $\widehat{x}$  (the terminal state of  $\widehat{\tau}$ ) requires access to values  $\widetilde{p}(y_k = \widehat{y}_k | \widehat{x})$ , which are not directly available. However, if the terminal classifier is learned as described above, the estimates  $w_i(\widehat{x}; \phi) = \widetilde{Q}_{\phi}(y_1 = i | x = \widehat{x})$  can be used instead. In this case, the loss and the sampling procedure for the non-terminal classifier rely on the outputs of the terminal classifier. In order to train two classifiers simultaneously, and avoid the instability due to a feedback loop, we employ the "target network" technique developed in the context of deep Q-learning [55]. We introduce a "target network" parameter vector  $\overline{\phi}$  which is used to produce the estimates  $w_i(\widehat{x}; \overline{\phi})$  for the non-terminal loss. We update  $\overline{\phi}$  as the exponential moving average of the recent iterates of  $\phi$ .
After putting all components together the training loss for the non-terminal state classifier is
$$
\mathcal {L} _ {N} (\phi , \bar {\phi}) = \underset {(\hat {\tau}, \hat {y} _ {1}) \sim \widetilde {p} (\tau , y _ {1})} {\mathbb {E}} \left[ \sum_ {\hat {y} _ {2} = 1} ^ {m} \dots \sum_ {\hat {y} _ {n} = 1} ^ {m} \left(\prod_ {k = 2} ^ {n} w _ {\hat {y} _ {k}} (\hat {x}; \bar {\phi})\right) \ell (\hat {\tau}, \hat {y} _ {1}, \dots , \hat {y} _ {n}; \phi) \right]. \tag {18}
$$
We refer the reader to Appendix D.4 for a more detailed derivation of the loss (18). Algorithm A.1 shows the complete classifier training procedure.
$p_1$
$p_1$
$p_2$
$p_2$
$p_3$
$p_1\otimes p_2$
$\widetilde{p}\left(x\bigg|_{y_1 = 2}^{y_1 = 1}\right)$
$p_1\mathbf{O}p_2$
$\widetilde{p}\left(x\begin{array}{c}|y_1 = 1\\ y_2 = 2\\ y_3 = 3 \end{array} \right)$
$p_1\mathbf{0}_{0.95}p_2$
$\widetilde{p}\left(x\bigg|_{\substack{y_1 = 2\\ y_2 = 2}}\right)$
$\widetilde{p}\left(x\begin{array}{c}y_1 = 2\\ y_2 = 2\\ y_3 = 2 \end{array}\right)$
Figure 2: Composed GFlowNets on  $32 \times 32$  grid domain. (Top) operations on two distributions. (Bottom) operations on three distributions. Cell probabilities are shown with color, darker is higher. Red circles indicate the high probability regions of  $p_1$ ,  $p_2$ ,  $p_3$ .
(a) Base at  $\beta = 32$
Figure 3: Reward distributions in the molecular generation domain. (a) Base GFlowNets at  $\beta = 32$ :  $p_{\mathrm{SEH}}$  and  $p_{\mathrm{SA}}$  are trained with  $R_{\mathrm{SEH}}(x)^{32}$  and  $R_{\mathrm{SA}}(x)^{32}$ . (b) harmonic mean of  $p_{\mathrm{SEH}}$  and  $p_{\mathrm{SA}}$  at  $\beta = 32$ . (c) contrasts at  $\beta = 32$ . (d) base GFlowNets at  $\beta = 96$ . (e) harmonic mean at  $\beta = 96$ . Lines show the contours of the level sets of the kernel density estimates in the  $(R_{\mathrm{SEH}}, R_{\mathrm{SA}})$  plane.
(b) Harmonic mean
(c) Contrasts
(d) Base at  $\beta = 96$
(e) Harmonic mean
2D distributions via GFlowNet. We tested our GFlowNet composition method on 2D grid [36], a controlled domain, where the ground-truth composite distributions can be evaluated directly.
In the 2D grid domain, the states are the cells of an  $H \times H$  grid. The starting state is the upper-left cell  $s_0 = (0,0)$ . At each state, the allowed actions are: 1) move right; 2) move down; 3) terminate the trajectory at the current position. We first trained GFlowNets  $p_i(x) \propto R_i(x)$  with reward functions  $R_i(x) > 0$ , and then trained classifiers and constructed compositions following Theorem 5.3.
Figure 2 (top row) shows the distributions obtained by composing two pre-trained GFlowNets (top row; left). The harmonic mean  $p_1 \otimes p_2$ , covers the regions that have high probability under both  $p_1$  and  $p_2$  and excludes locations where either of the probabilities is low.  $p_1 \oplus p_2$  resembles  $p_1$  but the relative masses of the modes of  $p_1$  are modulated by  $p_2$ : regions with high  $p_2$  have lower probability under contrast. The parameterized contrast  $p_1 \oplus_{0.95} p_2$  with  $\alpha = 0.05$  magnifies the contrasting effect: high  $p_2(x)$  implies very low  $(p_1 \oplus_{0.95} p_2)(x)$ . The bottom row of Figure 2 shows operations on 3 distributions. The conditional  $\widetilde{p}(x|y_1=1, y_2=2)$  is concentrated on the points that have high likelihood under both  $p_1$  and  $p_2$ . Similarly, the value  $\widetilde{p}(x|y_1=1, y_2=2, y_3=3)$  is high if  $x$  is likely to be observed under all three distributions at the same time. The conditionals  $\widetilde{p}(x|y_1=2, y_2=2)$  and  $\widetilde{p}(x|y_1=2, y_2=2, y_3=2)$  highlight the points with high  $p_2(x)$  but low  $p_1(x)$  and  $p_3(x)$ . Conditioning on three labels results in a sharper distribution compared to double-conditioning. We provide quantitative results and further details in Appendix F.1. The classifier learning curves are provided in Appendix G.4.
Molecule generation via GFlowNet. Next, we evaluated our method for GFlowNet composition on a large and highly structured data space, and assessed the effect that composition operations have
Table 1: Reward distributions of composite GFlowNets.
SEH SA QED	low	high
low	high	low	high
low	high	low	high	low	high	low	high
pSEH	0	0	0	0	62	9	24	5
pSA	0	0	73	4	0	0	18	5
pQED	0	40	0	26	0	21	0	13
(a) y= {SEH, SA}	1	0	16	2	6	3	54	18
(b) y= {SEH, QED}	0	11	0	4	1	48	4	32
(c) y= {SA, QED}	0	15	1	42	0	8	2	32
(d) y= {SEH, SA, QED}	0	7	2	11	2	19	10	49
(e) y= {SEH, SEH, SEH}	0	0	0	0	63	9	24	4
(f) y= {SA, SA, SA}	0	0	74	5	0	0	17	4
(g) y= {QED, QED, QED}	0	40	0	23	0	23	0	14
In each row, the numbers show the percentage of the samples from the respective model that fall into one of 8 bins according to rewards. The "low" and "high" categories are decided by thresholding SEH: 0.5, SA: 0.6, QED: 0.25.
Figure 4: 2D t-SNE embeddings of three base GFlowNets trained with  $R_{\mathrm{SEH}}(x)^{\beta}$ ,  $R_{\mathrm{SA}}(x)^{\beta}$ ,  $R_{\mathrm{QED}}(x)^{\beta}$  at  $\beta = 32$  and their compositions. The t-SNE embeddings are computed based on pairwise earth mover's distances between the distributions. Labels (a)-(g) match rows in Table 1.
on resulting data distributions in a practical setting. To that end, we conducted experiments with GFlowNets trained for the molecular generation task proposed by Bengio et al. [36]. We train base GFlowNets with 3 reward functions  $R_{\mathrm{SEH}}(x)$ ,  $R_{\mathrm{SA}}(x)$ , and  $R_{\mathrm{QED}}(x)$  measuring 3 distinct molecular properties (details in Appendix F.2). Following Bengio et al. [36], we introduced the parameter  $\beta$  which controls the sharpness (temperature) of the target distribution:  $p(x) \propto R(x)^{\beta}$ , increasing  $\beta$  results in a distribution skewed towards high-reward objects. We experimented with  $\beta = 32$  and  $\beta = 96$  (Figures 3(a), 3(d)). After training base GFlowNets with the respective reward functions, we trained classifiers with Algorithm A.1. In order to evaluate the base models and the composed GFlowNet policies, we generated 5000 samples with each policy and analyzed the samples. Further details are in Appendix F.2. Classifier learning curves are provided in Appendix G.4. Sample diversity statistics of base GFlowNets at different values of  $\beta$  are provided in Appendix G.5.
Figure 3 shows reward distributions of base GFlowNets (trained with rewards  $R_{\mathrm{SEH}}(x)^{\beta}$ ,  $R_{\mathrm{SA}}(x)^{\beta}$ ,  $\beta \in \{32,96\}$ ) and their compositions. Base GFlowNet distributions are concentrated on examples that score high in their respective rewards. For each model, there is considerable variation in the reward that was not used for training. The harmonic mean operation (Figures 3(b), 3(e)) results in distributions that are concentrated on the samples scoring high in both rewards. The contrast operation (Figure 3(c)) has the opposite effect: the distributions are skewed towards the examples scoring high in only one of the original rewards. Note that the tails of the contrast distributions are retreating from the area covered by the harmonic mean.
We show reward distribution statistics of three GFlowNets (trained with SEH, SA, and QED at  $\beta = 32$ ) and their compositions in Table 1. Each row of the table gives a percentage breakdown of the samples from a given model into one of  $2^{3} = 8$  bins according to rewards. For all three base models, the majority of the samples fall into the "high" category according to the respective reward, while the rewards that were not used for training show variation. Conditioning on two different labels (e.g.  $y = \{\mathrm{SEH}, \mathrm{QED}\}$ ) results in concentration on examples that score high in two selected rewards, but not necessarily scoring high in the reward that was not selected. The conditional  $y = \{\mathrm{SEH}, \mathrm{QED}, \mathrm{SA}\}$  shifts the focus to examples that have all three properties.
Figure 4 shows 2D embeddings of the distributions appearing in Table 1. The embeddings were computed with t-SNE based on the pairwise earth mover's distances (details in Appendix F.2, complete summary of distribution distances in Table G.5). The configuration of the embeddings provides insight into the relative positions of the base models and conditionals in the distribution space. The points corresponding to the pairwise conditionals lie in between the two base models selected for conditioning. The conditional  $y = \{\mathrm{SEH},\mathrm{SA},\mathrm{QED}\}$  appears to be near the centroid of the triangle  $(p_{\mathrm{SEH}},p_{\mathrm{SA}},p_{\mathrm{QED}})$  and lies close the the pairwise conditionals. The distributions obtained by repeated conditioning on the same label (e.g.  $y = \{\mathrm{SEH},\mathrm{SEH},\mathrm{SEH}\}$ ) are spread out to the boundary, lying closer to the respective base distributions and relatively far from pairwise conditionals.
Colored MNIST generation via diffusion models. Finally, we empirically tested our method for the composition of diffusion models on an image generation task. In this experiment, we composed
Figure 5: Composed diffusion models on colored MNIST. Samples from 3 pre-trained diffusion models and their various compositions.
three diffusion models trained to generate MNIST [56] digits  $\{0,1,2,3\}$  in two colors: cyan and beige. Each model was trained to generate digits with a specific property:  $p_1$  generated cyan digits,  $p_2$  generated digits less than 2, and  $p_3$  generated even digits.
We built the composition iteratively via the factorization  $\widetilde{p}(x|y_1,y_2,y_3) \propto \widetilde{p}(x)\widetilde{p}(y_1,y_2|x)\widetilde{p}(y_3|x,y_1,y_2)$ . To this end, we first trained a classifier  $\widetilde{Q}(y_1,y_2|x_t)$  on trajectories sampled from the base models. This allowed us to generate samples from  $\widetilde{p}(x|y_1,y_2)$ . We then trained an additional classifier  $\widetilde{Q}(y_3|x_t,y_1,y_2)$  on trajectories from compositions defined by  $(y_1,y_2)$  to allow us to sample from  $\widetilde{p}(x|y_1,y_2,y_3)$ . Further details can be found in Appendix F.3.
Figure 5 shows samples from the pre-trained models and from selected compositions. The negating effect of not conditioning on observations is clearly visible in the compositions using two variables. For example,  $\widetilde{p}(x|y_1 = 1, y_2 = 1)$  only generates cyan  $3$ 's. Because there we do not condition on  $p_2$  or  $p_3$ , the composition excludes digits that have high probability under  $p_2$  or  $p_3$ , i.e. those that are less than 2 or even. In  $\widetilde{p}(x|y_1 = 1, y_2 = 3)$ , cyan even digits have high density under both  $p_1$  and  $p_3$ , but because  $p_2$  is not conditioned on, the composition excludes digits less than two (i.e. cyan  $0$ 's). Finally,  $\widetilde{p}(x|y_1 = 1, y_2 = 2, y_3 = 3)$  generates only cyan  $0$ 's, on which all base models have high density.
We introduced Compositional Sculpting, a general approach for composing iterative generative models. Compositions are defined through "observations", which enable us to emphasize or de-emphasize the density of the composition in regions where specific base models have high density. We highlighted two binary compositions, harmonic mean and contrast, which are analogous to the product and negation operations defined on EBMs. A crucial feature of the compositions we have introduced is that we can sample from them directly. By extending classifier guidance we are able to leverage the generative capabilities of the base models to produce samples from the composition. Through empirical experiments, we validated our approach for composing diffusion models and GFlowNets on toy domains, molecular generation, and image generation.
Broader impact. We proposed a mathematical framework and methods for the composition of pre-trained generative models. While the primary emphasis of our work is on advancing foundational research on generative modeling methodology and principled sampling techniques, our work inherits ethical concerns associated with generative models such as creation of deepfake content and misinformation dissemination, as well as reproduction of biases present in the datasets used for model training. If not carefully managed, these models can perpetuate societal biases, exacerbating issues of fairness and equity. Our work further contributes to research on the reuse of pre-trained models. This research direction promotes eco-friendly AI development, with the long-term goal of reducing energy consumption and carbon emissions associated with large-scale generative model training.
TG and TJ acknowledge support from the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, DARPA Accelerated Molecular Discovery program, the NSF Expeditions grant (award 1918839) "Understanding the World Through Code", and from the MIT-DSTA collaboration.
SK and SDP were supported by the Technology Industries of Finland Centennial Foundation and the Jane and Aatos Erkko Foundation under project Interactive Artificial Intelligence for Driving R&D, the Academy of Finland (flagship programme: Finnish Center for Artificial Intelligence, FCAI; grants 328400, 345604 and 341763), and the UKRI Turing AI World-Leading Researcher Fellowship, EP/W002973/1.
VG acknowledges support from the Academy of Finland (grant decision 342077) for "Human-steered next-generation machine learning for reviving drug design", the Saab-WASP initiative (grant 411025), and the Jane and Aatos Erkko Foundation (grant 7001703) for "Biodesign: Use of artificial intelligence in enzyme design for synthetic biology".
GY acknowledges support from the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, http://iaifi.org).
We thank Sammie Katt and Pavel Izmailov for the helpful discussions and assistance in making the figures.
We thank NeurIPS 2023 anonymous reviewers for the helpful feedback on our work.