TabNAS: Rejection Sampling for Neural Architecture Search on Tabular Datasets

Introduction

To make a machine learning model better, one can scale it up. But larger networks are more expensive as measured by inference time, memory, energy, etc, and these costs limit the application of large models: training is slow and expensive, and inference is often too slow to satisfy user requirements.

Many applications of machine learning in industry use tabular data, e.g., in finance, advertising and medicine. It was only recently that deep learning has achieved parity with classical tree-based models in these domains . For vision, optimizing models for practical deployment often relies on Neural Architecture Search (NAS). Most NAS literature targets convolutional networks on vision benchmarks . Despite the practical importance of tabular data, however, NAS research on this topic is quite limited . (See Appendix A for a more comprehensive literature review.)

Weight-sharing reduces the cost of NAS by training a SuperNet that is the superset of all candidate architectures . This trained SuperNet is then used to estimate the quality of each candidate architecture or child network by allowing activations in only a subset of the components of the SuperNet and evaluating the model. Reinforcement learning (RL) has shown to efficiently find the most promising child networks for vision problems.

In our experiments, we show that a direct application of approaches designed for vision to tabular data often fails. For example, the TuNAS approach from vision struggles to find the optimal architectures for tabular datasets (see experiments). The failure is caused by the interaction of the search space and the factorized RL controller. To understand why, consider the following toy example with 2 layers, illustrated in Figure 1. For each layer, we can choose a layer size of $2$, $3$, or $4$, and the maximum number of parameters is set to 25. The optimal solution is to set the size of the first hidden layer to 4 and the second to 2. Finding this solution with RL is difficult with a cost penalty approach. The RL controller is initialized with uniform probabilities. As a result, it is quite likely that the RL controller will initially be penalized heavily when choosing option 4 for the first layer, since two thirds of the choices for the second layer will result in a model that is too expensive. As a result, option 4 for the first layer is quickly discarded by the RL controller and we get stuck in a local optimum.

This co-adaptation problem is caused by the fact that existing NAS methods for computer vision often use factorized RL controllers, which force all choices to be be made independently. While factorized controllers can be optimized easily and are parameter-efficient, they cannot capture all of the nuances in the loss landscape. A solution to this could be to use a more complex model such as an LSTM (e.g., ). However, LSTMs are often much slower to train and are far more difficult to tune.

Our proposed method, TabNAS, uses a solution inspired by rejection sampling. It updates the RL controller only when the sampled model satisfies the cost constraint. The RL controller is then discouraged from sampling poor models within the cost constraint and encouraged to sample the high quality models. Rather than penalizing models that violate the constraints, the controller silently discards them. This trick allows the RL controller to see the true constrained loss landscape, in which having some large layers is beneficial, allowing TabNAS to efficiently find global (not just local) optima for tabular NAS problems. Our contributions can be summarized as follows:

We identify failure cases of existing resource-aware NAS methods on tabular data and provide evidence this failure is due to the cost penalty in the reward together with the factorized space.

We propose and evaluate an alternative: a rejection sampling mechanism that ensures the RL controller only selects architectures that satisfy resource constraint. This extra rejection step allows the RL controller to explore parts of the search space that would otherwise be overlooked.

The rejection mechanism also introduces a systematic bias into the RL gradient updates, which can skew the results. To compensate for this bias, we introduce a theoretically motivated and empirically effective correction into the gradient updates. This correction can be computed exactly for small search spaces and efficiently approximated by Monte-Carlo sampling otherwise.

We show the resulting method, TabNAS, automatically learns whether a bottleneck structure is needed in an optimal architecture, and if needed, where to place the bottleneck in the network.

These contributions form TabNAS, our RL-based weight-sharing NAS with rejection-based reward. TabNAS robustly and efficiently finds a feasible architecture with optimal performance within the resource constraint. Figure 3 shows an example.

Notation and terminology

Math basics. We define $[n]=\{1,\ldots,n\}$ for a positive integer $n$. With a Boolean variable $\mathcal{X}$, the indicator function $\mathds{1}(\mathcal{X})$ equals 1 if $\mathcal{X}$ is true, and 0 otherwise. $|S|$ denotes the cardinality of a set $S$; $\text{stop\_grad}(f)$ denotes the constant value (with gradient 0) corresponding to a differentiable quantity $f$, and is equivalent to tensorflow.stop_gradient(f) in TensorFlow or f.detach() in PyTorch . $\subseteq$ and $\subset$ denote subset and strict subset, respectively. $\nabla$ denotes the gradient with respect to the variable in the context.

Weight, architecture, and hyperparameter. We use weights to refer to the parameters of the neural network. The architecture of a neural network is the structure of how nodes are connected; examples of architectural choices are hidden layer sizes and activation types. Hyperparameters are the non-architectural parameters that control the training process of either stand-alone training or RL, including learning rate, optimizer type, optimizer parameters, etc.

Neural architecture. A neural network with specified architecture and hyperparameters is called a model. We only consider fully-connected feedforward networks (FFNs) in this paper, since they can already achieve SOTA performance on tabular datasets . The number of hidden nodes after each weight matrix and activation function is called a hidden layer size. We denote a single network in our search space with hyphen-connected choices. For example, when searching for hidden layer sizes, in the space of 3-hidden-layer ReLU networks, 32-144-24 denotes the candidate where the sizes of the first, second and third hidden layers are 32, 144 and 24, respectively. We only search for ReLU networks; for brevity, we will not mention the activation function type in the sequel.

Loss-resource tradeoff and reference architectures. In the hidden layer size search space, the validation loss in general decreases with the increase of the number of parameters, giving the loss-resource tradeoff (e.g., Figure 3). Here loss and number of parameters serve as two costs for NAS. Thus there are Pareto-optimal models that achieve the smallest loss among all models with a given bound on the number of parameters. With an architecture that outperforms others with a similar or fewer number of parameters, we do resource-constrained NAS with the number of parameters of this architecture as the resource target or constraint. We call this architecture the reference architecture (or reference) of NAS, and its performance the reference performance. We do NAS with the goal of matching (the size and performance of) the reference. Note that the RL controller only has knowledge of the number of parameters of the reference, and is not informed of its hidden layer sizes.

Resource metric and number of parameters. We use the number of parameters, which can be easily computed for neural networks, as a cost metric in this paper. However, our approach does not depend on the specific cost used, and can be easily adapted to other cost metrics.

Methodology

Our NAS methodology can be decomposed into three main components: weight-sharing with layer warmup, REINFORCE with one-shot search, and Monte Carlo (MC) sampling with rejection.

As an overview, our method starts with a SuperNet, which is a network that layer-wise has width equal to the largest choice within the search space. We first stochastically update the weights of the entire SuperNet to “warm up” over the first 25% of search epochs. Then we alternate between updating the shared model weights (which are used to estimate the quality of different child models) and the RL controller (which focuses the search on the most promising parts of the space). In each iteration, we first sample a child network from the current layer-wise probability distributions and update the corresponding weights within the SuperNet (weight update). We then sample another child network to update the layerwise logits that give the probability distributions (RL update). The latter RL update is only performed if the sampled network is feasible, in which case we use rejection with MC sampling to update the logits with a sampling probability conditional on the feasible set.

To avoid overfitting, we split the labelled portion of a dataset into training and validation splits. Weight updates are carried out on the training split; RL updates are performed on the validation split.

In weight-sharing NAS, warmup helps to ensure that the SuperNet weights are sufficiently trained to properly guide the RL updates . With probability $p$, we train all weights of the SuperNet, and with probability $1-p$ we only train the weights of a random child model. When we run architecture searches for FFNs, we do warmup in the first 25% epochs, during which the probability $p$ linearly decays from 1 to 0 (Figure 5(a)). The RL controller is disabled during this period.

2 One-shot training and REINFORCE

We do NAS on FFNs with a REINFORCE-based algorithm. Previous works have used this type of algorithm to search for convolutional networks on vision tasks . When searching for $L$-layer FFNs, we learn a separate probability distribution over $C_{i}$ size candidates for each layer. The distribution is given by $C_{i}$ logits via the SoftMax function. Each layer has its own independent set of logits. With $C_{i}$ choices for the $i$th layer, where $i=1,2,\ldots,L$, there are $\prod_{i\in[L]}C_{i}$ candidate networks in the search space but only $\sum_{i\in[L]}C_{i}$ logits to learn. This technique significantly reduces the difficulty of RL and make the NAS problem practically tractable .

The alternation creates a positive feedback loop that trains the weights and updates the logits of the large-probability child networks; thus the layer-wise sampling probabilities gradually converge to more deterministic distributions, under which one or several architectures are finally selected.

Details of a resource-oblivious version is shown as Appendix B Algorithm 1, which does not take into account a resource constraint. In Section 3.3, we show an algorithm that combines Monte-Carlo sampling with rejection sampling, which serves as a subroutine of Algorithm 1 by replacing the probability in $J(y)$ with a conditional version.

3 Rejection-based reward with MC sampling

Only a subset of the architectures in the search space $S$ will satisfy resource constraints; $V$ denotes this set of feasible architectures. To find a feasible architecture, a resource target $T_{0}$ is often used in an RL reward. Given an architecture $y$, a resource-aware reward combines its quality $Q(y)$ and resource consumption $T(y)$ into a single reward. MnasNet proposes the rewards $Q(y)(T(y)/T_{0})^{\beta}$ and $Q(y)\max\{1,(T(y)/T_{0})^{\beta}\}$ while TuNAS proposes the absolute value reward (or Abs Reward) $Q(y)+\beta|T(y)/T_{0}-1|$. The idea behind is to encourage models with high quality with respect the resource target. In these rewards $\beta$ is a hyperparameter that needs careful tuning.

We find that on tabular data, RL controllers using these resource-aware rewards above can struggle to discover high quality structures. Figure 1 shows a toy example in the search space in Figure 5, in which we know the validation losses of each child network and only train the RL controller for 500 steps. The optimal network is 4-2 among architectures with number of parameters no more than 25, but the RL controller rarely chooses it. In Section 4.1, we show examples on real datasets.

This phenomenon reveals a gap between the true distribution we want to sample from and the distributions obtained by sampling from this factorized search space:

We only want to sample from the set of feasible architectures $V$, whose distribution is $\{\mathds{P}(y\>|\>y\in V)\}_{y\in V}$. The resources (e.g., number of parameters) used by an architecture, and thus its feasibility, is determined jointly by the sizes of all layers.

On the other hand, the factorized search space learns a separate (independent) probability distribution for the choices of each layer. While this distribution is efficient to learn, independence between layers discourages an RL controller with a resource-aware reward from choosing a bottleneck structure. A bottleneck requires the controller to select large sizes for some layers and small for others. But decisions for different layers are made independently, and both very large and very small layer sizes, considered independently, have poor expected rewards: small layers are estimated to perform poorly, while large layers easily exceed the resource constraints.

To bridge the gap and efficiently learn layerwise distributions that take into account the architecture feasibility, we propose a rejection-based RL reward for Algorithm 1. We next sketch the idea; detailed pseudocode is provided as Algorithm 2 in Appendix B.

In our new REINFORCE variant, motivated by rejection sampling, we do not update the logits when $y$ is infeasible. When $y$ is feasible, we replace the probability $\mathds{P}(y)$ in the REINFORCE update equation with the conditional probability $\mathds{P}(y\>|\>y\in V)=\mathds{P}(y)/\mathds{P}(y\in V)$. So $J(y)$ becomes

We can compute the probability of sampling a feasible architecture $\mathds{P}(V):=\mathds{P}(y\in V)$ exactly when the search space is small, but this computation is too expensive when the space is large. Instead, we replace the exact probability $\mathds{P}(y)$ with a differential approximation $\widehat{\mathds{P}}(y)$ obtained with Monte-Carlo (MC) sampling. In each RL step, we sample $N$ architectures $\{z^{(k)}\}_{k\in[N]}$ within the search space with a proposal distribution $q$ and estimate $\mathds{P}(V)$ as

For each $k\in[N]$, $p^{(k)}$ is the probability of sampling $z^{(k)}$ with the factorized layerwise distributions and so is differentiable with respect to the logits. In contrast, $q^{(k)}$ is the probability of sampling $z^{(k)}$ with the proposal distribution, and is therefore non-differentiable.

$\widehat{\mathds{P}}(V)$ is an unbiased and consistent estimate of $\mathds{P}(V)$; $\nabla\log[\mathds{P}(y)/\widehat{\mathds{P}}(V)]$ is a consistent estimate of $\nabla\log[\mathds{P}(y\>|\>y\in V)]$ (Appendix J). A larger $N$ gives better results (Appendix H); in experiments, we need smaller than the size of the sample space to get a faithful estimate (Figure 5(b), Appendix D and I) because neighboring RL steps can correct the estimates of each other. We set $q=\text{stop\_grad}(p)$ in experiments for convenience: use the current distribution over architectures for MC sampling. Other distributions that have a larger support on $V$ may be used to reduce sampling variance (Appendix J).

At the end of NAS, we pick as our final architecture the layer sizes with largest sampling probabilities if the layerwise distributions are deterministic, or sample from the distributions $m$ times and pick $n$ feasible architectures with the largest number of parameters if not. Appendix B Algorithm 3 provides the full details. We find $m=500$ and $n\leq 3$ suffice to find an architecture that matches the reference (optimal) architecture in our experiments.

In practice, the distributions often (almost) converge after twice the number of epochs used to train a stand-alone child network. Indeed the distributions are often useful after training the same number of epochs in that the architectures found by Algorithm 3 are competitive. Figure 1 shows TabNAS finds the best feasible architecture, 4-2, in our toy example, using $\widehat{\mathds{P}}(V)$ estimated by MC sampling.

Experimental results

Our implementation can be found at https://github.com/google-research/tabnas. We ran all experiments using TensorFlow on a Cloud TPU v2 with 8 cores. We use a 1,027-dimensional input representation for the Criteo dataset and 180 features for VolkertOur paper takes these features as given. It is worth noting that methods proposed in feature engineering works like and are complementary to and can work together with TabNAS.. The best architectures in our FFN search spaces already produce near-state-of-the-art results; details in Appendix C.2. More details of experiment setup and results in other search spaces can be found in Appendix C and D. Appendix E tabulates the performance of all RL rewards on all tabular datasets in our experiments. Appendix F shows a comparison with Bayesian optimization and evolutionary search in similar settings; Ablation studies in Appendix I show TabNAS components collectively deliver desirable results; Appendix H shows TabNAS has easy-to-tune hyperparameters.

Section 3.3 discussed the resource-aware RL rewards and highlighted a potential failure case. In this section, we show several failure cases of three resource-aware rewards, $Q(y)(T(y)/T_{0})^{\beta}$, $Q(y)\max\{1,(T(y)/T_{0})^{\beta}\}$, and the Abs Reward $Q(y)+\beta|T(y)/T_{0}-1|$, on our tabular datasets.

We use the 32-144-24 reference architecture (41,153 parameters). Figure 3 gives an overview of the costs and losses of all architectures in the search space. The search space requires us to choose one of 20 possible sizes for each hidden layer; details in Appendix D. The search has $1.7\times$ the cost of a stand-alone training run.

The RL controller is to blame. To verify that a low quality SuperNet was not the culprit, we trained a SuperNet without updating the RL controller, and manually inspected the quality of the resulting SuperNet. The sampling probabilities for the RL controller remained uniform throughout the search; the rest of the training setup was kept the same. At the end of the training, we compare two sets of losses on each of the child networks: the validation loss from the SuperNet (one-shot loss), and the validation loss from training the child network from scratch. Figure 7 shows that there is a strong correlation between these accuracies; Figure 7 shows RL that starts from the sufficiently trained SuperNet weights in 7 still chooses the suboptimal choice 64. This suggests that the suboptimal search results on Criteo are likely due to issues with the RL controller, rather than issues with the one-shot model weights. In a 3 layer search space we can actually find good models without the RL controller, but in a 5 layer search space, we found an RL controller whose training is interleaved with the SuperNet is important to achieve good results.

1.2 Volkert – 4 layer search space

We search for 4-layer and 9-layer networks on the Volkert dataset; details in Appendix D. For resource-aware RL rewards, we ran a grid search over the RL learning rate and $\beta$ hyperparameter. The reference architecture for the 4 layer search space is 48-160-32-144 with 27,882 parameters. Despite a hyperparameter grid search, it was difficult to find models with the right target cost reliably using the MnasNet rewards. Using the Abs Reward (Figure 8), searched models met the target cost but their quality was suboptimal, and the trend is similar to what has been shown in the toy example (Figure 1): a smaller $|\beta|$ gives an infeasible architecture that is beyond the reference number of parameters, and a larger $|\beta|$ gives an architecture that is feasible but suboptimal.

1.3 A common failure pattern

Apart from Section 4.1.1 and 4.1.2, more examples in search spaces of deeper FFNs can be found in Appendix D. In cases on Criteo and Volkert where where the RL controller with soft constraints cannot match the quality of the reference architectures, the reference architecture often has a bottleneck structure. For example, with a 1,027-dimensional input representation, the 32-144-24 reference on Criteo has bottleneck 32; with 180 features, the 48-160-32-144 reference on Volkert has bottleneck 48 and 32. As the example in Section 3.3 shows, the wide hidden layers around the bottlenecks get penalized harder in the search, and it is thus more difficult for RL with the Abs Reward to find a model that can match the reference performance. Also, Appendix C.2.1 shows the Pareto-optimal architectures in the tradeoff points in Figure 3 often have bottleneck structures, so resource-aware RL rewards in previous NAS practice may have more room for improvement than previously believed.

2 NAS with TabNAS reward

With proper hyperparameters (Appendix H), our RL controller with TabNAS reward finds the global optimum when RL with resource-aware rewards produces suboptimal results.

TabNAS does not introduce a resource-aware bias in the RL reward (Section 3.3). Instead, it uses conditional probabilities to update the logits in feasible architectures. We run TabNAS for 120 epochs with RL learning rate 0.005 and $N=3072$ MC samples.The 3-layer search space has $20^{3}=8000$ candidate architectures, which is small enough to compute $\mathds{P}(V)$ exactly. However, MC can scale to larger spaces which are prohibitively expensive for exhaustive search (Appendix D). The RL controller converges to two architectures, 32-160-16 (40,769 parameters, with loss 0.4457 $\pm$ 0.0002) and 32-144-24 (41,153 parameters, with loss 0.4455 $\pm$ 0.0003), after around 50 epochs of NAS, then oscillates between these two solutions (Figure 9). After 120-epochs, we sample from the layerwise distribution and pick the largest feasible architecture: the global optimum 32-144-24.

On the same hardware, the search takes $3\times$ the runtime of stand-alone training. Hence, as can be seen in Figure 3, the proposed architecture search method is much faster than a random baseline.

3 TabNAS automatically determines whether bottlenecks are needed

Previous NAS works like MnasNet and TuNAS (often or only on vision tasks) often have inverted bottleneck blocks in their search spaces. However, the search spaces used there have a hard-coded requirement that certain layers must have bottlenecks. In contrast, our search spaces permit the controller to automatically determine whether to use bottleneck structures based on the task under consideration. TabNAS automatically finds high-quality architectures, both in cases where bottlenecks are needed and in cases where they are not. This is important because networks with bottlenecks do not always outperform others on all tasks. For example, the reference architecture 32-144-24 outperforms the TuNAS-found 32-64-96 on Criteo, but the reference 64-192-48-32 (64,568 parameters, 0.0662 $\pm$ 0.0011) is on par with the TuNAS-and-TabNAS-found 96-80-96-32 (64,024 parameters, 0.0669 $\pm$ 0.0013) on Aloi. TabNAS automatically finds an optimal (bottleneck) architecture for Criteo, and automatically finds an optimal architecture that does not necessarily have a bottleneck structure for Aloi. Previous reward-shaping rewards like the Abs Reward only succeed in the latter case.

4 Rejection-based reward outperforms Abs Reward in NATS-Bench size search space

Although we target resource-constrained NAS on tabular datasets in this paper, our proposed method is not specific to NAS on tabular datasets. In Appendix G, we show the rejection-based reward in TabNAS outperforms RL with the Abs Reward in the size search space of NATS-Bench , a NAS benchmark on vision tasks.

Conclusion

We investigate the failure of resource-aware RL rewards to discover optimal structures in tabular NAS and propose TabNAS for tabular NAS in a constrained search space. The TabNAS controller uses a rejection mechanism to compute the policy gradient updates from feasible architectures only, and uses Monte-Carlo sampling to reduce the cost of debiasing this rejection-sampling approach. Experiments show TabNAS finds better architectures than previously proposed RL methods with resource-aware rewards in resource-constrained searches.

Many questions remain open. For example: 1) Can the TabNAS strategy find better architectures on other types of tasks such as vision and language? 2) Can TabNAS improve RL results for more complex architectures? 3) Is TabNAS useful for resource-constrained RL problems more broadly?

Acknowledgments and Disclosure of Funding

This work was done when Madeleine Udell was a visiting researcher at Google. The authors thank Ruoxi Wang, Mike Van Ness, Ziteng Sun, Xuanyi Dong, Lijun Ding, Yanqi Zhou, Chen Liang, Zachary Frangella, Yi Su, and Ed H. Chi for helpful discussions, and thank several anonymous reviewers for useful comments.

References

Appendix A Previous work

Neural architecture search (NAS) stems from the resurgence of deep learning. It focuses on tuning architectural hyperparameters in neural networks, like hidden layer sizes in feedforward networks or convolutional kernel sizes in convolutional networks. The earliest NAS papers trained thousands of network architectures from scratch for a single search. Due to the high costs involved, many works have proposed different methods to reduce the search cost. Most of these proposals are based around two complementary (but often intertwined) high-level strategies.

The first strategy is to reduce the time needed to evaluate each architecture seen during a search. For example, instead of training each network architecture from scratch, we can train a SuperNet – a single set of shared model weights that can be used to evaluate and rank any different candidate architecture in the search space . Other approaches include the use of network morphism to initialize weights of candidate networks that are close to previous instances.

The second strategy is to reduce the number of architectures we need to evaluate during a search. Proposed methods include reinforcement learning algorithms that learn a probability distribution over candidate architectures, evolutionary search that pursues more promising architectures from ancestors, and Bayesian optimization and parametric models that directly predict network performance.

Resource constraints are prevalent in deep learning. Finding architectures with outstanding performance and low costs are important to both NAS research and application. Apart from the surrogate models above that transfer knowledge across network candidates to avoid exhaustive search, specific techniques have been adopted to find networks with a good balance of performance and resource consumption. A popular method is to add regularizers that penalize expensive architectures. With hard resource constraints, greedy submodular maximization and heuristic scaling methods were used to grow or shrink networks during the search, so as to ensure the chosen candidate architecture obeys the constraint. TabNAS in this work operates under the hard resource constraint, and can find the global optima in the feasible set of architectures.

A.2 Deep learning on tabular datasets

Deep neural networks are gaining popularity on tabular datasets in academia and industry. In the pursuit of designing better architectures for tabular deep learning, an earlier line of work mimic the structure of tree-based models . Some other works use the attention mechanism or are based on feedforward or residual networks .

While automated machine learning (AutoML) on tabular datasets has been addressed from multiple perspectives and with different approaches , NAS on tabular datasets is less explored, partly due to insufficient understanding of promising architectures and a lack of benchmarks To disambiguate, the phrase “tabular NAS benchmark” in previous NAS benchmark literature often refers to tabulated performance of architectures on vision and language tasks.. Egele et al. interleaves NAS with aging evolution in a search space with multiple branches and hyperparameter tuning with Bayesian optimization. We show TabNAS can find architectures as simple as FFNs with a few layers that have outstanding performance and obey the resource constraint.

Appendix B Algorithm pseudocode

We show pseudocode of the algorithms introduced in Section 3.

Notice that in Algorithm 1, we show the weight and RL updates with the stochastic gradient descent (SGD) algorithm; in our experiments on the toy example and real datasets, we use Adam for both updates as in ProxylessNAS and TuNAS , since it synchronizes convergence across different layer size choices, and slows down the learning which would otherwise converge too rapidly.

Appendix C Details of experiment setup

We use the Adam optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.999$ and $\epsilon=0.001$ to update the logits. When we use the Abs Reward, the results are similar when $\eta\geq 0.05$, while the RL controller with $\eta<0.05$ converges too slow or is hard to converge. When we use the rejection-based reward, we use RL learning rate $\eta=0.1$; other $\eta$ values with which RL converges give similar results.

C.2 Real datasets

Table 1 shows the datasets we use. Datasets other than Criteo https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/ come from the OpenML dataset repository . For Criteo, we randomly split the labeled part (45,840,617 points) into 90% training (41,258,185 points) and 10% validation (4,582,432 points); for the other datasets, we randomly split into 80% training and 20% validationThe ranking of validation losses among architectures under such splits is almost the same as that of test losses under 60%-20%-20% training-validation-test splits.. The representations we use for Criteo are inspired by DCN-V2 .

Table 2 shows the hyperparameters we use for stand-alone training and NAS, found by grid search. With these hyperparameters, the best architecture in each of our search spaces (introduced in Appendix C.2.1) has performance that is within $\pm 5\%$ of the best performance in Kadra et al. Table 2, and we achieve these scores with FFNs that only have 5% parameters of the ones there. The Adam optimizer has hyperparameters $\beta_{1}=0.9$, $\beta_{2}=0.999$ and $\epsilon=0.001$. We use layer normalization for all datasets. We use balanced error (weighted average of classification errors across classes) for all other datasetsThe performance ranking of architectures under the balanced error metric is almost the same as under logistic loss. Also, the balanced error metric is only for reporting the final validation losses; both weight and RL updates use logistic loss. as in Kadra et al. , except for Criteo, on which we use logistic loss as in Wang et al. .

We use constant RL learning rates for NAS. The Connect-4https://www.openml.org/d/40668 and Higgshttps://www.openml.org/d/23512 datasets are easy for both the Abs Reward and rejection-based reward, in the sense that small FFNs with fewer than 5,000 parameters can achieve near-SOTA results ($\pm 5\%$ of the best accuracy scores listed in Kadra et al. Table 2, except that we do 80%-20% training-validation splits and use original instead of standardized features), and RL-based weight-sharing NAS with either reward can find architectures that match the Pareto-optimal reference architectures. The Aloi datasethttps://www.openml.org/d/42396 needs more parameters (more than 100k), but the other observations are similar to on Connect-4 and Higgs. Thus we omit the corresponding results.

The factorized search spaces we use for NAS are:

Criteo: Each layer has 20 choices {8, 16, 24, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 384, 512}.

Volkerthttps://www.openml.org/d/41166, 4-layer networks: Each layer has 20 choices {8, 16, 24, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 384, 512}.

Volkert, 9-layer networks: Each layer has 12 choices {8, 16, 24, 32, 48, 64, 80, 96, 112, 128, 144, 160}. This search space has fewer choices for each hidden layer than the 4-layer counterpart, but the size of the search space is over $3\times 10^{4}$ times larger.

Our goal is not to achieve state-of-the-art (SOTA) accuracy, but to find the best architecture that obeys a resource upper bound. This mimics a resource-constrained setting that is common in practice. Impressively, our method does nearly match SOTA performance, as our search space has architectures that are close to the best in previous literature. For example:

On Criteo: The best architecture in our search space achieves public and private scores 0.45284 and 0.45283 on Kaggle, ranking 20/717 on the leaderboardhttps://www.kaggle.com/competitions/criteo-display-ad-challenge/leaderboard.

On Volkert: The best architecture in our search space has balanced accuracy 0.695, within 2% of the best in Kadra et al. and better than most other works used for comparison in that work. The differences in settings are that we use original features instead of standardized, and we achieve this score with an FFN that only has 5% parameters of the one there.

On Aloi: The best architecture in our search space has balanced accuracy 0.957, within 2% of the best in Kadra et al. and better than most other works used for comparison in that work. Again, we use original features instead of standardized, and we achieve this score with an FFN that only has 5% parameters of the one used there. And the 0.957 balanced accuracy score is also within 1% of the best in Gorishniy et al. . The difference is that we use an FFN that only has $<$10% as many parameters as the one used there.

Each search space we use for exhaustive search and NAS has a fixed number of hidden layers. Resource-constrained NAS in a search space with varying number of hidden layers is an interesting problem for future studies. On each dataset, we randomly sample, train and evaluate architectures in the search space with the number of parameters fall within a range, in which there is a clear tradeoff between loss and number of parameters. These ranges are:

Volkert, 4-layer networks: 15,000 – 50,000

Volkert, 9-layer networks: 40,000 – 100,000

Figure 11 shows the tradeoffs between loss and number of parameters in these search spaces. When training each architecture 5 times, the standard deviation (std) across different runs is 0.0002 for CriteoOn Criteo, “a 0.001-level improvement (of logistic loss) is considered significant” . and 0.004 for Volkert, meaning that the architectures whose performance difference is larger than $2\times$ std are qualitatively different. We use Pareto-optimal architectures as the reference of resource-constrained NAS: we want an architecture that both matches (or even beatsNote that the Pareto optimality of the reference architecture is determined by only one round of random search. Thus because of the randomness across multiple training runs, the other architectures are likely to beat the reference architecture: a “regression toward the mean”.) the performance of the reference architecture and has no more parameters than the reference. Most Pareto-optimal architectures in Figure 11 have the bottleneck structure; Table 3 shows some examples.

C.2.2 More details on TPU implementation

When we run one-shot NAS on a TPU that has multiple TPU cores (for example, each Cloud TPU-v2 we use has 8 cores), each core samples an architectures independently, and we use the average loss and reward for weight and RL updates, respectively. This means our algorithm actually samples multiple architectures in each iteration and uses the tensorflow.tpu.cross_replica_sum() method to compute their average effect on the gradient. Since only a fraction of architectures are feasible in each search space, we set the losses and rewards given by the infeasible architectures to 0 before averaging, so that we are equivalently only averaging across the sampled architectures that are feasible. We then reweight the average loss or reward with number_of_cores / number_of_feasible_architectures to obtain an unbiased estimate.

C.2.3 More details on the NAS method comparison plot (Figure 3)

For each architecture below, we report its number of parameters and mean $\pm$ std logistic loss across 5 stand-alone training runs in brackets.

We have the reference architecture 32-144-24 (41,153 parameters, 0.4454 $\pm$ 0.0003) for NAS methods to match. In the search space with $20^{3}=8000$ candidate architectures:

TabNAS trials with no fewer than 2,048 Monte-Carlo samples and the RL learning rate $\eta$ among {0.001, 0.005, 0.01} consistently finds either the reference architecture itself, or an architectures that is qualitatively the same as the reference, like 32-112-32 (40,241 parameters, 0.4456 $\pm$ 0.0003).

NAS with the Abs Reward: After grid search over RL learning rate $\eta$ (among {0.0001, 0.0005, 0.001, 0.005, 0.01, 0.015, 0.02, 0.025, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.15, 0.2, 0.25, 0.3, 0.4, 0.5, 0.75, 1.0, 1.5, 2.0}) and $\beta$ (among {-0.0005, -0.001, -0.005, -0.01, -0.05, -0.1, -0.5, -0.75, -1.0, -1.25, -1.5, -2.0, -3.0}), the RL controller finds 32-64-96 (41,345 parameters, 0.4461 $\pm$ 0.0003) or 32-80-64 (40,785 parameters, 0.4459 $\pm$ 0.0002) among over $90\%$ trials that eventually find an architecture within $\pm 5\%$ of the target number of parameters 41,153.

C.3 Difficulty in using the MnasNet reward

With the MnasNet reward, only fewer than 1% NAS trials in our hyperparameter grid search (the ones with a medium $\beta$) can find an architecture whose number of parameters is within $\pm 5\%$ of the reference, and among which none or only one (out of tens) can match the reference performance. In contrast, TuNAS with the Abs Reward finds an architecture with number of parameters within $\pm 5\%$ of the reference among over 50% of the grid search trials described in Appendix C.2.3, and TabNAS with the rejection-based reward consistently finds such architectures at medium RL learning rates $\eta$ and decently large numbers of MC samples $N$. This means it is significantly more difficult to use the MnasNet reward than competing approaches in the practice of resource-constrained tabular NAS.

Appendix D More failure cases of the Abs Reward, and performance of TabNAS

For each architecture below, we report its number of parameters and mean $\pm$ std loss across 5 stand-alone training runs (logistic loss for Criteo, balanced error for the others) in brackets.

On Criteo, in the 4-layer search space. We have the reference architecture 48-128-16-112 (59,697 parameters, 0.4451 $\pm$ 0.0002) for NAS to match in the search space (shown as Figure 11(a)). Similar to Figure 3, we show similar results on NAS with rejection-based reward (TabNAS) and NAS with the Abs Reward (TuNAS) in Figure 12(a). In the search space with $20^{4}=1.6\times 10^{5}$ candidate architectures:

TabNAS with 32,768 Monte-Carlo samples and RL learning rate $\eta$ among {0.001, 0.005, 0.01} consistently finds architectures qualitatively the same as the reference. Example results include 48-128-24-32 (59,545 parameters, 0.4449 $\pm$ 0.0002), 48-144-16-48 (59,585 parameters, 0.4448 $\pm$ 0.0001), 48-112-16-144 (59,233 parameters, 0.4448 $\pm$ 0.0002) and the reference architecture itself.

NAS with the Abs Reward successfully finds the reference architecture 48-128-16-112 in 3 out of 338 hyperparameter settings on a $\beta$-$\eta$ grid. Other found architectures include 48-80-32-112 (59,665 parameters, 0.4452 $\pm$ 0.0002), 32-128-80-144 (59,249 parameters, 0.4453 $\pm$ 0.0003) and 48-160-8-48 (58,953 parameters, 0.4448 $\pm$ 0.0003), among which the first two are inferior to the TabNAS-found counterparts.

On Criteo, in the 5-layer search space. We have the reference architecture 48-240-24-256-8 (75,353 parameters, 0.4448 $\pm$ 0.0002) for NAS methods to match in the search space (shown as Figure 11(b)). Similar to Figure 3, we have similar results on the comparison among random sampling, NAS with rejection-based reward (TabNAS), and NAS with the Abs Reward as Figure 12(b). In the search space with $20^{5}=3.2\times 10^{6}$ candidate architectures:

TabNAS with 32,768 Monte-Carlo samples and the RL learning rate $\eta=0.005$ consistently finds architectures qualitatively the same as the reference. Example results include 48-176-64-16-256 (74,945 parameters, 0.4445 $\pm$ 0.0002), 48-208-48-48-64 (75,121 parameters, 0.4444 $\pm$ 0.0001), 48-256-32-80-24 (74,721 parameters, 0.4446 $\pm$ 0.0003) and 48-176-80-16-96 (75,153 parameters, 0.4445 $\pm$ 0.0002).

NAS with the Abs Reward finds 64-80-48-8-8 (75,353 parameters, 0.4448 $\pm$ 0.0001), 64-80-24-16-112 (75,353 parameters, 0.4447 $\pm$ 0.0001), 48-144-96-16-192 (75,329 parameters, 0.4446 $\pm$ 0.0001) and 64-96-8-32-64 (75,273 parameters, 0.4445 $\pm$ 0.0001) that are mostly inferior to the TabNAS-found architectures.

On Volkert, in the 4-layer search space. We have the reference architecture 48-160-32-144 (27,882 parameters, 0.3244 $\pm$ 0.0040) for NAS to match in the search space (shown as Figure 11(c)). Similar to Figure 3, we draw the comparison plot among random sampling, NAS with rejection-based reward (TabNAS), and NAS with the Abs Reward as Figure 12(c). In the search space with $1.6\times 10^{5}$ candidate architectures:

TabNAS with $10^{4}$ Monte-Carlo samples and the RL learning rate $\eta\in\{0.001,0.005,0.01,0.05\}$ consistently finds either the reference architecture itself or other architectures qualitatively the same. Examples include 64-128-48-16 (27,050 parameters, 0.3237 $\pm$ 0.0040), 80-48-112-32 (27,802 parameters, 0.3274 $\pm$ 0.0037), 64-96-80-24 (27,778 parameters, 0.3279 $\pm$ 0.0005), and 64-144-32-48 (27,658 parameters, 0.3204 $\pm$ 0.0038).

NAS with the Abs Reward finds 96-64-32-48 (27,738 parameters, 0.3302 $\pm$ 0.0042), 96-48-32-96 (27,738 parameters, 0.3305 $\pm$ 0.0047), 96-80-16-48 (27,738 parameters, 0.3302 $\pm$ 0.0050), 112-48-24-24 (27,722 parameters, 0.3301 $\pm$ 0.0034) and 80-80-48-48 (27,690 parameters, 0.3309 $\pm$ 0.0022) that are inferior.

On Volkert, in the 9-layer search space. We further do NAS on Volkert in the 9-layer search space to test the ability of TabNAS in searching among significantly deeper FFNs. The tradeoff between loss and number of parameters in the search space is shown in Figure 11(d). We have the reference architecture 144-128-112-16-16-48-144-24-160 (78,114 parameters, 0.3126 $\pm$ 0.0050) for NAS to match. We compare random sampling, NAS with rejection-based reward (TabNAS), and NAS with the Abs Reward in Figure 12(d). In the search space with $5.2\times 10^{9}$ candidate architectures (which is nearly impossible for exhaustive search):

TabNAS with $5\times 10^{6}$ Monte-Carlo samples and the RL learning rate $\eta\in\{0.002,0.005\}$ consistently finds architectures that are qualitatively the same as the reference. These architectures are found when the RL controller is far from converged and when $\mathds{P}(V)$ slightly decreases after RL starts. Example results include 144-144-112-64-24-16-128-8-128 (78,026 parameters, 0.3120 $\pm$ 0.0049), 128-160-96-32-24-64-64-32-160 (77,890 parameters, 0.3127 $\pm$ 0.0040), 128-144-112-32-64-64-80-16-128 (77,834 parameters, 0.3094 $\pm$ 0.0012), 160-128-96-32-48-64-48-24-112 (78,002 parameters, 0.3137 $\pm$ 0.0021), and 144-112-160-24-112-16-16-128-48 (77,986 parameters, 0.3119 $\pm$ 0.0029).

NAS with the Abs Reward finds 144-96-80-80-48-64-96-80-32 (78,170 parameters, 0.3094 $\pm$ 0.0039), 160-80-160-24-80-16-80-64-128 (78,114 parameters, 0.3158 $\pm$ 0.0020), 128-96-80-80-64-64-80-80-80 (78,106 parameters, 0.3128 $\pm$ 0.0020), and 144-128-80-16-16-16-96-160-24 (78,050 parameters, 0.3192 $\pm$ 0.0014). Interestingly, all architectures except 144-96-80-80-48-64-96-80-32 are inferior to the TabNAS-found architectures despite having slightly more parameters, and 144-96-80-80-48-64-96-80-32 does not have an evident bottleneck structure like the other architectures found here.

Appendix E Tabulated performance of different RL rewards on all datasets

Table 4 summarizes among results of RL with the rejection-based reward, the Abs Reward, two MNasNet rewards and the reward in RENA Equation 3. Bold results in each column indicate architectures that are on par with the best for the corresponding dataset. The rejection-based reward in TabNAS gets the best architectures overall.

Appendix F Comparison with Bayesian optimization and evolutionary search in one-shot NAS

train-then-search-with-RL: Train the SuperNet for $M$ epochs as in TabNAS, then do NAS for $0.75M$ epochs with rejection-based RL (with each iteration as Algorithm 2).

train-then-search-with-BO: Train the SuperNet for $M$ epochs as in TabNAS, then do BO (by Gaussian processes with expected improvement ) in the set of feasible architectures with a similar number of SuperNet forward passes for child network evaluation.

train-then-search-with-ES: Train the SuperNet for $M$ epochs as in TabNAS, then do ES in the set of feasible architectures as Algorithm 1 in Guo et al. with a similar number of SuperNet forward passes for child network evaluation: in each iteration, start with population size $P$, pick top-$k$ architectures, crossover and mutate these top-$k$ to each get $P/2$ architectures, then combine the crossover and mutation results to get the population for the next iteration.

On Criteo, the cost of forward passes for RL is comparable to evaluating 405 child networks on the validation set. The search space of 5-layer FFNs has 340,590 feasible architectures below the 75,353 parameters limit (corresponding reference architecture is 48-240-24-256-8) in Figure 11(b). In BO, we tune RBF kernel length scale, number of initially sampled architectures, and number of new architectures to sample in each step; in ES, we tune population size and $k$. We can see from the results in Table 5 that:

RL-based methods (TabNAS or train-then-search-with-RL) stably finds architectures that are qualitatively the best.

There is a large variance across architectures found by each hyperparameter setting of BO or ES. The search results are sensitive to initialization and are worse than those found by RL in over $2/3$ trials. We observe that each of the BO and ES searches quickly gets stuck at an architecture that is close to the best architecture in the initial set of samples.

The local optima that BO and ES get stuck at still often have bottleneck structures, but the number of parameters is often significantly below the limit (e.g., 63,817 parameters in the searched model vs. a limit of 75,353 parameters). Model performance suffers as a result.

Many interesting questions on BO and ES for one-shot NAS remain open for future work, including how to control the initialization randomness (for better exploration of the search space) and how to design methods that properly interleave weight training and NAS steps under such large exploration randomness (for better exploitation of promising architectures).

Appendix G Rejection-based reward outperforms Abs Reward in NATS-Bench size search space (full version)

In the vision domain, the NATS-Bench size search space has convolutional network architectures with a predefined skeleton and 8 candidate sizes {8, 16, 24, 32, 40, 48, 56, 64} for each of its 5 layers. In this search space with $8^{5}=32768$ candidate architectures, we use the true number of floating point operations (#FLOPs) as our cost metric, validation accuracy on CIFAR-100 as RL quality reward, and test error (1 - test accuracy) on CIFAR-100 as final performance metric (the use of error instead of accuracy is consistent with other results here). We use 75M #FLOPs as the resource limit, so that there are 13,546 (41.3%) feasible architectures in the search space. This experiment setting is the same as in the toy example (Figure 1), where we regard the network weights as given and only compare the RL controllers. We do grid search over NAS hyperparameters: {0.01, 0.05, 0.1, 0.5} for the RL learning rate for both the rejection-based reward and the Abs Reward, and {10, 5, 2, 1, 0.5, 0.2, 0.1} for $|\beta|$ in the Abs Reward. We show #FLOPs and test error statistics (mean $\pm$ std) of architectures found by the rejection-based RL reward and the Abs Reward across 500 NAS repetitions of each experiment setting.

Detailed results of #FLOPs and test errors of architectures found by either of the RL rewards at different NAS hyperparameters are listed in Table 6 and 7, respectively.

RL with the rejection-based reward finds architectures that are within the 75M #FLOPs limit.

When $|\beta|$ is large, the architectures found by RL with the Abs Reward are within the 75M #FLOPs limit, but they are inferior in quality (have larger test errors).

When $|\beta|$ is small, the architectures found by RL with the Abs Reward are similar or better in quality, but they exceed the 75M #FLOPs limit by >5%.

These observations are similar to what we saw in the toy example, showing the rejection-based reward outperforms in tasks from both tabular and vision domains.

Taking a closer look at the architectures that are Pareto-optimal and those found by each RL reward, we can see that:

The Pareto-optimal architectures often have bottleneck structures like their counterparts in tabular tasks; examples include (in NATS-Bench format, and same below) 8:32:16:40:48, 8:56:40:56:64, and 16:56:64:64:56 that have a bottleneck in their first channel.

Bottleneck structures occur less often in the architectures found by RL with the Abs Reward. Examples include 24:24:64:40:32, 24:40:56:40:48, and 32:48:40:32:32. The architectures found by RL with the rejection-based reward (our method) have more similar appearances as the Pareto-optimal architectures; examples include 8:56:64:64:64, 8:40:56:64:56, and 16:64:48:64:64.

This means the co-adaptation problem (described in Section 1 of the main paper) also occurs in RL-based NAS with resource-aware rewards in the image domain.

Appendix H Difficulty of hyperparameter tuning

Hyperparameter tuning has always been a headache for machine learning. In the design of NAS approaches, the hope is that the NAS hyperparameters are much easier to tune than the architectures NAS search over. We denote the RL learning rate and the number of MC samples by $\eta$ and $N$, respectively. The three resource-aware rewards (in MnasNet and TuNAS) have both $\eta$ and $\beta$ as hyperparameters; our TabNAS with the rejection-based reward has $\eta$ and $N$ to tune.

$\beta$ is difficult to tune in experiments: the best value varies by dataset and lies in the middle of its search space. Since $\beta<0$, we discuss its absolute value. In a NAS search space, the architecture that is feasible and can match the reference performance often has the number of parameters that is more than $98\%$ of the reference. A too small $|\beta|$ is not powerful enough to enforce the resource constraint, in which case NAS finds an architecture that is far from the target number of parameters and makes the search nearly unconstrained (e.g., the Abs Reward with $|\beta|=1$ in the toy example, shown in Figure 1 and towards the left end in Figure 13(a)). A too large $|\beta|$ severely penalizes the violation of the resource constraint, in which case the RL controller would always give an architecture close to the reference, with much bias (e.g., the Abs Reward with $|\beta|=2$ in Figure 1, and towards the right end in Figure 13(a)). Thus practitioners seek a medium $|\beta|$ in hyperparameter tuning to both obey the resource constraint and achieve a better result. In our experiments, such “appropriate” medium values vary largely across datasets: 1 on Criteo with the 32-144-24 reference architecture (41,153 parameters), 2 on Volkert with the 48-160-32-144 reference architecture (27,882 parameters), and 25 on Aloi with the 64-192-48-32 reference architecture (64,568 parameters).

H.2 RL learning rate η𝜂\eta

The RL learning rate $\eta$ is easier to tune and more generalizable across datasets than $\beta$. With a large $\eta$, the RL controller quickly converges right after the first $25\%$ epochs of layer warmup; with a small $\eta$, the RL controller converges slowly or may not converge, although there may still be enough signal from the layerwise probabilities to get the final result. It is thus straightforward to tune $\eta$ by observing the convergence behavior of sampling probabilities. In our experiments, the appropriate value of $\eta$ does not significantly vary across tasks: a constant $\eta\in[0.001,0.01]$ is appropriate for all datasets and all number of parameter limits.

H.3 Number of MC samples N𝑁N

The number of MC samples $N$ is also easier to tune than $\beta$. Resource permitting, $N$ is the larger, the better (Figure 13(b)), so that $\mathds{P}(V)$ can be better estimated. When $N$ is too small, the MC sampling has a high chance of missing the valid architectures in the search space, and thus incurs large bias and variance for the estimate of $\nabla\log[\mathds{P}(y\>|\>y\in V)]$. In such cases, $\widehat{\mathds{P}}(V)$ may miss all valid architectures at the beginning of RL and quickly converge to 0. $\widehat{\mathds{P}}(V)$ being equal or close to 0 is a bad case for our rejection-based algorithm: the single-step RL objective $J(y)$ that has a $-\log(\widehat{\mathds{P}}(V))$ term grows extremely large and gives an explosive gradient to stuck the RL controller in the current choice. Consequently, the criterion for choosing $N$ is to choose the largest that can afford, and hopefully, at least choose the smallest that can make $\widehat{\mathds{P}}(V)$ steadily increase during RL. Figure 14 shows the changes of $\widehat{\mathds{P}}(V)$ on Criteo with the 32-144-24 reference in the search space of 8,000 architectures at three $N$ values. The NAS succeeds when $N\geq 2048$, same as the threshold that makes $\widehat{\mathds{P}}(V)$ increase.

Overall, the RL controller with our rejection-based reward has hyperparameters that are easier to tune than with resource-aware rewards in MnasNet and TuNAS.

Appendix I Ablation studies

We do the ablation studies on Criteo with the 32-144-24 reference. The behavior on other datasets with other reference architectures are similar.

Whether to use $\widehat{\mathds{P}}(V)$ instead of $\mathds{P}(V)$. The Monte-Carlo (MC) sampling estimates $\mathds{P}(V)$ with $\widehat{\mathds{P}}(V)$ to save resources. Such estimations are especially efficient when the sample space is large. Empirically, the $\widehat{\mathds{P}}(V)$ estimated with enough MC samples (as described in Appendix H) enables the RL controller to find the same architecture as $\mathds{P}(V)$, because the $\widehat{\mathds{P}}(V)$ estimated with a large enough number of samples is accurate enough (e.g., Figure 5(b) and 9(d)).

Whether to skip infeasible architectures in weight updates. In each iteration of one-shot training and REINFORCE (Appendix B Algorithm 1) with the rejection mechanism (Appendix B Algorithm 2), we train the weights in the sampled child network $x$ regardless of whether $x$ is feasible. Instead, we may update the weights only when $x$ is feasible, in a similar rejection mechanism as the RL step. We find this mechanism may mislead the search because of insufficiently trained weights: the rejection-based RL controller can still find qualitatively the best architectures on Criteo with the 32-144-24 or 48-240-24-256-8 reference, but fails with the 48-128-16-112 reference. In the latter case, although the RL controller still finds architectures with bottleneck structures (e.g., 32-384-8-144), the first layer sizes of the found architectures are much smaller, leading to suboptimal performance.

Whether to differentiate through $\widehat{\mathds{P}}(V)$. Recall that REINFORCE with rejection has the objective

To update the RL controller’s logits, we compute $\nabla J(y)$, which requires a differentiable approximation of $\mathds{P}(V)$. From a theoretical standpoint, omitting the extra term $\mathds{P}(V)$ – or using a non-differentiable approximation – will result in biased gradient estimates. Empirically, we ran experiments with multiple variants of our algorithm where we omitted the term $\mathds{P}(V)$, but found that the quality of the searched architectures was significantly worse.

In the case that we do not skip infeasible architectures in weight updates, the largest hidden layer sizes may gain and maintain the largest sampling probabilities soon after RL starts. This is because most architectures in the 3-layer Criteo search space are above the number of parameters limit 41,153. When RL starts, the sampled feasible architectures underperform the moving average, thus their logits are severely penalized, making the logits of the infeasible architectures (which often have wide hidden layers) quickly dominate (Figure 15(a)). Accordingly, the (estimated) valid probability $\mathds{P}(V)$ (or $\widehat{\mathds{P}}(V)$) quickly decrease to 0 (Figure 15(b)), and the RL controller gets stuck (as described in Appendix H.3) in these large choices for hidden layer sizes.

In the case that we skip infeasible architectures in both weight and RL updates, the RL controller eventually picks feasible architectures with bottleneck structures, but the found architectures are almost always suboptimal: when RL starts, the controller severely boosts the logits of the sampled feasible architectures without much exploration in the search space, and quickly gets stuck there. For example, the search in Figure 15(c)) finds 24-384-16 (40,449 parameters) that is feasible but suboptimal; $\mathds{P}(V)$ and $\widehat{\mathds{P}}(V)$ quickly increase to 1 after RL starts (Figure 15(d)).

Strategy for choosing the final architecture after search. When RL finishes, instead of biasing towards architectures with more parameters (Appendix B Algorithm 3), we may also bias towards those that are feasible and have larger sampling probabilities. We find that when the final distributions are less deterministic, the architectures found by the latter strategy to perform worse: for example, the top 3 feasible architectures found with the final distribution in Figure 9 are 32-128-16, 32-160-16 and 32-128-8, and they are all inferior to 32-144-24.

Appendix J Proofs

Within the search space $S$, recall the definitions of $\mathds{P}(V)$ and $\widehat{\mathds{P}}(V)$:

$\mathds{P}(V)=\sum\limits_{z^{(i)}\in S}p^{(i)}\mathds{1}(z^{(i)}\in V)$

$\widehat{\mathds{P}}(V)=\frac{1}{N}\sum\limits_{k\in[N],z^{(k)}\in V}\frac{p^{(k)}}{q^{(k)}}=\frac{1}{N}\sum\limits_{k\in[N]}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in V)$

Unbiasedness. With $N$ architectures sampled from the proposal distribution $q$, we take the expectation with respect to $N$ sampled architectures:

Consistency. We first show the variance of $\mathds{P}(V)$ converges to 0 as the number of MC samples $N\rightarrow\infty$. Because of independence among samples,

which goes to 0 as $N\rightarrow\infty$. It worths noting that when we set $q=\text{stop\_grad}(p)$, the single-summand variance (Equation 3) becomes $\mathds{P}(V)-\mathds{P}(V)^{2}$, which is the variance of a Bernoulli distribution with mean $\mathds{P}(V)$.

Since $\mathds{P}(y\>|\>y\in V)=\frac{\mathds{P}(y)}{\mathds{P}(V)}$, we show $\operatorname*{plim}\limits_{N\rightarrow\infty}\nabla\log\widehat{\mathds{P}}(V)=\nabla\log\mathds{P}(V)$ below to prove consistency, in which $\operatorname*{plim}\limits_{N\rightarrow\infty}$ denotes convergence in probability.

Recall $p^{(i)}$ is the probability of sampling the $i$-th architecture $z^{(i)}$ within the search space $S$, and the definitions of $\mathds{P}(V)$ and $\widehat{\mathds{P}}(V)$ are:

$\mathds{P}(V)=\sum\limits_{z^{(i)}\in S}p^{(i)}\mathds{1}(z^{(i)}\in V)$,

Together with the condition that $\mathds{P}(V)>0$ (the search space contains at least one feasible architecture), we have the desired result for consistency as $\operatorname*{plim}\limits_{N\rightarrow\infty}\nabla\log\widehat{\mathds{P}}(V)=\operatorname*{plim}\limits_{N\rightarrow\infty}\frac{\nabla\widehat{\mathds{P}}(V)}{\widehat{\mathds{P}}(V)}=\frac{\operatorname*{plim}\limits_{N\rightarrow\infty}\nabla\widehat{\mathds{P}}(V)}{\operatorname*{plim}\limits_{N\rightarrow\infty}\widehat{\mathds{P}}(V)}=\frac{\nabla\mathds{P}(V)}{\mathds{P}(V)}=\nabla\log\mathds{P}(V)$, in which the equalities hold due to the properties of convergence in probability.