Paper: Aux-NAS: Exploiting Auxiliary Labels with Negligibly Extra Inference Cost (ICLR-2024)

Reviewer comment (Experiments):
Is it possible to use Normalization and Activation operations other than BatchNorm and ReLU in Eqs. 13 and 14?

Author response:
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.

Reviewer comment (Presentation):
In Fig. 3, should the cut-off dash line be between the 1x1 conv and the add operations?

Author response:
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.

Reviewer comment (Presentation):
I suggest the authors to move the supplementary material into the Appendix of the main text for better readability.

Author response:
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.

Reviewer comment (Reproducibility):
I 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?

Author response:
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.

Reviewer comment (Presentation):
I 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.

Author response:
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.

Reviewer comment (Presentation):
It 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.

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

Reviewer comment (Experiments):
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?

Author response:
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.

| NYU v2, Primary: Normal | mean | median | rmse | within 11.25 | within 22.5 | within 30 |
| ------ | ------ | ------ | ------ | ------ | ------ | ------ | 
| Init. Aux-to-Prim 0.5 (paper) | **12.4553** | **10.3271** | **15.5775** | **53.8358** | **85.9042** | 94.4133
| Init. Aux-to-Prim 0 | 12.5091 | 10.4472 | 15.6169 | 53.3136 | 85.6822 | **94.4466** |

Reviewer comment (Reproducibility):
Missing details: Are all auxiliary-to-primary connections are pruned at the end of training?

Author response:
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.

Reviewer comment (Novelty):
Method might be a bit too complex / cumbersome to be practically implemented widely -- especially given the size of the gains.

Author response:
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.

Reviewer comment (Experiments):
Method also significantly increases memory / compute overhead at training time.

Author response:
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.

Reviewer comment (Evaluation):
The experimental results have no error-bars. It's thus hard to judge the significance of the results.

Author response:
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:

| NYU v2, Primary: Normal | mean | median | rmse | within 11.25 | within 22.5 | within 30 |
| ------ | ------ | ------ | ------ | ------ | ------ | ------ | 
| mean | 12.5656 | 10.4620 | 15.6895 | 53.3011 | 85.6956 | 94.2751 | 
| std | 0.0479 | 0.0323 | 0.0702 | 0.1341 | 0.2018 | 0.1475 | 

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.

Reviewer comment (Writing):
The introduction has *a lot* of italicized text, many of which I think are unnecessary and distracting.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

| NYU v2, Primary: Normal | mean | median | rmse | within 11.25 | within 22.5 | within 30 |
| ------ | ------ | ------ | ------ | ------ | ------ | ------ | 
| Org results (paper) | 12.4553 | 10.3271 | 15.5775 | 53.8358 | **85.9042** | 94.4133
| Further finetuned results | **12.3787** | **10.3074** | **15.4840** | **53.9111** | 85.8241 | **94.5066** |

Reviewer comment (Experiments):
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.

Author response:
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:

| NYU v2, Primary: Seg | Aux-Head | Aux-Head + PCGrad-Aux | Aux-Head + GCS | Aux-FG-NAS | Aux-FG-NAS + PCGrad-Aux | Aux-FG-NAS + GCS |
| ------ | ------ | ------ | ------ | ------ | ------ | ------ | 
| mIoU | 34.7 | 35.2 | 35.3 | 36.0 | 36.2 | 36.3 |
| PAcc | 65.4 | 65.7 | 65.9 | 66.1 | 66.3 | 66.5 |

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.

Reviewer comment (Theory):
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?

Author response:
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].