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

Thank you for your constructive comments. We have thoroughly addressed each of the raised questions with detailed responses. We hope those well resolved your concerns. Please kindly let us know if any further discussion or clarification is needed.

Q1. Is it possible to use Normalization and Activation operations other than BatchNorm and ReLU in Eqs. 13 and 14?

It is true that we have the flexibility to leverage different Normalization and Activation operations other than BatchNorm and ReLU. As we are inserting new fusion operations to connect multiple well-trained single-task networks, the key design principle behind our fusion operations in Eqs. 13 and 14 is to ensure that their insertion does not significantly alter the original well-trained single-task networks. Therefore, any fusion operations, along with any associated Normalizations and Activations, are applicable as long as they adhere to this fundamental design principle.

Q2. In Fig. 3, should the cut-off dash line be between the 1x1 conv and the add operations?

During the NAS optimization, we cut-off the auxiliary-to-primary connections by regularizing $alpha^P$'s in Eq. 13 to 0. Therefore, in Fig. 3, we illustrated the cut-off dash line between the auxiliary input and the concatenation operation. Having said that, it is equivalent to draw the cut-off dash line be between the 1x1 conv and the add operations as suggested, because the output of 1x1 conv is also 0 when the input is 0.

Q3. I suggest the authors to move the supplementary material into the Appendix of the main text for better readability.

We have moved the supplementary material, originally in a separate file, to the Appendix of the main text. Additionally, we have included the appropriate references in the main text. Please kindly refer to our updated manuscript.

Thank you for your constructive comments. We have thoroughly addressed each of the raised questions with detailed responses. We hope those well resolved your concerns. Please kindly let us know if any further discussion or clarification is needed.

Q1. Details of NAS optimization.

It is true that the model weight $w$ and the architecture weight $\alpha$ are updated iteratively during our training of Aux-FG-NAS. Compared to training a fixed neural networks, our NAS optimization requires only one additional forward-backward pass to update the architecture weights in each iteration. This process is straightforward to implement, as indicated in our pseudo code provided above. We will release the complete training and evaluation codes upon acceptance.

Q2. Indicating the network backbone in the legends of Tables 3 and 4.

Indeed adding the network backbone into the table legends improves the readability. We followed this suggestion and updated our tables accordingly, please kindly refer to our updated manuscript.

Q3. Use vector images, remove the v-spacing.

We thank this suggestion. We have utilized vector images and minimized the v-spacings, please kindly refer to our updated manuscript.

The raw data of converged primary-to-auxiliary architecture weights for Question 4. The maximum value of all the auxiliary-to-primary architecture weights is 0 at the precision of 10^-2.

Q4. What is the final architecture produced by NAS? How consistent is the final performance across different random seeds and trials?

It is indeed interesting to see the final convergence of the learned architecture.

As discussed in the previous question, the auxiliary-to-primary connections are pruned out with the maximum architecture weights less than 0.01.

Regarding the primary-to-auxiliary connections, as we do not impose any regularization on their architecture weights (and those connections will be removed to maintain a single-task inference regardless how they converge), those connections typically converge to soft values between 0 and 1. Therefore, it is difficult to draw them clearly for illustration, given the amount of those connections and many replicate runs. Consequently, we instead gather their statistics of min, max, mean, median of 10 replicate runs with random seeds, and show the standard deviation of those statistics.

Specifically, we conducted the same experiments as those in Table 5 using ViTBase on NYUv2, with random seeds for 10 replicate runs. The mean and standard deviation of min, max, mean, median statistics of the primary-to-auxiliary architecture weights are shown below:

The above table demonstrates negligible standard deviation, which further produces subtle differences in the final performances, as shown in Question 3 of Reviewer CQEV. We believe that the subtle standard deviation in both the converged architectures and the final performances is ensured by: i) for the model weights, our network backbones are initialized by the converged single-task networks, and ii) for the architectures, we only search the cross-task connections, rather than altering the network backbones.

We include the raw primary-to-auxiliary architecture statistics in the following thread. We note that all the results presented here in the rebuttal and in our main text are obtained without specifying any seeds. We will release our training codes also without the specification of seeds.

We thank this discussion again, and please kindly let us know if any further discussion or clarification is needed.

Thank you for your constructive comments. We have thoroughly addressed each of the raised questions with detailed responses. We hope those well resolved your concerns. Please kindly let us know if any further discussion or clarification is needed.

Q1. The rationale behind using bi-directional NAS initialization. Ablation of the initialization with only primary-to-auxiliary connections.

The rationale behind initializing all the bi-directional connections is to exploit both features and gradients from the auxiliary tasks. During the training, the auxiliary-to-primary connections, which supply auxiliary features, are gradually pruned out during the NAS training to ensure a single-task inference.

We expect that the performance of bi-directional initialization suppresses that of only initializing primary-to-auxiliary connections. This is because the L1 regularization applied to auxiliary-to-primary connections leads to difficulty in re-establishing them if we do not initialize them. As a consequence, the model capacity, i.e., the valid search space, in the primary-to-auxiliary initialization only encompasses a subset of that in the bi-directional initialization.

To validate the above discussion, we conducted experiments by initializing only primary-to-auxiliary connections as instructed, mirroring the setup in Table 5 using ViTBase on NYUv2. As anticipated, the ablation results align with our expectations, as detailed below.

Q2. Ablation of the search space that only contains primary-to-auxiliary connections.

In addition to the ablation in the previous question where the search space contains bi-directional connections but auxiliary-to-primary links were initialized to 0, we also perform another ablation following your guidance, where the search space only contains primary-to-auxiliary connections. The results are listed below, which exhibit inferior performance w.r.t. our original Aux-FG-NAS. Also notably, the single-side search space setting here performs very much similarly to Q1 where auxiliary-to-primary links were initialized to 0. This outcome comes from similar reasons as discussed in the previous question, i.e., the model capacity, and the search space, of the single-directional setting are just a subset of our original bi-directional Aux-FG-NAS.

Q3. Are all auxiliary-to-primary connections pruned at the end of training?

We have collected the statistics of the auxiliary-to-primary architecture weights for all of our experiments, revealing that the maximum value of all the auxiliary-to-primary architecture weights is 0 at the precision of 10^-2.

Notably, all the results reported in our manuscript are the final single-task inference performance, where we manually set all the auxiliary-to-primary architecture weights to 0 before conducting the evaluation. This guarantees our objective of ensuring the single-task inference cost.

Thanks for the clarification. My concerns have been addressed if the discussions are included in the revised paper. Also, please compare the weights of aux-to-prim connections with those of prim-to-prim to show better contrast.

Dear Reviewer,

We sincerely apologize for the typos in our initial table entry for Q4 and its corresponding raw data table. The correct term should be Prim-to-Aux Arch Weight instead of Aux-to-Prim Arch Weight. We have rectified these errors. Our original text response of Q4 does not have those typos.

We gathered the maximum value of all the Aux-to-Prim architecture weights, which is 0 at the precision of 10^-2, please also refer to a more detailed discussion about the Aux-to-Prim architecture weights in Q3.

We re-evaluated our Aux-FG-NAS model of Table 5, without manually removing the Aux-to-Prim connections, the differences are subtle as shown below:

We apologize again for the confusion and appreciate the opportunity for further discussion. Please kindly let us know if any additional clarification is needed.

Following-up with our previous response, we would like to clarify that our search space only includes the cross-task prim-to-aux and aux-to-prim connections, we leave the architectures of the single-task backbones intact (equivalent to fix the prim-to-prim and aux-to-aux architecture weights to 1). This ensures the adaptability of our methods to diverse single-task backbones.

As acknowledge by Reviewer CQEV, our Aux-FG-NAS can be further extended to search for a better backbone, which can optionally exploits the priors of specific input single-task backbones and/or the characteristic of input tasks. But this falls outside the scope of our current focus, which is to design a general-purpose auxiliary learning method that adapts to diverse tasks, single-task backbones, datasets, and integrates with various Multi-Task Optimization methods.

We hope this well addresses your concerns, please let us know if further discussion is needed and we are happy to engage into that.

Q4. The unnecessary italicized text in the introduction.

We thank for pointing this out. We have removed unnecessary italicized text, please see our updated manuscript.

Q5. Relevant papers [4,5,6].

We appreciate for guiding us to those relevant papers, this certainly makes our taxonomy in Table 1 more comprehensive and self-contained. We have cited [4,6] in MTL & AL Methods - Optimization-based - For AL, and [5] in MTL & AL Methods - Optimization-based - For MTL, in Table 1. We also cite and discuss them in our related work section. Thanks again for pointing them out!

Q6. Further finetune the final model on the primary task after NAS converges.

As instructed, we continue to finetune the Aux-FG-NAS weights of Table 5 (i.e., the ViTBase NYUv2 experiment). We use the same learning rate and the same optimizer as those used to further train the converged single-task initializations. The results are presented below, showing a slight (but not significant) improvement over the original results. We thank this advice again.

References

[1] Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In ICLR, 2019.

[2] Sirui Xie, Hehui Zheng, Chunxiao Liu, Liang Lin. SNAS: Stochastic Neural Architecture Search. In ICLR, 2019.

[3] Guohao Li, Guocheng Qian, Itzel C. Delgadillo, Matthias Müller, Ali Thabet, Bernard Ghanem. SGAS: Sequential Greedy Architecture Search. In CVPR, 2020.

[4] Lucio M. Dery, Yann Dauphin, David Grangier. Auxiliary task update decomposition: the good, the bad and the neutral. In ICLR, 2021.

[5] Amelie Royer, Tijmen Blankevoort, Babak Ehteshami Bejnordi. Scalarization for Multi-Task and Multi-Domain Learning at Scale. In NeurIPS, 2023.

[6] Lucio M. Dery, Paul Michel, Mikhail Khodak, Graham Neubig, Ameet Talwalkar. AANG: Automating Auxiliary Learning. In ICLR, 2023.

Thank you for your constructive comments. We have thoroughly addressed each of the raised questions with detailed responses. We hope those well resolved your concerns. Please kindly let us know if any further discussion or clarification is needed.

Q1. Method might be a bit too complex/cumbersome to be practically implemented widely -- especially given the size of the gains.

We proposed two general-purpose auxiliary learning methods. The first one, Aux-G, is very simple, involving only the direct integration of the primary and auxiliary networks through primary-to-auxiliary connections. The combined network weights can then be trained directly using gradients, akin to training most neural networks with fixed architecture.

The second method, Aux-FG-NAS, leverages Neural Architecture Search for further improved performance. Notably, the chosen NAS optimization method, DARTS, is also easy to implement, requiring only one additional forward-backward pass in each iteration to train the architecture. As shown in our pseudo codes provided above, such an additional step can be implemented with just 7 lines of codes in PyTorch. We will release the complete training and evaluation codes upon acceptance.

Regarding the assessment of gains which can be subjective, it is important to acknowledge the challenge of designing a general-purpose auxiliary learning framework that seamlessly adapts to diverse tasks, network backbones, datasets, and integrates with various Multi-Task Optimization methods. Notably, in many of our experiments, the improvement achieved by our method versus the previous SOTA surpasses that of the previous SOTA versus the simple Aux-Head baseline. We also appreciate your acknowledgment in Strength 3 regarding the potential of our method to further exploit the task prior in network backbones for additional enhancements.

Q2. Method also significantly increases memory / compute overhead at training time.

The training memory / compute complexity for both of our methods scales to multiple auxiliary tasks linearly, which aligns with the Soft-Parameter Sharing Multi-Task Architectures (SPS-MTA) and most of Multi-Task Optimizations (MTO) that employ gradient manipulation (Note that MTO with gradient manipulation has linear instead of constant scalability because their training procedure requires to retain K copies of backward graph for storing and manipulating the gradients from K tasks).

Specifically, given K auxiliary tasks, denote model weights as $w$ and architecture weights as $\alpha$, our Aux-G method has the same training complexity as that of SPS-MTA methods and most MTO methods.

While for our Aux-FG-NAS, the training complexity is $O((K+1)|w| + K|\alpha|)$. This linear scalability is achieved by the efficient NAS optimization of DARTs [1], as detailed in Eq. 8 of [1], and is confirmed by our pseudo codes provided in the above response.

Further details of our Aux-FG-NAS complexity can be found in Sect. 4.2.4 and Appendix A of our updated manuscript (Appendix A was originally in a separate supplementary file at the time of the initial submission).

It's noteworthy that our method is primarily designed to address practical scenarios where the inference complexity takes precedence over the training complexity. Such scenarios are prevalent in customer-facing web applications where the model is only trained for a few times, then deployed online and subjected to million to billion inferences. While for the scenarios where the training resource is also constrained, we can at least use the Aux-G method with the same training complexity as SPS-MTA and most MTO methods, so as to ensure the improvement leveraging auxiliary tasks.

Q3. Error bar in the results.

Following your guidance, we perform 10 replicate runs with random seeds of the ViTBase experiments on NYU v2, as those in Table 5 of our manuscript. We did not observe a substantial standard deviation in the results:

We believe that the subtle standard deviation in the replicate runs is ensured by: i) for the model weights, our network backbones are initialized by the converged single-task networks, and ii) for the architectures, we only search the cross-task connections, without altering the network backbones. Notably, all the results presented in our manuscripts and here in the rebuttal are obtained with random seeds. Additionally, our training codes will also be released without any specification of seeds. We put the raw results of the 10 replicate runs in the following thread. We thank this suggestion again and hope this addressed your concern.

Thank you for your constructive comments. We have thoroughly addressed each of the raised questions with detailed responses. We hope those well resolved your concerns. Please kindly let us know if any further discussion or clarification is needed.

Q1. The baseline of sharing a single backbone, while projecting the gradient of the auxiliary task according to the main task on all (or selective) layers.

The suggested method is indeed a stronger baseline by using Multi-Task Optimization techniques. In fact, the suggested method is exactly what we refer to as PCGrad-Aux in Table 2. Specifically, we extend the original PCGrad [1] to auxiliary learning in PCGrad-Aux exactly by the way as suggested, i.e., fixing the primary gradient and projecting all auxiliary gradients w.r.t. the primary gradient, instead of randomly choosing a task for projection in the original PCGrad [1].

Moreover, we also validate our method using another Multi-Task Optimization method specifically designed for auxiliary learning named GCS [2]. In GCS, the auxiliary gradients are applied only if they exhibit a non-negative cosine similarity with the primary gradients; otherwise, GCS discards the auxiliary gradients.

The experimental results w.r.t. the auxiliary learning-based Multi-Task Optimization PCGrad-Aux and GCS from Table 2 are:

Our results, as presented above and in Table 2, demonstrate that i）Our vanilla method surpasses those Multi-Task Optimization methods, i.e., Aux-FG-NAS vs. Aux-Head + PCGrad-Aux & Aux-Head + GCS. ii) Importantly, our method is seamlessly incorporable with these Multi-Task Optimization methods, yielding further improvement in results, as demonstrated by Aux-FG-NAS vs. Aux-FG-NAS + PCGrad-Aux & Aux-FG-NAS + GCS.

We note that the second point was also acknowledged by Reviewers V6nv, Qzsf, fJaD. Therefore, we focus on comparing our method with other Multi-Task Architecture methods in the subsequent experiments to avoid potential distractions. We thank this discussion again and hope this resolved your concern.

Q2. Why "stitching" two backbones with different weights and objectives is logical? Are there any supporting theories or references?

Our method stitches the primary and the auxiliary networks with different weights and objectives, where the auxiliary networks can be regarded as a "loss function" to provide gradients, as well as a feature generator to offer features from another view (though those auxiliary features are pruned out after convergence). Moreover, during the training, we refrain from imposing any regularization that enforces similarity between the weights of the primary and the auxiliary networks. This deliberate choice allows us to stitch the primary and auxiliary networks with different weights and objectives, enabling the auxiliary networks to contribute unique "losses" and "features" distinct from the primary network.

We also note that the practice of stitching two backbones with different weights and objectives has a longstanding history in architecture-based MTL [3,4,5,6], where one of the seminal works happened to be termed as cross-stitching networks [3].

Q3. The marginal improvement of Table 3 and Table 4 given the loss in training efficiency.

While the assessment of the improvement can be subjective, we would like to highlight that our primary goal is to establish a general-purpose auxiliary-learning framework. This framework is designed to be seamlessly incorporable with diverse tasks, network backbones, datasets, and various Multi-Task Optimization methods, in a plug-and-play manner. In pursuit of this generality, we specifically concentrate on the cross-task connections while leaving the network backbone intact. As acknowledged by Reviewer CQEV (Strength 3), we have the potential to leverage the task priors, optionally by searching the task backbones, to further enhance the performance.

Moreover, it is worth noting that in Tables 3 and 4, the improvement of the previous SOTA over the simple Aux-Head baseline is notably less significant than that of our method over the previous SOTA. This observation also underscores the inherent challenges in designing a general-purpose auxiliary learning framework.

As for the training efficiency, please also kindly refer to the above pseudo codes and Question 2 of Reviewer CQEV.