Aux-NAS: Exploiting Auxiliary Labels with Negligibly Extra Inference Cost

TL;DR

Aux-NAS introduces a novel architecture-based NAS method that leverages auxiliary labels to improve primary task performance without increasing inference cost, by evolving networks with only primary-to-auxiliary links during inference.

Contribution

It proposes a NAS-based approach to exploit auxiliary labels through architecture design, enabling different training and inference networks with negligible extra inference cost.

Findings

Achieves improved primary task performance across six tasks.

Effectively removes auxiliary connections during inference.

Compatible with various backbones and auxiliary learning methods.

Abstract

We aim at exploiting additional auxiliary labels from an independent (auxiliary) task to boost the primary task performance which we focus on, while preserving a single task inference cost of the primary task. While most existing auxiliary learning methods are optimization-based relying on loss weights/gradients manipulation, our method is architecture-based with a flexible asymmetric structure for the primary and auxiliary tasks, which produces different networks for training and inference. Specifically, starting from two single task networks/branches (each representing a task), we propose a novel method with evolving networks where only primary-to-auxiliary links exist as the cross-task connections after convergence. These connections can be removed during the primary task inference, resulting in a single-task inference cost. We achieve this by formulating a Neural Architecture Search (NAS) problem, where we initialize bi-directional connections in the search space and guide the NAS optimization converging to an architecture with only the single-side primary-to-auxiliary connections. Moreover, our method can be incorporated with optimization-based auxiliary learning approaches. Extensive experiments with six tasks on NYU v2, CityScapes, and Taskonomy datasets using VGG, ResNet, and ViT backbones validate the promising performance. The codes are available at https://github.com/ethanygao/Aux-NAS.