# TabNAS: Rejection Sampling for
Neural Architecture Search on Tabular Datasets

Chengrun Yang1, Gabriel Bender1, Hanxiao Liu1, Pieter-Jan Kindermans1,
Madeleine Udell2, Yifeng Lu1, Quoc V. Le1, Da Huang1
{chengrun, gbender, hanxiaol<EMAIL_ADDRESS>
<EMAIL_ADDRESS>{yifenglu, qvl<EMAIL_ADDRESS>
1 Google Research, Brain Team 2 Stanford University

###### Abstract

The best neural architecture for a given machine learning problem depends on
many factors: not only the complexity and structure of the dataset, but also
on resource constraints including latency, compute, energy consumption, etc.
Neural architecture search (NAS) for tabular datasets is an important but
under-explored problem. Previous NAS algorithms designed for image search
spaces incorporate resource constraints directly into the reinforcement
learning (RL) rewards. However, for NAS on tabular datasets, this protocol
often discovers suboptimal architectures. This paper develops TabNAS, a new
and more effective approach to handle resource constraints in tabular NAS
using an RL controller motivated by the idea of rejection sampling. TabNAS
immediately discards any architecture that violates the resource constraints
without training or learning from that architecture. TabNAS uses a Monte-
Carlo-based correction to the RL policy gradient update to account for this
extra filtering step. Results on several tabular datasets demonstrate the
superiority of TabNAS over previous reward-shaping methods: it finds better
models that obey the constraints.

##  1 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 [18, 26].
For vision, optimizing models for practical deployment often relies on Neural
Architecture Search (NAS). Most NAS literature targets convolutional networks
on vision benchmarks [34, 9, 21, 55]. Despite the practical importance of
tabular data, however, NAS research on this topic is quite limited [15, 13].
(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 [6]. 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 [40, 9, 7] 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 [7] 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., [40, 8]). However, LSTMs are often much slower to train
and are far more difficult to tune.

Figure 1: A toy example for tabular NAS in a 2-layer search space with a
2-dimensional input and a limit of 25 parameters. Each cell represents an
architecture. The left half shows the number of parameters and loss of each
candidate in the search space. Infeasible architectures have red-striped
cells. The bottom left cell (bold border) is the global optimum with size 4
for the first hidden layer and size 2 for the second. The right half shows the
change of sampling probabilities in NAS with different RL rewards. Sampling
probabilities are shown both as percentages in cells, and intensity indicated
by the right colorbar. Orange bars on the top and right sides show the
(independent) sampling probability distributions of size candidates for
individual layers. With the Abs Reward, the sampling probability of each
architecture is the product of sampling probabilities of each layer; with the
rejection-based reward, the probability of an infeasible architecture is 0,
and probabilities of feasible architectures are normalized to sum to 1. At
epoch 500, the cell squared in bold shows the architecture picked by the
corresponding RL controller. RL with the Abs Reward $Q(y)+\beta|T(y)/T_{0}-1|$
proposed in TuNAS [7] either converges to a feasible but suboptimal
architecture ($\beta=-2$, middle row) or violates the resource constraint
($\beta=-1$, top row). Other latency-aware reward functions show similar
failures. In contrast, TabNAS converges to the optimum (bottom row).

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.

Figure 2: TabNAS reward distributionally outperforms random search and
resource-aware Abs Reward on the Criteo dataset within a 3-layer search space.
All error bars and shaded regions are 95% confidence intervals. The x axis is
the time relative to train time for a single architecture. The y axis is the
validation loss. More details in Appendix C.2.3.  Figure 3:  Validation loss
(logistic) vs. number of parameters on Criteo with a 3-layer search space. The
standard deviation (std) of architecture performance for different runs is
0.0002, so architectures whose performance difference is larger than $2$std
are qualitatively different with high probability. The search space and
Pareto-optimal architectures are shown in Appendix C.2.1.

##  2 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
[1] or f.detach() in PyTorch [39]. $\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 [26]. 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.

Search space. When searching $L$-layer networks, we use capital letters like
$X=X_{1}\mathchar 45\relax\dots\mathchar 45\relax X_{L}$ to denote the random
variable of sampled architectures, in which $X_{i}$ is the random variable for
the size of the $i$-th layer. We use lowercase letters like $x=x_{1}\mathchar
45\relax\dots\mathchar 45\relax x_{L}$ to denote an architecture sampled from
the distribution over $X$, in which $x_{i}$ is an instance of the $i$-th layer
size. When there are multiple samples drawn, we use a bracketed superscript to
denote the index over samples: $x^{(k)}$ denotes the $k$-th sample. The search
space $S=\\{s_{ij}\\}_{i\in[L],j\in[C_{i}]}$ has $C_{i}$ choices for the
$i$-th hidden layer, in which $s_{ij}$ is the $j$-th choice for the size of
the $i$-th hidden layer: for example, when searching for a one-hidden-layer
network with size candidates {5, 10, 15}, we have $s_{13}=15$.

Reinforcement learning. The RL algorithm learns the set of logits
$\\{\ell_{ij}\\}_{i\in[L],j\in[C_{i}]}$, in which $\ell_{ij}$ is the logit
associated with the $j$-th choice for the $i$-th hidden layer. With a fully
factorized distribution of layer sizes (we learn a separate distribution for
each layer), the probability of sampling the $j$-th choice for the $i$-th
layer $p_{ij}$ is given by the SoftMax function:
$p_{ij}=\exp(\ell_{ij})/\sum_{j\in[C_{i}]}\exp(\ell_{ij})$. In each RL step,
we sample an architecture $y$ to compute the single-step RL objective $J(y)$,
and update the logits with $\nabla J(y)$: an unbiased estimate of the gradient
of the RL value function.

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.

##  3 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.

###  3.1 Weight sharing with layer warmup

The weight-sharing approach has shown success on various computer vision tasks
and NAS benchmarks [40, 6, 9, 7]. To search for an FFN on tabular datasets, we
build a SuperNet where the size of each hidden layer is the largest value in
the search space. Figure 5 shows an example. When we sample a child network
with a hidden layer size $\ell_{i}$ smaller than the SuperNet, we only use the
first $\ell_{i}$ hidden nodes in that layer to compute the output in the
forward pass and the gradients in the backward pass. Similarly, in RL updates,
only the weights of the child network are used to estimate the quality reward
that is used to update logits.

In weight-sharing NAS, warmup helps to ensure that the SuperNet weights are
sufficiently trained to properly guide the RL updates [7]. 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.

###  3.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
[47, 9, 7]. 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 [9, 7].

The REINFORCE-based algorithm trains the SuperNet weights and learns the
logits $\\{\ell_{ij}\\}_{i\in[L],j\in[C_{i}]}$ that give the sampling
probabilities $\\{\ell_{ij}\\}_{i\in[L],j\in[C_{i}]}$ over size candidates by
alternating between weight and RL updates. In each iteration, we first sample
a child network $x$ from the SuperNet and compute its training loss in the
forward pass. Then we update the weights in $x$ with gradients of the training
loss computed in the backward pass. This weight update step trains the weights
of $x$. The weights in architectures with larger sampling probabilities are
sampled and thus trained more often. We then update the logits for the RL
controller by sampling a child network $y$ that is independent of the network
$x$ from the same layerwise distributions, computing the quality reward $Q(y)$
as $1-\textit{loss}(y)$ on the validation set, and then updating the logits
with the gradient of $J(y)=\text{stop\\_grad}(Q(y)-\bar{Q})\log\mathds{P}(y)$:
the product of the advantage of $y$’s reward over past rewards (usually an
exponential moving average) and the log-probability of the current sample.

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.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
[47] proposes the rewards $Q(y)(T(y)/T_{0})^{\beta}$ and
$Q(y)\max\\{1,(T(y)/T_{0})^{\beta}\\}$ while TuNAS [7] 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.

REINFORCE optimizes a set of logits $\\{\ell_{ij}\\}_{i\in[L],j\in[C_{i}]}$
which define a probability distribution $p$ over architectures. In the
original algorithm, we sample a random architecture $y$ from $p$ and then
estimate its quality $Q(y)$. Updates to the logits $\ell_{ij}$ take the form
$\ell_{ij}\leftarrow\ell_{ij}+\eta\frac{\partial}{\partial\ell_{ij}}J(y)$,
where $\eta$ is the learning rate, $\overline{Q}$ is a moving average of
recent rewards, and
$J(y)=\text{stop\\_grad}(Q(y)-\overline{Q})\cdot\log\mathbb{P}(y)$. If $y$ is
better (worse) than average, then $Q(y)-\overline{Q}$ will be positive
(negative), so the REINFORCE update will increase (decrease) the probability
of sampling the same architecture in the future.

Figure 4: Illustration of weight-sharing on two-layer FFNs for a binary
classification task. Edges denote weights; arrows at the end of lines denote
ReLU activations; circles denote hidden nodes; the square in the output layer
denotes the output logit. The size of each hidden layer can be one of {2, 3,
4}, thus the SuperNet is a two-layer FFN with size 4-4. At this moment, the
controller picks the child network 3-2, thus only the first 3 hidden nodes in
the first hidden layer and the first 2 hidden nodes in the second hidden
layer, together with the connected edges (in red), are enabled to compute the
output logits.

(a) Warmup probability

(b) Valid probability

Figure 5: Examples of layer warmup and valid probabilities. Figure 5(a) shows
our schedule: linearly decay from 1 to 0 in the first 25% epochs. Figure 5(b)
shows an example of the change of true and estimated valid probabilities
($\mathds{P}(V)$ and $\widehat{\mathds{P}}(V)$) in a successful search, with
8,000 architectures in the search space and the number of MC samples $N=1024$.
Both probabilities are (nearly) constant during warmup before RL starts, then
increase after RL starts because of rejection sampling.

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

$J(y)=\text{stop\\_grad}(Q(y)-\overline{Q})\cdot\log\left[\mathds{P}(y)/\mathds{P}(y\in
V)\right].$ (1)

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

$\widehat{\mathds{P}}(V)=\frac{1}{N}\sum_{k\in[N]}\frac{p^{(k)}}{q^{(k)}}\cdot\mathds{1}(z^{(k)}\in
V).$ (2)

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.

##  4 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 Volkert111Our paper takes these features as given. It is worth
noting that methods proposed in feature engineering works like [30] and [32]
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.

###  4.1 When do previous RL rewards fail?

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.

####  4.1.1 Criteo – 3 layer search space

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.

(a) Layer 1 (finally 32)

(b) Layer 2 (finally 64)

(c) Layer 3 (finally 96)

(d) retrain performance

Figure 6: Failure case of the Abs Reward on Criteo in a search space of
3-layer FFNs. The change of sampling probabilities and comparison of retrain
performance between the 32-144-24 reference and the 32-64-96 architecture
found with the $Q(y)+\beta|T(y)/T_{0}-1|$ Abs Reward, the target for the
reward was 41,153 parameters. Repeated runs of the same search find the same
architecture. Figure 6(d) shows the retrain validation losses of 32-64-96
(NAS-found) and 32-144-24 (reference).

Failure of latency rewards. Figure 6 shows the sampling probabilities from the
search when using the Abs Reward, and the retrain validation losses of the
found architecture 32-64-96. In Figures 6(a) – 6(c), the sampling
probabilities for the different choices are uniform during warmup and then
converge quickly. The final selected model (32-64-96) is much worse than the
reference model (32-144-24) even though the reference model is actually less
expensive. We also observed similar failures for the MnasNet rewards. With the
MnasNet rewards, the RL controller also struggles to find a model within $\pm
5\%$ of the constraint despite a grid search of the RL parameters (details in
Appendix C). In both cases, almost all found models are worse than the
reference architecture.

Figure 7: Left: 3-layer Criteo SuperNet calibration after 60 epochs (search
space in Appendix C): Pearson correlation is 0.96. The one-shot loss is
validation loss of each child network with weights taken from a SuperNet
trained with the same hyperparameters as in Figure 6 but with no RL in the
first 60 epochs; the stand-alone loss of each child network is computed by
training the same architecture with the same hyperparameters from scratch, and
has std 0.0003. Right: change in probabilities in layer 2 after 60 epochs of
SuperNet training and 40 of RL. Note the rapid changes due to RL.

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.

(a) $\beta=-10$, RL learning rate $\eta=0.001$

(b) $\beta=-25$, RL learning rate $\eta=0.001$

Figure 8: Abs Reward misses the global optimum on Volkert. Figure shows the
retrain validation losses of two architectures found by the Abs Reward vs. the
48-160-32-144 reference.

####  4.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.

####  4.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.

###  4.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.222The 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.

(a) Layer 1 (finally 32)

(b) Layer 2 (finally 144)

(c) Layer 3 (finally 24)

(d) valid probabilities

Figure 9: Success case: on Criteo in a search space of 3-layer FFNs, Monte-
Carlo sampling with rejection eventually finds 32-144-24, the reference
architecture, with RL learning rate 0.005 and number of MC samples 3,072.
Figure 9(d) shows the change of true and estimated valid probabilities.

###  4.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 [44] 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.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 [12], a NAS benchmark on vision
tasks.

##  5 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

  * Abadi et al. [2015] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.  TensorFlow: Large-scale machine learning on heterogeneous systems, 2015\.  URL https://www.tensorflow.org/.  Software available from tensorflow.org.
  * Abutbul et al. [2020] Ami Abutbul, Gal Elidan, Liran Katzir, and Ran El-Yaniv.  Dnf-net: A neural architecture for tabular data.  _arXiv preprint arXiv:2006.06465_ , 2020.
  * Arik and Pfister [2021] Sercan O Arik and Tomas Pfister.  Tabnet: Attentive interpretable tabular learning.  2021\.
  * Awad et al. [2020] Noor Awad, Neeratyoy Mallik, and Frank Hutter.  Differential evolution for neural architecture search.  _arXiv preprint arXiv:2012.06400_ , 2020.
  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Bender et al. [2018] Gabriel Bender, Pieter-Jan Kindermans, Barret Zoph, Vijay Vasudevan, and Quoc Le.  Understanding and simplifying one-shot architecture search.  In _International Conference on Machine Learning_ , pages 550–559. PMLR, 2018.
  * Bender et al. [2020] Gabriel Bender, Hanxiao Liu, Bo Chen, Grace Chu, Shuyang Cheng, Pieter-Jan Kindermans, and Quoc V Le.  Can weight sharing outperform random architecture search? An investigation with TuNAS.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 14323–14332, 2020.
  * Cai et al. [2018a] Han Cai, Tianyao Chen, Weinan Zhang, Yong Yu, and Jun Wang.  Efficient architecture search by network transformation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 32, 2018a.
  * Cai et al. [2018b] Han Cai, Ligeng Zhu, and Song Han.  Proxylessnas: Direct neural architecture search on target task and hardware.  _arXiv preprint arXiv:1812.00332_ , 2018b.
  * Chang et al. [2021] Chun-Hao Chang, Rich Caruana, and Anna Goldenberg.  Node-gam: Neural generalized additive model for interpretable deep learning.  _arXiv preprint arXiv:2106.01613_ , 2021.
  * Chen and Guestrin [2016] Tianqi Chen and Carlos Guestrin.  Xgboost: A scalable tree boosting system.  In _Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining_ , pages 785–794, 2016.
  * Dong et al. [2021] Xuanyi Dong, Lu Liu, Katarzyna Musial, and Bogdan Gabrys.  Nats-bench: Benchmarking nas algorithms for architecture topology and size.  _IEEE transactions on pattern analysis and machine intelligence_ , 2021.
  * Egele et al. [2021] Romain Egele, Prasanna Balaprakash, Isabelle Guyon, Venkatram Vishwanath, Fangfang Xia, Rick Stevens, and Zhengying Liu.  Agebo-tabular: joint neural architecture and hyperparameter search with autotuned data-parallel training for tabular data.  In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ , pages 1–14, 2021.
  * Elsken et al. [2018] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter.  Efficient multi-objective neural architecture search via lamarckian evolution.  _arXiv preprint arXiv:1804.09081_ , 2018.
  * Erickson et al. [2020] Nick Erickson, Jonas Mueller, Alexander Shirkov, Hang Zhang, Pedro Larroy, Mu Li, and Alexander Smola.  Autogluon-tabular: Robust and accurate automl for structured data.  _arXiv preprint arXiv:2003.06505_ , 2020.
  * Feurer et al. [2020] Matthias Feurer, Katharina Eggensperger, Stefan Falkner, Marius Lindauer, and Frank Hutter.  Auto-sklearn 2.0: The next generation.  _arXiv preprint arXiv:2007.04074_ , 24, 2020.
  * Fusi et al. [2018] Nicolo Fusi, Rishit Sheth, and Melih Elibol.  Probabilistic matrix factorization for automated machine learning.  _Advances in neural information processing systems_ , 31, 2018.
  * Gorishniy et al. [2021] Yury Gorishniy, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko.  Revisiting deep learning models for tabular data.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Guerrero-Viu et al. [2021] Julia Guerrero-Viu, Sven Hauns, Sergio Izquierdo, Guilherme Miotto, Simon Schrodi, Andre Biedenkapp, Thomas Elsken, Difan Deng, Marius Lindauer, and Frank Hutter.  Bag of baselines for multi-objective joint neural architecture search and hyperparameter optimization.  _arXiv preprint arXiv:2105.01015_ , 2021.
  * Guo et al. [2020] Zichao Guo, Xiangyu Zhang, Haoyuan Mu, Wen Heng, Zechun Liu, Yichen Wei, and Jian Sun.  Single path one-shot neural architecture search with uniform sampling.  In _European Conference on Computer Vision_ , pages 544–560. Springer, 2020.
  * Howard et al. [2019] Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, et al.  Searching for mobilenetv3.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 1314–1324, 2019.
  * Hu et al. [2020] Yibo Hu, Xiang Wu, and Ran He.  Tf-nas: Rethinking three search freedoms of latency-constrained differentiable neural architecture search.  In _European Conference on Computer Vision_ , pages 123–139. Springer, 2020.
  * Huang et al. [2020] Xin Huang, Ashish Khetan, Milan Cvitkovic, and Zohar Karnin.  Tabtransformer: Tabular data modeling using contextual embeddings.  _arXiv preprint arXiv:2012.06678_ , 2020.
  * Jin et al. [2019] Haifeng Jin, Qingquan Song, and Xia Hu.  Auto-keras: An efficient neural architecture search system.  In _Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining_, pages 1946–1956, 2019.
  * Jones et al. [1998] Donald R Jones, Matthias Schonlau, and William J Welch.  Efficient global optimization of expensive black-box functions.  _Journal of Global optimization_ , 13(4):455–492, 1998.
  * Kadra et al. [2021] Arlind Kadra, Marius Lindauer, Frank Hutter, and Josif Grabocka.  Well-tuned simple nets excel on tabular datasets.  In _Thirty-Fifth Conference on Neural Information Processing Systems_ , 2021.
  * Kandasamy et al. [2018] Kirthevasan Kandasamy, Willie Neiswanger, Jeff Schneider, Barnabas Poczos, and Eric P Xing.  Neural architecture search with bayesian optimisation and optimal transport.  _Advances in neural information processing systems_ , 31, 2018.
  * Ke et al. [2017] Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, Wei Chen, Weidong Ma, Qiwei Ye, and Tie-Yan Liu.  Lightgbm: A highly efficient gradient boosting decision tree.  _Advances in neural information processing systems_ , 30, 2017.
  * Ke et al. [2018] Guolin Ke, Jia Zhang, Zhenhui Xu, Jiang Bian, and Tie-Yan Liu.  Tabnn: A universal neural network solution for tabular data.  2018\.
  * Khawar et al. [2020] Farhan Khawar, Xu Hang, Ruiming Tang, Bin Liu, Zhenguo Li, and Xiuqiang He.  Autofeature: Searching for feature interactions and their architectures for click-through rate prediction.  In _Proceedings of the 29th ACM International Conference on Information & Knowledge Management_, pages 625–634, 2020.
  * Krizhevsky et al. [2009] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Liu et al. [2020] Bin Liu, Chenxu Zhu, Guilin Li, Weinan Zhang, Jincai Lai, Ruiming Tang, Xiuqiang He, Zhenguo Li, and Yong Yu.  Autofis: Automatic feature interaction selection in factorization models for click-through rate prediction.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pages 2636–2645, 2020.
  * Liu et al. [2018a] Chenxi Liu, Barret Zoph, Maxim Neumann, Jonathon Shlens, Wei Hua, Li-Jia Li, Li Fei-Fei, Alan Yuille, Jonathan Huang, and Kevin Murphy.  Progressive neural architecture search.  In _Proceedings of the European conference on computer vision (ECCV)_ , pages 19–34, 2018a.
  * Liu et al. [2018b] Hanxiao Liu, Karen Simonyan, and Yiming Yang.  Darts: Differentiable architecture search.  _arXiv preprint arXiv:1806.09055_ , 2018b.
  * Lyu et al. [2021] Bo Lyu, Hang Yuan, Longfei Lu, and Yunye Zhang.  Resource-constrained neural architecture search on edge devices.  _IEEE Transactions on Network Science and Engineering_ , 9(1):134–142, 2021.
  * Mehta et al. [2022] Yash Mehta, Colin White, Arber Zela, Arjun Krishnakumar, Guri Zabergja, Shakiba Moradian, Mahmoud Safari, Kaicheng Yu, and Frank Hutter.  Nas-bench-suite: Nas evaluation is (now) surprisingly easy.  _arXiv preprint arXiv:2201.13396_ , 2022.
  * Močkus [1975] Jonas Močkus.  On bayesian methods for seeking the extremum.  In _Optimization techniques IFIP technical conference_ , pages 400–404. Springer, 1975.
  * Olson and Moore [2016] Randal S Olson and Jason H Moore.  Tpot: A tree-based pipeline optimization tool for automating machine learning.  In _Workshop on automatic machine learning_ , pages 66–74. PMLR, 2016\.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32:8026–8037, 2019.
  * Pham et al. [2018] Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean.  Efficient neural architecture search via parameters sharing.  In _International Conference on Machine Learning_ , pages 4095–4104. PMLR, 2018.
  * Popov et al. [2019] Sergei Popov, Stanislav Morozov, and Artem Babenko.  Neural oblivious decision ensembles for deep learning on tabular data.  _arXiv preprint arXiv:1909.06312_ , 2019.
  * Rasmussen [2003] Carl Edward Rasmussen.  Gaussian processes in machine learning.  In _Summer school on machine learning_ , pages 63–71. Springer, 2003\.
  * Real et al. [2019] Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V Le.  Regularized evolution for image classifier architecture search.  In _Proceedings of the aaai conference on artificial intelligence_ , volume 33, pages 4780–4789, 2019.
  * Sandler et al. [2018] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 4510–4520, 2018.
  * Shi et al. [2020] Han Shi, Renjie Pi, Hang Xu, Zhenguo Li, James Kwok, and Tong Zhang.  Bridging the gap between sample-based and one-shot neural architecture search with bonas.  _Advances in Neural Information Processing Systems_ , 33:1808–1819, 2020.
  * Somepalli et al. [2021] Gowthami Somepalli, Micah Goldblum, Avi Schwarzschild, C Bayan Bruss, and Tom Goldstein.  Saint: Improved neural networks for tabular data via row attention and contrastive pre-training.  _arXiv preprint arXiv:2106.01342_ , 2021.
  * Tan et al. [2019] Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, and Quoc V Le.  Mnasnet: Platform-aware neural architecture search for mobile.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 2820–2828, 2019.
  * Vanschoren et al. [2013] Joaquin Vanschoren, Jan N. van Rijn, Bernd Bischl, and Luis Torgo.  Openml: Networked science in machine learning.  _SIGKDD Explorations_ , 15(2):49–60, 2013.  doi: 10.1145/2641190.2641198.  URL http://doi.acm.org/10.1145/2641190.2641198.
  * Wang et al. [2021] Ruoxi Wang, Rakesh Shivanna, Derek Cheng, Sagar Jain, Dong Lin, Lichan Hong, and Ed Chi.  Dcn v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems.  In _Proceedings of the Web Conference 2021_ , pages 1785–1797, 2021\.
  * Wen et al. [2020] Wei Wen, Hanxiao Liu, Yiran Chen, Hai Li, Gabriel Bender, and Pieter-Jan Kindermans.  Neural predictor for neural architecture search.  In _European Conference on Computer Vision_ , pages 660–676. Springer, 2020.
  * White et al. [2019] Colin White, Willie Neiswanger, and Yash Savani.  Bananas: Bayesian optimization with neural architectures for neural architecture search.  _arXiv preprint arXiv:1910.11858_ , 1(2):4, 2019\.
  * Wu et al. [2019] Bichen Wu, Xiaoliang Dai, Peizhao Zhang, Yanghan Wang, Fei Sun, Yiming Wu, Yuandong Tian, Peter Vajda, Yangqing Jia, and Kurt Keutzer.  Fbnet: Hardware-aware efficient convnet design via differentiable neural architecture search.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 10734–10742, 2019.
  * Xiong et al. [2019] Yunyang Xiong, Ronak Mehta, and Vikas Singh.  Resource constrained neural network architecture search: Will a submodularity assumption help?  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 1901–1910, 2019.
  * Yang et al. [2019] Chengrun Yang, Yuji Akimoto, Dae Won Kim, and Madeleine Udell.  OBOE: Collaborative filtering for AutoML model selection.  In _Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pages 1173–1183, 2019.
  * Ying et al. [2019] Chris Ying, Aaron Klein, Eric Christiansen, Esteban Real, Kevin Murphy, and Frank Hutter.  Nas-bench-101: Towards reproducible neural architecture search.  In _International Conference on Machine Learning_ , pages 7105–7114. PMLR, 2019.
  * Zhang et al. [2018] Chris Zhang, Mengye Ren, and Raquel Urtasun.  Graph hypernetworks for neural architecture search.  _arXiv preprint arXiv:1810.05749_ , 2018.
  * Zhao et al. [2021] Yiyang Zhao, Linnan Wang, Yuandong Tian, Rodrigo Fonseca, and Tian Guo.  Few-shot neural architecture search.  In _International Conference on Machine Learning_ , pages 12707–12718. PMLR, 2021.
  * Zhou et al. [2019] Hongpeng Zhou, Minghao Yang, Jun Wang, and Wei Pan.  Bayesnas: A bayesian approach for neural architecture search.  In _International conference on machine learning_ , pages 7603–7613. PMLR, 2019.
  * Zhou et al. [2018] Yanqi Zhou, Siavash Ebrahimi, Sercan Ö Arık, Haonan Yu, Hairong Liu, and Greg Diamos.  Resource-efficient neural architect.  _arXiv preprint arXiv:1806.07912_ , 2018.
  * Zoph and Le [2016] Barret Zoph and Quoc V Le.  Neural architecture search with reinforcement learning.  _arXiv preprint arXiv:1611.01578_ , 2016.
  * Zoph et al. [2018] Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoc V Le.  Learning transferable architectures for scalable image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 8697–8710, 2018.

##  Appendix A Previous work

###  A.1 Neural architecture search (NAS)

Neural architecture search (NAS) [60] 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 [60, 61]
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 [6, 34, 45, 57]. Other approaches include the
use of network morphism [27, 14] 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 [60, 9, 7] that learn a probability distribution over candidate
architectures, evolutionary search [33, 14, 20, 4, 43] that pursues more
promising architectures from ancestors, and Bayesian optimization [27, 24, 51,
58, 45] and parametric models [56, 50] 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 [52, 47, 7, 35] that
penalize expensive architectures. With hard resource constraints, greedy
submodular maximization [53] and heuristic scaling methods [22] 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 [29, 3, 41, 2, 10] mimic the structure
of tree-based models [11, 28]. Some other works use the attention mechanism
[23, 46, 18] or are based on feedforward [26, 18] or residual networks [18].

While automated machine learning (AutoML) on tabular datasets has been
addressed from multiple perspectives and with different approaches [16, 38,
17, 54, 15], NAS on tabular datasets is less explored, partly due to
insufficient understanding of promising architectures and a lack of benchmarks
333To disambiguate, the phrase “tabular NAS benchmark” in previous NAS
benchmark literature [55, 36] often refers to tabulated performance of
architectures on vision and language tasks.. Egele et al. [13] interleaves NAS
with aging evolution [43] 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.

Algorithm 1 (Resource-Oblivious) One-Shot Training and REINFORCE

1 search space $S$, weight learning rate $\alpha$, RL learning rate $\eta$

2 sampling probabilities $\\{p_{ij}\\}_{i\in[L],j\in[C_{i}]}$

3 initialize logits $\ell_{ij}\leftarrow 0$, $\forall i\in[L],j\in[C_{i}]$

4 initialize quality reward moving average $\bar{Q}\leftarrow 0$

5 layer warmup

6 for $\text{iter}=1$ to max_iter do

7  $p_{ij}\leftarrow\exp(\ell_{ij})/\sum_{j\in[C_{i}]}\exp(\ell_{ij})$,
$\forall i\in[L],j\in[C_{i}]$

8  $\triangleright$ weight update

9  for $i=1$ to $L$ do

10  $x_{i}\leftarrow$ the $i$-th layer size sampled from
$\\{s_{ij}\\}_{j\in[C_{i}]}$ with distribution $\\{p_{ij}\\}_{j\in[C_{i}]}$

11  $\textit{loss}(x)\leftarrow$ the (training) loss of $x=x_{1}\mathchar
45\relax\dots\mathchar 45\relax x_{L}$ on the training set

12  $w\leftarrow w-\alpha\nabla\textit{loss}(x)$, in which $w$ is the weights
of $x$ $\triangleright$ can be replaced with optimizers other than SGD

13  $\triangleright$ RL update

14  for $i=1$ to $L$ do

15  $y_{i}\leftarrow$ the $i$-th layer size sampled from
$\\{s_{ij}\\}_{j\in[C_{i}]}$ with distribution $\\{p_{ij}\\}_{j\in[C_{i}]}$

16  $Q(y)\leftarrow 1-\textit{loss}(y)$, the quality reward of
$y=y_{1}\mathchar 45\relax\dots\mathchar 45\relax y_{L}$ on the validation set

17  RL reward $r(y)\leftarrow Q(y)$ $\triangleright$ can be replaced with
resource-aware rewards introduced in Section 3.3

18  $J(y)\leftarrow\text{stop\\_grad}(r(y)-\bar{Q})\log\mathds{P}(y)$
$\triangleright$ can be replaced with Algorithm 2 when resource-constrained

19  $\ell_{ij}\leftarrow\ell_{ij}+\eta\nabla J(y)$, $\forall
i\in[L],j\in[C_{i}]$ $\triangleright$ can be replaced with optimizers other
than SGD

20
$\bar{Q}\leftarrow\frac{\gamma*\bar{Q}+(1-\gamma)*Q(y)}{\gamma*\bar{Q}+1-\gamma}$
$\triangleright$ update moving average with $\gamma=0.9$

Algorithm 2 Rejection with Monte-Carlo (MC) Sampling

1 Input: number of MC samples $N$, feasible set $V$, MC proposal distribution
$q$, quality reward moving average $\bar{Q}$, sampled architecture for RL in
the current step $y=y_{1}\mathchar 45\relax y_{2}\mathchar
45\relax\dots\mathchar 45\relax y_{L}$, current layer size distribution over
$\\{s_{ij}\\}_{j\in[C_{i}]}$ with probability $\\{p_{ij}\\}_{j\in[C_{i}]}$

2 Output: $J(y)$

3 if $y$ is feasible then

4  $Q(y)=$ the quality reward of $y$

5  $\mathds{P}(y):=\prod\limits_{i\in[L]}\mathds{P}(Y_{i}=y_{i})$

6  for $i=1$ to $L$ do

7  $\\{z_{i}^{(k)}\\}_{k\in[N]}\leftarrow$ $N$ samples of the $i$-th layer
size, sampled from $\\{s_{ij}\\}_{j\in[C_{i}]}$ with distribution
$\\{p_{ij}\\}_{j\in[C_{i}]}$

8  $p_{i}^{(k)}:=\mathds{P}(Z_{i}=z_{i}^{(k)})$, $\forall i\in[L],k\in[N]$

9  $p^{(k)}:=\prod\limits_{i\in[L]}p_{i}^{(k)}$, $\forall k\in[N]$

10
$\widehat{\mathds{P}}(V)\leftarrow\frac{1}{N}\sum\limits_{k\in[N],z^{(k)}\in
V}\frac{p^{(k)}}{q^{(k)}}$, in which $z^{(k)}:=z_{1}^{(k)}\mathchar
45\relax\dots\mathchar 45\relax z_{L}^{(k)}$

11
$J(y)\leftarrow\text{stop\\_grad}(Q(y)-\bar{Q})\log\frac{\mathds{P}(y)}{\widehat{\mathds{P}}(V)}$

12 else

13  $J(y)\leftarrow 0$

Figure 10: Illustration of the feasible set $V$ within the search space $S$.
Each green diamond or orange dot denotes a feasible or infeasible
architecture, respectively.  Algorithm 3 Sample to Return the Final
Architecture

1 Input: sampling probabilities $\\{p_{ij}\\}_{i\in[L],j\in[C_{i}]}$ returned
by Algorithm 1, number of desired architectures $n$, number of samples to draw
$m$

2 Output: the set of $n$ selected architectures $A$

3 for $i=1$ to $L$ do

4  $\\{x_{i}^{(k)}\\}_{k\in[m]}\leftarrow$ $m$ samples of the $i$-th layer
size, sampled from $\\{s_{ij}\\}_{j\in[C_{i}]}$ with distribution
$\\{p_{ij}\\}_{j\in[C_{i}]}$

5 $F:=\\{k\in[m]\mid x_{1}^{(k)}\mathchar 45\relax x_{2}^{(k)}\mathchar
45\relax\dots\mathchar 45\relax x_{L}^{(k)}\in V\\}$

6 $A\leftarrow$ $n$ unique architectures in $F$ with largest numbers of
parameters

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 [9]
and TuNAS [7], 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

###  C.1 Toy example

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
444https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/ come from
the OpenML dataset repository [48]. 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% validation555The 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 [49].

Table 1: Dataset details name | # points | # features | # classes |  embedding we use for each feature
---|---|---|---|---
|  | numerical | categorical |  |
Criteo | 51,882,752 | 13 | 26 | 2 |  original values for each numerical, 39-dimensional for each categorical
Volkert | 58,310 | 180 | 0 | 10 |  original values
Aloi | 108,000 | 128 | 0 | 1,000 |  original values
Connect-4 | 67,557 | 0 | 42 | 3 |  2-dimensional for each categorical
Higgs | 98,050 | 28 | 0 | 2 |  original values

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. [26] 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 [5] for all
datasets. We use balanced error (weighted average of classification errors
across classes) for all other datasets666The 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. [26], except for Criteo, on which we use logistic loss as in Wang et
al. [49].

Table 2: Weight training hyperparameter details name |  batch size |  learning rate |  learning rate schedule |  optimizer |  # training epochs | metric
---|---|---|---|---|---|---
Criteo |  512 |  0.002 |  cosine decay |  Adam |  60 | log loss
Volkert |  32 |  0.01 |  constant |  SGD with momentum 0.9 |  120 | balanced error
Aloi |  128 |  0.0005 |  constant |  Adam |  50 | balanced error
Connect-4 |  32 |  0.0005 |  cosine decay |  Adam |  60 | balanced error
Higgs |  64 |  0.05 |  constant |  SGD |  60 | balanced error

We use constant RL learning rates for NAS. The
Connect-4777https://www.openml.org/d/40668 and
Higgs888https://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. [26] 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
dataset999https://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}.

  * •

Volkert101010https://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
leaderboard111111https://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. [26] 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. [26] 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. [18]. The difference is that we
use an FFN that only has $<$10% as many parameters as the one used there.

####  C.2.1 More details on the tradeoff plot (Figure 3)

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:

  * •

Criteo: 0 – 200,000

  * •

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 Criteo121212On Criteo, “a
0.001-level improvement (of logistic loss) is considered significant” [49].
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 beats131313Note 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.

(a) Tradeoff on Criteo, in 4 layer search space

(b) Tradeoff on Criteo, in 5 layer search space

(c) Tradeoff on Volkert, in 4 layer search space

(d) Tradeoff on Volkert, in 9 layer search space

Figure 11: Tradeoffs between validation loss and number of parameters in four search spaces.  Table 3: Some Pareto-optimal architectures in Figure 11. All architectures shown here and almost all other Pareto-optimal architectures have the bottleneck structure.  | search space | Pareto-optimal architecture | number of parameters | loss
---|---|---|---|---
Figure 11(a) | Criteo 4-layer | 32-144-24-112 | 44,041 | 0.4454
Figure 11(a) | Criteo 4-layer | 48-112-8-80 | 56,537 | 0.4448
Figure 11(a) | Criteo 4-layer | 48-384-16-176 | 77,489 | 0.4441
Figure 11(a) | Criteo 4-layer | 96-144-32-240 | 125,457 | 0.4433
Figure 11(a) | Criteo 4-layer | 96-384-48-16 | 155,217 | 0.4430
Figure 11(b) | Criteo 5-layer | 32-240-16-8-96 | 45,769 | 0.4451
Figure 11(b) | Criteo 5-layer | 48-128-64-16-128 | 67,217 | 0.4446
Figure 11(b) | Criteo 5-layer | 48-256-16-8-384 | 69,977 | 0.4443
Figure 11(b) | Criteo 5-layer | 64-144-48-96-160 | 102,497 | 0.4437
Figure 11(b) | Criteo 5-layer | 96-512-24-256-48 | 179,449 | 0.4430
Figure 11(c) | Volkert 4-layer | 48-112-16-24 | 16,642 | 0.3314
Figure 11(c) | Volkert 4-layer | 32-112-24-224 | 20,050 | 0.3269
Figure 11(c) | Volkert 4-layer | 48-160-32-144 | 27,882 | 0.3149
Figure 11(c) | Volkert 4-layer | 48-256-24-112 | 31,330 | 0.3097
Figure 11(c) | Volkert 4-layer | 80-208-32-64 | 40,778 | 0.3054
Figure 11(d) | Volkert 9-layer | 64-64-160-48-16-144-16-8-48 | 40,482 | 0.3250
Figure 11(d) | Volkert 9-layer | 80-144-32-112-32-8-128-8-144 | 43,290 | 0.3238
Figure 11(d) | Volkert 9-layer | 112-144-32-32-24-24-24-128-32 | 51,890 | 0.3128
Figure 11(d) | Volkert 9-layer | 144-128-112-16-16-48-144-24-160 | 78,114 | 0.3019
Figure 11(d) | Volkert 9-layer | 160-144-144-32-112-32-48-32-144 | 94,330 | 0.3010

####  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.

(a) Criteo, 4-layer search space

(b) Criteo, 5-layer search space

(c) Volkert, 4-layer search space

(d) Volkert, 9-layer search space

Figure 12: Rejection-based reward distributionally outperforms random search
and resource-aware Abs Reward in a number of search spaces. The points and
error bars have the same meaning as in Figure 3. The time taken for each
stand-alone training run (the unit length for x axes) is 2.5 hours on Criteo
(Figure 3 in the main paper and 12(a), 12(b)), 10 minutes on Volkert with
4-layer FFNs (Figure 12(c)), and 22-25 minutes on Volkert with 9-layer FFNs
(Figure 12(d)).

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 [59] 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.

Table 4: Architectures found by RL with 5 reward functions on 5 tabular datasets in Table 1. Each reward-dataset combination has 3 repeated NAS runs. The bracket “(2 times)” indicates the same architecture was found by the corresponding reward function twice.  |  Criteo |  Volkert |  Aloi |  Connect-4 |  Higgs
---|---|---|---|---|---
reference architecture |  32-144-24 (41,153 parameters, 0.4454 $\pm$ 0.0003) |  48-160-32-144 (27,882 parameters, 0.3244 $\pm$ 0.0040) |  160-64-512-64 (162,056 parameters, 0.0470 $\pm$ 0.0004) |  32-22-4-20 (3,701 parameters, 0.2691 $\pm$ 0.0021) |  16-144-16 (5,265 parameters, 0.2794 $\pm$ 0.0013)
RL with rejection-based reward |  32-144-24 (2 times) (41,153 parameters, 0.4454 $\pm$ 0.0003) 32-112-32 (40,241 parameters, 0.4456 $\pm$ 0.0003) |  64-128-48-16 (27,050 parameters, 0.3237 $\pm$ 0.0040) 48-160-32-144 (27,882 parameters, 0.3244 $\pm$ 0.0040) 80-48-112-32 (27,802 parameters, 0.3274 $\pm$ 0.0037) |  176-144-144-80 (2 times) (161,672 parameters, 0.0458 $\pm$ 0.0007) 192-112-176-80 (161,432 parameters, 0.0461 $\pm$ 0.0012) |  20-32-14-48 (3,701 parameters, 0.2827 $\pm$ 0.0062) 20-30-24-22 (3,693 parameters, 0.2795 $\pm$ 0.0049) 28-10-12-56 (3,701 parameters, 0.2819 $\pm$ 0.0052) |  72-24-48 (2 times) (5,161 parameters, 0.2911 $\pm$ 0.0031) 80-32-8 (5,265 parameters, 0.2893 $\pm$ 0.0022) 64-16-128 (5,265 parameters, 0.2870 $\pm$ 0.0033)
RL with Abs Reward |  32-64-96 (2 times) (41,345 parameters, 0.4461 $\pm$ 0.0003) 32-80-64 (40,785 parameters, 0.4459 $\pm$ 0.0002) |  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) |  144-112-144-96 (3 times) (162,008 parameters, 0.0473 $\pm$ 0.0004) 128-96-96-112 (4 times) (162,072 parameters, 0.0488 $\pm$ 0.0012) 112-112-96-112 (161,816 parameters, 0.0502 $\pm$ 0.0007) |  30-22-12-12 (3,703 parameters, 0.2702 $\pm$ 0.0027) 26-20-18-26 (2 times) (3,703 parameters, 0.2807 $\pm$ 0.0029) 26-18-20-26 (3,703 parameters, 0.2765 $\pm$ 0.0046) |  104-16-24 (2 times) (5,233 parameters, 0.2896 $\pm$ 0.0036) 80-16-88 (2 times) (5,281 parameters, 0.2887 $\pm$ 0.0021) 88-16-64 (2 times) (5,217 parameters, 0.2874 $\pm$ 0.0027)
MNasNet (reward type 1: $Q(y)(T(y)/T_{0})^{\beta}$) |  32-64-16 (36,065 parameters, 0.4464 $\pm$ 0.0003) 32-64-8 (35,537 parameters, 0.4466 $\pm$ 0.0002) 32-24-64 (35,353 parameters, 0.4466 $\pm$ 0.0003) |  only found one feasible architecture: 32-32-224-24 (19,890 parameters, 0.3392 $\pm$ 0.0043) |  160-128-112-96 (163,544 parameters, 0.0469 $\pm$ 0.0009) 128-96-96-112 (162,072 parameters, 0.0488 $\pm$ 0.0012) 128-176-128-80 (153,192 parameters, 0.0473 $\pm$ 0.0007) |  24-18-16-44 (3,677 parameters, 0.2788 $\pm$ 0.0046) 28-12-32-14 (3,651 parameters, 0.2781 $\pm$ 0.0052) 22-32-16-22 (3,577 parameters, 0.2818 $\pm$ 0.0032) |  72-24-32 (4,745 parameters, 0.2898 $\pm$ 0.0019) 64-32-16 (4,545 parameters, 0.2888 $\pm$ 0.0018) 72-24-24 (2 times) (4,537 parameters, 0.2903 $\pm$ 0.0036)
MNasNet (reward type 2: $Q(y)\max\\{1,(T(y)/T_{0})^{\beta}\\}$) |  24-128-16 (29,953 parameters, 0.4468 $\pm$ 0.0005) 16-176-48 (27,985 parameters, 0.4478 $\pm$ 0.0006) 24-16-160 (27,953 parameters, 0.4479 $\pm$ 0.0003) |  only found one feasible architecture: 80-24-24-24 (17,874 parameters, 0.3521 $\pm$ 0.0013) |  112-144-240-64 (145,944 parameters, 0.0480 $\pm$ 0.0005) 160-112-128-64 (126,392 parameters, 0.0497 $\pm$ 0.0019) 256-64-224-32 (104,232 parameters, 0.0548 $\pm$ 0.0015) |  20-12-44-24 (3,679 parameters, 0.2871 $\pm$ 0.0069) 28-16-10-52 (3,745 parameters, 0.2776 $\pm$ 0.0061) 14-40-28-20 (3,581 parameters, 0.2886 $\pm$ 0.0047) |  16-8-16 (2 times) (777 parameters, 0.2929 $\pm$ 0.0020) 8-8-24 (553 parameters, 0.3003 $\pm$ 0.0031)
RENA [59] |  32-16-48 (34,289 parameters, 0.4472 $\pm$ 0.0002) 24-48-160 (33,873 parameters, 0.4468 $\pm$ 0.0001) 8-64-96 (15,137 parameters, 0.4523 $\pm$ 0.0004) |  only found one feasible architecture: 32-48-24-512 (26,482 parameters, 0.3389 $\pm$ 0.0039) |  didn’t find any (close to) feasible architecture |  24-26-20-16 (3,617 parameters, 0.2806 $\pm$ 0.0068) 30-18-10-32 (3,749 parameters, 0.2752 $\pm$ 0.0075) 26-26-18-8 (3,577 parameters, 0.2769 $\pm$ 0.0031) |  only found one feasible architecture: 80-8-152 (4,569 parameters, 0.2882 $\pm$ 0.0016)

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

Bayesian optimization (BO) and evolutionary search (ES) are popular strategies
for NAS (e.g., [24, 27, 51, 58, 45]; [33, 14, 4, 19]). We are not aware of any
work that successfully applies BO to one-shot NAS for tabular data; Guo et al.
[20] proposed an ES approach for one-shot NAS on vision datasets. BO and ES
omit the extra forward passes for RL controller training, but need extra
forward passes to evaluate child networks from SuperNet weights (Figure 5 in
the main paper). Thus we control the number of forward passes for a fair
comparison141414RL also does backward passes to optimize over the logits
$\\{\ell_{ij}\\}_{i\in[L],j\in[C_{i}]}$, but the number of logits in our
setting is much fewer than the size of the validation set (100 vs. 4,582,432),
which means the cost of forward passes dominates.. We design the following
one-shot NAS methods to compare with an $M$-epoch TabNAS that interleaves
weight and RL updates in its latter $75\%$ iterations:

  * •

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 [42] with expected improvement [37, 25]) 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.
[20] 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).

Table 5: TabNAS vs. train-then-search-with-RL/BO/ES for one-shot NAS on Criteo, among 5-layer FFNs. The reference architecture 48-240-24-256-8 has 75,353 parameters and validation loss 0.4448 $\pm$ 0.0002. Bold results below indicate architectures that are on par with the best validation loss 0.4444 $\pm$ 0.0002. We run each method 3 times, except for TabNAS, which was detailed in Appendix D. The square bracket “[3 times]” indicates that the same architecture was found by the corresponding method three times.  method |  found architectures (number of parameters, mean $\pm$ std loss) |  NAS cost
---|---|---
TabNAS ($\eta=0.005$ or $0.001$, $N=32768$) |  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-176-80-16-96 (75,153 parameters, 0.4445 $\pm$ 0.0002) |  $1.86\times 10^{9}$ forward passes (405 child network evaluations)
train-then-search-with-RL ($\eta=0.0005$ or $0.001$, $N=32768$) |  48-64-224-16-256 (75,249 parameters, 0.4446 $\pm$ 0.0002) 48-240-16-24-384 (75,353 parameters, 0.4445 $\pm$ 0.0002) 48-176-64-16-256 (74,945 parameters, 0.4445 $\pm$ 0.0002) |  $1.86\times 10^{9}$ forward passes (405 child network evaluations)
train-then-search-with-BO (RBF kernel L=10, white noise variance 1, 50 initial samples, 20 new samples per BO iteration) |  64-80-8-16-16 (72,073 parameters, 0.4447 $\pm$ 0.0002) 48-384-16-16-16 (74,881 parameters, 0.4444 $\pm$ 0.0002) 48-256-16-8-240 (68,537 parameters, 0.4447 $\pm$ 0.0002) |  153 child network evaluations, then gets stuck
train-then-search-with-BO (RBF kernel L=1, white noise variance 1, 100 initial samples, 10 new samples per BO iteration) |  48-80-48-96-96 (71,265 parameters, 0.4451 $\pm$ 0.0003) 48-128-8-64-8 (57,753 parameters, 0.4451 $\pm$ 0.0002) 64-96-16-32-16 (74,673 parameters, 0.4446 $\pm$ 0.0002) |  420 child network evaluations
train-then-search-with-BO (RBF kernel L=1, white noise variance 1, 50 initial samples, 20 new samples per BO iteration) |  64-112-16-8-16 (75,177 parameters, 0.4445 $\pm$ 0.0001) 48-32-256-8-240 (63,817 parameters, 0.4454 $\pm$ 0.0003) 64-96-16-24-48 (75,241 parameters, 0.4447 $\pm$ 0.0002) |  430 child network evaluations
train-then-search-with-BO (RBF kernel L=0.1, white noise variance 1, 50 initial samples, 20 new samples per BO iteration) |  [3 times] 64-96-16-24-48 (75,241 parameters, 0.4447 $\pm$ 0.0002) |  430 child network evaluations
train-then-search-with-BO (RBF kernel L=0.1, white noise variance 1, 100 initial samples, 10 new samples per BO iteration) |  [3 times] 64-96-16-24-48 (75,241 parameters, 0.4447 $\pm$ 0.0002) |  430 child network evaluations
train-then-search-with-ES (population $50$, $k=10$) |  48-96-32-224-16 (68,161 parameters, 0.4450 $\pm$ 0.0003) 48-48-8-160-64 (63,897 parameters, 0.4455 $\pm$ 0.0003) 48-96-32-24-64 (59,609 parameters, 0.4448 $\pm$ 0.0001) |  424 child network evaluations
train-then-search-with-ES (population $25$, $k=20$) |  48-144-32-16-144 (64,161 parameters, 0.4447 $\pm$ 0.0002) 48-80-16-128-64 (65,057 parameters, 0.4453 $\pm$ 0.0002) 48-96-32-48-64 (61,937 parameters, 0.4451 $\pm$ 0.0001) |  469 child network evaluations

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

In the vision domain, the NATS-Bench [12] 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
[31] 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.

Table 6: #FLOPs (M) of architectures found by RL with the rejection-based reward and the Abs Reward in NATS-Bench size search space. RL learning rate |  0.01 |  0.05 |  0.1 |  0.5
---|---|---|---|---
rejection-based reward, $N$=200 |  74.7 $\pm$ 0.3 (median 74.7) |  74.7 $\pm$ 0.2 (median 74.7) |  74.6 $\pm$ 0.3 (median 74.7) |  70.7 $\pm$ 4.5 (median 72.4)
Abs Reward, $\beta$=-10 |  77.0 $\pm$ 12.7 (median 76.4) |  75.1 $\pm$ 4.5 (median 74.8) |  75.2 $\pm$ 1.7 (median 75.1) |  75.7 $\pm$ 3.8 (median 75.2)
Abs Reward, $\beta$=-5 |  78.1 $\pm$ 13.1 (median 77.0) |  75.6 $\pm$ 4.3 (median 75.5) |  75.4 $\pm$ 1.7 (median 75.2) |  76.0 $\pm$ 4.4 (median 75.3)
Abs Reward, $\beta$=-2 |  80.6 $\pm$ 13.8 (median 78.5) |  76.3 $\pm$ 4.2 (median 76.1) |  75.7 $\pm$ 1.8 (median 75.5) |  75.9 $\pm$ 5.4 (median 75.4)
Abs Reward, $\beta$=-1 |  83.0 $\pm$ 13.3 (median 81.9) |  77.4 $\pm$ 4.4 (median 77.1) |  76.0 $\pm$ 2.1 (median 75.8) |  76.7 $\pm$ 5.3 (median 75.6)
Abs Reward, $\beta$=-0.5 |  87.1 $\pm$ 12.7 (median 86.7) |  78.8 $\pm$ 4.5 (median 78.7) |  77.0 $\pm$ 2.6 (median 76.6) |  76.9 $\pm$ 6.0 (median 76.1)
Abs Reward, $\beta$=-0.2 |  95.5 $\pm$ 12.8 (median 95.3) |  80.6 $\pm$ 5.2 (median 80.1) |  78.5 $\pm$ 3.6 (median 77.9) |  78.5 $\pm$ 9.2 (median 76.8)
Abs Reward, $\beta$=-0.1 |  101.8 $\pm$ 13.3 (median 103.0) |  84.2 $\pm$ 7.6 (median 83.0) |  81.1 $\pm$ 5.8 (median 80.5) |  80.1 $\pm$ 10.6 (median 78.6)
Table 7: Test errors of architectures found by RL with the rejection-based reward and the Abs Reward in NATS-Bench size search space. RL learning rate | 0.01 | 0.05 | 0.1 | 0.5
---|---|---|---|---
rejection-based reward, $N$=200 |  0.468 $\pm$ 0.050 |  0.435 $\pm$ 0.023 |  0.425 $\pm$ 0.014 |  0.420 $\pm$ 0.008
Abs Reward, $\beta$=-10 |  0.483 $\pm$ 0.056 |  0.468 $\pm$ 0.034 |  0.473 $\pm$ 0.046 |  0.490 $\pm$ 0.074
Abs Reward, $\beta$=-5 |  0.474 $\pm$ 0.049 |  0.461 $\pm$ 0.029 |  0.463 $\pm$ 0.039 |  0.481 $\pm$ 0.061
Abs Reward, $\beta$=-2 |  0.462 $\pm$ 0.039 |  0.450 $\pm$ 0.022 |  0.455 $\pm$ 0.028 |  0.468 $\pm$ 0.046
Abs Reward, $\beta$=-1 |  0.449 $\pm$ 0.027 |  0.442 $\pm$ 0.016 |  0.445 $\pm$ 0.020 |  0.456 $\pm$ 0.035
Abs Reward, $\beta$=-0.5 |  0.436 $\pm$ 0.019 |  0.434 $\pm$ 0.014 |  0.435 $\pm$ 0.015 |  0.444 $\pm$ 0.023
Abs Reward, $\beta$=-0.2 |  0.422 $\pm$ 0.013 |  0.424 $\pm$ 0.011 |  0.426 $\pm$ 0.012 |  0.433 $\pm$ 0.016
Abs Reward, $\beta$=-0.1 |  0.414 $\pm$ 0.009 |  0.416 $\pm$ 0.008 |  0.418 $\pm$ 0.009 |  0.426 $\pm$ 0.014

We can see that:

  * •

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.

(a) Tuning $\beta$ in Abs Reward

(b) Tuning $N$ in rejection

[1pt]

Figure 13: Tuning $\beta$ and $N$ on the toy example (Figure 1): the number of
MC samples $N$ in rejection-based reward is easier to tune than $\beta$ in Abs
Reward, and is easier to succeed. The lines and shaded regions are mean and
standard deviation across 200 independent runs, respectively.

Figure 14: Tuning $N$ on Criteo: the change of $\widehat{\mathds{P}}(V)$ when
the number of Monte-Carlo samples $N$ is 256, 2,048 or 5,120, and the time
taken for each iteration. We show results with RL learning rate $\eta=0.005$;
those other $\eta$ values have similar failure patterns.

###  H.1 Resource hyperparameter $\beta$

$\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$

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

$J(y)=\text{stop\\_grad}(Q(y)-\overline{Q})\cdot\log\left[\mathds{P}(y)/\mathds{P}(V)\right].$

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.

Below are our experimental findings:

  * •

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)).

(a) Layer 2, not skipping infeasible in training

(b) Valid probabilities, not skipping infeasible in training

(c) Layer 2, skipping infeasible in training

(d) Valid probabilities, skipping infeasible in training

Figure 15: Failure cases in ablation when $\widehat{\mathds{P}}(V)$ is non-
differentiable. We show results with RL learning rate $\eta=0.005$; those
under other $\eta$ values are similar.

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

###  J.1 $\widehat{\mathds{P}}(V)$ is an unbiased and consistent estimate of
$\mathds{P}(V)$

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:

$\displaystyle\operatorname{\mathbb{E}}[\widehat{\mathds{P}}(V)]$
$\displaystyle=\frac{1}{N}\operatorname{\mathbb{E}}\Bigg{[}\sum_{k\in[N],z^{(k)}\in
V}\frac{p^{(k)}}{q^{(k)}}\Bigg{]}$
$\displaystyle=\frac{1}{N}\operatorname{\mathbb{E}}\Bigg{[}\sum_{k\in[N]}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)\Bigg{]}$
$\displaystyle=\frac{1}{N}\sum_{k\in[N]}\operatorname{\mathbb{E}}\Bigg{[}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)\Bigg{]},$

in which each summand

$\displaystyle\operatorname{\mathbb{E}}\Bigg{[}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)\Bigg{]}$ $\displaystyle=\sum\limits_{z^{(k)}\in
S}q^{(k)}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in V)$
$\displaystyle=\mathds{P}(V),$

Thus $\operatorname{\mathbb{E}}[\widehat{\mathds{P}}(V)]=\mathds{P}(V)$.

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,

$\displaystyle\mathrm{Var}[\widehat{\mathds{P}}(V)]$
$\displaystyle=\frac{1}{N}\sum_{k\in[N]}\mathrm{Var}\Bigg{[}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)\Bigg{]},$

in which each summand

$\displaystyle\mathrm{Var}\Bigg{[}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)\Bigg{]}$
$\displaystyle=\operatorname{\mathbb{E}}\Big{[}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)-\mathds{P}(V)\Big{]}$ (3) $\displaystyle=\sum\limits_{z^{(k)}\notin
V}q^{(k)}\mathds{P}(V)^{2}+\sum_{z^{(k)}\in
V}q^{(k)}\Big{[}\frac{p^{(k)}}{q^{(k)}}-\mathds{P}(V)\Big{]}^{2}$
$\displaystyle=-\mathds{P}(V)^{2}+\sum_{z^{(k)}\in
V}\frac{(p^{(k)})^{2}}{q^{(k)}},$

thus the variance

$\displaystyle\mathrm{Var}[\widehat{\mathds{P}}(V)]$
$\displaystyle=\frac{1}{N}\sum_{k\in[N]}\mathrm{Var}\Big{[}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)\Big{]}$ $\displaystyle\frac{1}{N}\Big{[}-\mathds{P}(V)^{2}+\sum_{z^{(k)}\in
V}\frac{(p^{(k)})^{2}}{q^{(k)}}\Big{]},$

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)$.

The Chebyshev’s Inequality states that for a random variable $X$ with
expectation $\mu$, for any $a>0$,
$\mathds{P}(|X-\mu|>a)\leq\frac{\mathrm{Var}(X)}{a^{2}}$. Thus
$\lim\limits_{N\rightarrow\infty}\mathrm{Var}(X)=0$ implies that
$\lim\limits_{N\rightarrow\infty}\mathds{P}(|X-\mu|>a)=0$ for any $a>0$,
indicating consistency.

###  J.2 $\nabla\log[\mathds{P}(y)/\widehat{\mathds{P}}(V)]$ is a consistent
estimate of $\nabla\log[\mathds{P}(y\>|\>y\in 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)$,

  * •

$\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)$, in which each $p^{(k)}$ is differentiable with respect to all logits
$\\{\ell_{ij}\\}_{i\in[L],j\in[C_{i}]}$.

Thus we have

$\displaystyle\operatorname*{plim}_{N\rightarrow\infty}\widehat{\mathds{P}}(V)$
$\displaystyle=\operatorname*{plim}_{N\rightarrow\infty}\frac{1}{N}\sum\limits_{k\in[N]}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in
V)$ $\displaystyle=\frac{1}{N}\sum_{z^{(k)}\in
S}\frac{p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in V)Nq^{(k)}$
$\displaystyle=\sum_{z^{(k)}\in S}p^{(k)}\mathds{1}(z^{(k)}\in
V)=\mathds{P}(V),$

and

$\displaystyle\operatorname*{plim}_{N\rightarrow\infty}\nabla\widehat{\mathds{P}}(V)$
$\displaystyle=\operatorname*{plim}_{N\rightarrow\infty}\frac{1}{N}\sum\limits_{k\in[N]}\frac{\nabla
p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in V)$
$\displaystyle=\frac{1}{N}\sum_{z^{(k)}\in S}\frac{\nabla
p^{(k)}}{q^{(k)}}\mathds{1}(z^{(k)}\in V)Nq^{(k)}$
$\displaystyle=\sum_{z^{(k)}\in S}\nabla p^{(k)}\mathds{1}(z^{(k)}\in
V)=\nabla\mathds{P}(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.