GO4Align: Group Optimization for Multi-Task Alignment

We sincerely thank # Reviewer sLBS for their insightful comments. The following mainly addresses the concerns and answers questions.

---

1. Clustering methods.

In our main experiments, we employed standard K-means for instantiation. K-means is a widely used clustering approach that worked well in our experiments, so we did not explore this part further (as it was not the main focus).

We appreciate your suggestion and have conducted additional experiments to evaluate the impact of using alternative clustering algorithms. We thank you for bringing the reference [a]; however, due to the absence of open-source code for [a], we explored another SDP-based clustering method [b].

Specifically, we conduct the experiments on NYUv2 by substituting the K-means clustering in our method with SDP-based clustering [b] and spectral clustering [c]. As demonstrated in the table below, these alternative clustering methods also outperform state-of-the-art approaches (FAMO, -4.10%), particularly by enhancing the performance of each task over STL. For training efficiency, we implemented K-means using the >kmeans_pytorch package, which offers GPU support.

Interestingly, our experiments show that the K-means clustering algorithm we deployed outperforms both the spectral and SDP-based clustering methods. This is because the hyperparameters for the latter algorithms have yet to be thoroughly investigated. We have presented promising initial results and will include this ablation study and related discussions in the main manuscript.

---

2. Choice of the group indicators.

It is also plausible to use the gradient information as clustering indicators. As shown in Table 4, we tried to replace the proposed risk-guided group indicator with gradient-guided task weights(MGDA and NashMTL). Our work outperforms these alternatives. The main reason could be that our group indicator yields better representations of learning information by capturing the differences in the per-task risk scale and exploring the learning dynamics over time.

Moreover, from the perspective of training efficiency, the task-specific gradients need to back-propagate the shared architecture for $M$ times, where $M$ is the number of tasks, which increases computational cost linearly with the number of tasks.

---

3. Comparable time with linear scalarization and scale-up number of tasks.

We follow representative MTO works (RLW[39] and FAMO[14]) in choosing linear scalarization (LS) as the relative baseline for training time. This is a good choice because LS is a common MTL baseline, where each task has equal weights without extra loss-oriented or gradient-oriented techniques.

As shown in Fig. 4, when the number of tasks scales up from 2 to 40, the advantages of reducing computation cost with our method become increasingly significant compared to other gradient-oriented methods, e.g., NashMTL=2.07 versus Ours=1.01 with 2 tasks, NashMTL1=2.49 versus Ours=1.01 with 40 tasks. We will add this discussion in Line 299.

---

4. Questions.

Convergence difference in Figure 2: We evaluate the convergence difference by the standard deviation of the task-specific epoch numbers to reach convergence. We will polish this description in Line 90.

Definition of $\Delta m$: $\Delta m$ is the average per-task performance drop relative to STL as representative MTO works (NashMTL and FAMO). We will polish this in Line 268.

---

Reference:

[a] Awasthi, P., Bandeira, A. S., Charikar, M., Krishnaswamy, R., Villar, S., & Ward, R. (2015, January). Relax, no need to round: Integrality of clustering formulations. In Proceedings of the 2015 Conference on Innovations in Theoretical Computer Science (pp. 191-200).

[b] Tepper, M., Sengupta, A. M., & Chklovskii, D. (2017). The surprising secret identity of the semidefinite relaxation of k-means: manifold learning. arXiv preprint arXiv:1706.06028.

[c] Damle, A., Minden, V., & Ying, L. (2019). Simple, direct and efficient multi-way spectral clustering. Information and Inference: A Journal of the IMA, 8(1), 181-203.

---

Thank you for your time and efforts. We hope our experimental results and clarifications have addressed your concerns. Please don’t hesitate to reach out if you have further questions or need more information.

We sincerely thank # Reviewer 9PzS for their insightful comments. The following mainly addresses the concerns and answers questions.

---

1. The phenomenon of task imbalance and application scenarios.

Thanks for your kind suggestion. This phenomenon of task imbalance describes that some tasks are severely under-optimized during multi-task training. This can be observed in Fig. 2 and Table 1. (i) In the left subplot of Fig. 2, task “normal” is under-optimized, leading to convergence earlier than other tasks. (ii) In Table 1, most baselines achieve comparable performance to STL on the segmentation and depth estimation tasks but significantly sacrifice the performance on the surface normal estimation task. We note that on NYUv2, our work is the only method that improves each task’s performance relative to the corresponding STL results, especially for the surface normal estimation task. This demonstrates that our method performs better in alleviating task imbalance.

Meanwhile, we would like to specify one typical application scenario of GO4Align. It lies in vision-based autonomous driving, such as Autopilot, which needs to simultaneously handle several tasks (instance segmentation and depth estimation). Significant differences in the improvements across tasks could exist due to diverse scales, different task difficulties, or asynchronous learning dynamics. In this case, the proposed method can be deployed to balance their learning process, particularly without extra training costs.

---

2. What is the optimization method for $\omega$ and $\mathcal{G}$?

We use standard K-means with “Euclidean” distance (import kmeans_pytorch) on $\gamma$ to optimize the $\omega$ and $\mathcal{G}$.

---

3. Why is it possible to invert the matrix $\mathcal{G}$?

Special thanks for pointing out this. The assignment matrix $\mathcal{G}$ is a full-row rank. For simplicity, we formulate its generalized inverse, especially one-sided right inverse, $\mathcal{G}_{R}^{-1}=\mathcal{G}^{\top}(\mathcal{G}\mathcal{G}^{\top})^{-1}.$ We will add a detailed description in Line 172.

---

4. Imbalance phenomenon in datasets.

The task-imbalance phenomenon in the MTO literature [1, 27] mostly refers to the imbalanced optimization process rather than data distributions in the task space. We will clarify this in the paper, and specifically, we will add the description of task imbalance in Line 257.

---

Thank you for your time and efforts. We hope this has addressed your concerns and answered your questions. Please don’t hesitate to reach out if you have further questions or need more information.

We sincerely thank # Reviewer AXsM for their insightful comments. The following addresses their concerns and provides answers to their questions.

---

1. GO4Align versus FAMO.

We thank the reviewer for the comment. GO4Align and FAMO are both strong candidates for handling task imbalance in the MTO literature. However, on NYUv2, our work is the only method that improves each task’s performance relative to the corresponding STL results, especially for the challenging surface normal estimation task. This demonstrates that our method performs better in alleviating task imbalance.

---

2. Questions.

Performance change with only Eq.4: Eq.4 is a bi-level optimization process, where lower-level optimization takes the group indicators (Eq.5 and Eq.6) as inputs to update the assignment matrix and group weights. Thus, the optimization process can not be performed individually without specifying the group indicator. More details can be found in Code Line 6-10 of Algorithm 1.

Computational cost: Each method’s training time is computed relative to a common MTL baseline (LS, linear scalarization), the same as used in representative MTO works (RLW[39] and FAMO[14]). LS weights each task equally without extra loss-oriented or gradient-oriented techniques. Moreover, our method has the advantage of computational efficiency over STL. STL takes o(M) the space or time cost, where M is the number of tasks. In contrast, our work uses o(1) space and time due to the shared architecture and loss-oriented mechanism.

---

Thanks again for your time and efforts. We hope this has answered your questions. If you have any other questions, we are happy to discuss them and provide further clarification.

We sincerely thank # Reviewer FqZ5 for their insightful comments. The following addresses their concerns and provides answers to their questions.

---

1. Motivations.

Thank you for acknowledging the two technical contributions to our paper. We will clarify the motivations for both in the main manuscript as follows:

Group indicator: The motivation of the group indicator derivations is to fully utilize the risk information to explore the relationships among tasks, where applying risk information can avoid the computation cost of gradients. Compared with other loss-oriented methods (such as RLW, DWA, UW, and FAMO), our group indicator can capture the differences in the per-task risk scale and fully utilize the learning dynamics over time, yielding better representations of risk information.

Empirically, Table 4 in Line 340 shows comparisons with other task-specific weights (MGDA, NashMTL and FAMO) + the proposed AGRM principle; our group indicators result in superior performance over other alternatives.

Group-specific weights versus task-specific weights: The motivation for group-specific weights is that tasks with close group indicators have similar behaviors on risk, and similar tasks can benefit from training together by sharing parameters as much as possible [21], even weights in the multi-task objective. With the group-specific weights, our method can align tasks with similar risk behaviors during joint training, alleviating the task imbalance issue.

As shown in Table 3 in Line 307, our work with group-specific weights outperforms its variants with task-specific weights (without Eq.4). This further demonstrates that our group-specific weights can effectively alleviate the under-optimization of some tasks.

---

2. Caption of Figure 5.

Thanks for your suggestion. We've taken your advice and added the description to the caption of Figure 5.

The x-axis in the subplots denotes the epoch, and the intensity of the color indicates the weight value.

---

3. Datasets with a larger number of tasks.

In MTL literature, it is common to evaluate datasets with a small number of tasks, such as NYUv2 (3 tasks) and CityScapes (2 tasks) [6, 13, 14, 16]. As task grouping may be beneficial for larger numbers of tasks, we also included QM9 (11 tasks) and CelebA (40 tasks). On these last two benchmarks, our method achieves better overall performance than other loss-oriented methods, with higher training efficiency than gradient-oriented baselines. The overall results are in Table 2 in Line 276 and details in Appendix Table 5 & 6.

---

4. Experimental analysis on Table 2.

We appreciate the reviewer's careful observations of Table 2. We want to provide clarifications on two points:

Weight calculations on CityScape: GO4Align achieves comparable performance, securing the third position ($\Delta m$) on CityScape. It’s critical to note that this benchmark, comprising only two tasks, offers limited grouping options, which constrains the effectiveness of the proposed weight calculation (risk-guided group indicator). Furthermore, in contrast to other weight calculation methods on NYUv2, such as Table 4 in Line340, the proposed risk-guided group indicator surpasses alternative weight calculations.

MTL methods underperform STL in terms of the overall evaluation: This is a common phenomenon in MTL literature (CAGrad, NashMTL, and FAMO), caused by the task-imbalance issue, where some tasks are under-optimized during joint training. In particular, we note that: (i) In Appendix Table 5 & 6 with detailed task-specific performance, we can observe that most baselines achieve comparable performance over STL on some tasks but significantly sacrifice the performance on the other tasks. (ii) MTL still offers advantages such as improved computational efficiency and reduced training time, which significantly contribute to real-world systems.

---

5. Our model with grouping versus without grouping.

We refer the reviewer to the ablations in Table 3 in Line307. As the grouping is performed by Eq.(4), the first two rows in Table 3 are the variants of our method without grouping, and the last row is our method with grouping. Table 3 empirically examines the performance gains of task grouping over models without grouping. We will make this more clear in Line 307.

---

6. Detailed experimental setup and open-source plan.

This work follows the same experimental setting used in NashMTL [13] and FAMO [14], including the dataset partition for training, validation, and testing. The benchmark partition is attached. We will add this table in Appendix Sec.B.1. we also note that NYUv2 and Cityscapes do not have validation sets. Following the protocol in [13-14], we report the test performance averaged over the last ten epochs. Importantly, we will release our code to facilitate MTO research after the final decision.

---

Thank you for your feedback. We greatly appreciate the time and effort you put into reviewing our work. We have carefully considered your comments and made improvements based on your suggestions. We hope you will reconsider your evaluation of our work. Thank you once again.