Transforming Vision Transformer: Towards Efficient Multi-Task Asynchronous Learner

Thank you for your detailed review, we will try to address any concerns encountered in the paper.

Q1: Is the router truly necessary?

It is an insightful question regarding the role of the router in our framework. Initially, we followed to the conventional MoE framework which includes a router. This comment prompted us to undertake a investigation and conduct additional experiments. From our analysis and experimental findings, we draw the following conclusions:

1. When using MoE for multi-task learning, the router plays a crucial role in adaptively weighting different experts for each token. During the training process, for the same input token, calculating different weights for different experts allows for a nuanced and fine-grained optimization of each expert. This results in a well-specialized ensemble of experts, enhancing the model's overall performance and expertise in handling diverse tasks.
2. Our experiments with a Router Fade strategy demonstrated a consistent improvement, with a mean accuracy increase of 0.34. This enhancement is non-trivial and underscores the utility of the router in achieving better performance across tasks.

Based on these observations, we conclude that the router is indeed necessary for optimizing performance in our multi-task MoE framework.

Q2: So how many clusters we should have?

In response to your query about the number of clusters and their impact:

1. Regarding the optimal number of clusters, our approach balances between too many and too few clusters. Having too many clusters would result in each expert having too few channels, limiting their individual expressiveness. Conversely, too few clusters make each expert's low-rank characteristics less pronounced, which dilutes the distinct functionalities of each expert. We recommend that the number of experts similar to the number of attention heads.
2. Our hyperparameter ablation studies, specifically looking at different numbers of clusters, showed optimal performance of 90.27 with cluster size of 16.

These findings underscore the importance of carefully selecting the number of clusters to maintain the effectiveness of each expert.

Q2: What if we don't cluster them?

The MOELoRA method in Table 1 is the scenario where MoE without clustering. In this setup, experts are directly employed in the FFN without forming clusters, resulting in a Mean Accuracy of 88.04. Moreover, we perform additional experiments in the table above without clustering (the clusters is set to 1).

Q3: Applying QR loss on other MTL models?

Indeed, the QR loss can be integrated into any MTL model to potentially enhance its performance. To illustrate this, we conducted experiments using the MOELoRA framework, which is detailed in a recent study. Our findings indicate that incorporating QR loss directly into the MOELoRA model results in a performance improvement. This demonstrates the versatility and effectiveness of QR loss in boosting the capabilities of MTL models, making it a valuable addition to existing and future MTL architecture design methods.

Q4: Fair comparison of the MTO methods and the MTL architectures.

In response to your valuable suggestion, we incorporate the content from Table 3 and extend experiment on QR loss within the same model designs, into the Table A of 'Rebuttal.pdf'. This allow for direct and convenient comparisons under the respective categories of MTO and MTL architectures.

Q5: Extension to pixel-to-pixel prediction tasks.

To demonstrate the extensibility of our method, we conducted experiments on the NYUv2 dataset using ViT-B due to the time constraints.

Specifically, we made the following attempts, integrating our method with the existing open-source method TaskPrompter. For MTL structure, we clustered and divided the FFN in the backbone and used the Router fade strategy, keeping the rest of the architecture unchanged. For MTO, We directly applied QR to the semantic segmentation task head for experiments. The results are as follows:

Specifically, our method significantly enhanced the performance on semantic segmentation and depth estimation tasks, while achieving comparable results on surface normal estimation and object boundary detection tasks. Overall, it resulted in an average relative improvement of 1.57%, confirming its effectiveness for pixel-level prediction tasks.

Q6: The inference time.

In addition to the ablation study in Table 4, we add a column to the Table A of 'Rebuttal.pdf', to highlight the efficiency of our unified model during inference.

Thanks a lot for your insightful comments and suggestions. Our responses are summarized as below:

W1. On grouping strategies.

  A1. To the best of our knowledge, our work is the first one that employs the idea of grouping similar weights to establish the LoRA experts, and there lack grouping strategies that can be directly used for comparisons. In our opinion, this design incurs the following two advantages:

1). The experts built by grouping have more specialized functions as described in [28], as well as low-rank properties, thus being natural for updating through LoRA.

2). LoRA learns low-rank updates [32] for each expert, restraining the rank of the experts' weights from increasing, which in turn promotes the specialization of the experts.

As aforementioned, there lack grouping strategies that can be directly used for comparison. Inspired by existing works, we implement the following two methods for comparison: 1) Co-activation similar to [28] that clusters weights based on the activations for each channels; 2) Gradient-cluster inspired by SPT [47] that clusters weights based on the cumulative gradients.

As displayed below, our method reaches the best results, clearly demonstrating its effectiveness.  

W2: Generalizability to the heterogeneous task settings.

Thanks for your suggestion. We mainly follow the work [6], and evaluate on FGVC for fair comparisons. However, we do agree that FGVC adopts a homogeneous task setting. The other reviewer also mentioned this issue, and suggests more results on other heterogeneous benchmarks (e.g. NYUv2). Due to the time limitation, we conduct experiments on NYUv2, as it contains four substantially heterogeneous tasks including Semantic Segmentation, Monocular Depth Estimation, Surface Normal Estimation, and Object Boundary Detection. For comparison, as the released code of the suggested MLoRE approach has some unfixed issues and fails to reproduce the reported results after contacting the authors, we finally select another open-source SOTA method TaskPrompter [Ref.1] for comparison. In implementation our method, we follow the framework of the existing open-source SOTA method TaskPrompter [Ref.1]. For the architecture, we cluster and divide the FFN into experts and apply router fade. For the multi-task optimization, we apply QR to the classification branch of the Semantic Segmentation task head, and omit it in other tasks.

The comparison results are summarized below:    

As displayed, our method promotes TaskPrompter by 1.57% on average, showing the effectiveness of our method for the heterogeneous task settings.

[Ref.1] Ye, Hanrong, et al. "Taskprompter: Spatial-channel multi-task prompting for dense scene understanding". In ICLR. 2023.

W3&W4: Applicability of Quality Retaining (QR), and comparison to previous multi-task optimization (MTO) methods.

Intuitively, QR can be integrated to any task (e.g. detection and segmentation) that contains the classification branch, thus being applicable to typical computer vision tasks including the segmentation. Due to the time limitation, we evaluate our overall method on the segmentation task on NYUv2, by comparing to the SOTA approach TaskPrompter with the ViT-Base backbone. As shown in the second column of the table above, the results clearly displays the applicability and effectiveness of our method on the segmentation task beyond classification.

We have also add comparison results between QR with the representative MTO methods including the loss based approaches GLS [24] and ATML [27], and the gradient based approaches Nash-MTL [22] and Aligned-MTL[23]. For fair comparisons, we use the same architecture design for all compared methods. The results in Table A in "Rebuttal.pdf" clearly show the effectiveness of QR.

We will add the discussion and results in the final version, and further study the extension of QR to more tasks in future.

W5: Comparison to training from scratch.  

As shown in the table below, training a ViT from scratch usually yields extremely poor results, due to overfitting on the limited training data in downstream tasks. In contrast, our method reparametrizes a pre-trained transformer into an MoE-like structure, providing a meaningful initialization, thus leading to better performance.  

*indicates the best results after carefully tuning the hyperparameters.    

W6: On unified model.  

Most existing approaches including MOELoRA and MLoRE reparameterize the LoRA experts based on the task-driven router's weights, by preserving a specialized model or a bypass branch for each task during inference, thus failing to generate a unified model for different tasks. In contrast, by using the reparameterization and the proposed router fade strategy, our method generates a single model without employing additionally task-relevant models or branches during inference, thus being dubbed as a unified model across tasks.

Thanks for your insightful suggestions. Our point-to-point responses are summarized as below.

W1. On ablation study.

Firstly, the setting alpha=0 from the very beginning is not equivalent to the vanilla LoRA without clustering. We refer to this as Cluster+Finer LoRA, and provide extra ablation results by adding the router and the router fade strategy, summarized as below:

The results clearly display the effectiveness of the router.

W2. On the backbone.

The ViT-B/16 backbone [1] used in our experiments is supervised pre-trained on ImageNet-21K, which is widely used in previous works, and is adopted in our paper for comparisons. However, we do agree that exploring other backbones is necessary. Therefore, we conducted additional experiments using the backbone pre-trained by MAE [2] on the Multi-task FGVC benchmark. The results are summarized below:

The experimental results indicate that our method is effective compared with the SOTA MTL methods by using the self-supervised pre-trained backbones, but fails to surpass that of the supervised ones.

W3. On more complex tasks.

We acknowledge the importance of evaluating our method on tasks with greater domain shifts. Several reviewers have raised this point, and we agree that it is crucial to validate our approach on a broader range of benchmarks, including NYUv2, WILDS, PASCAL-Context, and Taskonomy benchmark. Since time constraints limit the comprehensive evaluation on these benchmarks, we perform evaluation the NYUv2 benchmarks to further validate our method's performance on tasks with significant domain differences. The NYUv2 dataset includes tasks such as Semantic Segmentation, Monocular Depth Estimation, Surface Normal Estimation, and Object Boundary Detection. These tasks exhibit substantial domain differences, making NYUv2 a suitable benchmark for evaluating our method's robustness to domain shifts.

The results indicate that our method performs well across tasks with significant domain differences, demonstrating its robustness and generalization capabilities.

W4. On missing comparison of PEFT methods.

Table 2 is designed not only to validate the extensibility of our method across more benchmarks but also to compare our approach with existing single-task efficient fine-tuning paradigms. This highlights the effectiveness of our multi-task PEFT method. Therefore, we additionally included comparisons with public results of existing PEFT methods in Table 2.

BOFT, GPS [Ref.1] and LoSA [Ref.2] are published after or very close to the submission date of NeurIPS, which are missed in our submission. But we do agree it will make the comparison more compherehensive to add them for comparison. However, unlike most existing works, BOFT employs a larger DINOv2-Large backbone instead of the standard supervised ViT-B/16, making it unfair to directly compare. Due to the time limitation, we mainly add GPS and LoSA for fair comparisons as below. | Method | Reference | Patch Camelyon | EuroSAT | Resisc45 | Retinopathy | Mean Acc | Params. (M) |

[Ref. 1] Zhang, Zhi, et al. "Gradient-based Parameter Selection for Efficient Fine-Tuning". In CVPR, 2024.

[Ref. 2] Otniel-Bogdan Mercea, et al. "Time- memory- and parameter-efficient visual adaptation". In CVPR, 2024.

The results demonstrate the effectiveness of our method compared to existing single-task efficient fine-tuning methods. Our approach not only achieves the state-of-the-art performance but also avoids selecting from multiple specialized models trained separately, thereby improving the inference speed.

Q1-Q3. On presentations in L64-69, Figure 2, L129.

Thanks for your suggestions on the presentations. We will carefully improve the presentation and description in the final version.

Q4&Q6&Q7. On the notation r, PoolFT, Eqs. (14) and (15).

Thanks for pointing out these typos. PoolFT should be Union FT that fine-tunes on the joint feature space. i is indeed the iteration step and L_t shoud be $L _ {CE,t}$. We will correct them in the revision.

Q5. On reduction of inference time.

Indeed, the LoRA experts are merged back into the pre-trained model. As shown in Table 4, the inference time is reduced from 13.87ms to 7.15ms.

[1]: ViT-B/16 supervised pre-trained on ImageNet-21K. https://storage.googleapis.com/vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz

[2]: ViT-B/16 Self-supervised MAE. https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth