SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tung Nguyen, Sudhanshu Agrawal, Aditya Grover
{tungnd,adityag}@cs.ucla.edu,sudhanshuagr27@g.ucla.edu
Experimental design for optimizing black-box functions is a fundamental problem in many science and engineering fields. In this problem, sample efficiency is crucial due to the time, money, and safety costs of real-world design evaluations. Existing approaches either rely on active data collection or access to large, labeled datasets of past experiments, making them impractical in many real-world scenarios. In this work, we address the more challenging yet realistic setting of few-shot experimental design, where only a few labeled data points of input designs and their corresponding values are available. We introduce Experiment Pretrained Transformers (ExPT), a foundation model for few-shot experimental design that combines unsupervised learning and in-context pretraining. In ExPT, we only assume knowledge of a finite collection of unlabelled data points from the input domain and pretrain a transformer neural network to optimize diverse synthetic functions defined over this domain. Unsupervised pretraining allows ExPT to adapt to any design task at test time in an in-context fashion by conditioning on a few labeled data points from the target task and generating the candidate optima. We evaluate ExPT on few-shot experimental design in challenging domains and demonstrate its superior generality and performance compared to existing methods. The source code is available at https://github.com/tung-nd/ExPT.git.
The design of experiments to optimize downstream target objectives is a ubiquitous challenge across many science and engineering domains, including materials discovery [27], protein engineering [7, 49, 2], molecular [22] design, mechanical design [4, 38], and neural architecture optimization [66]. The key criterion of interest in experimental design (ED) is sample-efficiency, as the target objectives are often black-box functions and evaluating these objectives for any candidate design often involves expensive real-world experiments. A standard class of approaches learn a surrogate to approximate the target objective and actively improve the approximation quality through online experiments [53]. However, online data acquisition may be infeasible in the real world due to high costs, time constraints, or safety concerns. As an alternate, recent works have proposed offline ED [7, 58, 36, 57, 12, 34], wherein a model learns to perform optimization from a fixed dataset of past experiments. While this is more practical than the online setting, current offline methods and benchmarks assume access to large experimental datasets containing thousands of data points, which are hard or even impossible to obtain in high-stake and emerging science problems. Even when these datasets exist, the past experiments might be of very poor quality resulting in poor surrogate learning and optimization.
In this paper, we aim to overcome these limitations for hyper-efficient experimental design that does not require large experimental datasets. To this end, we introduce few-shot experimental design, a more challenging setting that better resembles real-world scenarios. We describe few-shot ED as a two-phased pretraining-adaptation paradigm. In the pretraining phase, we only assume access to
Figure 1: Experiment Pretrained Transformers (ExPT) follow a pretraining-adaptation approach for few-shot experimental design. During pretraining (left), the model has access to unlabeled designs from domain  $\mathcal{X}$  without their corresponding scores. For adaptation (right), the model conditions on a small set of (design, score) pairs and the desired score  $y^{\star}$  to generate the optimal design  $x^{\star}$ .
unlabeled data, i.e., input designs without associated function values. During the adaptation phase, we have access to a few labeled examples of past experiments to adapt the model to the downstream task. This setup offers several advantages. First, it alleviates the requirement for costly annotated data and relies mainly on unlabeled inputs that are easily accessible. Second, unsupervised pretraining enables us to utilize the same pretrained backbone for adapting to multiple downstream optimization tasks within the same domain. For example, in molecule design, one may want to optimize for multiple properties, including drug-likeness, synthesizeability, or similarity to target molecules [8, 17].
The key question in this setup is how to make use of the unlabeled data to facilitate efficient generalization to downstream tasks during optimization. Our intuition here is that, while the objective function is unknown, we can use the unlabeled inputs to generate pretraining data from other synthetic functions. If a model can few-shot learn from a diverse and challenging set of functions, it should be able to generalize quickly to any target objective during the adaptation phase, in line with recent foundation models for language [9] and vision [3]. This insight gives rise to our idea of synthetic pretraining, wherein we pretrain the model on data generated from a rich family of synthetic functions that operate on the same domain as the target task. Specifically, for each function drawn from this family, we sample a set of points by using the unlabeled data as inputs. We divide these points into a small context set and a target set, and train the model via in-context learning to perform conditional generation of the target input  $x$  given the context points and the target value  $y$ . A model that works well on this task should be able to efficiently capture the structures of the underlying function, i.e., how different regions of the input space influence the function value, from a small context set. By explicitly training the model to perform this task on a diverse set of functions, the model can generalize efficiently to downstream functions during adaptation requiring only limited supervision. After pretraining, we can perform optimization by conditioning the model on a few labeled examples from the downstream task and generating an input that achieves the optimum  $y^{\star}$ .
Inspired by recent advances in few-shot learning in language [9,42] and other domains [44,45,19,37], we instantiate a novel foundation model with a transformer-based architecture [59], which we call Experiment Pretrained Transformers (ExPT). ExPT is an encoder-decoder architecture, in which the encoder is a transformer [59] network that encodes the context points and the target value, and the decoder is a VAE [32] model that predicts the high-dimensional target input. The transformer encoder allows ExPT to perform few-shot generation and optimization purely through in-context learning in a gradient-free fashion. We compare the performance of ExPT and various baselines on 2 few-shot settings created from Design-Bench [58], a standard database benchmark for ED. The two settings allow us to examine how different methods perform with respect to different quantities and qualities of few-shot data. In both these settings, results show that ExPT achieves the highest average score and the highest average ranking with respect to median performance, mean performance, and best-achieved performance. Especially in the more challenging setting, ExPT outperforms the second-best method by  $70\%$  in terms of the mean performance. Additionally, we explore the potential of using the same pretrained ExPT for multiple objectives, and conduct extensive ablation studies to validate the effectiveness of our design choices for synthetic data generation and ExPT architecture.
2 Experiment Pretrained Transformers
Let  $f: \mathcal{X} \to \mathbb{R}$  be a black-box function that operates on a  $d$ -dimensional domain  $\mathcal{X} \subseteq \mathbb{R}^d$ . In experimental design (ED), the goal is to find the input  $x^\star$  that maximizes  $f$ :
$$
x ^ {\star} \in \underset {x \in \mathcal {X}} {\arg \max } f (x). \tag {1}
$$
Typically,  $f$  is a high-dimensional and complex function that often involves expensive physical experiments. Existing approaches either assume the ability to actively query  $f$  to collect data [53] or access to a large dataset of past experiments [58]. Both assumptions are too strong in many real-world applications where data collection is hard or even impossible [11]. Therefore, we propose few-shot ED, a more challenging yet realistic setting to overcome these limitations. In few-shot ED, the goal is to optimize any objective function in the domain  $\mathcal{X}$  given only a handful of examples. We approach this problem with a pretraining-adaptation pipeline. In the pretraining phase, we assume access to an unlabeled dataset  $\mathcal{D}_{\text{unlabeled}} = \{x_i\}_{i=1}^{|D_{\text{unlabeled}}|}$  from the optimization domain  $\mathcal{X} \subseteq \mathbb{R}^d$ . We note that  $\mathcal{D}_{\text{unlabeled}}$  only contains potential design inputs without their corresponding scores, for example, potential molecules in molecule optimization or different combinations of hyperparameters in neural architecture search. This means the objective function  $f$  is unspecified during pretraining.
During the adaptation phase, one can use the pretrained model to optimize any objective function  $f$  in the same domain  $\mathcal{X}$ . We now have access to a few-shot labeled dataset that the model can use to adapt to the downstream function  $\mathcal{D}_{\mathrm{few - shot}} = \{(x_1,y_1),\ldots ,(x_n,y_n)\}$ , in which  $y_{i} = f(x_{i})$  and  $n = |\mathcal{D}_{\mathrm{few - shot}}|\ll |\mathcal{D}_{\mathrm{unlabeled}}|$ . After adaptation, we evaluate a few-shot optimization method by allowing it to propose  $Q$  input  $x^{\prime}s$  and query their scores using the black-box function  $f$ , where  $Q$  is often called the optimization budget [58, 36, 57, 12, 34]. The performance of a black-box optimizer is then measured by computing the max, median, and mean of the  $Q$  evaluations. This setup provides two key benefits. First, it resembles many real-world scenarios, where the potential design inputs are cheap and easy to obtain while their target function values are expensive to evaluate. For example, in molecular optimization, we have databases of millions of molecules [30, 5, 48] but only the properties of a handful are known [28, 47, 23]. Second, unsupervised pretraining allows us to train a general backbone that we can adapt to multiple optimization tasks in the same domain.
2.2 Synthetic Pretraining and Inverse Modeling for Scalable Experimental Design
Intuitively, the adaptation phase in §2.1 resembles a few-shot learning problem, in which a model is tasked to produce the optimal input  $x^{\star}$  by conditioning on a few labeled examples in  $D_{\mathrm{few-shot}}$ . To perform well in this task, a model has to efficiently capture the structure of a high-dimension function  $f$ , i.e., what regions of the function lead to higher values and vice versa, from very few examples in  $D_{\mathrm{few-shot}}$ . Given this perspective, the question now is how to make use of the unlabeled dataset  $D_{\mathrm{unlabeled}}$  to pretrain a model that achieves such efficient generalization to the objective function  $f$ . Our key insight is, if a model learns to perform in-context learning on a diverse and challenging set of functions, it should be able to adapt quickly to any objective function at test time. While the function values are unknown during pretraining, we can use the unlabeled inputs  $x^{\prime}s$  to generate pretraining data from other functions. This gives rise to our idea of synthetic pretraining, wherein we pretrain the model to perform few-shot learning on a family of synthetic functions  $\tilde{F}$  that operate on the same input domain  $\mathcal{X}$  of the objective  $f$ . We discuss in detail our mechanism for synthetic data generation in Section 2.3. For each function  $\tilde{f}$  generated from  $\tilde{F}$ , we sample a set of function evaluations  $\{(x_i,y_i)\}_{i = 1}^N$  that we divide into a small context set  $\{(x_i,y_i)\}_{i = 1}^m$  and a target set  $\{(x_j,y_j)\}_{j = m + 1}^N$ . We train the model to predict the target points conditioning on the context set.
There are two different approaches to pretraining a model on this synthetic data. The first possible approach is forward modeling, where the model is trained to predict the target outputs  $y_{m + 1:N}$  given the context points and the target inputs  $x_{m + 1:N}$ . This is similar to the approach followed by TNPs [45], a model recently proposed in the context of meta-learning. During adaptation, we can condition the model on the labeled examples in  $\mathcal{D}_{\mathrm{few-shot}}$  and perform gradient ascent updates to improve an existing design input  $x_{t}$ . However, as commonly observed in previous works [58, 36, 57], this approach is susceptible to producing highly suboptimal inputs. This is because performing gradient ascent with respect to an imperfect forward model may result in points that have high values
Figure 2: The pretraining-adaptation phases for ExPT. We sample synthetic data from  $\tilde{F}$  and pretrain the model to maximize  $\log p(x_{m+1:N} \mid x_{1:m}, y_{1:m}, y_{m+1:N})$ . At adaptation, the model conditions on  $\mathcal{D}_{\mathrm{few-shot}}$  and  $y^{\star}$  to generate candidates. ExPT employs a transformer encoder that encodes the context points and target outputs and a relevant decoder that predicts the target inputs.
under the model but are poor when evaluated using the real function. Instead, we propose to perform inverse modeling, where the model learns to predict the inputs  $x_{m+1:N}$  given the output values  $y_{m+1:N}$  and the context points. As the model learns to directly generate input  $x' s$ , it is less vulnerable to the aforementioned problem. Another advantage of inverse modeling is after pretraining, we can simply condition on  $\mathcal{D}_{\mathrm{few-shot}}$  and the optimal value  $y^{\star}$  to generate the candidate optima. Our loss function for pretraining the model is:
$$
\begin{array}{l} \theta = \underset {\theta} {\arg \max } \mathbb {E} _ {\tilde {f} \sim \tilde {F}, x _ {1: N} \sim \mathcal {D} _ {\mathrm {u n l a b e l e d}}, y _ {1: N} = \tilde {f} (x _ {1: N})} \left[ \log p (x _ {m + 1: N} \mid x _ {1: m}, y _ {1: m}, y _ {m + 1: N}) \right] \\ = \arg \max  _ {\theta} \mathbb {E} _ {\tilde {f} \sim \tilde {F}, x _ {1: N} \sim \mathcal {D} _ {\text {u n l a b e l e d}}, y _ {1: N} = \tilde {f} \left(x _ {1: N}\right)} \left[ \sum_ {i = m + 1} ^ {N} \log p \left(x _ {i} \mid x _ {1: m}, y _ {1: m}, y _ {i}\right) \right], \tag {2} \\ \end{array}
$$
where we assume the target points are independent given the context set and the target output. Figure 2 illustrates the proposed pretraining and adaptation pipeline. Typically, we use a small context size  $m$  during pretraining to resemble the test scenario.
After pretraining, ExPT can adapt to any objective  $f$  in the domain in a gradient-free fashion. Samples in the few-shot dataset  $\mathcal{D}_{\mathrm{few - shot}}$  become the context points and the model conditions on only one target  $y^{\star}$ , which is the optimal value of  $f$ , to generate candidate optima. Note that we only assume the knowledge of  $y^{\star}$  and not  $x^{\star}$ . This assumption is common in many prior works [34, 46, 13, 14]. In practice,  $y^{\star}$  might be known based on domain knowledge. For example, in molecule design, there are physical limits on the value of certain properties such as relaxed energy, in robot applications, the optimal performance can be computed from the cost function, and in neural architecture search, we know the theoretical limits on the highest possible accuracy for classifiers.
Next, we present the details of synthetic data generation and our proposed model architecture, the two components that constitute our proposed foundation model, which we refer to as Experiment Pretrained Transformers (ExPT).
2.3 Data generation
We need a family of functions to generate synthetic data for pretraining ExPT. A good family of functions should be easy to sample from and should be capable of producing diverse functions. Many possible candidates exist for synthetic function families, such as Gaussian Processes (GPs), randomly constructed Gaussian Mixture Models, or randomly initialized or pretrained neural networks. Among these candidates, we choose to generate synthetic data from Gaussian Processes with an RBF kernel. This is for several reasons. First, they are a natural choice as they represent distributions over functions. Second, it is easy and cheap to sample data from prior GPs. And third, a GP with an RBF
kernel is a universal approximator to any function [41]. Specifically,  $\tilde{f}$  is sampled as follows,
$$
\tilde {f} \sim \mathcal {G P} (0, \mathcal {K}), \quad \mathcal {K} (x, x ^ {\prime}) = \sigma^ {2} \exp \left(- \frac {(x - x ^ {\prime}) ^ {2}}{2 \ell^ {2}}\right), \tag {3}
$$
in which  $\sigma$  and  $\ell$  are the two hyperparameters of the RBF kernel. The variance  $\sigma$  scales the magnitudes of the covariance matrix. A larger variance results in a wider range of function values, while a smaller variance restricts the function values to a narrower range. On the other hand, the length scale  $\ell$  determines how strongly the covariance matrix varies with respect to the distance between  $x$  and  $x'$ . A very small length scale  $\ell$  means the kernel is sensitive to the distance between  $x$  and  $x'$ , leading to sharp transitions between neighboring points and a lack of smoothness in the functions. In contrast, if  $\ell$  is too large, the covariance between points will be similar for both near and distant points, leading to function values that are very similar. In other words, too large a length scale reduces the diversity of the synthetic functions. In practice, we randomize both  $\sigma$  and  $\ell$  to increase the diversity of the pretraining data. Appendix C.1 demonstrates the empirical importance of these hyperparameters.
2.4 Model architecture
To optimize the loss function in Equation (2), we need a model architecture that can condition on a few examples drawn from an underlying function to make predictions for other points. This resembles the idea of in-context learning that has proven very successful in language [9, 42] and other domains [44, 45, 19, 37]. The key to success in these works is a transformer architecture that performs in-context learning efficiently via the attention mechanism [59]. Inspired by this, we instantiate ExPT with a transformer-based architecture. Figure 2 illustrates the ExPT overall architecture. Specifically, ExPT employs a transformer encoder that encodes the context points  $\{(x_i, y_i)\}_{i=1}^m$  and the target inputs  $y_{m+1:N}$ , and outputs hidden vectors  $h_{m+1:N}$ . To inform the model that  $x_i$  and  $y_i$  are the input and the corresponding output from  $\tilde{f}$ , we concatenate them to form a token. This results in the sequence  $\{(y_1, x_1), \ldots, (y_m, x_m), y_{m+1}, \ldots, y_N\}$ . We then embed these tokens using two 1-layer MLP networks, one for the pairs and one for the target inputs, before feeding the sequence to the transformer layers. We implement a masking mechanism that prevents the context points from attending the target points, as they do not contain information about the underlying function  $\tilde{f}$ .
Each hidden vector  $h_i$  output by the transformer encoder encompasses the information of the context points and the target input  $y_i$ . Therefore, given  $h_i$ , the conditional probability  $p_{\theta}(x_i \mid x_{1:m}, y_{1:m}, y_i)$  reduces to  $p_{\theta}(x_i \mid h_i)$ . As  $x_i$  is high-dimensional, we can utilize existing generative models to model the conditional distribution  $p(x_i \mid h_i)$ . In this work, we train a conditional VAE model [32] alongside the transformer encoder because of its training stability, light hyperparameter tuning, and good empirical performance. For discrete tasks, we follow the same procedure as Trabucco et al. [58] that emulates logit values by interpolating between a uniform distribution and the one hot values. We train the entire model by maximizing the lower bound of the conditional likelihood  $\log p(x_i \mid h_i)$ :
$$
\log p _ {\theta} \left(x _ {i} \mid h _ {i}\right) \geq \mathbb {E} _ {q _ {\phi} \left(z \mid x _ {i}, h _ {i}\right)} \left[ \log p _ {\theta} \left(x _ {i} \mid z, h _ {i}\right) \right] - \mathrm {K L} \left(q _ {\phi} \left(z \mid x _ {i}, h _ {i}\right) | | p (z)\right), \tag {4}
$$
in which  $q_{\phi}(z \mid x_i, h_i)$  is the encoder of the conditional VAE and  $p(z)$  is a standard Gaussian prior.
3.1 Synthetic experiments
We first evaluate the performance of ExPT in a synthetic experiment, where we train ExPT on data generated from Gaussian Processes (GPs) with an RBF kernel and test the model on four out-of-distribution functions drawn from four different kernels: Matern, Linear, Cosine, and Periodic. Figure 3 shows the performance of ExPT on four test functions through the course of training. The model performs well on all four functions, achieving scores that are much higher than the max value in the few-shot dataset, and approaching the true optimal value, even for kernels that are significantly different from RBF like Cosine and Periodic. Moreover, the performance improves consistently as we pretrain, showing that pretraining facilitates generalization to out-of-distribution functions with very few examples. See Appendix A for a detailed setup of this experiment.
Figure 3: The performance of ExPT on 4 out-of-distribution synthetic tasks through the pretraining phase. We average the performance across 3 seeds.
3.2 Design-Bench experiments
Tasks We consider 4 tasks from Design-Bench [58]. D'Kitty and Ant are continuous tasks with input dimensions of 56 and 60, respectively. In D'kitty and Ant, the goal is to optimize the morphological structure of two simulated robots, Ant [6] to run as fast as possible, and D'kitty [1] to reach a fixed target location. TF Bind 8 and TF Bind 10 are two discrete tasks, where the goal is to find the length-8 and length-10 DNA sequence that has a maximum binding affinity with the SIX6_REF_R1 transcription factor. The design space in these two tasks consists of sequences of one of four categorical variables, corresponding to four types of nucleotide. For each task, Design-Bench provides a public dataset, a larger hidden dataset which is used to normalize the scores, and an oracle. We have an exact oracle to evaluate the proposed designs in all 4 tasks we consider.
Few-shot settings We create 2 few-shot settings from the above tasks, which we call random and poorest. In random, we randomly subsample  $1\%$  of data points in the public set of each task as the few-shot dataset  $\mathcal{D}_{\mathrm{few - shot}}$ . The poorest setting is more challenging, where we use  $1\%$  of the data points which have the lowest scores. The two settings examine how sensitive different methods are to the quantity and quality of the data. In both settings, we use  $x^{\prime}s$  in the public dataset as  $\mathcal{D}_{\mathrm{unlabeled}}$ .