SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Eeshaan Jain $^{\S}$ , Tushar Nandy $^{\S}$ , Gaurav Aggarwal $^{\dagger}$ , Ashish Tendulkar $^{\dagger}$ , Rishabh Iyer $^{*}$  and Abir De $^{\S}$
$^{\S}$ IIT Bombay, †Google, *UT Dallas
{eeshaan,tnandy,abir}@cse.iitb.ac.in,
{gauravaggarwal,ashishvvt}@google.com, rishabh.iyer@utdallas.edu
Existing subset selection methods for efficient learning predominantly employ discrete combinatorial and model-specific approaches which lack generalizability. For an unseen architecture, one cannot use the subset chosen for a different model. To tackle this problem, we propose SUBSELNET, a trainable subset selection framework, that generalizes across architectures. Here, we first introduce an attention-based neural gadget that leverages the graph structure of architectures and acts as a surrogate to trained deep neural networks for quick model prediction. Then, we use these predictions to build subset samplers. This naturally provides us two variants of SUBSELNET. The first variant is transductive (called as Transductive-SUBSELNET) which computes the subset separately for each model by solving a small optimization problem. Such an optimization is still super fast, thanks to the replacement of explicit model training by the model approximator. The second variant is inductive (called as Inductive-SUBSELNET) which computes the subset using a trained subset selector, without any optimization. Our experiments show that our model outperforms several methods across several real datasets.
In the last decade, neural networks have drastically enhanced the performance of state-of-the-art ML models. However, they often demand massive data to train, which renders them heavily contingent on the availability of high-performance computing machineries such as GPUs and RAM. Such resources entail heavy energy consumption, excessive  $\mathrm{CO}_{2}$  emission, and maintenance cost.
Driven by this challenge, a recent body of work focuses on suitably selecting a subset of instances so that the model can be trained quickly using lightweight computing infrastructure [4, 23, 51, 32, 54, 37, 18-21, 36]. However, these methods are not generalizable across architectures—the subset selected by such a method is tailored to train only one specific architecture and thus need not be optimal for training another architecture. Hence, to select data subsets for a new architecture, they need to be run from scratch. However, these methods rely heavily on discrete combinatorial algorithms, which impose significant barriers against scaling them for multiple unseen architectures. Appendix C contains further details about related work.
Responding to the above limitations, we develop SUBSELNET, a trainable subset selection framework. Specifically, we make the following contributions.
Novel framework on subset selection that generalizes across models. SUBSELNET is a subset selector that generalizes across architectures. Given a dataset, once SUBSELNET is trained on a set of model architectures, it can quickly select a small optimal training subset for any unseen (test) architecture, without any explicit training of this test model. SUBSELNET is a non-adaptive method 37th Conference on Neural Information Processing Systems (NeurIPS 2023).
since it learns to select the subset before the training starts for a new architecture, instead of adaptively selecting the subset during the training process. Our framework has several applications in the context of AutoML [35, 68, 30, 43, 61, 9, 2, 24, 22, 3]. For example, Network Architecture Search (NAS) can have a significant speed-up when the architectures during selection can be trained on the subsets provided by our method, as compared to the entire dataset. In hyperparameter selection, such as the number and the widths of layers, learning rates or scheduler-specific hyperparameters, we can train each architecture on the corresponding data subset obtained from our method to quickly obtain the trained model for cross-validation.
Design of neural pipeline to eschew model training for new architecture. We initiate our investigation by writing down a combinatorial optimization problem instance that outputs a subset specifically for one given model architecture. Then, we gradually develop SUBSELNET, by building upon this setup. The key blocker in scaling up a model-specific combinatorial subset selector across different architectures is the involvement of the model parameters as optimization variables along with the candidate data subset. We design the neural pipeline of SUBSELNET to circumvent this blocker specifically. This neural pipeline consists of the following three components: (1) GNN-guided architecture encoder: This converts the architecture into an embedded vector space. (2) Neural model approximator: It approximates the predictions of a trained model for any given architecture. Thus, it provides the accuracy of a new (test) model per instance without explicitly training it. (3) Subset sampler: It uses the predictions from the model approximator and an instance to provide a selection score of the instance. Due to the architecture encoder and the neural approximator, we do not need to explicitly train a test model for selecting the subset since the model approximator directly provides the predictions the model will make.
Transductive and Inductive SUBSELNET. Depending on the functioning of the subset sampler in the final component of our neural pipeline, we design two variants of our model.
Transductive-SUBSELNET: The first variant is transductive in nature. For each new architecture, we utilize the model approximator's predictions for replacing the model training step in the original combinatorial subset selection problem. However, the candidate subset still remains involved as an optimization variable. Thus, we still solve a fresh optimization problem with respect to the selection score provided by the subset sampler every time we encounter a new architecture. However, the direct predictions from the model approximator allow us to skip explicit model training, making this strategy extremely fast in terms of memory and time.
Inductive-SUBSELNET: In contrast to Transductive-SUBSELNET, the second variant does not require to solve any optimization problem and instead models the selection scores using a neural network. Consequently, it is extremely fast.
We compare our method against six state-of-the-art methods on five real world datasets, which show that SUBSELNET provides the best trade-off between accuracy and inference time as well as accuracy and memory usage, among all the methods.
Setting. We are given a set of training instances  $\{(\pmb{x}_i, y_i)\}_{i \in D}$  where we use  $D$  to index the data. Here,  $\pmb{x}_i \in \mathbb{R}^{d_x}$  denotes the features, and  $y_i \in \mathcal{V}$  denotes the labels. In our experiments, we consider  $\mathcal{V}$  as a set of categorical labels. However, our framework can also be used for continuous labels. We use  $m$  to denote a neural architecture and represent its parameterization as  $m_\theta$ . We also use  $\mathcal{M}$  to denote the set of neural architectures. Given an architecture  $m \in \mathcal{M}$ ,  $G_m = (V_m, E_m)$  provides the graph representation of  $m$ , where the nodes  $u \in V_m$  represent the operations and the  $e = (u_m, v_m)$  indicates an edge, where the output given by the operation represented by the node  $u_m$  is fed to one of the operands of the operation given by the node  $v_m$ . Finally, we use  $H(\cdot)$  to denote the entropy of a probability distribution and  $\ell(m_\theta(\pmb{x}), y)$  as the cross entropy loss hereafter.
2.1 Combinatorial subset selection for efficient learning
Given a dataset  $\{(\pmb{x}_i, y_i)\}_{i \in D}$  and a model architecture  $m \in \mathcal{M}$  with its neural parameterization  $m_\theta$ , the goal of a subset selection algorithm is to select a small subset of instances  $S$  with  $|S| = b < < |D|$  such that training  $m_\theta$  on the subset  $S$  gives nearly the same accuracy as training on the entire dataset  $D$ . Existing works [20, 47, 19] adopt different strategies to achieve this goal, but all of them aim to simultaneously optimize for the model parameters  $\theta$  as well as the candidate subset  $S$ . At the outset,
Figure 1: Illustration of SUBSELNET. (a) Overview: Given a model architecture  $m \in \mathcal{M}$ , SUBSELNET takes its graph  $G_{m}$  as input to the architecture encoder  $\mathrm{GNN}_{\alpha}$  to compute the architecture embedding. This, together with  $\pmb{x}$  is fed into the model approximator  $g_{\beta}$  which predicts the output of the trained model  $m_{\theta^{*}}(\pmb{x})$ . Then this is fed as input to the subset sampler  $\pi$  to obtain the training subset  $S$ . (b) Neural architecture of different components:  $\mathrm{GNN}_{\alpha}$  consists of recursive message passing layer. The model approximator  $g_{\beta}$  performs a BFS ordering on the embeddings  $H_{m} = \{h_{u}\}$  and feeds them into a transformer. Subset sampler optimizes for  $\pi$  either via direct optimization for  $\pi$  (Transductive) or via a neural network  $\pi_{\psi}$  (Inductive).
we may consider the following optimization problem.
$$
\underset {\theta , S \subset D: | S | = b} {\text {m i n i m i z e}} \sum_ {i \in S} \ell \left(m _ {\theta} \left(\boldsymbol {x} _ {i}\right), y _ {i}\right) - \lambda \text {D I V E R S I T Y} \left(\left\{\boldsymbol {x} _ {i} \mid i \in S \right\}\right), \tag {1}
$$
where  $b$  is the budget, DIVERSITY( $\{\pmb{x}_i \mid i \in S\}$ ) measures the representativeness of  $S$  with respect to the whole dataset  $D$  and  $\lambda$  is a regularizing coefficient. One can use submodular functions [11, 17] like Facility Location, Graph Cut, or Log-Determinant to model DIVERSITY( $\{\pmb{x}_i \mid i \in S\}$ ). Here,  $\lambda$  trades off between training loss and diversity.
Bottlenecks of the combinatorial optimization (1). For every new architecture  $m$ , one needs to solve a fresh version of the optimization (1) problem from scratch to find  $S$ . Therefore, this is not generalizable across architectures. Moreover, the involvement of both combinatorial and continuous optimization variables, prevents the underlying solver from scaling across multiple architectures.
We address these challenges by designing a neural surrogate of the objective (1), which would lead to the generalization of subset selection across different architectures.
3 Overview of SUBSELNET
Here, we give an outline of our proposed model SUBSELNET that leads to substituting the optimization (1) with its neural surrogate, which would enable us to compute the optimal subset  $S$  for an unseen model, once trained on a set of model architectures.
3.1 Components
At the outset, SUBSELNET consists of three key components: (i) the architecture encoder, (ii) the neural approximator of the trained model, and (iii) the subset sampler. Figure 1 illustrates our model.
GNN-guided encoder for neural architectures. Generalizing any task across the different architectures requires the architectures to be embedded in vector space. Since a neural architecture is essentially a graph between multiple operations, we use a graph neural network (GNN) [59] to achieve this goal. Given a model architecture  $m \in \mathcal{M}$ , we first feed the underlying DAG  $G_{m}$  into a GNN  $(\mathrm{GNN}_{\alpha})$  with parameters  $\alpha$ , which outputs the node representations for  $G_{m}$ , i.e.,  $H_{m} = \{h_{u}\}_{u \in V_{m}}$ .
Approximator of the trained model  $m_{\theta^{*}}$ . To tackle lack of generalizability of the optimization (1), we design a neural model approximator  $g_{\beta}$  which approximates the predictions of any trained model for any given architecture  $m$ . To this end,  $g_{\beta}$  takes input as  $H_{m}$  and the instance  $x_{i}$  and compute  $g_{\beta}(H_{m},x_{i}) \approx m_{\theta^{*}}(x_{i})$ . Here,  $\theta^{*}$  is the set of learned parameters of the model  $m_{\theta}$  on dataset  $D$ .
Subset sampler. We design a subset sampler using a probabilistic model  $\mathrm{Pr}_{\pi}(\bullet)$ . Given a budget  $b$ , it sequentially draws instances  $S = \{s_1, \dots, s_b\}$  from a softmax distribution of the logit vector  $\pi \in \mathbb{R}^{[D]}$  where  $\pi(x_i, y_i)$  indicates a score for the element  $(x_i, y_i)$ . We would like to highlight that we use  $S$  as an ordered set of elements, selected in a sequential manner. However, such an order does not affect the trained model, which is inherently invariant of permutations of the training data; it only affects the choice of  $S$ . Now, depending on how we compute  $\pi$  during test, we have two variants of SUBSELNET: Transductive-SUBSELNET and Inductive-SUBSELNET.
Transductive-SUBSELNET: During test, since we have already trained the architecture encoder  $\mathrm{GNN}_{\alpha}$  and the model approximator  $g_{\beta}$ , we do not have to perform any training when we select a subset for an unseen architecture  $m'$ , since the trained model can then be replaced with  $g_{\beta}(\pmb{H}_{m'}, \pmb{x}_i)$ . Thus, the key bottleneck of solving the combinatorial optimization (1)—training the model simultaneously with exploring for  $S$ —is ameliorated. Now, we can perform optimization over  $\pi$ , each time for a new architecture. However, since no model training is involved, such explicit optimization is fast enough and memory efficient. Due to explicit optimization every time for an unseen architecture, this approach is transductive in nature.
Inductive-SUBSELNET: Here, we introduce a neural network to approximate  $\pi$ , which is trained together with  $\mathrm{GNN}_{\alpha}$  and  $g_{\beta}$ . This allows us to directly select the subset  $S$  without explicitly optimizing for  $\pi$ , unlike Transductive-SUBSELNET.
3.2 Training and inference
Training objective. Using the approximation  $g_{\beta}(\pmb{H}_m, \pmb{x}_i) \approx m_{\theta^*}(\pmb{x}_i)$ , we replace the combinatorial optimization problem in Eq. (1) with a continuous optimization problem, across different model architectures  $m \in \mathcal{M}$ . To that goal, we define
$\Lambda (S;m;\pi ,g_{\beta},\mathrm{GNN}_{\alpha}) = \sum_{i\in S}\ell (g_{\beta}(\pmb {H}_m,\pmb {x}_i),y_i) - \lambda H(\operatorname *{Pr}_{\pi}(\bullet))$  with,  $H_{m} = \mathrm{GNN}_{\alpha}(G_{m})$  (2) and seek to solve the following problem:
$$
\min  _ {\pi , \alpha , \beta} \sum_ {m \in \mathcal {M}} \underset {S \sim \Pr_ {\pi}} {\mathbb {E}} \left[ \Lambda (S; m; \pi , g _ {\beta}, \operatorname {G N N} _ {\alpha}) + \sum_ {i \in S} \gamma K L \left(g _ {\beta} \left(\boldsymbol {H} _ {m}, \boldsymbol {x} _ {i}\right), m _ {\theta^ {*}} \left(\boldsymbol {x} _ {i}\right)\right) \right] \tag {3}
$$
Here, we use entropy on the subset sampler  $H(\operatorname{Pr}_{\pi}(\bullet))$  to model the diversity of samples in the selected subset. We call our neural pipeline, which consists of architecture encoder  $\mathrm{GNN}_{\alpha}$ , the model approximator  $g_{\beta}$ , and the subset selector  $\pi$ , as SUBSELNET. In the above,  $\gamma$  penalizes the difference between the output of the model approximator and the prediction made by the trained model, which allows us to generalize the training of different models  $m \in \mathcal{M}$  through the model  $g_{\beta}(H_m, x_i)$ .
4 Design of SUBSELNET
Bottlenecks of end-to-end training and proposed multi-stage approach. End-to-end optimization of the above problem is difficult for the following reasons. (i) Our architecture representation  $H_{m}$  only represents the architectures and thus should be independent of the parameters of the architecture  $\theta$  and the instances  $\mathbf{x}$ . End-to-end training can make them sensitive to these quantities. (ii) To enable the model approximator  $g_{\beta}$  accurately fit the output of the trained model  $m_{\theta}$ , we need explicit training for  $\beta$  with the target  $m_{\theta}$ .
In our multi-stage training method, we first train the architecture encoder  $\mathrm{GNN}_{\alpha}$ , then the model approximator  $g_{\beta}$  and then train our subset sampler  $\operatorname{Pr}_{\pi}$  (resp.  $\operatorname{Pr}_{\pi_{\psi}}$ ) for the transductive (inductive) model. In the following, we describe the design and training of these components in details.
4.1 Design of architecture encoder using graph neural network
Architectures can be represented as directed acyclic graphs with forward message passing. During forward computation, at any layer for node  $v$ , the output  $a(v)$  can be represented as  $a(v) = \mathsf{Act}\left(\sum_{u\in \mathrm{InNbr}(v)}0\mathsf{p}_v(a(u))\right)$  with the root output as the input. Here, Act is the activation function and  $0\mathsf{p}_{\bullet}$  are operation on a node of the network. Given a GNN has a similar computation process, the permutation-equivariant node representations generated are good representations of the operations within the architecture. This allows further coupling with transformer-based architectures since they are universal approximators of permutation equivariant functions [63].
Neural parameterization. Given a model  $m \in \mathcal{M}$ , we compute the representations  $H_{m} = \{h_{u}|u \in V_{m}\}$  by using a graph neural network  $\mathrm{GNN}_{\alpha}$  parameterized with  $\alpha$ , following the proposal of Yan et al. [59]. We first compute the feature vector  $f_{u}$  for each node  $u \in V_{m}$  using the one-hot encoding of the associated operation (e.g., max, sum) and then feeding it into a neural network to compute an initial node representation  $h_{u}[0] = \mathrm{INITNODE}_{\alpha}(f_{u})$ . Then, we use a message-passing network, which collects signals from the neighborhood of different nodes and recursively computes the node representations [59, 58, 12]. Given a maximum number of recursive layers  $K$  and the node
$u$ , we compute the node embeddings  $\pmb{H}_{m} = \{\pmb{h}_{u}|u\in V_{m}\}$  by gathering information from the  $k < K$  hops using  $K$  recursive layers as follows.
$$
\boldsymbol {h} _ {(u, v)} [ k ] = \operatorname {E D G E E M B} _ {\alpha} \left(\boldsymbol {h} _ {u} [ k ], \boldsymbol {h} _ {v} [ k ]\right), \quad \boldsymbol {h} _ {u} [ k + 1 ] = \operatorname {U P D A T E} _ {\alpha} \left(\boldsymbol {h} _ {u} [ k ], \sum_ {v \in \operatorname {N b r} (u)} \boldsymbol {h} _ {(u, v)} [ k ]\right) \tag {4}
$$
Here,  $\mathrm{Nbr}(u)$  is the set of neighbors of  $u$ . EDGEEMB is injective mappings, as used in [58]. Note that trainable parameters from EDGEEMB and UPDATE are decoupled. They are represented as the set of parameters  $\alpha$ . Finally, we obtain our node representations as  $\pmb{h}_u = [\pmb{h}_u[0],..,\pmb{h}_u[K - 1]]$ .
Parameter estimation. We perform unsupervised training of  $\mathrm{GNN}_{\alpha}$  using a variational graph autoencoder (VGAE). This ensures that the architecture representations  $\pmb{H}_{m}$  remain insensitive to the model parameters. We build the encoder and decoder of our GVAE by following existing works on graph VAEs [59, 46]. Given a graph  $G_{m}$ , the encoder  $q(\mathcal{Z}_m|G_m)$ , which takes the node embeddings  $\{\pmb{h}_u\}_{u\in V_m}$  and maps it into the latent space  $\mathcal{Z}_m = \{\pmb{z}_u\}_{u\in V_m}$ . Specifically, we model the encoder  $q(\mathcal{Z}_m|G_m)$  as:  $q(\pmb{z}_u|G_m) = \mathcal{N}(\mu (\pmb{h}_u),\Sigma (\pmb{h}_u))$ . Here, both  $\mu$  and  $\Sigma$  are neural networks. Given a latent representation  $\mathcal{Z}_m = \{\pmb{z}_u\}_{u\in V_m}$ , the decoder models a generative distribution of the graph  $G_{m}$  where the presence of an edge is modeled as Bernoulli distribution BERNOULLLI  $(\sigma (\pmb{z}_u^\top \pmb{z}_v))$ . Thus, we model the decoder as  $p(G_m|\mathcal{Z}) = \prod_{(u,v)\in E_m}\sigma (\pmb{z}_u^\top \pmb{z}_v)\cdot \prod_{(u,v)\notin E_m}[1 - \sigma (\pmb{z}_u^\top \pmb{z}_v)]$ . Here,  $\sigma$  is a parameterized sigmoid function. Finally, we estimate  $\alpha ,\mu ,\Sigma$  and  $\sigma$  by maximizing the evidence lower bound (ELBO):  $\max_{\alpha ,\mu ,\Sigma ,\sigma}\mathbb{E}_{\mathcal{Z}\sim q(\bullet |G_m)}[p(G_m|\mathcal{Z})] - \mathrm{KL}(q(\bullet |G_m)||p_{\mathrm{pr}}(\bullet))$
4.2 Design of model approximator
Neural parameterization. Having computed the architecture representation  $\mathbf{H}_m = \{h_u \mid u \in V_m\}$ , we next design the model approximator, which leverages these embeddings to predict the output of the trained model  $m_{\theta^*}(\boldsymbol{x}_i)$ . To this aim, we developed a model approximator  $g_\beta$  parameterized by  $\beta$  that takes  $\mathbf{H}_m$  and  $\boldsymbol{x}_i$  as input and attempts to predict  $m_{\theta^*}(\boldsymbol{x}_i)$ , i.e.,  $g_\beta(\mathbf{H}_m, \boldsymbol{x}_i) \approx m_{\theta^*}(\boldsymbol{x}_i)$ . It consists of three steps. In the first step, we generate an order on the nodes. Next, we feed the representations  $\{\boldsymbol{h}_u\}$  in this order into a self-attention-based transformer layer. Finally, we combine the output of the transformer and  $\boldsymbol{x}_i$  using a feedforward network to approximate the model output.
Node ordering using BFS order. We first sort the nodes using breadth-first-search (BFS) order  $\rho$ . Similar to You et al. [62], this sorting method produces a sequence of nodes and captures subtleties like skip connections in the network structure  $G_{m}$ .
Attention layer. Given the BFS order  $\rho$ , we pass the representations  $\pmb{H}_{m} = \{\pmb{h}_{u} | u \in V_{m}\}$  in the sequence  $\rho$  through a self-attention-based transformer network. Here, the Query, Key, and Value functions are realized by linear networks on  $h_{\bullet}$ . We compute an attention-weighted vector  $\zeta_{u}$  as:
$$
\operatorname {A t t} _ {u} = \boldsymbol {W} _ {c} ^ {\top} \sum_ {v} a _ {u, v} \operatorname {V a l u e} \left(\boldsymbol {h} _ {v}\right) \text {w i t h}, \quad a _ {u, v} = \operatorname {S O F T M A X} _ {v} \left(\operatorname {Q u e r y} \left(\boldsymbol {h} _ {u}\right) ^ {\top} \operatorname {K e y} \left(\boldsymbol {h} _ {v}\right) / \sqrt {k}\right) \tag {5}
$$
Here  $k$  is the dimension of the latent space, and the softmax operation is over the node  $v$ . Subsequently, for each node  $u$ , we use a feedforward network, preceded and succeeded by layer normalization operations to obtain an intermediate representation  $\zeta_{u}$  for each node  $u$ . We present additional details in Appendix D. Finally, we feed  $\zeta_{u}$  for the last node  $u$  in the sequence  $\rho$ , i.e.,  $u = \rho(|V_{m}|)$ , along with the feature  $x_{i}$  into a feedforward network parameterized by  $W_{F}$  to model the prediction  $m_{\theta^{*}}(x_{i})$ . Thus, the final output of  $g_{\beta}(H_{m}, x_{i})$  is
$$
\boldsymbol {o} _ {m, \boldsymbol {x} _ {i}} = \operatorname {F F} _ {\boldsymbol {W} _ {F}} \left(\zeta_ {\rho \left(| V _ {m} |\right)}, \boldsymbol {x} _ {i}\right) \tag {6}
$$
Here,  $W_{\bullet}$ , parameters of Query, Key and Value and layer normalizations form  $\beta$ .
Parameter estimation. We train our model approximator  $g_{\beta}$  by minimizing the KL-Divergence between the approximated prediction  $g_{\beta}(\pmb{H}_m, \pmb{x}_i)$  and the ground truth prediction  $m_{\theta^*}(\pmb{x}_i)$ , where both these quantities are probabilities across different classes. The training problem is as follows:
$$
\operatorname {m i n i m i z e} _ {\beta} \sum_ {i \in D, m \in \mathcal {M}} \mathrm {K L} \left(m _ {\theta^ {*}} \left(\boldsymbol {x} _ {i}\right) \| g _ {\beta} \left(\boldsymbol {H} _ {m}, \boldsymbol {x} _ {i}\right)\right) \tag {7}
$$