Aux-NAS: Asymmetric Auxiliary Learning

• Aux-NAS is an auxiliary-learning method that uses asymmetric cross-task connections, allowing auxiliary supervision during training without adding inference cost.
• It employs Neural Architecture Search to identify optimal primary-to-auxiliary connections, thus reducing negative transfer by distinguishing useful gradient signals from dispensable auxiliary features.
• Experimental results on datasets like NYU v2, CityScapes, and Taskonomy demonstrate significant performance gains over single-task models across diverse architectures.

Aux-NAS is an auxiliary-label learning method that uses labels from an independent auxiliary task to improve a designated primary task while preserving a primary-task inference cost essentially identical to that of a single-task model. It is explicitly architecture-based rather than optimization-based: instead of relying only on loss-weight adaptation, gradient manipulation, or objective selection, it searches for an asymmetric cross-task connectivity pattern in which auxiliary information is useful during training but dispensable at test time. The resulting training network and inference network therefore differ by design, with the final deployed model reduced to the primary branch alone (Gao et al., 2024).

1. Problem formulation and conceptual scope

Aux-NAS is defined for a setting with one primary task that matters at test time and one or more auxiliary tasks whose labels are available during training. Typical examples in the reported experiments include semantic segmentation as the primary task and depth or surface normal prediction as auxiliary tasks. The stated objective has two parts: improving primary-task performance by exploiting auxiliary labels, and preserving single-task inference cost for the primary task (Gao et al., 2024).

This formulation differs from standard multi-task learning. In standard multi-task learning, the objective is to improve all tasks jointly, and inference typically retains the full multi-task network. Aux-NAS is not intended to optimize all tasks equally; it is designed for the asymmetric case in which only the primary task matters at deployment. It also differs from conventional auxiliary-learning approaches that keep a fixed shared architecture and attempt to manage negative transfer through optimization procedures such as adaptive loss weighting, gradient projection, gradient clipping, or auxiliary objective selection. Aux-NAS shifts the central question from how gradients should be combined to what inter-task pathway should exist at all (Gao et al., 2024).

A common misconception is to treat Aux-NAS as a variant of ordinary hard-parameter-sharing or soft-parameter-sharing multi-task learning. The paper’s formulation instead separates the task branches and introduces asymmetric connectivity so that the auxiliary task can regularize or supervise the primary branch during training without creating an inference-time dependency on auxiliary computation. This suggests that the method is better understood as a deployment-constrained auxiliary-learning framework than as a symmetric multi-task architecture.

2. Asymmetric architecture and branch-level decomposition

The starting point is two independent single-task backbones: a primary branch $P$ and an auxiliary branch $A$. Each is a single-task network for its own task. The architectural asymmetry lies in allowing interaction during training while guaranteeing that the primary branch does not require auxiliary features in its forward path at inference (Gao et al., 2024).

The paper contrasts this with a symmetric soft-sharing baseline in which both branches consume features from both tasks. In that symmetric case, feature sharing is effective, but the auxiliary branch cannot be removed because the primary branch depends on auxiliary features. Aux-NAS instead enforces an asymmetric structure:

$P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$

$$A_i^{\text{fea}} = \mathcal{O}^{AP} P_{i-1}^{\text{fea}} + \mathcal{O}^{AA} A_{i-1}^{\text{fea}}$$

with backward propagation:

$$P_{i-1}^{\text{grad}} = P_i^{\text{grad}} + \frac{d\mathcal{O}^{AP}}{d\theta} A_i^{\text{grad}}$$

$$A_{i-1}^{\text{grad}} = \frac{d\mathcal{O}^{AA}}{d\theta} A_i^{\text{grad}}.$$

The practical consequence is precise. The primary branch does not consume auxiliary features in its forward computation, but the auxiliary branch can still propagate auxiliary gradients back into the primary branch. Auxiliary supervision therefore influences optimization without imposing an inference-time computational dependency. The paper characterizes this as the core asymmetry of the method and attributes its utility to the ability to reduce negative transfer by separating task branches while retaining controlled cross-task coupling (Gao et al., 2024).

3. Aux-G and Aux-NAS

The paper presents two related methods: Aux-G and Aux-NAS. Aux-G is the simpler construction and uses only primary-to-auxiliary connections. At auxiliary layer $i$,

$$A_i = \mathcal{O}^{A}\Big(A_{i-1}, \alpha_{0,i} P_{0,i}, \ldots, \alpha_{i-1,i} P_{i-1,i}\Big),$$

where $\alpha_{j,i} \in \{0,1\}$ indicates whether a connection from primary layer $j$ to auxiliary layer $A$0 exists. In this design, auxiliary labels help mainly through auxiliary gradients rather than through auxiliary features in the primary forward pass (Gao et al., 2024).

Aux-NAS enlarges the search space by initializing bi-directional cross-task connections: primary-to-auxiliary and auxiliary-to-primary. The training network can therefore exploit both auxiliary features and auxiliary gradients. Neural Architecture Search is then used to drive the architecture toward a converged pattern in which auxiliary-to-primary connections are pruned away, leaving only primary-to-auxiliary connections. The method thus operates in three distinct regimes: a richer bi-directional network during training, a converged architecture with only primary-to-auxiliary links, and a final inference configuration in which all auxiliary-related computation is removed (Gao et al., 2024).

This distinction is central to the paper’s contribution. Aux-G already demonstrates that auxiliary gradients alone can be beneficial. Aux-NAS adds a search procedure that identifies where cross-task links should be placed and, during training, can exploit auxiliary features before eliminating the removable direction. A plausible implication is that Aux-G isolates the benefit of asymmetric supervision, whereas Aux-NAS tests whether learned sparse connectivity and temporary bidirectionality yield further gains.

4. NAS formulation, fusion operator, and training–inference asymmetry

The search formulation separates standard model weights from architecture weights. Let $A$1 denote model parameters, $A$2 architecture weights for auxiliary-to-primary edges, and $A$3 architecture weights for primary-to-auxiliary edges. The optimization problem is

$A$4

with

$A$5

Only $A$6 is regularized, so auxiliary-to-primary connections are pushed toward zero, whereas primary-to-auxiliary connections are not regularized because they are removed at inference anyway (Gao et al., 2024).

The corresponding layerwise fused features are

$A$7

and

$A$8

This search space spans all possible cross-task links between layers of the two fixed backbones (Gao et al., 2024).

The implemented fusion operator uses concatenation, $A$9 convolution, normalization, and activation:

$P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$0

and

$P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$1

The efficiency argument is that once $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$2 is pruned to zero, the expensive $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$3 convolution for auxiliary-to-primary fusion disappears. The paper describes the remaining overhead as negligible and states that, at inference, the entire auxiliary branch and all cross-task connections are cut, so the deployed model is just the primary single-task network (Gao et al., 2024).

Training and inference are therefore intentionally different. During training, both branches are active, auxiliary labels are used, and $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$4 and $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$5 are trained alternately using two non-overlapped batches, following standard one-shot NAS practice. During evaluation, the paper states that auxiliary-to-primary weights are manually set to zero before evaluation to guarantee single-task inference behavior (Gao et al., 2024).

5. Empirical scope and reported performance

The evaluation covers three datasets—NYU v2, CityScapes, and Taskonomy—and six tasks across experiments: semantic segmentation, surface normal prediction, depth estimation, monocular disparity estimation, object classification, and scene classification. The reported backbones are VGG-16, ResNet-50, and ViT-Base. Baselines include Single, Aux-Head, Adashare, Adashare-Aux, Aux-G-Stage, Aux-G-Layer, and Aux-NAS; optimization-based comparisons include Uncertainty, DWA, PCGrad, PCGrad-Aux, CAGrad, and GCS (Gao et al., 2024).

Depending on the task, the metrics include mIoU and pixel accuracy for semantic segmentation; mean, median, and RMSE angular error plus thresholded angular percentages such as $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$6, $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$7, and $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$8 for surface normals; Top-1 and Top-5 accuracy for classification; and corresponding regression metrics for depth and disparity estimation described in the appendix (Gao et al., 2024).

The paper reports that the Aux-NAS family outperformed optimization-based auxiliary-learning baselines on NYU v2 segmentation with surface normals as auxiliary, and that combining Aux-NAS with optimization-based methods improved results further. It also states that Aux-NAS consistently beat strong architecture-based baselines such as Adashare and Aux-G (Gao et al., 2024).

Setting	Single	Aux-NAS
NYU v2, primary segmentation, VGG-16	33.5 mIoU / 64.1 PAcc	36.0 / 66.1
NYU v2, primary normal prediction, VGG-16	14.6 mean angular error	12.5 mean angular error
CityScapes, primary segmentation	68.3 mIoU	71.1 mIoU
Taskonomy, primary object classification	34.3 Top-1 / 65.9 Top-5	39.8 / 70.7

The reported gains also extend across task direction and backbone family. The method worked when segmentation was primary and normals auxiliary, when normals were primary and segmentation auxiliary, and with VGG-16, ResNet-50, and ViT-Base. The paper further notes that ViT-Base on NYU v2 also showed improved normal estimation over baselines. This suggests that the method is not tied to a single task ordering or to a specific CNN architecture family (Gao et al., 2024).

6. Ablations, compatibility, and limitations

The ablation study isolates three factors: auxiliary gradients, auxiliary features, and NAS search. The reported conclusion is hierarchical: using auxiliary gradients alone already helps; adding NAS further improves performance by finding better connection locations; and using both auxiliary features and NAS yields the best result. The paper interprets this as evidence that cross-task connections matter, their placement matters, and sparse asymmetric connectivity discovered by NAS is beneficial (Gao et al., 2024).

The appendix reports convergence behavior for the architecture weights. Auxiliary-to-primary architecture weights converge to very small values, with maximum values around $P_i^{\text{fea}} = P_{i-1}^{\text{fea}}$9, whereas primary-to-auxiliary weights remain soft because they are not penalized and are discarded at inference anyway. The paper also reports stable statistics over multiple runs, indicating search robustness (Gao et al., 2024).

A further characteristic is orthogonality to optimization-based auxiliary learning. The paper explicitly lists Uncertainty weighting, DWA, PCGrad, CAGrad, and GCS as compatible because those methods modify optimization dynamics, whereas Aux-NAS modifies architecture. The claimed workflow is therefore compositional: the architecture search determines helpful cross-task structure, and gradient surgery or loss weighting can still be applied during training (Gao et al., 2024).

The limitations are also clearly delimited. Training is more complex than for a plain single-task model, the NAS stage adds optimization overhead, auxiliary labels must be available during training, and the method is designed for cases where only the primary task matters at inference rather than for standard multi-task settings that seek to optimize all tasks equally. Within those constraints, the method is most applicable when related auxiliary annotations are available during training but deployment cannot tolerate inference cost beyond that of a primary-only model (Gao et al., 2024).