SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Angelica Chen*
angelica.chen@nyu.edu
David M. Dohan†
david@ddohan.com
David R. So†
Jane Street
david.r.so.ai@gmail.com
Given the recent impressive accomplishments of language models (LMs) for code generation, we explore the use of LMs as general adaptive mutation and crossover operators for an evolutionary neural architecture search (NAS) algorithm. While NAS still proves too difficult a task for LMs to succeed at solely through prompting, we find that the combination of evolutionary prompt engineering with soft prompt-tuning, a method we term EVOPROMPTING, consistently finds diverse and high performing models. We first demonstrate that EVOPROMPTING is effective on the computationally efficient MNIST-1D dataset, where EVOPROMPTING produces convolutional architecture variants that outperform both those designed by human experts and naive few-shot prompting in terms of accuracy and model size. We then apply our method to searching for graph neural networks on the CLRS Algorithmic Reasoning Benchmark, where EVOPROMPTING is able to design novel architectures that outperform current state-of-the-art models on 21 out of 30 algorithmic reasoning tasks while maintaining similar model size. EVOPROMPTING is successful at designing accurate and efficient neural network architectures across a variety of machine learning tasks, while also being general enough for easy adaptation to other tasks beyond neural network design.
Scaling of Transformers (Vaswani et al., 2017) has produced language models (LM) with impressive performance. Beyond achieving state-of-the-art results on conventional natural language processing tasks, these LMs demonstrate breakthrough technical capabilities, such as learning how to code (Chen et al., 2021), doing math (Noorbakhsh et al., 2021), and solving reasoning problems (Wei et al., 2022). Yet, despite these strides, several works have noted LMs' current limitations in solving complex problems and creating novel solutions (Qian et al., 2022; Dakhel et al., 2022). In this work, we improve upon a base LM's ability to propose novel and diverse solutions to complex reasoning problems by iteratively evolving in-context prompts and prompt-tuning the LM. We call this technique EVOPROMPTING and demonstrate its success on the difficult task of deep learning architecture design. Our key finding is that, while LMs perform poorly at designing novel and effective neural architectures via naive few-shot prompting, EVOPROMPTING enables LMs to create novel and effective deep neural architectures, particularly when combined with prompt-tuning methods.
*Work done while a Student Researcher at Google DeepMind.
$^{\dagger}$ Work done while at Google DeepMind.
Figure 1: An overview of EvOPROMPTING. After initializing the search with a handful of manually designed program seeds, the meta-learning loop begins. First, our code-pretrained LM uses the seeds as in-context prompt examples to generate candidate architectures. Those candidate architectures are then trained on the task training data and evaluated on the task validation set. Next, the most fit members of the population are selected as in-context examples for the next meta-learning loop and all evaluated individuals are used as training data for prompt-tuning the LM. From there, the meta-learning loop begins again.
EVOPROMPTING is based on the recently popularized practice of in-context prompting. Prompting is the technique of conditioning a LM's decoded output on a custom prefix known as a prompt, which can include natural language task instructions or a few input-output examples. The prompt is used only at inference time and requires no gradient updates (Brown et al., 2020). In past work, prompting has been demonstrated to elicit impressive performance on a wide variety of tasks without requiring task-specific fine-tuning (Sanh et al., 2021; Wei et al., 2022; Kojima et al., 2022). Here, we leverage LM prompting for the task of designing improved deep learning architectures.
To engineer adequately powerful prompts, we draw inspiration from existing ideas in the field of neural architecture search. There, evolution has long been used to search over discrete spaces to efficiently discover improved deep learning architectures (Yao, 1999; Real et al., 2017). However, evolutionary approaches typically require careful manual design of a discrete search space (e.g. a small set of known convolutional neural network components, as in Real et al. (2017) or TensorFlow primitives, as in So et al. (2021)). As a result, the performance of the evolutionary algorithm is then sensitive to and possibly limited by the design of the search space. In EvOPROMPTING the LM's vocabulary replaces the search space, which both increases the flexibility of the search and reduces reliance on manual design. The LM is also an adaptive mutation/crossover operator, in the sense that it can be improved round over round via prompt-tuning. Furthermore, EvOPROMPTING also improves on naive few-shot prompting by using an evolutionary search approach to iteratively improve the in-context examples for few-shot prompting.
To demonstrate the effectiveness of this method, we first do extensive testing and analyses on the relatively low-compute problem of MNIST-1D (Greydanus, 2020). The key finding of these experiments is that EVOPROMPTING is capable of producing conventional convolutional architectures superior to published manually designed models (Section 4.1). In Section 4.2 we then apply our method to the more challenging task of designing graph neural networks using problems from the CLRS Algorithmic Reasoning Benchmark (Veličković et al., 2022), where EVOPROMPTING generates novel architectures that outperform state-of-the-art models on 21 out of 30 algorithmic reasoning tasks (Appendix 3).
1. We propose EvOPROMPTING, a method that utilizes evolutionary search to create and curate data to improve LM in-context prompting examples. Although this work focuses on the specific task of neural architecture design to develop this method, EvOPROMPTING is generally applicable to LM tasks that rely on in-context learning (ICL) or prompt-tuning.
2. A study applying LMs to code-level neural architecture design. Our experiments demonstrate that applying few-shot prompting alone to neural architecture design is unsuccessful, but few-
shot prompting with EVOPROMPTING enables LMs to create architectures that outperform those designed by human experts.
3. Novel graph neural network architectures that were discovered using EVOPROMPTING. These architectures outperform the current state-of-the-art architecture, TripleGMPNN (Ibarz et al., 2022), on 21 out of 30 CLRS Algorithmic Reasoning Benchmark tasks (Appx. 3).
LMs for code generation Scaling Transformers (Vaswani et al., 2017) is currently a popular route for reliably creating state-of-the-art natural language systems (Brown et al., 2020; Du et al., 2021; BigScience Workshop et al., 2022; Zhang et al., 2022; Thoppilan et al., 2022; Chowdhery et al., 2022). Many works have observed that large LMs are capable of performing technical tasks such as writing code (Chen et al., 2021), doing math (Noorbakhsh et al., 2021), and solving complex reasoning problems (Wei et al., 2022). Our work is most closely related to efforts that have applied LMs to coding tasks (Chen et al., 2021; Odena et al., 2021; Xu et al., 2022; Wang et al., 2021; Ahmad et al., 2021; Feng et al., 2020), since our technique proposes architectures in code.
Prompting Brown et al. (2020) demonstrated that LMs can be prompted with in-context examples to steer LM decoding towards solving problems in-context without gradient updates. Numerous works have utilized this prompting to further boost LM abilities (Sanh et al., 2021; Wei et al., 2022; Kojima et al., 2022). Others have focused on optimizing these prompts (Min et al., 2022; Liu et al., 2021) as via approaches such as augmentation with retrieval systems (Rubin et al., 2021), permutations of few-shot examples (Lu et al., 2021; Zhao et al., 2021), generating prompts via LMs (Zhou et al., 2022), and instruction-tuning (Wei et al., 2021; Ouyang et al., 2022; Sanh et al., 2021). From the perspective of Dohan et al. (2022), prompts are parameters that can be tuned using probabilistic inference techniques. Brooks et al. (2022) proposes using few-shot prompts to implement both the rollout policy and world model of a policy iteration algorithm. Our EVO PROMPTING method extends these efforts by proposing evolutionary search as a means to both better design prompts for ICL and tune the base LM to use the prompt more effectively.
**Evolutionary Algorithms** Our method is closely related to evolutionary neural architecture search (NAS) (Real et al., 2017, 2018; Elsken et al., 2018; So et al., 2019; Liu et al., 2020), in which architectures are represented as discrete DNAs, and evolved and filtered based on fitness metrics that assess architecture performance. However, our method can search over arbitrary strings of code, whereas conventional evolutionary NAS algorithms rely on hand-crafted search spaces that can strongly bias and constrain the search (Li & Talwalkar, 2019; Sciuto et al., 2019; Bender et al., 2020; Real et al., 2020; So et al., 2021). A work close to ours is Lehman et al. (2022), in which an LM is fine-tuned to produce Python code diffs given one of three fixed messages that describe what should be changed, and then used as the mutation operator in an evolutionary algorithm. Their work is validated on the Sedarace domain. Our work differs in that we use an LM as a crossover operator, without specifying the class of changes to make, which may offer greater flexibility. Furthermore, we evaluate our approach on the real-world task of NAS, rely on mixed temperature sampling of the LM for diversity instead of using a QD algorithm, and also use prompt-tuning in our algorithm. We choose not to use a QD algorithm such as MAP-Elites since this approach requires the design and discretization of a descriptor space, which is complex and difficult to hand-design for the space of all possible neural networks.
Another concurrent work is Meyerson et al. (2023), which uses an LM as a crossover operator to produce variations of text-based genotypes in the domains of symbolic regression, text sentiment, images, and Sedaracer programs. Like Lehman et al. (2022), they use MAP-Elites to trade off quality with diversity in two of the domains and demonstrate that their overall algorithm reliably produces a diverse range of outputs. They additionally demonstrated performance comparable to state-of-the-art approaches on the toy task of symbolic regression. Their study varies from ours in a number of ways – we apply our algorithm to the real-world task of NAS, we optimize for a tradeoff between state-of-the-art task performance and model size, we condition on target performance in our prompts, we do not use MAP-Elites, and we use prompt-tuning to iteratively improve the LM's crossover abilities instead.
3 EVOPROMPTING Method
3.1 Architecture search problem formulation
Let our target task be denoted by  $\mathcal{T}$  and  $\mathcal{D}$  be a dataset consisting of input-output pairs  $(x,y)\in \mathcal{D}$  for task  $\mathcal{T}$ . Define the probability distribution  $\pi_{\theta}:\mathcal{V}\to \{0,1\}$  over vocabulary  $\mathcal{V}$  as a language/code model parameterized by  $\theta$ , from which we can sample code segments  $c\in \mathcal{V}^*$  (for  $\mathcal{V}^*$  the Kleene closure of  $\mathcal{V}$ , i.e. the set of all concatenations of symbols in  $\mathcal{V}$ ). We also have an evaluation function  $\mathrm{Eval}_{\mathcal{T}}(c,\mathcal{D}): \mathcal{V}^* \times \mathcal{D} \to \mathbb{R}$  that trains the model architecture given by code  $c$  on  $\mathcal{D}$  and outputs some real-valued fitness score  $s\in \mathbb{R}$ , which can be a function of model accuracy and other model characteristics. Our ultimate goal is to identify some set of code samples  $c\sim \mathcal{V}^*$  that define neural network architectures that, when trained on  $\mathcal{D}$ , maximize the reward  $\mathrm{Eval}_{\mathcal{T}}(c,\mathcal{D})$ .
3.2 LMs for evolutionary crossover and mutation
The goal of our algorithm is to generate a set  $C$  consisting of  $k$  neural network architectures that maximize the reward  $\mathrm{Eval}_{\mathcal{T}}(c,\mathcal{D})$  for arbitrary pairs of  $(\mathcal{D},\mathcal{T})$ :
$$
\arg \max  _ { \begin{array}{c} C = \left\{c \mid c \sim \pi_ {\theta} \right\} \\ | C | = k \end{array} } \mathbb {E} _ {c \in C} \mathbb {E} _ {(x, y) \in \mathcal {D}} \left[ \operatorname {E V A L} _ {\mathcal {T}} (c, \mathcal {D}) \right] \tag {1}
$$
Since this optimization problem is generally intractable, we turn to a black-box evolutionary approach for iteratively generating, scoring, and selecting the best neural network architectures. Indeed, evolution has been demonstrated to perform particularly well in this domain because of how sparse high quality solutions tend to be (Real et al., 2017, 2018). Although evolution has been used for architecture search many times before (Real et al., 2017, 2018; Elsken et al., 2018; So et al., 2019), we improve upon this approach by using an LM for crossover and mutation operations.
Using an LM in this manner has multiple appealing properties. While past evolutionary approaches for neural architecture search have required careful design and specification of a discrete search space (e.g. the space of high level modules (Real et al., 2018; So et al., 2019), TensorFlow statements (So et al., 2021), or basic mathematical operations (Real et al., 2020)), our algorithm's search space includes any neural network architecture that can be represented in Python. This allows for greater flexibility and diversity of the output architectures, and reduces the amount of manual design and human bias involved in the algorithm. Furthermore, modern pre-trained LMs are typically trained on massive datasets containing a significant number of source code files. This pre-training process encodes useful knowledge about code structure and functionality that is not otherwise available in evolutionary algorithms. Lastly, LMs can also be used as self-adaptive crossover operators, in which the crossover operator is incrementally trained round after round to generate higher reward crossovers.
3.3 EVOPROMPTING meta-learning algorithm
Our complete algorithm is described in Algorithm 1. At the core of our algorithm is a scoring function, which describes the general "fitness" of a model on the task at hand. Since higher accuracy can often be achieved simply by increasing the number of parameters in a model, we use the negative product of the validation error and the model size as the fitness (see step 6 in Algorithm 3). More complicated objective functions have previously been used for dual objective neural architecture search (Bender et al., 2020), but we find this simple product works best in our case and requires minimal tuning. Generally the higher the fitness, the better (with some caveats, noted in our description of fitness-based selection below).
The end-to-end meta-learning algorithm has several stages, which we describe below:
Initialization We start by setting our global historical population  $G$  to the empty list and initializing our current population  $P$  with a few seed architectures that are known to be well-designed (step 3 in Algorithm 1), which warm-starts the search (So et al., 2019). These seed models are evaluated using the same  $\mathrm{Eval}_{\mathcal{T}}(c, \mathcal{D})$  function that is used to evaluate new candidate models (see below).
Algorithm 1 Complete meta-learning evolutionary algorithm using  $p_{\theta}$  as a crossover and mutation operator.
1: Input: LM  $\pi_{\theta_0}$ , dataset  $\mathcal{D}$ , task  $\mathcal{T}$ ,  $T$  number of rounds,  $m$  number of few-shot prompts per round,  $n$  number of samples to generate per prompt,  $k$  number of in-context examples per prompt,  $p$  number of survivors to select per generation,  $\alpha$  the upper threshold for the test error  
2:  $G \gets []$   
3:  $P \gets \text{INITIALIZEPOPULATION}(p)$   
4:  $t \gets 0$   
5: while  $t < T$  do  
6:  $C \gets \text{CROSSPUT}(\pi_{\theta_t}, P, m, k, n)$   
7:  $C_{\text{EVALED}} \gets \text{FILTERANDEVAL}(C, \mathcal{T}, \mathcal{D}, \alpha)$   
8:  $G \gets G + C_{\text{EVALED}}$   
9: if  $t < T - 1$  then  
10:  $P \gets \text{GETTOP}(G, p)$   
11:  $\theta_{t+1} \gets \text{TRAIN}(\theta_t, C_{\text{EVALED}} \setminus P)$   
12: end if  
13:  $t \gets t + 1$   
14: end while  
15: Return GETTOP $(G, p)$
Algorithm 2 The crossover and mutation algorithm, CROSSmut  $(\pi_{\theta_t}, P, m, k, n)$ , where  $\mathrm{Uniform}(P)$  denotes the uniform distribution over the set  $P$ . The set of potential parents  $P$  consists of the top examples from the previous round.
1: Input: LM  $\pi_{\theta}$ , population of code samples and fitnesses  $P = \{(c,s) | c \in \mathcal{V}^*$ , EVAL $\tau(c,\mathcal{D}) = s\}$ ,  $m$  number of few-shot prompts to create,  $k$  number of in-context examples in each prompt, and  $n$  number of samples to sample per prompt.  
2:  $C \gets []$   
3:  $i \gets 0$   
4: while  $i < m$  do  
5:  $E \gets \{x_j\}_{j=1}^k$ , where  $x_j \stackrel{i.i.d.}{\sim} \text{Uniform}(P)$   
6:  $p \gets \text{MAKEFEWSHOTPROMPT}(E)$   
7:  $C_i \gets \{c_j\}_{j=1}^n$ , where  $c_j \stackrel{i.i.d.}{\sim} \pi_{\theta}(\cdot|p)$   
8:  $C \gets C + C_i$   
9:  $i \gets i + 1$   
10: end while  
11: Output:  $C$
Algorithm 3 The algorithm for filtering and scoring child models, FILTERANDEVAL  $(C,\mathcal{T},\mathcal{D},\alpha)$
1: Input: set of code samples  $C$ , task  $\mathcal{T}$ , dataset  $\mathcal{D}$ , evaluation function  $\mathrm{Eval}_{\mathcal{T}}(c,\mathcal{D})$ , upper threshold for error  $\alpha$   
2:  $C_{\mathrm{EVALED}} \gets []$   
3: for c in  $C$  do  
4: c(error)  $\leftarrow \mathrm{Eval}_{\mathcal{T}}(c,\mathcal{D})$   
5: if c.error <  $\alpha$  then  
6: s  $\leftarrow -c.$  model_size  $\times c.$  error  
7:  $C_{\mathrm{EVALED}} \gets C_{\mathrm{EVALED}} + [(c,s)]$   
8: end if  
9: end for  
10: Output:  $C_{\mathrm{EVALED}}$
Crossing over and mutating the parent models To mutate and apply crossover to the parents  $P$  selected in the last step, we use both the source code and the evaluation metrics of each model in  $P$  to create few-shot prompts.
In the last line of the prompt, we create a target set of metrics to condition  $\pi_{\theta}$ 's generations on that indicate the desired validation accuracy and model size of the proposed architecture. We set the target
model size as  $90\%$  of the minimum model size of the parent models, rounded to the nearest 100 parameters, and the target validation accuracy as  $102\%$  of the maximum validation accuracy of the parent models, rounded to the nearest tenth of a percent. We create  $m$  such prompts per round, each with  $k$  in-context examples selected uniformly at random from  $P$ . An example of a prompt is shown in Listing 1.
1 ""Metrics:   
2 {'num_parameters': '4800', 'val.accuracy': '0.865'}   
3   
4 class Model(nnModule):   
5 @nn.compact   
6 def __call__(self, x):   
7 x = nn(Dense features=10)(x)   
8 return x   
9   
10 ""Metrics:   
11 {'num_parameters': '4300', 'val.accuracy': '0.880'}   
12   
13 class Model(nnModule):
Listing 1: The format of our few-shot prompts. In practice we use 2-shot prompts but we omit the second in-context example here for brevity.
Finally, we use  $\pi_{\theta}$  to generate  $n$  samples per prompt, yielding a total of  $n\times m$  child samples per round of evolution. We denote this portion of the algorithm as CROSSMut  $(\pi_{\theta_t},P,m,k,n)$  (Algorithm 2 and step 6 of Algorithm 1).
Filtering and scoring child samples To score and filter child samples  $c$  generated by  $\pi_{\theta}$ , we use the evaluation function  $\mathrm{Eval}_{\mathcal{T}}(c,\mathcal{D})$ , which trains the model encoded by  $c$  on the dataset  $\mathcal{D}$  and returns the lowest validation error encountered during training. All child models are trained for the same number of steps, with the same optimizer hyperparameters. Since our fitness function can potentially be gamed by generating arbitrarily small models, we also add a validation error threshold  $\alpha$ , which is the upper limit of the validation error that a model can incur without being removed from  $G$ , the global population. We refer to this function as FILTERANDEVAL  $(C,\mathcal{T},\mathcal{D},\alpha)$  (Algorithm 3 and step 7 of Algorithm 1). Lastly, we add the remaining trainable models and their associated fitness scores into  $G$  (step 8 of Algorithm 1).
Fitness-based selection After evaluating all child models in the current round, we apply fitness-based selection to identify top candidate models for crossover (step 10 of Algorithm 1). We denote this as  $\mathrm{GETTOP}(G,p)$ , which refers simply to selecting the  $p$  models with the highest fitness scores from  $G$ . Once these models have been selected, they are permanently removed from the population and cannot be used again as parents for crossover.
Training  $\pi_{\theta_t}$  Lastly, all child models generated in the current round that were not previously selected for crossover (i.e.  $C_{\mathrm{EVALED}} \setminus P$ ) are used to prompt-tune  $\pi_{\theta}$  for the next round (step 11 of Algorithm 1).
We evaluate our meta-learning algorithm on two datasets – MNIST-1D (Greydanus, 2020) and the CLRS algorithmic reasoning benchmark (Veličković et al., 2022). While the former benchmark is lightweight and permits us to do a more thorough analysis of our algorithm, the latter is a newer benchmark that covers 30 different algorithms with more headroom for discovering novel architectures with better performance.
In all of our experiments, our  $\pi_{\theta_0}$  (i.e. the crossover operator) is a 62B parameter PALM model (Chowdhery et al., 2022) pre-trained on 1.3T tokens of conversational, web, and code documents. It was additionally fine-tuned on a corpus of 64B tokens containing near-deduplicated, permissively-licensed Python source code files from Github. We always sample from  $\pi_{\theta_0}$  with
mixed temperature sampling, in which the sampling temperature is selected uniformly from  $[0.2, 0.6, 0.8, 1.0]$ . Between each round, the model is prompt-tuned (Lester et al., 2021) for 5 epochs with a soft prompt length of 16, batch size of 16, and learning rate of 0.1 (as described in Section 3.3 and Step 11 of Algorithm 1). Unless stated otherwise, we run 10 rounds of evolution with 10 prompts per round and 16 samples generated per prompt, yielding a total of 160 models generated per round and 1600 models generated during the entire search. Duplicate models and un-trainable models are not scored, but do count into the 1600. All other EVOPROMPTING hyperparameters are listed in Appendix A.1.
4.1 MNIST-1D
Dataset We apply our method first to MNIST-1D (Greydanus, 2020), a one-dimensional, scaled-down version of the MNIST-1D dataset containing examples that are 20 times smaller than the original MNIST dataset. Each example is only 40-dimensional, with 4000 examples in the training dataset and 1000 in test. Since there is no validation dataset, we randomly set aside 500 examples from the training dataset to use as the validation dataset. Despite being more lightweight, MNIST-1D distinguishes more between different architecture types (Greydanus, 2020) than its larger counterpart MNIST (LeCun et al., 1998).
Meta-learning set-up Throughout the model search we use the AdamW optimizer (Loshchilov & Hutter, 2019) to train each child model on a single NVIDIA Tesla P100 GPU for 8000 steps, with learning rate 0.01 and batch size 128. We score child models according to the best validation accuracy achieved during training. We also seed the search with 4 seed models - the 3 hand-designed neural baselines from the original MNIST-1D paper (Greydanus, 2020) (GRU, CNN, and MLP) and a fourth, larger CNN model of our own design. All four are implemented with Flax (Heek et al., 2020). We refer the reader to Appendix A.2 for the source code of these seed models.
Baselines We compare EVOPROMPTING with the following baselines: