TabNAS: Rejection Sampling for Neural Architecture Search on Tabular Datasets

TL;DR

TabNAS introduces a rejection sampling-based reinforcement learning approach for neural architecture search on tabular datasets, effectively handling resource constraints and outperforming previous methods in finding optimal models.

Contribution

It proposes a novel rejection sampling method integrated with RL for resource-constrained NAS on tabular data, improving model quality and constraint adherence.

Findings

Outperforms previous reward-shaping NAS methods

Finds better models that satisfy resource constraints

Demonstrates effectiveness across multiple tabular datasets

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.