NTKMTL: Mitigating Task Imbalance in Multi-Task Learning from Neural Tangent Kernel Perspective

We appreciate your recognition of our method's theoretical innovation, clear algorithmic design, and broad experimental evaluation. Thank you for your careful and valuable comments! We hope to address your concerns by answering your questions below.

Q1: Incompatibility of NTKMTL-SR with Non-Task-Specific Architectures.

A1: We acknowledge that NTKMTL-SR is an approximate algorithm specifically designed for models with task-specific layers, which might limit its direct applicability to settings like MT10 or large language model (LLM) training. However, NTKMTL supports non-task-specific architectures and has demonstrated superior performance in multi-task reinforcement learning. Additionally, NTKMTL achieves a speed advantage of over 2x compared to some gradient-oriented MTL methods like PCGrad and Nash-MTL, which require additional gradient projection or iterative optimization.

Furthermore, our proposed extended NTK theory in MTL still holds potential for generating new algorithms in these scenarios. For instance, a promising direction involves balancing training losses from different data sources (e.g., code, math) during LLM pre-training by computing the NTK matrix with respect to deep-layer parameters for each data source. We welcome any insights you may have on this!

Q2: The weight in NTKMTL does not strictly ensure that the task weights sum to k (the number of tasks).

A2: Yes, early MTL methods (e.g. DWA [1]) often constrained the sum of weights to k (the number of tasks). However, the vast majority of recent MTL algorithms (e.g., Nash-MTL [2], FAMO [3], FairGrad [4], GO4Align [5]) do not adhere to this constraint. They operate under the premise that the scale information inherent in the calculated weights is beneficial for training, and that the optimal weight scales differ across various training stages. Forcibly normalizing these weights to k would lead to a loss of this crucial information.

Especially for algorithms like Nash-MTL and FairGrad, the task weights are derived by solving an optimization problem. In such cases, forced normalization would invalidate them as solutions to the optimization problem. Therefore, most contemporary MTL algorithms only require that the weights be bounded. Their specific scale is determined by the algorithm's hyperparameters, which are often tuned to find the optimal settings.

[1] Liu S, Johns E, Davison A J. End-to-end multi-task learning with attention[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019: 1871-1880.

[2] Navon A, Shamsian A, Achituve I, et al. Multi-task learning as a bargaining game[J]. arXiv preprint arXiv:2202.01017, 2022.

[3] Liu B, Feng Y, Stone P, et al. Famo: Fast adaptive multitask optimization[J]. Advances in Neural Information Processing Systems, 2023, 36: 57226-57243.

[4] Ban H, Ji K. Fair resource allocation in multi-task learning[J]. International Conference on Machine Learning, 2024.

[5] Shen J, Wang Q, Xiao Z, et al. Go4align: Group optimization for multi-task alignment[J]. Advances in Neural Information Processing Systems, 2024.

Q3: Aside from Table 1, most other tables do not report the results of all baseline methods.

A3: NYUv2, featured in Table 1, is one of the most widely used benchmarks in the MTL field, so nearly all MTL methods conduct experiments on it. Except for QM9, where we re-implemented almost all methods for a fair comparison, results for other MTL methods on other benchmarks were obtained from their original papers.

Since some works have conducted experiments on only two or three benchmarks, we cannot obtain the results from their papers. Many of them overlook CelebA (40-task), a benchmark crucial for evaluating MTL performance with a large number of tasks. This explains why the number of methods shown in Table 2 is fewer than in Table 1. For the CityScapes benchmark, due to space limitations in the main text, we presented its results alongside CelebA in Table 2, showing only 14 MTL methods. We report the detailed results for this benchmark, encompassing 18 methods, in Appendix Table 7.

Thank you for your suggestion. We will consider to re-implement the missing results for some MTL methods on CelebA and integrate them into the latest version of the manuscript.

Q4: Are there more ablation studies on hyperparameter n for NTKMTL?

A4: For NTKMTL-SR, we have conducted a comprehensive ablation study on n , with the detailed results presented in the supplementary material. At your request, we also conducted an ablation study on the hyperparameter n for NTKMTL. We performed experiments CelebA benchmark, and the results are presented in the table below. We maintain consistency with the settings of other MTL methods discussed in the main text, training for 15 epochs on CelebA. For each value of n, we conduct three repeated experiments and report the mean of $\Delta m \\%$.

The table shows that increasing n does yield modest performance gains; however, the computational overhead grows almost linearly with n, in sharp contrast to the negligible cost incurred by NTKMTL-SR. Because the results with n=1 are already highly competitive, and to ensure a fair comparison with other MTL methods, we recommend setting n=1 as the default for NTKMTL. In this case, it aligns with other gradient-oriented methods, requiring k gradient backpropagations for shared parameters per iteration.

Q5: Further Analysis for Off-Diagonal NTK Components.

A5: Thank you for showing this point. Our motivation is to balance the convergence speeds of different tasks, and thus we mainly require the maximum eigenvalue of the respective NTK matrix of each task to capture the convergence speed. Therefore, we primarily used the diagonal blocks of $\tilde{K}$ for analysis.

We also recognize that the off-diagonal blocks of $\tilde{K}$, specifically $\tilde{K}_{ij}$, might contain rich information of task similarity. These information can potentially lead to new MTL methods based on task grouping. However, as this represents a technical direction orthogonal to our paper's focus of mitigating task imbalance by balancing convergence speeds, we did not elaborate extensively on it here. Instead, we included it in the limitations section for future work.

 

These discussions have further enriched the quality of the manuscript. We will incorporate the ablation studies for both NTKMTL and NTKMTL-SR into the appendix, and simultaneously refine parts of the main text for enhanced clarity. If you have any further questions, please don't hesitate to contact us. Thank you again for your valuable time and effort during the review process!

We appreciate your recognition of our method's new theoretical perspective and extensive experimental validation. Thank you for your careful and valuable comments! We hope to address your concerns by answering your questions below.

Q1: Visualization experiments for the NTK eigenvalues during training.

A1: Thanks for your suggestion. We agree that experimental visualization is essential to support the proposed theory. Hence, we conducted experiments on NYUv2, visualizing the change in the maximum eigenvalue of the NTK matrix for the three tasks during training under linear scalarization (i.e., equal weighting). As images cannot be uploaded, we present the results in the table below.

'Avg over xx' denotes the mean of the maximum NTK eigenvalues during the first xx iterations. On the NYUv2 dataset, the difficulty levels of the three tasks show significant variation. Previous methods generally outperform the Single Task Learning (STL) baseline in segmentation and depth estimation tasks, but almost all of them consistently underperform STL on the surface normal prediction task, leading to a significant task imbalance in the overall results.

As evident from the table, throughout training, the largest eigenvalue corresponding to surface normal prediction remains substantially smaller than those of segmentation and depth estimation. Given that the maximum NTK eigenvalue reflects a task's convergence speed, this observation aligns with the empirical finding that many existing MTL algorithms struggle to converge on surface normal prediction. Although some prior methods (e.g., MGDA) attempt to prioritize the most difficult tasks, their performance in surface normal prediction tasks remains unsatisfactory due to the challenge in accurately quantifying the "difficulty" and "convergence speed" of different tasks. In contrast, by leveraging NTK theory to accurately characterize and balance the convergence speed of each task during training, NTKMTL delivers SOTA results on surface normal prediction and is one of only two methods that outperform single-task learning on all three tasks.

Q2: Question about "different tasks exhibit vastly different loss scales and heterogeneous ultimate loss minima".

A2: In MTL, simple equal weighting fails to balance task losses precisely because different tasks inherently possess distinct natures and may utilize distinct loss calculation methods (e.g., cross-entropy loss for classification vs. L2 loss for regression). This leads to vastly different loss scales and different ultimate loss minima, a fact widely acknowledged by many MTL works (e.g., UW, FairGrad, GO4Align). Directly applying equal weighting often results in the training process being dominated by tasks with larger loss scales.

For instance, on the NYUv2 benchmark, the losses for semantic segmentation and depth prediction tasks are typically an order of magnitude higher than that for surface normal prediction. On the even more heterogeneous QM9 dataset, the loss discrepancies among different tasks can span up to three orders of magnitude.

Thanks for your suggestion. We will rephrase this statement in lines 32-39 of the manuscript to enhance clarity.

Q3: Further elaboration for the key challenges of extending NTK theory from STL to MTL.

A3: In STL, the NTK expression can be directly derived from the gradient flow. However, MTL explicitly divides the objective function into k components, each associated with a specific task. Since all tasks share parameters, the gradient flow of one task also influences the losses of other tasks, which explicitly increases the complexity of the analysis. In MTL, analyzing only the overall loss using STL's NTK theory (Theorem 3.1) fails to capture fine-grained, per-task training dynamics. To fully understand the differences in convergence speed among tasks, it's essential to extend STL's NTK theory, which poses a new theoretical challenge.

We analyzed the specific scenario of MTL and proposed Theorem 3.2 and Proposition 3.3 (proofs detailed on pages 15-17). The most intuitive physical interpretation these theories provide is that with the introduction of weighting coefficients $w_i$, the new NTK matrix for each task becomes $w_i^2K_{ii}$, and its maximum eigenvalue is accordingly scaled to $w_i^2 \lambda_i$. This ultimately led to our final method design.

Thanks for your suggestion. We will further clarify the necessity of our proposed theory in lines 175-197 of the method section and provide a smoother transition to the method design.

Q4: Missing References.

A4: Thank you for the additions! We have previously read both of these two papers. Dynamic Task Prioritization (DTP), in particular, is a significant work in the MTL field. It was an oversight on our part to have omitted these two papers from the related work section. We've now included them in the manuscript.

Q5: Do all methods train for the same number of epochs in the experiment?

A5: Yes, to ensure fairness, all methods are trained for the same number of epochs across all benchmarks. Therefore, the training time per epoch can essentially represent the comparison of total training duration.

Q6: On CITYSCAPES, the MR of NTKMTL is worse than some baseline methods. How do different datasets affect the effectiveness of NTKMTL?

A6: While our method's performance varies across different benchmarks, a consistent pattern emerges: it consistently optimizes more difficult tasks more effectively, leading to more balanced optimization and achieving superior average performance.

When evaluating MTL methods, $\Delta m\\%$ is typically the most intuitive metric as it relies solely on the objective performance of the MTL method on a specific benchmark. Mean Rank (MR), however, depends on the number and performance of other MTL methods used for comparison, and thus usually serves as a supplementary criterion when $\Delta m\\%$ values are similar.

Why does NTKMTL exhibit suboptimal MR performance on the CityScapes dataset? We report the detailed performance of 17 different MTL methods on the CityScapes in Appendix Table 7 (page 19). The CityScapes task comprises two subtasks: segmentation and depth prediction. Due to the disparity in task difficulty, most prior methods excelled at the segmentation task but performed significantly worse than the STL baseline on the depth prediction task, exhibiting severe task imbalance. Our method effectively balances the performance of different tasks, achieving significant progress in the depth task. However, this also implies a slight decrease in performance on the segmentation task, placing our method at a disadvantage in MR calculation.

In summary, NTKMTL's suboptimal MR performance on CityScapes does not indicate a failure of the method on this dataset. In fact, it effectively mitigates task imbalance and achieves the optimal $\Delta m\\%$.

Q7: Why does NTKMTL-SR perform better than NTKMTL in certain scenarios?

A7: In Table 2, NTKMTL-SR does not outperform NTKMTL in terms of $\Delta m\\%$ (0.23 vs. -0.77), achieving only a slightly better Mean Rank (MR) (4.33 vs. 4.35). As stated in A6, the MR metric is susceptible to changes based on the inclusion or exclusion of other comparative MTL methods.

However, on the QM9 dataset, NTKMTL-SR indeed surpasses NTKMTL and achieves state-of-the-art performance. We also analyzed this phenomenon, attributing it to the relatively small number of model parameters used in this regression task; consequently, even with approximation, the training dynamics of different tasks can be captured effectively. As a result, a larger n produces a more accurate estimate of convergence speed. To further validate this, we conducted detailed ablation experiments in the supplementary material, where we set n to [1, 2, 3, 4, 6] for NTKMTL-SR and performed three experiments for each setting, reporting the mean and error bars. As shown in Fig. 2 of the supplementary material, when n=1, the performance of NTKMTL-SR is comparable to that of NTKMTL. However, when increasing n from 1 to 2 or more, NTKMTL-SR shows a noticeable improvement in performance, accompanied by a reduction in the variance of performance across repeated experiments.

Q8: How does the task-relatedness/similarity influence task dominance and NTK singular values?

A8: During our experiments, we observed that tasks exhibiting clear similarities (e.g., tasks 7-10 on QM9) not only have similarly scaled maximum eigenvalues of their NTK matrices but also show consistent trends in their evolution throughout training. Consequently, NTKMTL assigns essentially equal weights to these tasks to facilitate balanced optimization. (As shown in Table 3, these highly correlated tasks exhibit very close final MAE values.) We have also visualized this phenomenon, and the images and corresponding analysis will be included in the latest version of the appendix.

Furthermore, we acknowledge that the off-diagonal blocks of $\tilde{K}$, specifically $\tilde{K}_{ij}$, contain rich information. These blocks can, to some extent, represent task similarity and potentially lead to new MTL methods based on task grouping. However, due to the orthogonality of the task grouping technical approach to our current method, we include it in the limitations section for future work.

 

These discussions have further enriched the quality of the manuscript. We will integrate all rebuttal content, adding visualizations as figures to the appendix, and simultaneously revise and refine certain statements in the main text. If you have any further questions, please don't hesitate to contact us. Thank you again for your valuable time and effort during the review process!

We appreciate your recognition of our method's novelty, strong theoretical foundation, effectiveness, and efficiency. Thank you for your careful and valuable comments! We hope to address your concerns by answering your questions below.

Q1: Ablation Study on Hyperparameter n

A1: We also noticed that exploring the impact of n as a hyperparameter on experimental results is meaningful, especially for NTKMTL-SR. Therefore, we conducted a comprehensive ablation study on n for NTKMTL-SR, with the detailed results presented in the supplementary material. We set n to $[1, 2, 3, 4, 6]$ and conduct an ablation study on the QM9 (11-task) benchmark. For each value of n, we perform 3 repeated experiments with different random seeds and calculate the mean and error bar for $\Delta m\\%$. Detailed results and analysis are available in the supplementary material, and we have also excerpted them here for convenience.

The results indicate that increasing n from 1 to 2 or more shows a noticeable improvement in NTKMTL-SR's performance, accompanied by a reduction in the error bar of performance across repeated experiments. We attribute this to the fact that increasing the number of mini-batches leads to a larger NTK matrix dimension, which in turn reduces stochastic error and allows our method to more accurately characterize the convergence speed of the tasks. However, since further increasing n (i.e., from 4 to 6) does not yield significant additional benefits and introduces further computational overhead, we posit that the selection of hyperparameter n is a trade-off between performance and training speed, and n=4 generally presents a favorable compromise.

Additionally, at your request, we also conducted an ablation study on the hyperparameter n for NTKMTL. Due to limited rebuttal time, these experiments were performed on the smaller CelebA benchmark, and the results are presented in the table below. We maintain consistency with the settings of other MTL methods discussed in the main text, training for 15 epochs on CelebA. For each value of n, we conduct three repeated experiments and report the mean of $\Delta m \\%$.

The table shows that increasing n does yield modest performance gains; however, the computational overhead grows almost linearly with n, in sharp contrast to the negligible cost incurred by NTKMTL-SR. Because the results with n=1 are already highly competitive, and to ensure a fair comparison with other MTL methods, we recommend setting n=1 as the default for NTKMTL. In this case, it aligns with other gradient-oriented methods, requiring k gradient backpropagations for shared parameters per iteration.

Q2: Further analysis of the full extended NTK matrix $\tilde{K}$, particularly the off-diagonal $\tilde{K}_{ij}$ blocks.

A2: Thank you for showing this point. Our motivation is to balance the convergence speeds of different tasks, and thus we mainly require the maximum eigenvalue of the respective NTK matrix of each task to capture the convergence speed. Therefore, we primarily used the diagonal blocks of $\tilde{K}$ for analysis.

We also recognize that the off-diagonal blocks of $\tilde{K}$, specifically $\tilde{K}_{ij}$, might contain rich information about task similarity. These information can potentially lead to new MTL methods based on task grouping. However, as this represents a technical direction orthogonal to our paper's focus of mitigating task imbalance by balancing convergence speeds, we did not elaborate extensively on it here. Instead, we included it in the limitations section for future work.

Q3: About typographical errors.

A3: Thank you for your careful review! We've corrected these typos, and they will be eliminated in the latest version of the paper.

 

These discussions have further enriched the content of this paper. We will integrate the relevant ablation studies into the main text and Appendix, while also carefully reviewing the manuscript to eliminate any typos. Thank you again for your high appreciation of our work and for the valuable time and effort during the review process!

We appreciate your recognition of our method's rationality, the comprehensiveness and diversity of our experiments, and the broad adaptability of the proposed method. Thank you for your careful and valuable comments! We hope to address your concerns by answering your questions below.

Q1: Why does the proposed method achieve better results compared to other approaches (e.g. other methods that balance tasks based on gradients)?

A1: The superior performance of the proposed NTKMTL can be attributed to its ability to effectively balance the convergence rates of individual tasks, which fosters more equitable optimization and ultimately yields enhanced overall performance.

Why balanced optimization counts: Recent works (e.g., FairGrad[1] and GO4Align[2]) have demonstrated that more balanced cross-task performance typically yields better overall results (See Fig.2 in [2]). To more clearly elucidate the reason for NTKMTL's superior performance, we analyze it using the NYUv2 benchmark as an example. On the NYUv2 dataset, the difficulty levels of the three tasks show significant variation. Previous methods generally outperform the Single Task Learning (STL) baseline in the tasks of semantic segmentation and depth estimation, but almost all of them consistently underperform STL on the surface normal prediction task, leading to a significant task imbalance in the overall results. Although some prior methods (e.g., MGDA) attempt to prioritize the most difficult tasks, their performance in surface normal prediction tasks remains unsatisfactory due to the challenge in accurately quantifying the "difficulty" and "convergence speed" of different tasks. In contrast, by leveraging NTK theory to accurately characterize and balance the convergence speed of each task during training, NTKMTL demonstrates strong performance in the surface normal prediction task. Notably, among all existing methods, only NTKMTL and GO4Align consistently outperform the STL baseline across all three tasks, achieving a more balanced optimization.

[1] Ban H, Ji K. Fair resource allocation in multi-task learning[J]. International Conference on Machine Learning, 2024.

[2] Shen J, Wang Q, Xiao Z, et al. Go4align: Group optimization for multi-task alignment[J]. Advances in Neural Information Processing Systems, 2024.

Experimental validation: To further substantiate this, we conducted additional experiments and visualized: 1) the maximum eigenvalue of the NTK matrix for different tasks during NYUv2 training, and 2) the relationship between cross-task performance variance and $\Delta m\\%$ for various methods on NYUv2. As images cannot be uploaded, we present these findings in the tables below.

1). The maximum eigenvalue of the NTK matrix for different tasks during NYUv2 training.

In the table above, we report the evolution of the largest eigenvalue of the NTK matrix for each task when utilizing conventional linear scalarization method. ‘Avg over xxxx’ denotes the mean of the maximum NTK eigenvalues during the first xxxx iterations. Throughout training, the largest eigenvalue corresponding to surface-normal prediction remains substantially smaller than those of segmentation and depth estimation. Since the maximum NTK eigenvalue reflects the convergence speed of a task, this observation aligns with the empirical finding that many existing MTL algorithms struggle to converge on surface-normal prediction. By accurately quantifying these differing convergence rates and re-allocating weights accordingly, NTKMTL attains more balanced performance, delivers state-of-the-art results on surface-normal prediction, and is one of only two methods that outperform single-task learning on all three tasks.

2). The relationship between cross-task performance variance and $\Delta m\\%$ on NYUv2.

The table above reports $\Delta m\\%$ (average performance drop) and Var[$\Delta m\\%$] (variance of the performance drop across tasks) on NYUv2 for NTKMTL and other 10 existing MTL methods. Notably, the performance $\Delta m\\%$ of different MTL methods exhibits a striking positive correlation with cross-task performance variance, and methods with smaller performance variance tend to achieve better overall performance. This is precisely why “mitigating task imbalance” is of paramount importance in the field of multi-task learning, and it also supports why NTKMTL can achieve a superior $\Delta m\\%$ through more balanced optimization.

Q2: Are there any modules or parameter choices in the proposed method that are particularly critical (Have any ablation studies been conducted)?

A2: Since the weights are derived by computing the maximum eigenvalue of the NTK matrix for different tasks, the samples within a batch need to be split into n mini-batches. While a larger n can lead to a more precise NTK representation, it also incurs greater computational overhead.

We also noticed that exploring the impact of n as a hyperparameter on experimental results is meaningful, especially for NTKMTL-SR. Therefore, we conducted a comprehensive ablation study on n for NTKMTL-SR, with the detailed results presented in the supplementary material. We set n to $[1, 2, 3, 4, 6]$ and conduct an ablation study on the QM9 (11-task) benchmark. For each value of n, we perform 3 repeated experiments with different random seeds and calculate the mean and error bar for $\Delta m\\%$. Detailed results and analysis are available in the supplementary material, and we have also excerpted them here for convenience.

The results indicate that increasing n from 1 to 2 or more shows a noticeable improvement in NTKMTL-SR's performance, accompanied by a reduction in the error bar of performance across repeated experiments. We attribute this to the fact that increasing the number of mini-batches leads to a larger NTK matrix dimension, which in turn reduces stochastic error and allows our method to more accurately characterize the convergence speed of the tasks. However, since further increasing n (i.e., from 4 to 6) does not yield significant additional benefits and introduces further computational overhead, we posit that the selection of hyperparameter n is a trade-off between performance and training speed, and n=4 generally presents a favorable compromise.

For NTKMTL, due to the memory and time cost associated with fully computing the gradients of n mini-batches for shared parameters to construct the NTK matrix, we set n=1 to ensure fairness when comparing with other methods. In this case, it aligns with other gradient-oriented methods, requiring k gradient backpropagations for shared parameters per iteration. Additionally, during the rebuttal period, we conducted further ablation experiments on the hyperparameter n for NTKMTL. You can refer to A1 in our response to Reviewer 2 (yMmP) for details.

Q3: Why does the proposed method exhibit varying performance across different tasks?

A3: In the field of MTL, achieving optimal performance across all subtasks on a given benchmark simultaneously is challenging, as it essentially means directly extending the Pareto frontier. While our method's performance varies across different benchmarks, a consistent pattern emerges: it consistently optimizes more difficult tasks more effectively, leading to more balanced optimization, effectively mitigating task imbalance, and achieving superior average performance.

For instance, on NYUv2, NTKMTL achieved SOTA results on the most challenging surface normal prediction task. It also obtained the optimal MR and $\Delta m\\%$. On the CityScapes benchmark, most prior methods excelled at the segmentation task but performed significantly worse than the Single-Task Learning (STL) baseline on depth prediction. NTKMTL, however, better balanced the performance of these two tasks and achieved the optimal $\Delta m\\%$.

 

These discussions have further enriched the content of this paper. We will integrate the rebuttal content and update the manuscript accordingly. Specifically, the visualizations from A1 regarding the maximum eigenvalue of the NTK matrix for different tasks during training and the relationship between cross-task performance variance and $\Delta m\\%$ on various benchmarks will be added as figures to the Appendix. The corresponding analysis from A1 will also be incorporated into the main experimental section. The results of the ablation studies will be integrated into both the main text and the Appendix. The discussion from A3 will be added around lines 252-270 in Section 4.2 to enhance the analysis of the experimental results.

If you have any further questions, please do not hesitate to contact us. Thank you again for your valuable time and effort during the review process!