Aux-NAS: Exploiting Auxiliary Labels with Negligibly Extra Inference Cost
Authors: Yuan Gao
WEIZHONG ZHANG
Wenhan Luo
Lin Ma
Jin-Gang Yu
Gui-Song Xia
Jiayi Ma
Venue: ICLR 2024 poster (2024)
Keywords: Auxiliary Learning; Neural Architecture Search; Soft Parameter Sharing; Multi-Task Learning; Single Task Inference Cost
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We propose a novel soft-parameter sharing architecture-based method optimized by Neural Architecture Search, which exploits auxiliary task labels to boost the primary task performance without increasing the inference cost for the primary task.

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.

Published as a conference paper at ICLR 2024
AuX-NAS: EXPLOITING AUXILIARY LABELS WITH NEGLIGIBLY EXTRA INFERENCE COST
Yuan Gao', Weizhong Zhang', Wenhan Luo", Lin Jin-Gang Yu', Gui-Song Xia'; Jiayi Ma' Wuhan University. FFudan University, *HKUST, "Meituan, South China University of Technology {ethan. y. gao, whluo. china, forest linma, jyma2010J6gmaal com, wei zhongzhange fudan. edu. cn, jingangyu@scut. edu. .n, guisong xww@whu. cn
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 learn- ing methods are optimization-base tion, our method is architecture-basse with 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 rep- resenting a task), we propose a novel method with evolving networks where only links exist as the cross-task connections after convergence. These connections can removed during the primary task inference, resulting single-task inference cost. We achieve this by formulating Neural Architec- ture Search (NAS) problem, where we initialize bi-directional connections in the search space and guide the NAS optimization converging connections. Moreover, our method can
relying on loss weights/gradients manipula-
architecture with
only the single-side primary-to-auxilinn be incorporated with optimization-ban based auxiliary learning approaches. Extensive experiments with six tasks on NYU v2, CityScapes, and Taskonomy datasets us- ing VGG, ResNet, and VIT backbones validate the promising performance. The codes are available at https//ggthub..om//hhaaaaa//////
INTRODUCTION
In this paper, we tackle the practical issue of auxiliary learning, which involves improving the perfor- mance of a specific task (i.e., the primary task) while incorporating additional auxiliary labels from different tasks (i.e., the auxiliary tasks). We aim to efficiently leverage these auxiliary labels to en- hance the primary task's 's performance while maintaining a comparable computational and parameter cost to single-task network when evaluating the primary task.
Our problem is closely related to Multi-Task Learning (MTL) but has two distinct practical require- ments: 1) only the performance of the primary task. rather than that of all the tasks, is of our interest, and 2) we aim to maintain a single-task cost during the inference. For example, we may primarily concerned with the semantic segmentation performance but have auxiliary depth labels available. The question is that can we leverage those auxiliary labels to improve the primary task, while pre- serving the inference cost for the primary task similar to single-task network?
The primary obstacle to achieving these goals is the presence of the inherent negative transfer, which caused by the conflict gradients from different tasks that flow to the shared layer Ruder (2017):
Ruder et al. (2019); Vandenhende et al. (2000b). To alleviate the negative transfer. most existing auxiliary learning methods are optimization-
which formulates a shared network and deals
with the conflicted gradients to the shared layer by modifying the auxiliary loss weights gradients Du et al. (2020); [Lu et al.] (2022); Navon et al. (202T): Verboven et al. (2020); [Shi al. (2020). However, two most recent studies independently show that it is challenging
solve the negative transfer solely by manipulating the loss weights or gradients Xin et al. (2022); Kurin et al. (2022).
"Corresponding Author.

---

Published as a conference paper at ICLR 2024
Primary Features
Auxiliary Features
Fusion Operation
Fixed
Searchable Connections w/ L1 Regu.
Searchable Connections w/o Regu.
Discovered Connections
Training Network:
Inference Network:
Training Network:
Aux-Grudient Method
Single-Tack Cost
Aux-Feature & Gradient Melhod with NAS
Discovered Arch.
Search Space
Opt. Arch. Regu.
Input
Exput
Ingont
Inpor
Figure 1: Overview of the proposed methods. Our methods are based on an asymmetric architec- ture that employs different networks for training and inference, where exploit gradients and/or features from the auxiliary task during the training, and preserve a single-task cost for evaluating the primary task. Our first method (Leftmost) leverages the auxiliary gradients. Our second method (Rightmost) exploits both auxiliary features and gradients, where the auxiliary-to-prime primary connec- tions (green dash lines) are gradually pruned by NAS, resulting only primary-to-aue toauiiiiary connections (the line widths indicate the converged architecture weights). Finally, the primary-to to-auxiliary connections, as well the auxiliary branch, can be safely removed to obtain single task network (Middle) to inference the primary task. The network arrows indicate directiinnsinverse directions the feature/gradient flow. (Best view colors.)
converged architecture with
Instead, the architecture-based methods with soft parameter sharing assign separated task-specific model parameters, which avoids the conflicting gradients from the root al, (2023). Therefore, we aim to better solve the negative transfer using the soft parameter sharing based architectures. In other words, we propose to learn separate set of parameters/features for each task avoid the negative transfer, where different primary and auxiliary features interact with each other. However, is challenging to achieve single-task inference cost with separate parameters for different tasks, as the primary and auxiliary networks rely each other input to generate higher-level features.
The above analysis inspires us to design an asymmetric network architecture that produces change- able networks between the training and the inference phases, i.e., the training network can be more complex to better exploit the auxiliary labels, while those auxiliary-relatte computations can re- moved during the inference. Starting with multiple single-task branches (each for task), key design is to ensure that the network is asymmetric and only includes directed acyclic primary- to-auxiliary forward connections. This allows primary task inference, resulting in a multi-task level performance and a single-task inference cost.
safely remove inter-task connections during
Motivated by this, first method exploits the gradients from the auxiliary task extra regulariza- tion to the primary task, while those auxiliary computations can be removed during the inference as the gradients are no longer required. We implement this by establishing multiple layerwise forward connections from the primary to the auxiliary tasks, where the auxiliary task leverages the features from, and thus pack-propagates the gradients to, the primary task, as shown Fig. (Left).
The follow-up question is that can we harness both the features and gradients from the auxiliary task during the training, while still maintaining a single-task cost in the inference? Fortunately, this can be achieved by training an evolving network with novel Neural Architecture Search (NAS) algorithm that employs asymmetric constraints for different architecture weights. Specifically, the NAS search space is initialized to include all bi-directional primary-to-auxilaary and auxiliary-to- primary connections. We then impose l constrains on the auxiliary-to- -orimmar architecture weights to gradually prune them out during the search procedure. We illustrate this in Fig. (Right).
Both of our proposed methods are -eneral-applicable to various auimlry- task combina- tions" with different single task backbones. Moreover, the proposed methods are orthogonal to, auxiliary learning methods
and can be incorporated with, most existing optimization-
'Specifically, given the primary and auxiliary task(s), we do not assume how much they related. Instead, methods designed to automatically leam what share hetween the tasks.
2

---

Published as a conference paper at ICLR 2024
Chen et aL 2018:
2020
Guo et al. 2018 Kendall et al.
20118: Karin [2022] : Royer
(2023 Sener
For MTL
Knltun 2018
Suren & Gua y (219 Lau et al. (219
202 2022 et [2022] 2019 Xin nl. (2022
Optimization-basa
Chen et al. [2022 Dery ct al. (202] [2023 | Du et nl.
For AL
2022 et al. 2020); Sun et al. (2020): Verboven el al. 2020
[2020] [
nl. (2021] Shh
Bruggemann al.
2016 2020): Guo al. (2020 (Hashimoto
Fur MTL
Kokkanos (2017) Lau el al. (2015 e
Manims al. 2019:
Van
MTL & AL Methods
denhende 20200 Xu 2018
Yang Hospodales
201
Hard
For AL
Arch based
For MTL
Gao 2020 2019
Misra
Ruder 2019
Soft
For
Our Method
Virtual Aux Label Gen.
For Al.
Liu
2022
et
2021
Gao K2019 2020
et al. 2019d
Ours
Sun al. 2020
Arch-based Methods
Misra et al. 2016 Kuder 2019
Manmus et al. 20119
(Soft)
(Solt)
(Hard)
(Hard)
Inference FLOPs
+ IJ + + 1/K
N A7
S N
Table The taxonomy of our method the MTL and AL areas (Top), and the inference FLOPs of our method and the representative architecture-based ones (Bottom). AL. Soft/Hard mean Auxiliary Learning, Soft/Hard Parameter Sharing. N the FLOPs for a single task, is the number of auxiliary tasks, and M is the fusion FLOPs for each task pair in Gao et al. (219 [Mssra
MMisra et al.
2016) Ruder et al. (2019), ) or the extra attention FLOOPs Liu et al. (2119d) Maninis et al. (2019).
Du al. 2020 Liu et al. 2022 Navon al. 202 1: Verboven al. (2020); Shi et al. (2020). We validate our methods with 6 tasks (see Sect. 4 using VGG- 16, ResNet-50, and ViTBase backbones,
on the NYU v2, CityScapes, and Taskonomy datasets. Our contributions three-fold:
We tackle the auxiliary learning problem with novel asymmetric architecture, which pro- duces different networks for training and inference, facilitating multi-task level perfor- mance with single-task level computations/parance
We implement the above idea with a novel training architecture with layerwise primary- to-auxiliary forward connections, where the auxiliary task computations provide additional gradients and can be safely removed during the inference.
We propose a more advanced method with evolving architectures to leverage both the aux- iliary features and gradients, where the auxiliary-to to--riaorrray connections can be gradually cut off by a NAS algorithm with a novel search space and asymmetric regularizations.
1.1 TAXONOMY OF OUR METHOOSS
The taxonomy of our methods in both MTL and auxiliary learning areas is illustrated in Table 1 Our methods fall under the category of architecture soft-parameter sharing scheme. We note that the virtual aux-label generation methods operate in distinct context from ours without auxiliary labels, and thus are beyond our scope.
methods for auxiliary learning with
In contrast all the existing auxiliary learning methods which focus on designing effective opti- mization strategies, we instead to explore novel auxiliary learning architecture design, for which its objective can be freely integrated with any optimization-based methods, as validated in Sect. 4.1
While architecture-base methods are generally better at migrating negative transfer Shi et al. (202), there few (if any) approach designed for auxiliary learning except for ours Moreover, methods leverage soft parameter sharing, which further alleviates the negative transfer compared with its hard parameter sharing counterpart Ruuder (2017), by implementing independent model pa- rameters. While independent model parameters often result in an increased inference cost outlined in Table (Bottom), our methods are different from them by preserving single-task inference cost.
2 RELATED WORK
Multi-Task Learning. MTL aims to improve the performance of all input tasks Long et al. (2017);
Kokkinos (2077; Zamir al. (2018), which can be categorized into Multi-Task Optimizations
3

---

Published as a conference paper at ICLR 2024
(MTO) and Multi-Task Architectures (MTA) Ruder (2017); Vandenhende et al. (2020b). The multi-
task optimization methods manipulate the task gradients/loss weights to tackle the negative transfer
Kendall et al.
(2018); Chen et al. (2018) ; Liu et al. (2019d) Lin et al. (2022); Chen et al. 22020) Guo
et al.
2018); Sener & Koltun (2018); Lin et al. (019)) Liu et I 2021b); Suteu & | (2019) 9) Liu
al. a 2022); Royer al. (2023). D. Our methods are orthogonal to MTO methods and follow
the .TAA category, i.e., learning better features for different tasks via elaborated Hard or Soft Param-
eter Sharing (HPS or SPS) network architectures Misra et al. (2016); Ruder (207)); Xu et (2018)
Ruder et al. (2019) Gao et al. 2019): Maninis et al. (2019); :) Gao et al. (2020); Liu et al. (2019d);
Vandenhende et al. (2020b a): Guo et al. (2020) Bruggemann et 2020); Yang & Hospedales
(2017); Kokkinos )); Hashimoto et al. (2016); Liu et al. (2019e). Our methods leverage SPS
scheme Misra et al. (2016); Ruder et al. (2019), : Gao et al. (2019 2020), which uses separate model
weights for each task to better tackle the negative transfer with single-task inference cost.
Auxiliary Learning. Most existing auxiliary learning methods are optimization-ba
which use shared feature set with auxiliary gradients/loss weights manipulation Liebel & Körner (2018); Du
cct al. 2020); Navon et al. (202T) Verboven et al. (2020); Shi et al. (2020); Liu et al. (2022); Chen
ee al. y Sun et al. (2020): Dery et al. (2021). Most recently, Dery et al. (2023) proposed to search for appropriate auxiliary objectives from a candidate pool and achieves remarkable perfor-
mance. Instead, our methods learn unique feature set for each task, which can integrate with those
methods. Recent works generate virtual auxiliary labels Liu et al. (2019c); Navon et al. (2021)
al. (022), which have a very different setting from ours and is thus beyond our scope.
Network Pruning & Neural Architecture Search. Network pruning aims at removing unimportant layers without severely deteriorating the performance Han et al. (2016); He et al. (2017); Liu et al.
(2177: Luo et al. (2017); Ye et al. Gordon et al. (2018). As the pruning process is crucial for the final performance Frankle & Carbin (2019). (2019) our algorithm gradually prunes hamaar-to-auxilian
connections by a single-shot gradient based NAS method Guo et al. (2019): Pham et al. (2018);
Saxena & Verbeek (2016); Bender et al. (2018); Liu et al. (2019)); Xie et al. (2019); Akimoto et al. (2019): el 2019: Zhang et al. (2019) iu et al. (2019a): Mei et (2020); Gao et al. (2200)
3 METHODS
In this section, we propose two novel methods exploiting the auxiliary labels to enhance our primary task, while keeping single task inference cost. Based the soft parameter sharing architecture, our key design is an asymmetric architecture which creates different networks for training and inference, i.e. the training architecture is more complex to exploit the auxiliary labels, while those auxiliary- related computations/parameerss can be safely removed during the primary task inference.
In the following, we first discuss our asymmetric architecture design. Then, we implement two novel algorithms, where the first method exploits the auxiliary gradients, and the second method leverages both the auxiliary features and gradients in an evolving network trained by novel NAS algorithm. After that, we implement the feature fusion operations. Finally, we provide a taxonomy of our methods within the areas of both MTL and auxiliary learning.
3.1 THE ASYMMETRIC ARCHITECTURE WITH SOFT PARAMETER SHARING
We tackle auxiliary learning by the
Primury Features
Amailayy Fentures
architecture- methods. Due to its aaeeneeeaaeee
Fniio Uperation
Eorward Connections
merit of better migrating the negative trans- fer, our methods follow the soft parameter sharing architecture, where different tasks exhibit separated network branches (i.e., in- dependent/unssaaeed network parameters) with feature fusion connections across them Ruder 2017); Vandenhende et al. (2020b).
Symmeerti Saft Parametec Sharling Strecture
Proposed Avyunetric Seft Purameter Shuring, Structure for Am--Learnig
ples/and
A
reloted compotatinus
Make auailiary.
pit//und
removeable daring the inferenee
OPp
O
04
plai/and
A
Ita
plea/arl
Bo/gnd
Given a primary task and an auxiliary task, the widely used soft parameter sharing structure in MTL is shown in Fig. (Left). Let plea and Afsa be the primary and auxiliary features for
Figure 2: The asymmetric primary-auxiliary ar- chitecture with soft parameter sharing. (Best view in colors.)
4

---

Published as a conference paper at ICLR 2024
the i-th layer, and Parad and A be the corresponding gradients, then the forward and the back-
ward for Fig. 2 (Left) are Eqs. 1 4 where OP [OPP, (PA] and O = [OAP, OAA] are the learnable fusion operations parameterized by the model weights and dO /d0 are the corre- sponding derivatives. Equation [ shows that the primary feature from the higher layer i takes both the primary and the auxiliary features from the i 1th layer. Thus, the auxiliary branch cannot be removed when inferencing the primary task in this case.
Symmetric, Fig. 2 (Lfft):
Asymmetric, Fig. (Right) (Ours):
pl 1
+
(1)
p 1
(5)
A 1 Aam I oAP p
(2)
Alt ( =01P
(6)
[-1 40"P pumu
+ 40 A
(3)
p P
60 APT
(7)
4 a pura
(4)
A 4o AA A
(8)
Since the gradients rameter sharing structure, which enables removing the auxiliary computations during the inference. As shown in Fig. 2 (Right), propose only exploit the auxiliary gradients (rather than features) as additional regularization for the primary task. According to the corresponding forward (Eqs. 5 6 and backward (Eqs. [) Eq. 7 shows the auxiliary gradients are used to train the primary task while Eq. enables to maintain single-task cost during the primary task inference.
longer required during the inference, design an asymmetric soft pa-
Remark The above analysis indicates that the structure for our problem should follow Fig. (Right). We implement two methods in Sects. and 33. Sect. directly applies Fig. 2 (Right), while Sect. 333 establishes evolving architecture that is initialized with bi-directional inter-task connections as Fig. (Left), then the eexiliary-t--rrmmne connections are gradually cut off using NAS during training, resulting in converged structure as Fig. 2 (Right).
3.2 THE AUXILIARY GRADIENT METHOD
Starting from two independent/unshared single task networks like Misra et al. (2016); Gao et al. (2019); Ruder et al. (2019), our first method implement Eqs. 5 g with moderate extension by inserting multiple layerwise primary-to-auxilian connections (representing the forward feature fu- sion) between the two branches. We denote this method as the Auxiliary Gradient method (Aux-G).
Our training architecture is shown in Fig. D (Left). We use fusion operation each layer of the auxiliary branch, which takes multiple features as input and produces a single output feature. As result, it enables multiple inter-task connections pointing to the same sink node of the auxiliary branch, which in turn allows multiple gradients routing to the primary branch.
Denoting the auxiliary task feature from the (i - 1)-th layer A; and the primary feature from the j-th layer as P3 where j I i, the fused auxiliary feature at i-th layer A₃ is:
A; = 0 A 00,; P 01.11
0
where O is the fusion operation, whose implementation will be discussed in Sect. 3.4 aj, is binary indicator representing the location of the there exists connection from j-th primary layer i-th auxiliary layer, otherwise, aj, is O.
connection, i.e., aj,i is if
3.3 THE AUXILIARY FEATURE AND GRADIENT METHOD WITH NAS
We further extend Aux-G to exploit both auxiliary features and gradients via evolving architec- ture trained by a novel NAS algorithm. Being initialized with bi-directional connections like Fig. 2 (Left), our NAS method guides the architecture converging to Fig, 2 (Right) with only primary- to-auxiliary connections, which can be removed during the primary task inference. We denote this method as the Auxiliary Feature and Gradient method with NAS (Aux-NAS).
Remark 2 Our search space applies bi-directional connections auxillary-to-primaryy at every layer between the two fixed single lask backbones.
cinimarr-to-uuxilion and
Such search space design enables us to exploit both the features and the gradients from the auxiliary task, which also exhibits two additional merits: the proposed search space is general-applicabe to
5

---

Published as a conference paper at ICLR 2024
any primary and auxiliary task combinations as it searches the general feature fusing scheme, and ii) it efficiently exploits the layerwise features of each task without introducing negative transfer.
We also propose a novel search algorithm that facilitates the evolving architectures con- verging to a model where only the crimary-to-axiilinn connections exist between the tasks. There- fore, those connections, as well as the auxiliary branch, can be safely removed without affecting the primary task performance during the inference.
Remark 3
By such design, we implicitly assume that besides the discovered discrete architecture, it is also important to exploit the mixed models (where the architecture weights are between O and 1) in the training procedure during the search phase. We note that, however, such importance leveraging the mixed-model training was not fully exploited in the popular one-shot gradient based NAS algorithms (e.g.ARRTS) Liu et al.] (2019b). Specifically, the most widely used training procedure of a typical one-shot gradient based NAS algorithm includes search phase and a retrain phase which produces the inconsistent learning objectives between those two phases. As a consequence, performance gap or generalization issue is witnessed between the discovered mixed architecture and the retrained single architecture Xie et al. (2019); Li et al. 2019).
Our method does not suffer from this issue, because all the primary-to -o-uxiiiyy connections cut off in our evaluation, regardless of a mixture model or a single model they converge to. The only requirement of our method is to restrict the architecture weights associated with the auxiliary- to-primary connections converging to a small value. Our NAS algorithm without a retrain phase makes it meaningful exploit features and gradients from the auxiliary branch during the search phase, as the searched model (with the auxiliary-related computations removed) is directly used for evaluation. We achieve this using l regularization on the auxiliary-to-ppl to-immaryy architecture weights, which gradually cuts off all the auxiliary-to-primas to-primary during the search phase. We do not impose constraints on the primary-to- to-aaryy---uxiiiaryyyyr connections. This method is shown in Fig. (Right).
Formally, let w be the model weights, denote the architecture weights for the seviliar-too-rrmaae connections as a {al. ji. with S j} where af for the connection from the auxiliary layer to the primary layer j. denote similarly the architecture weights for the primary-to-auxiliary connections as a {ac). J(i,j) with i S j}. our optimization problem becomes:
min (P w)) + (A(a w)) + R(aP), with R(a AIIaP111, (10) a _.,aa ,00
where L is the loss function for the primary task, L is the loss function for the auxiliary task, R is the regularization term on a with À as the regularization weight, P = (P,, Vi} are all the
fused primary features with P, as that of the i-th layer, and A {A,, Vi) are all the fused auxiliary features with A, as that of the i-th layer. Similar to Eq. 9 P. and A, are:
P. w -0 (P-1.,
of 1 A A
(11)
A,(a A w) = O ( 0b Po
(A
a 1 1, Pi-1)
(12)
Note that, being different from Au NAS is necessary for this method to exploit the auxiliary features, it is used to gradually cut off the ciiliiarytt--priman connections.
3.4 FUSION OPERATION
We design a unique fusion operation for both of our methods, so as to better illustrate the merit of the proposed novel network connections. There are two principles design the fusion operation: i) the fusion operation can take an arbitrary number of input features, and ii) there is a negligible cost on both parameters and computations in P. Pa ur) of Eq. 11 with a as all O.
Existing researches discussed the fusion operations on several input features, such as those based on weighted-sum Misra et al. (2016), attention [iu et al.) (2019d), and neural discriminative dimen- sionality reduction (NDDR) Gao et al. (2019). where both Liu al. (2019d) and Gao et al. (2019) leverage Ix] convolution for feature transformation. Besides the extension of taking arbitrary input features, method also integrates the advantages of those methods, where we implement heavier
'The two- pho-p training procedure is needed because the search phase usually converges to a mixed model with the architecture weights hetween 0 and 1, thus the retrain phase comes up prune the mixed model and retrain single model with fixed architecture for evaluation.

---

Published as a conference paper at ICLR 2024
computations (i.e, 1xl convolution) on the inference-removable features and negligible compu- tations (e.g., BatchNorm and ReLU) on the feature that remains during the inference
Add
Actuntion
Veaggned
AU2% VIR
As show in Fig. eration in Eqs. nation, 1xl convolution, summation, normalization, and activation, where [-] demotes a feature con- catenation along the channel dimension, Activ and Norm can be ReLU and BatchNorm. Additionally, Bill neaneaaain Interp applied each input feature when necessary before concatenation, it resizes the input features to the output spatial resolution, we omit it from Eqs. 13 and [4 for simplicity.
implement the fusion op- and [2 with feature concate-
Figure 3: The illustration of the proposed fu- sion operator. Note that the auxiliary features (in orange color) are concatenated by its ar- chitecture weights. The dash line indicates that the regularized NAS objective in Eq. I0 enables to cut off the whole auxiliary compu- tations (also the following 1x1 conv due to input). (Best view in colors.)
A A.-1))) Pr , w) Activ Activ (NorrNNNonPPP.. + 1x1 Lxl.coov((abb.aal. A
(13)
A:(aA w) = Activ Norm(A:-1 + 1x1.conv([og, Po... 1,,, Pi-I))
(14)
Equation 13 enables to discard the heavier Ixl convolution when a are all 0 (as constrained by R(a in Eq. 100, which introduces negligible computation from a BatchNorm and a ReLUU Specifically, when introducing fusion operations in n layers, there are only additionally 2n parame- ters (i.e., 3, y from BatchNorm), 2n summations and 2n productions (both from BatchNorm), and truncations (from ReLU). Those cost is negligible because at most is the number of layers in a network (e.g., at most n is 16 for VGG-16, 50 for ResNet-50, and 12 for ViTBase). Equation [4] does not introduce additional inference cost as it can be removed as whole during the inference.
EXPERIMENTS
We fully assess our methods following the taxonomy in Table We first show that our methods can be incorporated with the optimization-bases methods. Then, we evaluate our methods against the architecture- based methods across various datasets, network backbones, and task combinations:
Network Backbones. We evaluate methods on both CNNs and Transformers, i.e., VGG-16 Simonyan Zisserman (015), ResNet-50 He et al. (2016), and ViTBase Dosovitskiy et al. 2021).
Datasets. We perform our experiments the NYU v2 Silberman et al. (0112). the CityScapes Cordts et al. (206)), and the Taskonomy Zamir et al. 2018) datasets.
Tasks. Six tasks performed including semantic segmentation, surface normal prediction, depth estimation, monocular disparity estimation, object classification, and scene classification. We detail the task losses and the evaluation metrics in Appendix B
We implement Aux-G with two granularities: Aux-G-Stage with linked connections at the last layer of each stage, and Aux-G-Layer with linked connections every layer within each stage. We allow the bi-directional links for Aux-NAS. We restrict the fusion operation to only receive the features within 3 layers, i.e., the i-th layer fusion operation only receive features of 3, i 2, i layers from the other task, where (a.P aAA) and w are trained alternatively by sampling two non-overlapped batches Gao et al. (2020). We give more details in Appendix and analyze the architecture convergence in Appendix
4.1 COMPARISON WITH OPTIMIZATION BITTIIZZAIIOO-BAEED METHODS
Our architecture-based methods are mathematically orthogonal to, and can be incorporated with, the optimization-basee based methods. In this section, we compare and incorporate our methods with four MTL methods: Uncertainty Kendall et al. 2018). DWA Liu et al. (2019d), PCGrads Yu et al. 2020), CAGrads Liu et al.] (2021a), and one existing auxiliary learning GCS Du et al. (2020) We also implement an auxiliary learning variant of PCGrads, i.e., PCGrads-Aux, which always
Note that it is direct include multiple candidate fusion operations in the search space using NAS, hut chasing the fusion-op of each layer is prohlem-depend (thus distracting) and heyond the scope of this paper.

---

Published as a conference paper at ICLR 2024
Primary: Seg
Primary: Normal
CityScapes
(%) (f)
NYU v2
(%) (1)
Err (4)
Within 1 (%) (1)
Primary: Seg
mloU
PAcc
mloU
PAcc
Mean
Med.
RMSE
11.25
22.5
Single
68.3
94.5
Single
46.4
Aux-Head
70.0
94.6
Aux-Head
34.7
65.4
15.3
11.7
20.0
48.4
75.9
Adashare
70.3
Adushare
35.0
65.6
14.8
11.5
18.8
50.3
77.2
Adashare-Aus
70.1
94.8
Adushare-Aux
35.0
65.7
14.4
19.1
51.2
78.3
Aux-G-Stage
70.1
94.8
Aux-G-Stage
35.4
65.9
Aux-G-Layer
70.2
94.8
Aux-G-Layer
35.6
65.9
14.4
11.0
18.8
50.8
78.7
Aux-NAS
36,0
66.1
14.2
10.6
18.9
52.4
79.0
Table Semantic seg. on CityScapes using VGG-16
Table 3: Semantic segmentation and Surface normal prediction
on NYU v2 using VGG-16, with the other task the auxiliary task.
with disparity as aux. task.
projects the auxiliary gradients to the primary task. Next, first show the (vanilla) performance of those meth- ods on the basic Aux-Head network with fully shared encoder and separated heads for different tasks, then we report the performance integrating our networks. We con- duct semantic segmentation experiments on NYU v2 with surface normal estimation as the auxiliary task.
(t)
NYU x2, Primary: Seg
Aux-Head
34.7
65.4
Aux-Head + Uncertainty
35.2
65.6
Aux-Head DWA
35.3
65.7
Aux-Head + CAGrad
Aux- + PCGrad
35.0
65.5
Aux- Head + PCGrad Aux
35.2
65.7
The results are shown in Table 2 which illustrates
Aux-Head GCS
35.3
65.9
that all of the proposed networks without incorporating optimization-base
Aux-G-Stage
35,4
65.9
approaches have already outper- formed the best vanilla optimization- -based methods. This might imply the superior of the soft parameter sharing ar- chitecture methods in dealing with the negative transfer over the optimization- based methods Xin et al.] (2022); Shi al. (2023). Moreover, performance be further improved by incorporating with the optimization- based approaches.
Aux-G-Layer
65.9
Aux-NAS
66.1
Aux-NAS + Uncertainty
36.0
66.3
Aux-NAS DWA
36.0
66.1
Aux NAS CAGrad
35,9
66.0
Aux-NAS PCGrad-Aux
36.2
66.3
Aux-NAS GCS
36.3
66.5
Table 2: Comparison with the opt-based methods on NYU v2 semantic seg., with surface normal prediction as auxiliary.
4.2 COMPARISON WITH ARCHITECTURE-BASED METHODS
Note that methods exhibit single task inference cost, which is not ensured in many state-of-the- art architecture-baset methods as shown in Table [ Therefore, besides the Single and Aux- Aex-Head baselines, we compare with Adashare (Snn al. (220)) its inference cost is equal that of a single task. For more fair comparisons, we implement an auxiliary learning variant of Adashare, where for the primary task, we fix the select-or-skip policy as all I's and remove its sparsity regularization. This produces Adashare-Auu with the same inference cost of ours.
less than
4.2.1 DIFFERENT PRIMARY-AUXILIAS TASK COMBINATIONS
This section is show that our methods consistently outperforms the state-of-the-art methods across different primary- auxiiiyyy combinations. We use NYU v2 Silberman et al. (2012) dataset to conduct our experiments. We carry out two different primary-aume auxillary semantic segmentation and surface normal estimation combinations, where each of them serves the primary task and the other the auxiliary. We use the same network backbone and the same losses as Sect. [4. Table show show
that methods outperforms the basclines for all the experiments. And our Aux-NAS outperforms Aux-G, as Aux-NAS leverages the additional auxiliary features, which formulates a much larger capacity in the NAS search space, therefore it achieves better convergence.
4.2.2 DIFFERENT DATASETS
To further assess the applicability of our methods across datasets, Table show that our methods consistently outperform other baselines, when using auxiliary monocular disparity estimation to assist the primary semantic segmentation on CityScapes dataset. We also perform object and scene classification tasks Taskonomy dataset in Appendix D
8

---

Published as a conference paper at ICLR 2024
4.2.3 DIFFERENT BACKBONES
NYU v2, Pri-
Err (4)
Within t" (%) (t)
mary: Normal
Mean
Med.
RMSE
11.25
22.5
This section is to show the robustness of our methods across single-task backbones. We validates this on both ViTBase with DPT decoder Ranfti et al. (202T) Table S and
Single
14.6
12.9
17.7
43.2
80.8
Aux -eeal
14.8
13.2
17.9
41.9
80.1
Adushure
13.2
11.4
16.8
49.7
82.2
Adashare-Aun
12.9
11.0
16.7
51.9
85.5
G Layer
12.6
15.7
52.3
85.9
ResNet-50 in Appendix E which demon-
Aux NAS
12.5
10.3
15.6
53.8
85.9
strate that the proposed methods consistently outperform other state-of-the-art.
Table 5: Surface normal prediction on NYU v2 as- sisted by segmentation using ViTBase. We do not have Aux-G-Stage as no Stage concept in ViTBase.
4.2.4 SCALABILITY
TO MORE AUXILIARY TASKS
Our method maintains a consistent single- task inference cost with more auxiliary tasks.
NYU Pri- mary: Normal
Err (U)
Within ( (f)
Mod.
RMSE
11.25
22.5
Moreover, our training complexity scales
Aux GaL--G (1)
12.6
10.7
15.7
52.3
85.9
linearly with the number of auxiliary tasks.
Aux-NAS (1)
12.5
10.3
15.6
53.8
85.9
This is because i) linear network Growth: as focus solely
Aux-G-Layer
12.5
10.6
15.5
52.7
86.1
the primary task perfor- mance, given K auxiliary tasks, we only link K inter-task connections from each auxil- iary task to the primary task. This is in con-
Aux-NAS (2)
12.2
10.2
15.3
54.5
86.7
Table 6: Scalability to more auxiliary tasks. (1) and
(2) denote the number of auxiliary tasks.
trast to (K + K)K22 inter-task connections for all the task pairs in the soft parameter sharing MTL methods Gao al. [2119] 2020); Misra et al.) (2016); [Ruder et al. (2019); ii) linear NAS complexity w.r.l the network size: the single-shot gradient-based NAS algorithm has a linear complexity, as proved in Sect. 2.3 of Liu et al. (2019b). We discuss in details in Appendix A
We perform surface normal estimation as the primary task on NYU v2 using ViTBase, where we use i) semantic segmentation, and ii) semantic segmentation and depth estimation, as the auxiliary task(s). Table [6] shows that the performance of our method can be further improved with more auxiliary tasks. The average barward-backwand tasks (i.e., and 3 tasks including the primary task) are 0.337s and 0.556s (roughly 3/2 times of 0.337s), demonstrating the linear scalability of training complexity more auxiliary tasks.
time of each 10-samlle batch for 1 and 2 auxiliary
5 ABLATION ANALYSIS
In this section, we first show the effects of our design, i.e.the effects of the Auxiliary Gradients, the Auxiliary Features, with and without the NAS training. We perform ablations using VGG-16 on NYU v2 with semantic seg- mentation the primary task, to study the effects of the auxiliary gradients, the auxiliary features, and the NAS training, respectively. Our results in Table demonstrate that i) NAS improves the method which only exploits the auxiliary gradients (i.e., Aux-r-Layer) by discovering the best locations of the primary-to-auxiliar connections, and ii) our Aux-NAS further boost the per- formance of the Aux-G method with NAS, it also leverages the auxiliary features.
Seg. (%) (f)
Gradient
Feature
NAS
mloU
PAce
35.4
65.9
35.7
36.0
66.1
Table 7: Effects of the auxiliary gradi- ents, the auxiliary features, and the NAS training.
6 CONCLUSION
We aim to exploit the additional auxiliary labels to improve the primary task performance, while keeping the inference as a single task network. We design novel asymmetric structure which produces changeable architectures for training and inference. We implement two novel methods by exploiting the auxiliary gradients solely, or by leveraging both the auxiliary gradients and fea- tures with neural architecture search. Both of our methods converges to an architecture with only crimary-to-auxilley -uxiiiiiary connections, which can be safely removed during the inference to achieve a tangle-task inference cost for the primary task. Extensive experiments show that method can be incorporated with the optimization based approaches. and generalizes to different task combinations, different datasets, and different single task backbones.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank Haoping Bai for constructive discussion. This work was supported by the National Nat- ural Science Foundation of China (62306214, 62325111, 62372480, 62076099), Natural Science Foundation of Hubei Province of China (2023AFB19)), and Knowledge Innovation Program of Wuhan-Shugung Project (20230102010202588)
REFERENCES
Youhei Akimoto, Shinichi Shirakawa, Nozomu Yoshinari, Kento Uchida, Shota Saito, and Kouhei Nishida. Adaptive stochastic natural gradient method for one-shot neural architecture search. In ICML, pp. 171-180, 2019.
Gabriel Bender, Pieter-Jan Kindermans, Barret Zoph, Vijay Vasudevan, and Quoc Le. Understand- ing and simplifying one-shot architecture search. In ICML, 2018.
David Bruggemann, Menelaos Kanakis, Stamatios Georgoulis, and Luc Van Gool, Automated search for resource-efficient branched multi-task networks. In BMVC, 2020.
Hong Chen, Xin Wang, Chaoyu Guan, Yue Liu, and Wenwu Zhu. Auxiliary learning with joint task and data scheduling. ICML, 2022.
Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In ICML, 2018.
Zhao Chen, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov. Just pick sign: Optimizing deep multitask models with gradient sign dropout. In NeurlPS, 2020.
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The cityscapes dataset for semantic urban scene understanding. In CVPR, 2016.
Lucio M Dery, Yann Dauphin, and David Grangier. Auxiliary task update decomposition: The good, the bad and the neutral. In ICLR, 2021.
Lucio M Dery, Paul Michel, Mikhail Khodak, Graham Neubig, and Ameet Talwalkar. Aang: Au- tomating auxiliary learning. In ICLR, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x 16 words: Transformers for image recognition at scale. ICLR, 2021.
Yunshu Du, Wojeiech M Czarnecki, Siddhant M Jayakumar, Razvan Pascanu, and Balaji Lakshmi- narayanan. Adapting auxiliary losses using gradient similarity. arXiv preprint arXiv: / / 0.02224, 2020.
David Eigen and Rob Fergus. Predicting depth, surface normals and semantic labels with common multi-scale convolutional architecture. In CVPR, 2015.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In ICLR, 2019.
Yuan Gao, Jiayi Ma, Mingbo Zhao, Wei Liu, and Alan L Yuille. NDDR-CNN: fusing in multi-task cnns by neural discriminative dimensionality reduction. In CVPR, pp. 3205- 3214, 2019.
Layerwise feature
Yuan Gao, Haoping Bai, Zequn Jie, Jiayi Ma, Kui Jia, and Wei Liu. MTL-NAS: Task- agnostic neural architecture search towards ceneral-purposs multi-task learning. arXiv preprint arXiv:2003. /4058, 2020.
Ariel Gordon, Elad Eban, Ofir Nachum, Bo Chen, Hao Wu, Tien-Ju Yang, and Edward Choi. Mor- phnet: Fast & simple resource--oosstraiee structure learning of deep networks. In CVPR, pp. 1586-1595, 2018.
10

---

Published as a conference paper at ICLR 2024
Michelle Guo, Albert Haque, De-An Huang. Serena Yeung, and Li Fei-Fei. Dynamic task prioriti- zation for multitask learning. ECCV, 2018.
Pengsheng Guo, Chen-Ye Yu Lee, and Daniel Ulbricht. Learning to branch for multi-task learning. In ICML, 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. arXiv preprint arXiv:1904.00420, 2019.
Song Han, Huizi Mao, and William J. Dally. Deep compression: Compressing deep neural network with pruning, trained quantization and huffman coding. ICLR, 2016.
Kazuma Hashimoto, Caiming Xiong, Yoshimasa Tsuruoka, and Richard Socher. A joint many-task model: Growing neural network for multiple nlp tasks, In EMNLP, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, pp. 770-778, 2016.
Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural net- works. In ICCV, pp. 1398-1406, 2017.
Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. CVPR, 2018.
lasonas Kokkinos. Ubernet: Training auniversal convolutional neural network for low-, mid-, and high-level vision using diverse datasets and limited memory. In CVPR, 2017.
Vitaly Kurin, Alessandro De Palma, Ilya Kostrikov, Shimon Whiteson, and M. Pawan Kumar. In defense of the unitary scalarization for deep multi-task learning. In NeurlPS, 2022.
Guohao Li, Guocheng Qian, Itzel Delgadillo, Matthias Müller, Ali Thabet, and Bernard Ghanem. Sgas: Sequential greedy architecture search. arXiv preprint /9Xiv:/912.000 2019.
Lukas Liebel and Marco Körner. Auxiliary tasks in multi-task learning. arXiv: 2018.
arXiv preprint
Baijiong Lin, Feiyang Ye, Yu Zhang, and Ivor Tsang. Reasonable effectiveness of random weight- ing: litmus test for multi-task learning. In ICLR, 2022.
Xi Lin, Hui-Ling Zhen, Zhenhua Li, Qingfu Zhang, and Sam Kwong. Pareto multi-task learning. In NeurlPS, 2019.
Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-task learning. In NeurIPS, 2021a.
Chenxi Liu, Liang-Chieh Chen, Florian Schroff, Hartwig Adam, Wei Hua, Alan L. Yuille, and Li Fei-Fei. Auto-deeplab: Hierarchical neural architecture search for semantic image segmenta- tion. In CVPR, June 2019a.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In ICLR, 2019b.
Liyang Liu, Yi Li, Zhanghui Kuang, Jing-Hao Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Towards impartial multi-task learning. In ICLR, 2021b.
Shikun Liu, Andrew J. Davison, and Edward Johns. Self-supervised generalisation with meta aux- iliary learning. In NeurlPS, 2019c.
Shikun Liu, Edward Johns, and Andrew J Davison. End-to-end multi-task learning with attention. In CVPR, pp. 1871-1880, 2019d.
Shikun Liu, Stephen James, Andrew Davison, and Edward Johns. Auto-lambda: Disentangling dynamic task relationships. Transactions on Machine Learning Research, 2022.

---

Published as a conference paper at ICLR 2024
Yifan Liu, Bohan Zhuang. Chunhua Shen, Hao Chen, and Wei Yin. Auxiliary learning for deep multi-task learning. arXiv preprint arXiv:1909.02214, 2019c.
Zhuang Liu, Jianguo Li, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang. Learn- ing efficient convolutional networks through network slimming. In ICCV, pp. 2755-2763, 2017.
Mingsheng Long. Zhangjie Cao, Jianmin Wang, and Philip S. Yu. Learning multiple tasks with deep relationship networks. In NewrlPS, 2017.
Jian-Hao Luo, Jianxin Wu, and Weiyao Lin. Thinet: A filter level pruning method for deep neural network compression. ICCV, pp. 5068-5076, 2017.
Kevis-Kokitty Maninis, Ilja Radosavovic, and lasonas Kokkinos. Attentive toagle-tskknngg of multi- ple tasks. In CVPR, pp. 1851-1860, 2019.
Jieru Mei, Yingwei Li, Xiaochen Lian, Xiaojie Jin, Linjie Yang, Alan Yuille, and Jianchao Yang. AtomNAS: Fine- tine-graieed end-to-end neural architecture search. In ICLR, 2020.
Ishan Mísra, Abhinav Shrivastava, Abhinav Gupta, and Martial Hebert. Cross-stitch networks for multi-task learning. In CVPR, pp. 3994-4003, 2016.
Aviv Navon, Idan Achituve, Haggai Maron, Gal Chechik, and Ethan Fetaya. Auxiliary learning by implicit differentiation. In ICLR, 2021.
Hieu Pham, Melody Y Guan, Barret Zoph, Quoe V Le, and Jeff Dean. Efficient neural architecture search via parameter sharing. arXiv preprint arXiv:/802.03226. 2018.
René Ranftl, Alexey Bochkovskiy, and Vladlen Koltun. Vision transformers for dense prediction. In ICCV, 2021.
René Ranftl, Katrin Lasinger, David Hafner, Konrad Schindler, and Vladlen Koltun. Towards robust monocular depth estimation: Mixing datasets for zero-shot cross-dataset transfer. TPAMI, 44(3), 2022.
Amelie Royer, Tijmen Blankevoort, and Babak Ehteshami Bejnordi. Scalarization for multi-task and multi-domain learning at scale. In NeurIPS, 2023.
Sebastian Ruder. An overview of multi-task learning in deep neural networks. arXiv preprint arXiv: 2017.
Sebastian Ruder, Joachim Bingel, Isabelle Augenstein, and Anders Søgaard. Latent multi-task ar- chitecture learning. In AAAI, 2019.
Shreyas Saxena and Jakob Verbeek. Convolutional neural fabrics. NIPS, pp. 2016.
Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. In NeurIPS, 2018.
Baifeng Shi, Judy Hoffman, Kate Saenko, Trevor Darrell, and Huijuan Xu. Auxiliary task reweight- ing for minimum-data learning. In NeurlPS, 2020.
Guangyuan Shi. Qimai Li, Wenlong Zhang. Jiaxin Chen, and Xiao-Ming Wu. Recon: Reducing conflicting gradients from the root for multi-task learning. In ICLR, 2023.
Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus. Indoor segmentation and sup- port inference from rgbd images. ECCV, pp. 746-760, 2012.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In ICLR, 2015.
Ximeng Sun, Rameswar Panda, Rogerio Feris, and Kate Saenko. Adashare: Learning what to share for efficient deep multi-task learning. In NewrIPS, 2020.
Mihai Suteu and Yike Guo. Regularizing deep multi-task networks using orthogonal gradients. arXiv preprint arXiv: /912v:9912.066 1912.0684, 2019.
12

---

Published as a conference paper at ICLR 2024
Simon Vandenhende, Stamatios Georgoulis, Bert De Brabandere, and Luc Van Gool. Branched multi-task networks: Deciding what layers to share. In BMVC, 2020a.
Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Mare Proesmans, Dengxin Dai, and Luc Van Gool. Multi-task learning for dense prediction tasks: A survey. arXiv preprint arXiv:2004. 13379, 2020b.
Sam Verboven, Muhammad Hafeez Chaudhary, Jeroen Berrevoets, and Wouter Verbeke. Hydalearm: Highly dynamic task weighting for multi-task learning with auxiliary tasks. arXiv preprint arXiv:2008. 11643, 2020.
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 CVPR, 2019.
Sirui Xie, Hehui Zheng, Chunxiao Liu, and Liang Lin. SNAS: stochastic neural architecture search. In ICLR, 2019.
Derrick Xin, Behrooz Ghorbani, Ankush Garg, Orhan Firat, and Justin Gilmer. Do current multi-task optimization methods deep learning even help? NeurIPS, 2022.
Dan Xu, Wanli Ouyang, Xiaogang Wang, and Nicu Sebe. Pad-net: Multi-tasks guided prediction- and-distillation network for simultaneous depth estimation and scene parsing. In CVPR, pp. 675- 684, 2018.
Yongxin Yang and Timothy Hospedales. Deep multi-task representation learning: A tensor factori- sation approach. In ICLR, 2017.
Jianbo Ye, Xin Lu, Zhe Lin, and James Z. Wang. Rethinking the smaller-oom--es--imfmmaaat assumption in channel pruning of convolution layers. In ICLR, 2018.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In NeurlPS, 2020.
Amir R. Zamir, Alexander Sax, William B. Shen, Leonidas J. Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2018.
Yiheng Zhang, Zhaofan Qiu, Jingen Liu, Ting Yao, Dong Liu, and Tao Mei. Customizable architec- ture search for semantic segmentation. In CVPR, 2019.
13

---

Published as a conference paper at ICLR 2024
A DETAILS OF THE TRAINING COST, THE NAS COMPLEXITY, AND THEIR LINEAR SCALABILITY TO MORE AUXILIARY TASKS
We first emphasize that the models for both of our Aux-G and Aux-NAS methods increase linearly with more auxiliary tasks, since we focus solely on the primary task performance. As shown in Fig. given K auxiliary tasks, we only have K inter-task connections from each auxiliary task to the primary task and do not need the interactions between the auxiliary tasks. This is in contrast to (K 1/ inter-task connections for all the task pains in the soft parameter sharing MTL methods Gao et al. (2019 2020): Misra et al. 20166; Ruder et al. 2019).
The Auxiliary Learning Architectures (ours):
The Multi-Tasi Learning Architectures: All the input tasks are pears, resulting in square inter-task connections wrt. the number of nuxiliary tasks, ic., the inter- comnections ure established for the task puirs.
It focus solely on the primary task performance and only connects the ptimary task to each of the auxiliary tasks, resulting in linear inter lusk connections w.r. the number muxiliary tasks.
Task
Task
Aux
Prim
Aux
Task
Figure 4: An illustration for the inter-task connections (i..,, the search space) of the auxiliary learn- ing (ours) and the multi-task learning architectures. We use 3 tasks (or primary task plus auxiliary tasks) as an example.
Our Aux-G model is simply trained by gradients, whose complexity is linear w.r.t. the model FLOPs and scales linearly with more auxiliary tasks.
Our Aux-NAS model is trained using the single-shot gradient-based NAS algorithm, as proofs in Sect. 2.3 and Eq. of Liu et al.] (2019b), the complexity of the single-shot gradient-based NAS algorithm is linear w.r.t. the amount of the model weights and that of the architecture weights az
O(w + lax).
(15)
As shown in Fig 4 given K auxiliary tasks, the amount of the model weights becomes (K + w) (including the primary task), and that of the architecture weights becomes K [ax] (as we only connect the primary task to each auxiliary task, rather than between the auxiliary task pairs), therefore, the training complexity of our Aux-NAS model is:
O(KK + 1)]w] + K|aa),,
(16)
which scales linear with K. Our experiments in Table the main text also verifies the training time for each batch increases from 0.337s for K 1 to 0.556s for K 2, where the ratio is between (2+1)/ / (1+1)= = 1.5 (for model weights complexity) and 2/1 = (for architecture weights complexity).
B DETAILS OF THE TASK LOSSES AND THE EVALUATION METRICS
Losses. We use the cross-entropy loss for semantic segmentation. Surface normal prediction is trained by the cosine similarity loss. The scale and shift invariant loss Ranftl et al. (2022) is used for depth estimation. We train disparity estimation with the mean squared error (MSE). For the object classification and scene classification on Taskonomy, we use the &g loss as the "groundtruth" given by the Taskonomy dataset is the "soft classification labels" predicted by large network.
Evaluation Metries. For semantic segmentation, we evaluate with the pixel-wise accuracy (PAcc) and mean intersection over union (mloU). We mean, median, and root square (RMSE)
14

---

Published as a conference paper at ICLR 2024
Taskonomy
(%)
(1)
Primary: Object Cls.
Top-1
Top-5
Single
34.3
65.9
Aux-Head
34.7
66.6
Adashare
35.9
67.1
Adashare-Aux
36.3
67.7
Aux-G-Stage
37.4
67.9
Aux-G-Layer
37.2
68.3
Aux-NAS
39.8
70.7
Table 8: Object classification on the Taskonomy dataset with scene classification as the auxiliary task using the ResNet-50 network.
angle difference, also the percentage of pixels within 11 and 22.5 w.r.t the ground truth for surface normal prediction. We report the Top-1 and Top-5 accuracy for object classification.
C MORE TRAINING AND EVALUATION DETAILS
Training Strategy. We train our Aux-G method simply by gradients. For the Aux-NAS method, we train (a ", A) and w alternatively by sampling two non-overlapped batches similar Gao et al. (2020).
Evaluation Strategy. We simply cut off all the cimmarytoo-uxxilinn connections and the auxiliary branch for inference, which remains a single-task model for the primary task without extra inference computations.
More Training Details. We use 321 321 image samples for the CNN backbones (i.e., ResNet-50 and 16)-16), and 224 224 image samples for transformer backbone (i.e., ViTBase Dosovitskii et al. (2021). The ViTBase backbone we used is vit wwt._aaeeepttteeeeeeeeeeeeeettee
6_224 of the huggingface timm package. We initialize the single task branches with the pretrained single task model weights. For the fusion operations, initialize the Ixl convolution with all 0. We gradually increase À of Eq. 10 from 100 during the training, and initialize all a' 'ss of Eqs. 13 and 14 to 0.5.
For the pixel-labeling tasks (i.e., semantic segmentation, surface normal prediction, depth estima- tion, and monocular disparity estimation) on CNN backbones (i.e., ResNet-50 and VGG-16), we Deeplab head with atrous convolutions Liu al. (2022). We DPT head Ranftl et al. (202T) for those tasks on Transformer backbone (i.e., ViT'aase.. For the classification tasks (i.e., object classification and scene classification), simple the feature dimensions to the number of classes.
Tully-Connected/MII.P layer map
We use the official train/val split for the NYU Eigen Fergus 2015) and the CityScapes Cordts eet al. (2016) datasets. For the Taskonomy Zamir et al. (2018) dataset, we use the official Tiny split of it.
We treat the convolutions layers of VGG-16, the residue blocks of ResNet-50, and the multi-head attention blocks of ViTBase the basic building elements to construct the inter-task connections.
D THE EXPERIMENTS ON THE TASKONOMY DATASET
In this section, we perform the object classification task, assisted by the scene classification, on the Taskonomy dataset, using the Res-50 network. The Top-1 and Top-5 recognition rates are reported Table 8 which, accompanied with Tables and in the main text, demonstrate that our method consistently outperforms the state-of-the-art various tasks and datasets. 33.8 63.0 37.8 70.5 Multiple 34.1 66.1 37.8 71.2
E THE EXPERIMENTS WITH THE RESNET-50 BACKBONE
This section is to further demonstrate the robustness of our methods across backbones. We validates this ResNet-50 network in Table which, accompanied with Tables 3 and 5 in the main text,
15

---

Published as a conference paper at ICLR 2024
NYU v2
(%)
(f)
Primary: Seg
mloU
PAce
Single
34.1
65.0
Aux-Head
32.9
64.6
Adashare
33.3
64.9
Adashare-Aux
33.7
Aux-G-Stage
33.4
64.8
Aux-G-Layer
32.9
64.5
Aux-NAS
36.8
66.7
Table 9: Semantic segmentation on the NYU v2 dataset with surface normal prediction as the aux- iliary task using the ResNet-50 network.
demonstrate that the proposed methods consistently improve the performance across different single task backbones.
F ANALYSIS OF THE ARCHITECTURE CONVERGENCE
In the following, we analyze the convergence of the searched architecture based on the Aux-NAS weights of Table 5
Auxiliary-to-Prina Primtar Connections: The auxiliary-to -orimmary architecture weights converges to very small values, with the maximum value around 0.01, which indicates that method success- fully pruned those auxiliary-to-primayy down. Notably, all the results reported in our manuscript are the final single-task inference performance, where manually set all the auxiliary-to- primary architecture weights to 0 before conducting the evaluation. This guarantees our ob- jective of ensuring the single-task inference cost.
Regarding the primary-to-auxiln connections, as we do not impose any regularization on their architecture weights (and those connections will be removed to maintain single-task inference regardless how they converge), those connections typically converge soft values between 0 and We conducted 10 replicate runs with random seeds using the same experimental setup as those in Table (i.e.. ViTBase on NYUv2), the mean and standard deviation of min, max, mean, median statistics of the primary-to-auxilary architecture weights are shown below:
Connections:
Prim-to-Aux Arch Weight
min
max
mean
median
mean std
0.1561
0.5041
0.3073
0.2725
0.0110
0.0120
0.0117
0.0150
Table 10: Statistics of the converged primary-to-auxillan weights. Those results are obtained by
performing 10 replicate runs of the experiment in Table 5
The table above illustrates a small standard deviation of the converged primary- to-ayy-toouuiiiaryy weights in replicate runs, indicating degree of stability in method. We attribute this small deviation in the converged architectures to: i) for the model weights, initializing our network backbones with converged single-task networks, and ii) for the architecture weights, focus our architecture search solely on cross-task connections without altering the network backbones.
16


Reviews:
[{"review_id": "uCTzTFo7Vf", "reviewer": "Reviewer_fF8U", "summary": "The paper studies how to harness additional auxiliary labels from an auxiliary task to elevate the performance of the main task without escalating the inference cost. To do so, the authors propose to employ individual networks for different tasks and only regularize the main task with the auxiliary task’s gradient. It’s understandable that this act allows the network trained on the auxiliary task to be completely pruned during inference. Furthermore, the authors propose to search for the most appropriate structure that satisfies the previously mentioned constraint with NAS. The paper accentuates its methodology's compatibility with prevailing optimization-based auxiliary learning techniques. The empirical validation, evident from experiments on NYU v2, CityScapes, and Taskonomy datasets using well-known backbones like VGG-16, ResNet-50, and ViT-B, demonstrates the efficacy of the proposed method.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "++ The paper is well-written with clearly motivated arguments and insights. The auxiliary learning task is also meaningful when we only seek to boost one task with another and aim at quick inference. \n\n++ Table 1 provides a comprehensive understanding and meticulous survey of the field. The authors offer an exhaustive overview of both Multi-Task Learning (MTL) and Auxiliary Learning (AL) methods. Authors have incorporated a wide range of references from multiple years, indicating a holistic survey of both seminal works and recent advancements. The inclusion of their method alongside existing techniques also provides clarity on its positioning within the broader research landscape. \n\n++ The proposed method is backbone- and task-agnostic that is applicable to multiple backbones and tasks.", "weaknesses": "-- I am not very familiar with auxiliary learning. However, I do think one baseline might be meaningful, which is to share a single backbone while projecting the gradient of the auxiliary task to the orthogonal direction of the main task on all (or selective) layers. This baseline also has no inference lag while exploiting the auxiliary objective signals. \n\n-- The authors use NAS to search for suitable architectures to optimize the main task's objective. However, distinct backbones are used for each task, and their weights can vary significantly. I'm uncertain why \"stitching\" two backbones with different weights and objectives is logical. Are there any supporting theories or references?\n\n-- The authors claim that their method achieves \"promising performance.\" However, based on Tables 3 and 4, it appears that the performance gain of the proposed method is only marginal. Considering the additional training costs in the NAS search and optimization, I am not sure whether the loss in training efficiency is worth it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "opp2nVn4RB", "reviewer": "Reviewer_CQEV", "summary": "This paper presents an architecture based take on auxiliary learning. They recognize that the asymmetry between the auxiliary and primary tasks can be exploited by learning architectures with constraints that favor transfer of information from the auxiliary to the primary, but in an indirect way so as to minimize the possibility / effect of negative transfer.", "questions": "1. Did you try further finetuning the final model on the primary task only after being done with the auxiliary-task based NAS ? This could result in extra performance boost", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea is novel and interesting.  I think the use of joint training, followed by the slow trimming of the aux-to-prim connections via L1 regularization is a clever way of more intimately introducing the auxiliary task  but remove it later to avoid needing it during inference.\n2. The paper is clearly written and easy to follow\n3. This has interesting implications for Auxiliary learning based architecture search -- since what was searched for in this paper were connections, there are expansions on this that can focus on other parts of the architecture space.", "weaknesses": "1. Method might be a bit too complex / cumbersome to be practically implemented widely -- especially given the size of the gains.\n2. Method also significantly increases memory / compute overhead at training time\n3. The experimental results have no error-bars. It's thus hard to judge the significance of the results\n\n### Nitpicks\n1. The introduction has *a lot* of italicized text, many of which I think are unnecessary and distracting.\n\n### Some relevant papers\nOn gradient conflict\n1. Dery, Lucio M., Yann Dauphin, and David Grangier. \"Auxiliary task update decomposition: The good, the bad and the neutral.\" arXiv preprint arXiv:2108.11346 (2021).\n2. Royer, Amelie, Tijmen Blankevoort, and Babak Ehteshami Bejnordi. \"Scalarization for Multi-Task and Multi-Domain Learning at Scale.\" arXiv preprint arXiv:2310.08910 (2023).\n\nOn NAS-like construction of auxiliary objectives\n1. Dery, Lucio M., et al. \"AANG: Automating Auxiliary Learning.\" arXiv preprint arXiv:2205.14082 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zU2HZqCAAq", "reviewer": "Reviewer_V6nv", "summary": "The paper introduces a new framework for auxiliary learning, in which the goal is to improve the performance on a task of interest (i.e., primary task) by utilizing auxiliary information. In particular, the proposed method aims to tackle auxiliary learning problems without introducing computational or parameter overhead during inference. To this end, the paper borrows inspiration from multi-task learning and neural architecture search to design an asymmetric network architectures, where the connections from primary-task network layers and auxiliary-task network-layers are directed (from primary to auxiliary), such that computations or parts of networks for auxiliary information can be removed during inference.", "questions": "Written in the weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method successfully tackles auxiliary learning without inducing extra computational overhead during inference, by utilizing NAS to design a network that has asymmetric connections directed from primary-task network parts to auxiliary-task network parts.\n\n- The proposed method is flexible in that it can be combined with different auxiliary learning methods\n\n- The paper is clearly written; easy to read and follow.", "weaknesses": "- Is there a need to initialize search space to include all bi-directional connections? why not start from networks with only primary-to-auxiliary connections right away?\n\n- Lack of ablation studies related with the question above: the performance change as the search space only contains primary-to-auxiliary connections.\n\n- Missing details: Are all auxiliary-to-primary connections are pruned at the end of training?\n\n- Missing details: What is the final architecture produced by NAS? How consistent is the final performance across different random seeds and trials?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xyqzjE5EaM", "reviewer": "Reviewer_Qzsf", "summary": "This paper introduces a learnable and flexible asymmetric network architecture designed for general-purpose auxiliary learning, where the auxiliary task plays a pivotal role in supporting the primary task's training process, and can be freely removed during the inference. As a result, the proposed method achieves a multi-task level performance while keeping a single-tasks level inference cost. The authors implement their design as adaptive layerwise feature fusion of multiple single-task branches, where the full network converges to an asymmetric architecture with only primary-to-auxiliary connections existed, enabling the removal of the auxiliary task during the inference. Two algorithms are developed to achieve this, where the more advanced one exploits a specifically designed NAS pruning to achieve an asymmetric architecture after convergence. The experiments are extensive across 6 tasks with 3 network backbone architectures, which sufficiently demonstrate the promising performance.", "questions": "The author claimed that they implement the tailored version of PCGrad and AdaShare specifically for the auxiliary task learning, i.e., PCGrad-Aux and AdaShare-Aux. What are the details of those auxiliary task learning variants?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1.\tThis paper tackles the general-purpose auxiliary learning towards a multi-task level performance and a single-tasks level inference cost. The proposed method can be applied to various tasks and network backbones mathematically and also validated experimentally.\n2.\tThe proposed method can also be freely combined with various multi-task or auxiliary task optimization methods listed in Table 1, which was also validated by the experiments.\n3.\tThe single-task level inference cost is assured through the resultant converged asymmetric network architecture. Furthermore, the training cost exhibits a linear increase when incorporating additional auxiliary tasks, which is enabled by the supernet architecture for NAS that only encompasses the connections between the primary task and each of auxiliary tasks. \n4.\tTable 1 present a very clear and comprehensive taxonomy about the position of the proposed method among the area of multi-task learning and the auxiliary task learning.\n5.\tThe experiments are extensive, validating the generalization on 6 diverse tasks within 3 datasets, and 3 network backbones including both CNNs and Transformers.", "weaknesses": "This paper is well written, and I do not see major weakness, but the clarification of the following minor issues would further improve the paper: \n1.\tI appreciate that the authors provide the full NAS objective in Eq. 10, but the details about how it is optimized need to be further elaborated. If I understand correctly, the model weight w and the architecture weight alpha should be updated iteratively?\n2.\tI suggest indicating the network backbone in the legends of Tables 3 and 4, as there are several tables in a similar shape that only differ from backbones.\n3.\tIt is suggested to replace the figures with vector images for a better resolution. The paper, in its current version, used a lot of v-spacing; it is also advised to remove them for better readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "okSymgdlDy", "reviewer": "Reviewer_fJaD", "summary": "This paper aims to harness the auxiliary tasks to enhance the performance of the primary task, while maintaining a single task inference cost for the primary task. The authors achieve this by designing an asymmetric network and further develops two algorithms: the first algorithm directly uses the asymmetric primary-to-auxiliary architecture, where the auxiliary tasks can be directly removed during the inference; the second algorithm initiates with an architecture with bi-directional connections, and subsequently exploits a tailored L1 constrained NAS optimization to prune all the auxiliary-to-primary connections, thereby enabling to remove the auxiliary task during inference. The proposed soft-parameter sharing architecture-based method can be integrate with existing optimization-based methods. The author validates their method with extensive experiments on 6 tasks with 3 CNN and transformer architectures.", "questions": "Please respond to those in the Weakness section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper formulates the auxiliary learning problem through a task-oriented adaptive feature fusion approach without the need of explicitly identifying the task similarity. Mathematically, such architecture-based method can be seamlessly integrated with a variety of multi-task/auxiliary optimization methods such as loss re-weighting and gradient manipulation. The paper is well written and easy to understand, with an extensive literature review in Table 1 clearly demonstrating the contribution of the proposed method.\n2. The evolving and asymmetric network design, coupled with a tailored NAS algorithm, ensures the converged network comprises only the primary-to-auxiliary connections, thereby guaranteeing a single-task inference cost for the learned architecture.\n3. Beyond the benefits in the single-task inference cost, the authors also show (in Sect. 4.2.4 and the supplementary) that the training complexity exhibits a linear scalability to multiple auxiliary tasks.\n4. The experiments are extensively performed on 6 highly diverse tasks with 3 base net architectures including CNN and transformers. The authors also checked the performance when the primary and the auxiliary tasks possess different architectures in the supplementary. The results of all those experiments are promising.", "weaknesses": "1. Is it possible to use Normalization and Activation operations other than BatchNorm and ReLU in Eqs. 13 and 14?\n2. In Fig. 3, should the cut-off dash line be between the 1x1 conv and the add operations?\n3. I suggest the authors to move the supplementary material into the Appendix of the main text for better readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)