SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(1) We acknowledge that there indeed is a computation time for pre-training the model approximator. However, as we mentioned in amortization analysis in Section 5.2, the one-time overhead is offset quickly by the speed and effectiveness of the following selection pipeline for the subset selection of unseen architectures. Since the prior training is only a one-time overhead, the overall cost is amortized by the number of unseen architectures during inference time training. In practice, when it is used to predict the subset for a large number of unseen architectures, the effect of training overhead quickly vanishes. As demonstrated by our experiments (Figures 2 and 3, Table 6), once the pipeline of all the neural networks is set up, the selection procedure is remarkably fast and can be easily adapted for use with unseen architectures.
To give an analogy, premier search engines invest a lot of resources in making fast inferences rather than training. They build complex models that are difficult and computationally intensive, but their inference is fast for several queries once trained. Thus, the cost is amortized by the large number of queries. Another example is locality-sensitive hashing. Researchers design trainable models for LSH whose purpose is to make fast predictions. Training LSH models can take a lot of time, but again this cost is amortized by the number of unseen queries.
Finally, we would like to highlight many efficient model training methods without complete training (running a few epochs via curriculum learning [66]), which one can easily explore and plug with our method for a larger dataset like Imagenet-1K.
(2) We use the space of neural architectures which comprises only of CNNs. We did not experiment with sequence models such as RNNs or transformers. However, we believe that our work can be extended with RNNs or transformer based architectures.
(3) If the distribution of network architectures varies widely from training to test, then there is significant room for performance improvement. In this context, one can develop domain adaptation methods for graphs to tackle different out-of-distribution architectures more effectively.
B Broader Impact
Our work can be used to provide significant compute efficiency by the trainable subset selection method we propose. It can be used to save a lot of time and power, that ML model often demands. Specifically, it can be used in the following applications in the context of AutoML.
Fast tuning of hyperparameters related to optimizer/training. Consider the case where we need to tune non-network hyperparameters, such as learning rate, momentum, and weight decay. Given the architecture, we can choose the subset obtained using our method to train the underlying model parameters for different hyperparameters, which can then be used for cross-validation. Note that we would use the same subset in this problem since the underlying model architecture is fixed, and we obtain a subset for the given architecture independent of the underlying non-network hyperparameters. We have shown utility of our method in our experiments in Section 5.2.
Fast tuning of model related hyperparameters. Consider the case where we need to tune network-related hyperparameters, such as the number of layers, activation functions, and the width of intermediate layers. Instead of training each instance of these models on the entire data, we can train them on the subset of data obtained from our method to quickly obtain the trained model, which can then be used for cross-validation.
Network architecture search. As we shown in our experiments, our method can provide speedup in network architecture search. Here, instead of training the network the entire network during
architecture exploration, we can restrict the training on a subset of data, which can provide significant speedup.
Note that, the key goal of our method is design a trainable subset selection method that generalizes across architectures. As we observed in our experiments, these methods can be useful in the above applications—however, our method is a generic framework and not tailored to any one of the above applications. Therefore, our method may need application specific modifications before directly deploying it practice. However, our method can serve as a base model for the practitioner who intends to speed up for one of the above applications.
We do not foresee any negative social impact of our work.
C Additional discussion on related work
Our work is closely related to representation learning for model architectures, network architecture search, data subset selection.
Representation learning for model architectures. Recent work in network representation learning use GNN based encoder-decoder to encapsulate the local structural information of a neural network into a fixed-length latent space [64, 40, 59, 34]. By employing an asynchronous message passing scheme over the directed acyclic graph (DAG), GNN-based methods model the propagation of input data over the actual network structure. Apart from encodings based solely on the structure of the network, White et al. [55], Yan et al. [60] produce computation-aware encodings that map architectures with similar performance to the same region in the latent space. Following the work of Yan et al. [59], we use a graph isomorphism network as an encoder but instead of producing a single graph embedding, our method produces a collection of node embeddings, ordered by breadth-first-search (BFS) ordering of the nodes. Our work also differs in that we do not employ network embeddings to perform downstream search strategies. Instead, architecture embeddings are used in training a novel model approximator that predicts the logits of a particular architecture, given an architecture embedding and a data embedding.
Machine learning on architecture space. We use NAS-Bench-101 in our method. This dataset was built in the context of network architecture search (NAS). The networks discovered by NAS methods often come from an underlying search space, usually designed to constrain the search space size. One such method is to use cell-based search spaces [35, 68, 30, 43, 61, 9]. Although we utilize the NAS-Bench-101 search space for architecture retrieval, our work is fundamentally different from NAS. In contrast to the NAS methods, which search for the best possible architecture from the search space using either sampling or gradient-descent based methods [1, 67, 44, 45, 31, 49], our work focuses on efficient data subset selection given a dataset and an architecture, which is sampled from a search space. Our work utilizes graph representation learning on the architectures sampled from the mentioned search spaces to project an architecture under consideration to a continuous latent space, utilize the model expression from the latent space as proxies for the actual model and proceed with data subset selection using the generated embedding, model proxy and given dataset.
Data subset selection. Data subset selection is widely used in literature for efficient learning, coreset selection, human centric learning, etc. Several works cast the efficient data subset selection task as instance of submodular or approximate-submodular optimization problem [19, 51-53, 20, 47]. Another line of work focuses on selecting coresets which are expressed as the weighted combination of subset of data, approximating some characteristics, e.g., loss function, model prediction [10, 37, 15, 4, 33]. Among other works, Toneva et al. [50] showed coresets can be selected by omitted several instances based on the forgetting dynamics at the time of training. Na et al. [39] selects proxy data based on entropy of the model. Coleman et al. [5] uses proxy model and then use it for coreset selection. Guo et al. [14] develop a library on coreset selection.
Our work is closely connected to simultaneous model learning and subset selection [7, 6, 47]. These existing works focus on jointly optimizing the training loss, with respect to the subset of instances and the parameters of the underlying model. Among them [7, 6] focus on distributing decisions between human and machines, whereas [47] aims for efficient learning. However, these methods adopt a combinatorial approach for selecting subsets and consequently, they are not generalizable across architectures. In contrast, our work focuses on differentiable subset selection mechanism, which can generalize across architectures.
D Additional details about experimental setup
D.1 Dataset
Datasets  $(D)$
Dataset	No. of Classes	Imbalanced	Train-Test Split	Shape	Transformations Applied
FMNIST	10	X	(60K,10K)	28x28x1	Normalize
CIFAR10	10	X	(50K,10K)	32x32x3	RandomHorizontalFlip, RandomCrop, Normalize
CIFAR100	100	X	(50K,10K)	32x32x3	RandomHorizontalFlip, RandomCrop, Normalize
Tiny-Imagenet	200	X	(100K,10K)	64x64x3	RandomHorizontalFlip, RandomVerticalFlip, Normalize
Caltech-256	257	✓	(24.5K,6.1K)	96x96x3	RandomHorizontalFlip, Resize, Normalize
Table 7: A brief description of the datasets used along with the transformations applied during training
Architectures  $(\mathcal{M})$ . We leverage the NASBench-101 search space as an architecture pool. It consists of 423, 624 unique architectures with the following constraints - (1) number of nodes in each cell is at most 7, (2) number of edges in each cell is at most 9, (3) barring the input and output, there are three unique operations, namely  $1 \times 1$  convolution,  $3 \times 3$  convolution and  $3 \times 3$  max-pool. We utilize the architectures from the search space in generating the sequence of embeddings along with sampling architectures for the training and testing of the encoder and datasets for the subset selector. As mentioned in the experimental setup, pre-training large number of models can be expensive.
Figure 5: Distribution of parameters of architectures in  $\mathcal{M}_{\mathrm{tr}}$  when  $|\mathcal{M}_{\mathrm{tr}}| = 423k$  (blue), and  $\mathcal{M}_{\mathrm{tr}}$  with the sampled set of 250 architectures (orange).
Therefore, we choose a diverse subset of architectures from  $\mathcal{M}_{\mathrm{tr}}$  of size 250 that ensures efficient representation over low-parameter and high-parameter regimes. The distributions of true and sampled architectures are given in Figure 5.
Note that the pre-training can be made faster by efficient model training methods without complete training (running a few epochs via curriculum learning [66]), which can be easily plugged with our method.
D.2 Implementation details about baselines
Facility Location (FL). We implemented facility location on all the three datasets using the apricot library. The similarity matrix was computed using Euclidean distance between data points, and the objective function was maximized using the naive greedy algorithm.
Pruning. It selects a subset from the entire dataset based on the uncertainty of the datapoints while partial training. In our setup, we considered ResNet-18 as a master model, which is trained on each dataset for 5 epochs. Post training, the uncertainty measure is calculated based on the probabilities of each class, and the points with highest uncertainty are considered in the subset. We train the master model at a learning rate of 0.025.
Proxy. It selects a subset
Glister and Grad-Match. We implemented GLISTER [20] and Grad-Match [19] using the CORDS library. We trained the models for 50 epochs, using batch size of 20, and selected the subset after every 10 epochs. The loss was minimized using SGD with learning rate of 0.01, momentum of 0.9 and weight decay with regularization constant of  $5 \times 10^{-4}$ . We used cosine annealing for scheduling
<https://github.com/jmschrei/apricot>
the learning rate with  $T_{max}$  of 50 epochs, and used  $10\%$  of the training data as the validation set. Details of specific hyperparameters for stated as follows.
Glider uses a greedy selection approach to minimize a bi-level objective function. In our implementation, we used stochastic greedy optimization with learning rate 0.01, applied on the data points of each mini-batch. Online-Glider approximates the objective function with a Taylor series expansion up to an arbitrary number of terms to speed up the process; we used 15 terms in our experiments.
Grad-Match applies the orthogonal matching (OMP) pursuit algorithm to the data points of each mini-batch to match gradient of a subset to the entire training/validation set. Here, we set the learning rate is set to 0.01. The regularization constant in OMP is 1.0 and the algorithm optimizes the objective function within an error margin of  $10^{-4}$ .
GraNd. This is an adaptive subset selection strategy in which the norm of the gradient of the loss function is used as a score to rank a data point. The gradient scores are computed after the model has trained on the full dataset for the first few epochs. For the rest of epochs, the model is trained only on the top- $k$  data points, selected using the gradient scores. In our implementation, we let the model train on the full dataset for the first 5 epochs, and computed the gradient of the loss only with respect to the last layer fully connected layer.
EL2N. When the loss function used to compute the GraNd scores is the cross entropy loss, the norm of the gradient for a data point  $\mathbf{x}$  can be approximated by  $\mathbb{E}||p(\mathbf{x}) - y||_2$ , where  $p(\mathbf{x})$  is the discrete probability distribution over the classes, computed by taking softmax of the logits, and  $y$  is the one-hot encoded true label corresponding to the data point  $\mathbf{x}$ . Similar to our implementation of GraNd, we computed the EL2N scores after letting the models train on the full data for the first 5 epochs.
D.3 Implementation details about our model
$\mathbf{GNN}_{\alpha}$ . As we utilize NASBench-101 space as the underlying set of neural architectures, each computational node in the architecture can comprise of one of five operations and the one-hot-encoded feature vector  $\mathbf{f}_u$ . Since the set is cell-based, there is an injective mapping between the neural architecture and the cell structure. We aim to produce a sequence of embeddings for the cell, which in turn corresponds to that of the architecture. For each architecture, we use the initial feature  $\mathbf{f}_u \in \mathbb{R}^5$  in as a five dimensional one-hot encoding for each operation. This is fed into INITNODE to obtain an 16 dimensional output. Here, INITNODE consists of a  $5 \times 16$  linear, ReLU and  $16 \times 16$  linear layers cascaded with each other. Each of EDGEMBED and UPDATE consists of a  $5 \times 128$  linear-BatchNorm-ReLU cascaded with a  $128 \times 16$  linear layer. Moreover, the symmetric aggregator is a sum aggregator.
We repeat this layer  $K$  times, and each iteration gathers information from  $k < K$  hops. After all the iterations, we generate an embedding for each node, and following [62] we use the BFS-tree based node-ordering scheme to generate the sequence of embeddings for each network.
The GVAE-based architecture was trained for 10 epochs with the number of recursive layers  $K$  set to 5, and the Adam optimizer was used with learning rate of  $10^{-3}$ . The entire search space was considered as the dataset, and a batch-size of 32 was used. Post training, we call the node embeddings collectively as the architecture representation.
To train the latent space embeddings, the parameters  $\alpha$  are trained in an encoder-decoder fashion using a variational autoencoder. The mean  $\mu$  and variance  $\sigma$  on the final node embeddings  $h_u$  are:
$$
\mu = \operatorname {F C N} \left(\left[ \boldsymbol {h} _ {u} \right] _ {u \in V _ {m}}\right) \text {a n d} \sigma = \exp \left(\operatorname {F C N} \left(\left[ \boldsymbol {h} _ {u} \right] _ {u \in V _ {m}}\right)\right)
$$
The decoder aims to reconstruct the original cell structure (i.e the nodes and the corresponding operations), which are one-hot encoded. It is modeled using single-layer fully connected networks followed by a sigmoid layer.
Model Approximator  $g_{\beta}$ . The model approximator  $g_{\beta}$  is essentially a single-head attention block that acts on a sequence of node embeddings  $H_{m} = \{h_{u}|u\in V_{m}\}$ . The Query, Key and Value are three linear networks with parameters:  $W_{query}$ ,  $W_{key}$  and  $W_{value}\in \mathbb{R}^{16\times 8}$ . Note that the matrix  $W_{C}\in \mathbb{R}^{8\times 16}$  in Eq. (5). As described in Section 4.2, for each node  $u$ , we use a feedforward network, preceded and succeeded by layer normalization operations, which are given by the following set of
(a) (a) FMNIST
(b) (b) CIFAR10
Figure 6: Kullback-Leibler divergence values  $(\mathrm{KL}(m_{\theta^{*}}(\pmb{x}_{i})||g_{\beta}(\pmb{H}_{m},\pmb{x}_{i}))$  computed during the training of the model encoder  $g_{\beta}$  over 80 epochs.
equations (where LN the denotes Layer-Normalization operation):
$$
\zeta_ {u, 1} = \operatorname {L N} \left(\operatorname {A t t} _ {u} + \boldsymbol {h} _ {u}; \gamma_ {1}, \gamma_ {2}\right),
$$
$$
\zeta_ {u, 2} = \boldsymbol {W} _ {2} ^ {\top} \operatorname {R E L U} \left(\boldsymbol {W} _ {1} ^ {\top} \zeta_ {u, 1}\right),
$$
$$
\zeta_ {u, 3} = \operatorname {L N} \left(\zeta_ {u, 1} + \zeta_ {u, 2}; \gamma_ {3}, \gamma_ {4}\right)
$$
The fully connected network acting on  $\zeta_{u,1}$  consists of matrices  $W_{1} \in \mathbb{R}^{16 \times 64}$  and  $W_{2} \in \mathbb{R}^{64 \times 16}$ . All the trainable matrices along with the layer normalizations were implemented using the Linear and LayerNorm functions in Pytorch. The last item of the output sequence  $\zeta_{u,3}$  is concatenated with the data embedding  $\pmb{x}_{i}$  and fed to another 2-layer fully-connected network with hidden dimension 256 and dropout probability of 0.3. The model approximator is trained by minimizing the KL-divergence between  $g_{\beta}(H_m, \pmb{x}_i)$  and  $m_{\theta^*}(\pmb{x}_i)$ . We used an AdamW optimizer with learning rate of  $10^{-3}$ ,  $\epsilon = 10^{-8}$ , betas = (0.9, 0.999) and weight decay of 0.005. We also used Cosine Annealing to decay the learning rate, and used gradient clipping with maximum norm set to 5. Figure 6 shows the convergence of the outputs of the model approximator  $g_{\beta}(H_m, \pmb{x}_i)$  with the outputs of the model  $m_{\theta^*}(\pmb{x}_i)$ .
Neural Network  $\pi_{\psi}$ . The inductive model is a three-layer fully-connected neural network with two Leaky ReLU activations and a sigmoid activation after the last layer. The input to  $\pi_{\psi}$  is the concatenation  $(H_{m}; o_{m,i}; x_{i}; y_{i})$ . The hidden dimensions of the two intermediary layers are 64 and 16, and the final layer is a single neuron that outputs the score corresponding to a data point  $x_{i}$ . While training  $\pi_{\psi}$  we add a regularization term  $\lambda' (\sum_{i \in D} \pi_{\psi}(H_{m}, o_{m,i}, x_{i}, y_{i}) - |S|)$  to ensure that nearly  $|S|$  samples have high scores out of the entire dataset  $D$ . Both the regularization constants  $\lambda$  (in equation 3) and  $\lambda'$  are set to 0.1. We train the model weights using an Adam optimizer with a learning rate of 0.001. During training, at each iteration we draw instances using  $\mathrm{Pr}_{\pi}$  and use the log-derivative trick to compute the gradient of the objective. During each computation step, we use one instance of the ranked list to compute the unbiased estimate of the objective in (3).
D.4 Hyperparameter Optimization
The hyperparameter search was done over optimizer and scheduler-based hyperparameters using Ray Tune [29]. We set the possible optimizers to be SGD, Adam and RMSprop, and the possible schedulers to be CosineAnnealing and StepLR. The search space parameters are given below:
- Optimizers
1. SGD: learning_rate  $\in$  (0.001, 0.1), momentum  $\in$  (0.7, 1.0), weight_Decay  $\in$  (0.01, 0.0001)
2. Adam: learning_rate  $\in$  (0.001, 0.1), weight Decay  $\in$  (0.01, 0.0001)
3. RMSprop: learning_rate  $\in$  (0.001, 0.1), momentum  $\in$  (0.7, 1.0), weight_Decay  $\in$  (0.01, 0.0001)
- Schedulers
1. StepLR: step_size  $\in$  [10, 20, 30, 40], gamma  $\in$  (0.05, 0.5)
2. CosineAnnealingLR
We employ TPE [2] as the hyperparameter search algorithm, and ASHA [28] as the hyperparameter scheduling algorithm. The hyperparameter search runs for 100 epochs in all cases. The random baseline is run for 5 runs, and we report the average speedup and test error.
E Additional experiments
E.1 Comparison with additional baselines
Here, we compare the performance of SUBSELNET against two baselines. They are the two variants of our method-Bottom- $b$ -loss and Bottom- $b$ -loss+gumbel.
In Bottom- $b$ -loss, we sort the data instances based on their predicted loss  $\ell(g_{\beta}(\pmb{H}_m, \pmb{x}), y)$  and consider those points with the bottom  $b$  values.
In Bottom- $b$ -loss+gumbel, we add noise sampled from the gumbel distribution with  $\mu_{\mathrm{gumbel}} = 0$  and  $\beta_{\mathrm{gumbel}} = 0.025$ , and sort the instances based on these noisy loss values, i.e.,  $\ell(g_{\beta}(\mathbf{H}_m, \mathbf{x}), y) + G_{\mathrm{gumbel}}(0, \beta_{\mathrm{gumbel}} = 0.025)$ .
Figure 7 compares the performance of the variants of SUBSELNET, Bottom- $b$ -loss, and Bottom- $b$ -loss+gumbel. We observe that Bottom- $b$ -loss and Bottom- $b$ -loss+gumbel do not perform that well in spite of being efficient in terms of time and memory.
Figure 7: Comparison of Transductive-SUBSELNET and Inductive-SUBSELNET with Bottom- $b$ -loss and Bottom- $b$ -loss+gumbel. In Bottom- $b$ -loss, we select top- $b$  instances in terms of their predicted loss  $\ell(g_{\beta}(\pmb{H}_m, \pmb{x}), y)$  computed using the model approximator. In Bottom- $b$ -loss+gumbel, we add gumbel noise Gumbel(0, 0.025) to the loss and sort the instances based on these noisy loss values.
E.2 Comparison of performance on ImageNet
Here, we compare the performance of SUBSELNET against two baselines: Selection via-Proxy [5] and Pruning [48] on ImageNet-1K (1.28M images) [8] for  $b \in (0.05|D|, 0.8|D|)$ . Figure 8 compares Transductive-SUBSELNET and Inductive-SUBSELNET with the two baselines on basis of speedup, memory and budget.
Figure 8: Comparison of performance of the top four non-adaptive subset selectors (Transductive-SUBSELNET, Inductive-SUBSELNET, Pruning, and Selection via-Proxy) on ImageNet-1K [8].
E.3 Analysis of compute efficiency in high accuracy regime
We analyze the compute efficiency of all the methods, given an allowance of reaching  $20\%$  and  $10\%$  of the relative accuracy reduction (RAR) ( $80\%$  and  $90\%$  of the accuracy) achieved by training on the full data. We make the following observations:
1. Transductive-SUBSELNET achieves the best speedup and consumes the least memory, followed by Inductive-SUBSELNET.
2. For CIFAR100, Tiny-Imagenet and Caltech-256, Bottom- $b$ -loss and Bottom- $b$ -loss+gumbel achieve better performance than the baselines which are able to reach the desired RAR milestones (10% or 20%).
3. For CIFAR100, Tiny-Imagenet and Caltech-256, most baselines could not achieve an accuracy of either  $10\%$  or even  $20\%$  of the RAR on full data.
FMNIST	CIFAR10	CIFAR100
Speedup	Memory	Speedup	Memory	Speedup	Memory
Method	10%	20%	10%	20%	10%	20%	10%	20%	10%	20%	10%	20%
GLISTER	5.64	7.85	98.51	116.36	1.52	2.12	515.96	365.05	0.54	1.02	1427.77	758.55
GradMatch	4.17	5.24	136.40	243.75	1.69	2.20	457.67	362.47	—	0.84	—	917.04
EL2N	6.50	16.42	77.63	139.89	1.93	4.78	413.90	170.03	—	—	—	—
GraNd	—	1.18	—	450.73	—	—	—	—	—	—	—	—
FacLoc	0.82	2.37	652.67	81.01	—	0.80	—	558.56	—	—	—	—
Pruning	3.12	4.68	559.44	19.55	3.54	5.53	221.10	139.41	—	1.71	—	452.09
Selection-via-Proxy	3.65	18.09	168.20	35.27	1.95	1.03	819.22	410.26	—	1.02	—	765.05
Bottom-b-loss	1.68	2.98	393.28	190.40	—	1.77	—	433.07	—	1.67	—	465.39
Bottom-b-loss+gumbel	2.70	10.18	203.83	59.37	1.63	2.04	489.30	363.07	—	1.78	—	446.95
Inductive-SUBSELNET	28.64	69.24	22.73	8.24	3.63	8.99	221.99	99.54	1.93	2.82	417.16	274.17
Transductive-SUBSELNET	28.63	68.36	21.25	8.24	5.61	16.52	142.45	53.67	2.35	3.47	331.45	222.91