SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Naive few-shot prompting: This baseline simply generates code samples  $c \sim \pi_{\theta_0}(\cdot | p)$ , where  $p$  is a 2-shot prompt constructed using in-context examples randomly selected from the seed models (Listing 1). This is essentially an ablation of steps 7-12 in Algorithm 1 with  $T = 1$ . We increase the number of samples generated per prompt for the naive prompting baseline such that the total number of samples generated by  $\pi_{\theta}$  matches that of the other baselines.
- EVOPROMPTING (- prompt-tuning): We run the entire algorithm as is, but without prompt-tuning between each round. This is an ablation of step 11 from Algorithm 1
- EVOPROMPTING (random parents): Instead of selecting the most fit models from the last round as parents for the next round, we select parents randomly. This is an ablation of Step 10 in Algorithm 1, which is the GETTOP  $(G,p)$  step.
EVOPROMPTING finds smaller and more accurate models Figure 2a shows a comparison of the test error and model size of the top 20 models discovered by EVOPROMPTING compared with those of our seed models and three baselines. The points approximate a Pareto frontier, below which each algorithm cannot improve on one dimension without hurting the other. EVOPROMPTING possesses the Pareto frontier closest to the origin, indicating that it finds more optimal models in terms of accuracy and size. In fact, many models in EVOPROMPTING's top 20 discovered models are orders of magnitude smaller than those of the other baselines, while still having lower test error.
We also note that – on this task in particular – EVOPROMPTING excels especially at optimizing convolutional architectures. Many of the top 20 models are narrower and deeper convolutional architectures, with smaller strides, less padding, and no dense layers. These models consistently perform better than the shallower, denser, and wider convolutional architectures seen in earlier rounds of the model search.
Another important aspect of a meta-learning algorithm is the relationship between the number of individuals evaluated and the maximum fitness observed so far, i.e. the sample efficiency. Neural architecture search can be an expensive process, with the most open-ended searches requiring the evaluation of trillions of individuals (Real et al., 2020). Thus, it is crucial to identify fit candidates
(a) Pareto frontiers of the model size versus test error of the top 20 experiments for each variation of the MNIST1D model search. Frontiers closer to the origin are considered more desirable.
(b) Number of child models generated versus maximum fitness in sample, as estimated using 100 bootstrap samples of size 20 for each point along the x-axis.
Figure 2: EVOPROMPTING discovers smaller and better performing architectures on MNIST-1D than alternative search methods.
Figure 3: Number of child models generated versus maximum fitness of top model seen so far (as estimated using 100 bootstrap samples of size 20 for each point along the x-axis) when searching over neural network models for three CLRS tasks. As mentioned in Section 4.2, these algorithms were selected because our preliminary analyses indicated that they had the most headroom for architectural improvements.
using as few samples as possible. Figure 2b compares how the fitness of the best-performing child model improves as a function of the number of child samples generated thus far. The random parents baseline plateaus the quickest, reaching a maximum fitness by the time approximately 200 individuals have been generated. Furthermore, the maximum fitness it reaches is significantly worse than that of the other experiments. On the other hand, EVOPROMPTING without prompt-tuning and normal EVOPROMPTING do not plateau until much later on. EVOPROMPTING's plateau is the highest and therefore fitter on average than the individuals discovered by any of the other experiments.
It is also evident from both Figure 2a and 2b that performance suffers when any individual component is removed. Interestingly, Figure 2a indicates that prompting with randomly selected parents combined with prompt-tuning is no more effective than naive prompting alone. This highlights the importance of selecting helpful in-context examples, particularly in a task for which we assume that less training signal exists in the pre-training data. However, selecting more fit models as in-context examples without prompt-tuning also does not perform nearly as well as our full method.
Trajectory over meta-learning rounds We also explored the trajectory of our meta-learning algorithm round over round, as shown in Appendix A.3. In general, we observe that EVOPROMPTING starts out further away from the origin (in round 0) and ends up closest to the origin in round 10, which signifies that it discovers – on average – the smallest and most accurate models in the last round. However, the search does not always yield improvements on both axes between consecutive rounds. In rounds 0-2 and 6-10, EVOPROMPTING improves test error while trading off model size. On the other hand, both dimensions are simultaneously improved upon in rounds 3-5.
4.2 CLRS
Although the MNIST-1D task offers an efficient and practical setting for evaluating a meta-learning algorithm, CNN architectures already perform fairly well on this task and neural image classification architectures have been extensively studied as a whole. There also exists the possibility that our LM has seen many convolutional architectures in its pre-training data. Instead, we turn to a different learning task and class of neural network architectures in order to assess whether our meta-learning framework generalizes to other tasks, datasets, and neural architectures.
Dataset The CLRS algorithmic reasoning benchmark (Veličković et al., 2022) evaluates the ability of neural networks to learn algorithmic reasoning across a set of 30 classical algorithms covered in the Introduction to Algorithms textbook by Cormen, Leiserson, Rivest and Stein (Cormen et al., 2009). This benchmark is useful not only as a difficult logical reasoning task for neural networks, but also as a measure of a neural network's algorithmic alignment (Xu et al., 2020). In brief, algorithmic alignment refers to a model's ability to reason like an algorithm (i.e. using the computation graph for a task), rather than relying upon memorization or other less sample efficient learning strategies. Although a model can approximate an algorithm by pattern-matching against similar inputs or relying on other shortcuts, it cannot generalize to arbitrarily long inputs or edge cases without learning the computation graph underlying the algorithm.
Accordingly, the CLRS benchmark represents the algorithms' inputs and outputs as graphs, and the steps of the algorithm as a trajectory of operations over the input graph. This problem setup can be straightforwardly processed by graph neural networks, which is explored in Ibarz et al. (2022). They find that a Triplet-GMPNN model (a message-passing neural network (Gilmer et al., 2017) with gating and triplet edge processing) exhibits the best performance when trained and evaluated across all 30 algorithms at once.
Table 1: A comparison of OOD accuracy and model size (in number of parameters) of models newly discovered by EVOPROMPTING on select CLRS tasks where EVOPROMPTING has discovered more accurate architectures without large increases in model size, compared with the baseline model (the Triplet-GMPNN from Ibarz et al. (2022)). OOD accuracy numbers for the baseline model are from Ibarz et al. (2022). For the full table of results on all CLRS tasks, including accuracies of our own implementation of the Triplet-GMPNN, see Appendix 3.
CLRS Task	Best Performing Model	Model Size ↓	OOD Accuracy ↑
Ours	Baseline	Ours	Baseline
Articulation Points	QUADNODEMINMAX	497969	531913	93.5 ± 1.8%	88.3± 2.0%
BFS	MAXMEAN	522931	523963	100.0 ± 0.0%	99.7± 0.0%
Bubble Sort	CONCATREP	568533	524477	88.9 ± 2.8%	67.7± 5.5%
DFS	DIV2MEAN	660158	661190	68.1 ± 1.4%	47.8± 4.2%
Floyd Warshall	CONCATREP	669145	625089	61.4 ± 0.8%	48.5± 1.0%
Heapsort	CONCATREP	703710	659654	69.9 ± 4.2%	31.0± 5.8%
Insertion Sort	DIV2MEAN	523445	524477	89.5 ± 2.6%	78.1± 4.6%
Quicksort	DIV2MEAN	524727	525759	85.2 ± 4.3%	64.6± 5.1%
Task Scheduling	TANHEXPANDTRIPLETS	262333	262333	88.2 ± 0.4%	87.3± 0.4%
Meta-learning set-up Similar to our MNIST-1D set-up, we use the AdamW optimizer to train each child model on a single NVIDIA Tesla P100 GPU. However, since most of the explored child models were much larger than the MNIST-1D models, we only trained each child model for 2000 steps. Anecdotally, we observed that the performance of different models often diverged by 2000 steps, which provided sufficient signal for the model search process. We otherwise followed the hyperparameters for single-task training in Ibarz et al. (2022) and evaluated models using validation accuracy.
Unlike our MNIST-1D set-up, we only search over the triplet representations of a Triplet-GMPNN model (see Ibarz et al. (2022) for more details), rather than the entire graph processor. We also seed the search with nine different seed models - each a variant of a Triplet-GMPNN model with a different triplet representation. Each seed triplet representation incorporates a minor tweak of a single component of the original triplet representation designed by Ibarz et al. (2022). These include a fully-connected output layer, a sum aggregation, fully-connected node/edge/graph representations,
a simple linear triplet representation, and a bilinear representation (Mnih & Hinton, 2007). All nine are implemented with Haiku (Hennigan et al., 2020), an object-oriented neural network library for Jax (see Appendix A.5 for the source code of the seed models.)
Generalizing beyond image classification models We search using EVOPROMPTING on 3 individual algorithms in the CLRS benchmark – the articulation points, Graham scan, and Kruskal's minimum spanning tree algorithms. We select these algorithms because our preliminary analyses with hand-designed architectures showed that they had the most headroom for improvement, although we found that the discovered architectures transfer well to other CLRS benchmark tasks as well (Appx. 3). Our search results are shown in Figure 3. EVOPROMPTING continues to find models that are more "fit" than our other two baselines, though we observed that the results also show more variation than our results for MNIST-1D did.
Analyzing newly discovered models Our search across triplet representations yielded several new designs that we sought to evaluate across all algorithms in the CLRS benchmark. Although these new models were discovered in model searches over single algorithms, they oftentimes generalized to other algorithms that were unseen during the model search. Figure 5 shows the trajectory of validation accuracy during training and Table 1 provides OOD accuracies for these models on a few select algorithms. (We defer the reader to Appendix A.4 for the full source code of each newly discovered model and Table A.6 for the full list of OOD accuracies for every algorithm in the CLRS benchmark.)
We note that the model search suggested several simple but effective changes. For example, instead of taking the maximum of the triplet representation, the QUADNODEMINMAX model uses quadruplet node representations instead of triplets, and it subtracts the minimum of the quad representation from the max instead. CONCATREP represents the node, edge, and graph representations as a concatenation of a projection feedforward layer, and MAXMEAN takes the maximum of the triplet representations prior to taking the mean and passing it through the output dense layer. DIV2MEAN scales each of the node representations by  $1/2$  and uses a mean aggregation of the triplet representations instead of the max aggregation. TANHEXPANDTRIPLETS applies additional dimension expansion to the triplet representations and applies a hyperbolic tangent function after the max aggregation. See Appx. A.4 for the full code of each discovered model.
Of the 5 newly discovered models that we chose to analyze, CONCATREP is the only one that increases model size. However, as shown in Table 1, CONCATREP frequently yielded improvements in OOD accuracy that far exceeded the percent increase in model size. For instance, on the heapsort algorithm CONCATREP increased OOD accuracy by  $125.19\%$  while only increasing model size by  $6.68\%$  over the baseline. The other four newly discovered models shown in Table 1 simultaneously improved OOD accuracy while decreasing model size on the articulation points, BFS, DFS, insertion sort, quicksort, and task scheduling algorithms. On the rest of the CLRS algorithms (Table A.6), our newly discovered models typically achieved OOD accuracy comparable to or better than the baseline, while maintaining similar model size.
We have shown that embedding a pre-trained LM in an evolutionary algorithm significantly improves the LM's performance on the task of neural architecture design. Our approach has demonstrated success at not only optimizing convolutional architectures for the MNIST-1D task, but also at developing new kinds of GNNs for the CLRS algorithmic benchmark. This demonstrates: 1) using evolutionary techniques can vastly improve the in-context capabilities of pre-trained LMs, and 2) EVOPROMPTING can discover novel and state-of-the-art architectures that optimize for both accuracy and model size. Furthermore, EVOPROMPTING is general enough to be easily adapted to search for solutions to other kinds of reasoning tasks beyond NAS. We leave the adaptation of EVOPROMPTING for other tasks to future work.
However, our study is limited by the lack of an extensive comparison against other standard NAS techniques because EVOPROMPTING was designed for open-ended search, whereas other techniques were not, which would introduce a potential confounder. We include one such comparison on NATS-Bench in Appendix A.7, as well as a discussion of the confounders thereof.
We thank Maarten Bosma, Kefan Xiao, Yifeng Lu, Quoc Le, Ed Chi, Borja Ibarz, Petar Velickovic, Chen Liang, Charles Sutton, and the Google Brain AutoML team for providing valuable discussions and feedback that influenced the direction of this project. We also thank the Google Student Researcher program for providing the resources and opportunities necessary for this project to take place.
Xu, F. F., Alon, U., Neubig, G., and Hellendoorn, V. J. A systematic evaluation of large language models of code. Proceedings of the 6th ACM SIGPLAN International Symposium on Machine Programming, 2022.
Xu, K., Li, J., Zhang, M., Du, S. S., ichi Kawarabayashi, K., and Jegelka, S. What can neural networks reason about? In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=rJxbJeHFPS.
Yao, X. Evolving artificial neural networks. Proc. IEEE, 87:1423-1447, 1999.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., Mihaylov, T., Ott, M., Shleifer, S., Shuster, K., Simig, D., Koura, P. S., Sridhar, A., Wang, T., and Zettlemoyer, L. Opt: Open pre-trained transformer language models, 2022. URL https://arxiv.org/abs/2205.01068.
Zhao, T., Wallace, E., Feng, S., Klein, D., and Singh, S. Calibrate before use: Improving few-shot performance of language models. ArXiv, abs/2102.09690, 2021.
Zhou, Y., Muresanu, A. I., Han, Z., Paster, K., Pitis, S., Chan, H., and Ba, J. Large language models are human-level prompt engineers. ArXiv, abs/2211.01910, 2022.
A.1 EVOPROMPTING Hyperparameters
Table 2: Values of hyperparameters used in EVOPROMPT.
HYPERPARAMETER	DESCRIPTION	VALUE
p	Num. parents to select in every generation	10
k	Num. in-context examples in prompt	2
T	Num. rounds of evolution	10
m	Num. prompts per round	10
n	Num. samples to generate per prompt	16
α	Lower threshold for test error	0.5
A.2 MNIST-1D Seed Models
Below we provide the source code for the four seed models used in the MNIST-1D model search.
1 class Model(nnModule):  
2 features: int = 32  
3 nlayer: int = 3  
4  
5 @nn.compact  
6 def __call__(self, x):  
7 x = x[..., None]  
8 x = nn.Conv_features = self_features, kernel_size = (3,) (x)  
9 x = nn.ReLU(x)  
10  
11 x = nn AVG_pool(x, window_shape = (2,) , strides = (2,) )  
12 for _ in range(self.nlayer - 1):  
13 xp = nn.Conv(  
14 features = self_features,  
15 kernel_size = (3,) ,  
16 ) (x)  
17 xp = nn.relu(xp)  
18 x = x + xp  
19  
20 x = nn_avg_pool(x, window_shape = (2,) , strides = (2,) )  
21 x = x.reshape((x.shape[0], -1)) # flatten  
22 x = nn(Dense_features=256)(x)  
23 x = nn.ReLU(x)  
24 x = nn(Dense_features=10)(x)  
25 return x
Listing 2: A hand-designed convolutional model.
class Model(nnModule):
    features: int = 25
    @nn.compact
    def __call__(self, x):
        x = x[...] None
        x = nn.Conv(
            features= self/features, kernel_size=(5), strides=(2),
            padding=(1,)
        )(x)
        x = nn.ReLU(x)
        for _ in range(2):
Listing 3: A Flax implementation of the convolutional baseline from Greydanus (2020).
12  $\mathbf{x} = \mathbf{nn}$  .Conv(   
13 features  $\equiv$  self_features，kernel_size=(3,)，strides=(2,)   
14 padding=(1,)   
15 )x  $\equiv$  nnrelu(x)   
16 x=x.reshape((x.shape[0]，-1))   
17 x  $\equiv$  nn.Dense(features=10)(x)   
18 return x
Listing 4: A Flax implementation of the GRU baseline from Greydanus (2020).
class Model(nn.Module):
    '''A simple GRU model'''  
hidden_size: int = 6  
seed: int = 42  
@nn.summary  
def __call__(self, x):  
    x = jnpexpand_dims(x, -1)  
rng =jax_random.PRNGKey(self.seed)  
gru = recurrent.GRU(
    hidden_size= self-hidden_size,
    num_layers=1,
    dropout_rate=0.0,
    bidirectional=True,
)
lengths = np.full([x.shape[0]], x.shape[1])
initialized.params = gru.init(rng, x, lengths)
params = initialized.params['params']
outputs, _ = gru.apply{'params': params}, x, lengths)
outputs = outputs.reshape((outputs.shape[0], -1))
x = nn(Dense/features=10)(outputs)  
return x
Listing 5: A Flax implementation of the fully connected baseline from Greydanus (2020).
1 class Model(nnModule):  
2 hidden_size: int = 100  
3  
4 @nn.compact  
5 def __call__(self, x):  
6 x = nn.Dense_featuresself-hidden_size(x)  
7 x = nnrelu(x)  
8 x = x + nnrelu(nn.Dense_featuresself-hidden_size(x))  
9 x = nn(Dense_features=10)(x)  
10 return x  
11  
12 return Model
A.3 Trajectory of search for MNIST-1D models
Figure 4: The average model size and test error of the child models produced in each round of the model search. Data points closer to the origin represent rounds that yielded more "fit" models.
A.4 Newly Discovered CLRS GNNs
Figure 5: Maximum fitness scores of five of the newly discovered models, compared against the baseline, on eight of the CLRS tasks.
Below we list the Python source code of five of the newly discovered GNNs.
def get triplet msgs Quad(z, edge_fts, graph_fts, nb triplet_fts, out_size):  
    node_reps = [hk.Linear(nb triplet_fts) for _ in range(4)]  
    triplet_node_reps = [node_rep(z) for node_rep in node_reps]  
    node_pair_inversions = [(1, 2), (1, 3), (2, 3), (3, 1)]  
    triplets = functools.reduce()  
    lambda x, y: x + y,  
7 [  
8 jnp.exp_dims(tri_node_rep, axis=perm)  
9 for tri_node_rep, perm in zip(
10 triplet_node_reps, node_pair_inversions
11 }
12 ], 
13 )
14 return jnp.max(triplets, axis=1) - jnp.min(triplets, axis=1)
Listing 6: The triplet representation that we refer to as QUADNODEMINMAX.
def get triplet msgss_catrep(z, edge_fts, graph_fts, nb triplet_fts, out_size): def rep_fn(x, size): proj = hk.nets.MLP([size]) ff = hk.nets.MLP([size * 4, size]) return jnp.catenate([[proj(x), ff(x), ], axis=-1)   
triplet_node_reps = [rep_fn(z, nb triplet_fts) for _ in range (3)]   
triplet_edge_reps = [rep_fn(edge_fts, nb triplet_fts) for _ in range(3)]   
triplet_graph_rep = rep_fn(graph_fts, nb triplet_fts)   
node_pair_permutations = [(2, 3), (1, 3), (1, 2)]   
triplets = functools.reduce( lambda x, y: x + y, [ jnp.exp_dims(tri_node_rep, axis=perm) for tri_node_rep, perm in zip( triplet_node_reps, node_pair_permutations ] ],   
triplets += functools.reduce( lambda x, y: x + y, [ jnp.exp_dims(tri_edge_rep, axis=i) for tri_edge_rep, i in zip(triplet_edge_reps, range(3, 0,-1)) ],   
triplets += jnp.exp_dims(triplet_graph_rep, axis=(1, 2, 3)) output_layer = hk.Linear(out_size)   
return output_layer(jnp.max(triplets, axis=1))
Listing 7: The triplet representation that we refer to as CONCATREP.
def get tripletmsgs_tanhexplandtriplets(z, edge_fts, graph_fts, nb triplet_fts, out_size):  
    node_reps = [hk.Linear(nb triplet_fts) for _ in range(3)]  
    edge_reps = [hk.Linear(nb triplet_fts) for _ in range(3)]  
    graph_rep = hk.nets.MLP([nb triplet_fts])  
    triplet_node_reps = [node_rep(z) for node_rep in node_reps]  
    triplet_edge_reps = [edge_rep(edge_fts) for edge_rep in edge_reps]  
    node_pair.permutations = [(2, 3), (1, 3), (1, 2)]  
    triplets = functools.reduce()  
    lambda x, y: x + y,  
10 [  
11 jnp.expand_dims(tri_node_rep, axis=perm)  
12 for tri_node_rep, perm in zip(
Listing 8: The triplet representation that we refer to as TANHEXPANDTRIPLETS.
triplet_node_reps, node_pair_permutations
) 
], 
) 
triplets += functools.reduce(
    lambda x, y: x + y,
    [ 
        jnp.expand_dims(tri_edge_rep, axis=i)
        for tri_edge_rep, i in zip(triplet_edge_reps, range(3, 0, -1))
        ], 
) 
triplets += jnp.expand_dims(graph_rep/graph_fts), axis=(1, 2, 3))
triplets += jnp.expand_dims.graph_rep/graph_fts), axis=(2, 3, 1))
triplets += jnp.expand_dims.graph_rep/graph_fts), axis=(3, 1, 2))
output_layer = hk.Linear(out_size)
return output_layer(jnp.tanh(jnp.max(tripllets, axis=1)))
Listing 9: The triplet representation that we refer to as DIV2MEAN.
def get triplet msgs_div2mean(z, edge_fts, graph_fts, nb triplet_fts, out_size):  
    node_reps = [hk.Linear(nb triplet_fts) for _ in range(3)]  
    edge_reps = [hk.Linear(nb triplet_fts) for _ in range(3)]  
    triplet_node_reps = [node_rep(z / 2) for node_rep in node_reps]  
    triplet_edge_reps = [edge_rep(edge_fts) for edge_rep in edge_reps]  
    node_pair.permutations = [(2, 3), (1, 3), (1, 2)]  
    triplets = functools.reduce()  
    lambda x, y: x + y,  
        [  
            jnp expand_dims(tri_node_rep, axis=perm)  
            for tri_node_rep, perm in zip(  
                triplet_node_reps, node_pair.permutations  
            )  
        ],  
    )  
    triplets += functools.reduce()  
    lambda x, y: x + y,  
        [  
            jnp expand_dims(tri_edge_rep, axis=perm)  
            for tri_edge_rep, perm in zip(triplet_edge_reps, range(3, 0, -1))  
            ],  
    )  
    output_layer = hk.Linear(out_size)  
    return output_layer(jnp.mean(tripllets, axis=1))
def get triplet msgs_maxmean(z, edge_fts, graph_fts, nb triplet_fts, out_size):  
    node_reps = [hk.Linear(nb triplet_fts) for _ in range(3)]  
    edge_reps = [hk.Linear(nb triplet_fts) for _ in range(3)]  
    graph_rep = hk.nets.MLP([nb triplet_fts])  
    triplet_node_reps = [node_rep(z) for node_rep in node_reps]
13   
14   
15   
16   
17   
18   
19   
20   
21   
22   
23   
24   
25   
26   
27   
28
triplet_edge_reps = [edge_rep(edge_fts) for edge_rep in edge_reps]
node_pair.permutations = [(2, 3), (1, 3), (1, 2)]
triplets = functools.reduce(
    lambda x, y: x + y,
    [ jnp.expand_dims(tri_node_rep, axis=perm)
        for tri_node_rep, perm in zip(
            triplet_node_reps, node_pair.permutations
        )
        ], 
) 
triplets += functools.reduce(
    lambda x, y: x + y,
    [ jnp.expand_dims(tri_edge_rep, axis=i)
        for tri_edge_rep, i in zip(triplet_edge_reps, range(3, 0, -1))
        ],
) 
triplets = jnp.maximun(triplets, -100.0)
output_layer = hk.Linear(out_size)
return output_layer(jnp.mean(triplets, axis=1))
Listing 10: The triplet representation that we refer to as MAXMEAN.
6   
7   
8   
9   
10   
11   
12   
13   
14   
15   
16   
17   
18   
19   
20   
21   
22   
23   
24   
25   
26
A.5 CLRS Seed Models
Below we provide the source code for the nine seed models used in the CLRS model search.