GatedMTL: Learning to Share, Specialize, and Prune Representations for Multi-task Learning

Comparison to recent baselines and parameter size: In response to the reviewers feedback regarding the selection of baselines, we have extensively expanded our experiments to include the suggested state-of-the-art methods, i.e. PCGrad [1], CAGrad [2], MGDA [3], and RDW [4]. The results are presented in the table below. As can be seen, GatedMTL substantially outperforms all competing baselines across both $\Delta_{MTL}$ and the newly added Mean Rank metric (suggested by R1). We additionally report the parameter counts of the models as suggested by the reviewer. Our method compares favorably against Cross-stitch and MTAN in terms of number of parameters and performance. Please note that as discussed above, our comparisons do not include MoE-based approaches as they are not suitable for our MTL setup for solving all tasks simultaneously.

Updated performance comparison on NYUD-v2 using HRNet-18 backbone. The new baselines are shown in bold.

---

Updated performance comparison on NYUD-v2 using ResNet-50 backbone.

The choice of $\omega_t$: The hyper-parameter $\omega_t$ denotes the scalarization weights. We use the weights suggested in prior work but also report numbers of uniform scalarization. For NYUD-v2, we use uniform scalarization as suggested in [5,6], and for PASCAL-Context, we similarly use the weights suggested in the paper [5,6].

Target computational cost: We agree with the reviewer and reformulated the wording of the claim in the abstract and additionally mentioned this as a limitation of our method. As correctly described by the reviewer, although both $\lambda_s$ and $\tau_t$ can control the trade-off between performance and computational cost and closely result in the desired FLOPs count, we still cannot guarantee landing on a particular target computational cost.

Minor comments: We thank the reviewer for their detailed and meticulous comments. We fixed all of them directly in the manuscript.

References:

[1] Gradient Surgery for Multi-Task Learning. NeurIPS, 2020.

[2] Conflict-Averse Gradient Descent for Multi-task Learning. NeurIPS, 2021.

[3] Multi-Task Learning as Multi-Objective Optimization. NeurIPS, 2018.

[4] Reasonable Effectiveness of Random Weighting: A Litmus Test for Multi-Task Learning. TMLR, 2022.

[5] Attentive Single-Tasking of Multiple Tasks, CVPR, 2019.

[6] Multi-task learning for dense prediction tasks: A survey, PAMI, 2021.

We thank the reviewer for their thorough review of our paper and insightful comments.

Open-source Code Availability: We acknowledge the critical role of open-sourcing in advancing research and reproducibility. We have initiated the process for obtaining legal approval to open-source our code. Additionally, to facilitate understanding and reproducibility in the interim, we reported all detailed hyperparameters and have included detailed pseudo-code for GatedMTL in Appendix E.

Clarification of our method: We apologize for any confusion caused and appreciate the opportunity to clarify. The encoder in GatedMTL consists of both shared parameters $\Psi$ and task-specific parameters $\Phi_t$. The shared $\psi$ and task-specific features $\varphi_t$ are computed using their respective parameters. In Figure 1, we illustrate the architecture of a single GatedMTL layer within the unified encoder, highlighting how task-specific gates blend shared and task-specific features to form $\varphi_t^{\prime}$, serving as the input to the following task-specific layers. The purple blocks represent the features extracted from the shared branch, while the green and yellow blocks represent features extracted by task $1$ and task $t$ (we show only two task-specific branches for simplicity). For each task, a corresponding task-specific gate (e.g. yellow round block) makes a selection between using shared features and task specific features resulting in $\varphi_t^{\prime}$ shown as a blended representation with two colors. Subsequently, $\varphi_t^{\prime}$ represents the input to the next task-specific non-linear transformations with parameters $\Phi_t$. The input to the next shared layer is constructed via a linear combination of the task-specific features of all tasks, followed by a non-linear transformation parameterized by $\Psi$. There are fundamental differences between GatedMTL and MoE-based approaches which we highlight in the next point.

Comparison to MoE-based approaches: A major difference between GatedMTL and MoE-based approaches is in how they operate at inference. GatedMTL is designed to solve all tasks simultaneously in one forward pass, a strict requirement in many real-world practical scenarios (like XR and autonomous driving applications requiring edge deployment for solving multiple tasks simultaneously). In the contrary, existing MoE-based approaches, can only solve one task at a time, requiring one forward pass per task. In our setting, this makes MoE approaches more expensive than conventional MTL approaches, considering the cost of the accumulated forward passes. Unlike GatedMTL which targets efficient inference for solving all tasks simultaneously, MoE-based approaches focus on improving accuracy via dynamic routing to expert modules while solving one task at a time. Therefore, MoE models are not suitable for the setup we are targeting in this work. We have made these clearer in the manuscript and highlighted the key differences in the related work section.

Using single-task weights and implications on training cost: Our work primarily aims to enhance inference efficiency, vital for deploying multi-task learning (MTL) models in edge computing environments. In numerous real-world applications, the gains in efficient inference over time, especially when scaled across millions of edge devices, can significantly offset the initial higher training costs. While we acknowledge the importance of training efficiency and have discussed GatedMTL’s training cost limitations in our manuscript, it’s important to highlight that competing methods like MTO approaches (PCGrad [1], CAGrad [2], MGDA [3], etc.) also bear considerable training costs, requiring multiple backward passes during training. Our GatedMTL significantly surpasses these in performance, demonstrating its advantage for practical applications that demand efficient inference on the edge.

In addition, in both MTL research and practice, establishing single-task baselines is a standard practice to gauge the relative performance of MTL models. Therefore, having access to single task models is not unreasonable from a practical perspective. In addition, pre-trained single task models for many applications are abundantly available which may obviate the need to separately train STL models.

Thanks to the authors for the rebuttal. I still have some concerns about my initial comments.

1. I still do not understand the architecture of GateMTL in Figure 1. Are there $T+1$ encoders where one is shared among different tasks and the remaining $T$ are task-specific? If so, why the parameter size of GateMTL is smaller than STL? If not, how do we obtain each layer's shared and task-specific features?

2. Figure 1 in [2] summarizes different MTL architectures. The blue block means the shared block and the red and green ones are task-specific. (1) The architecture of PLE is similar to GateMTL. (2) Why MoE-based methods can only solve a task at a time?

3. I appreciate that this paper aims to enhance inference efficiency, but is it related to initializing task-specific branches with their corresponding single-task weights? It is unfair to compare with the baselines that do not use the information from single-task models. I suggest training GateMTL without single-task weights initialization for a fair comparison.

4. The method name in [8] is RLW rather than RDW.

We thank the reviewer for this discussion.

Comparison to [4]: We will improve the visualization in Figure 1. It appears there are some misunderstandings about how GatedMTL works. In [4], the DSMTL-IL variant is extremely expensive (all encoders are executed) and incomparable to GatedMTL both methodologically and computationally. The DSMTL-AL, which focuses on saving compute at inference, has $T$ private encoders for each specific task. The main focus is how to learn to branch out from the shared encoder to the task-specific encoder. In other words, part of the trunk is shared among tasks, and tasks decide to exit the shared branch at some point and switch to their task-specific encoder blocks. We find this an interesting idea and will discuss it in the paper, but DSMTL-AL is not an MoE architecture. DSMTL is fundamentally different from GatedMTL. GatedMTL does not gate blocks but rather individual channels/filters making it far more flexible and fine-grained. Our work does not focus on when to branch out from a shared branch, but rather to decide how to effectively choose from every individual shared-parameter and task-specific ones, resulting in a more expressive architecture. Unfortunately, the paper [4] does not provide inference timing or source code which makes comparison difficult.

Regarding the comparison with MoE-based methods starting from single-task models, two primary challenges hinder their suitability in our multi-task learning (MTL) context. First, as we mentioned before, the MoE framework, by its nature, requires running tasks sequentially, leading to a cumulative inference cost equivalent to executing $T$ single tasks. This approach contradicts the efficiency principles of GatedMTL and limits the relevance of a direct comparison.

Furthermore, the initialization of single-task weights into MoE architectures presents unresolved complexities. Consider an MoE with 16 experts at each layer solving three tasks, where the router dynamically selects from these experts. How is it possible to load single task weights inside these 16 experts in an MoE architecture? Another challenge arrises when we match the number of experts and number of tasks. A fundamental challenge in MoEs is jointly learning the router and experts. The router needs to direct samples to the most appropriate experts, but these experts only become proficient if they receive the correct samples by the router. Initializing experts with single-task weights in this scenario can lead the router to trivial solutions selecting the corresponding expert for each task (mode collapse), thereby ignoring other experts. This pattern contradicts the intended learning dynamics of MoEs and undermines the model’s overall adaptability and learning potential.

We thank the reviewer for their thorough review of our paper and insightful comments.

The motivation of the gating design for MTL: We would like to highlight the core motivation for the design of our GatedMTL model. The core concept of our paper addresses a significant challenge in MTL: task interference [1]. This phenomenon is particularly evident in models with limited representation capacity. In such models, while certain features are universally beneficial across all tasks, task-specific features often compete for the remaining capacity. This interference can lead to dominance of certain tasks over others during training, influenced by the dynamics of training and the magnitudes of loss functions.

Our primary innovation lies in recognizing that a rigid or pre-defined allocation of feature capacity between task-specific and shared features is sub-optimal. Such static partitioning fails to adapt to the varying needs of different tasks and layers within the network. To overcome this, we propose a flexible gating mechanism. This mechanism is not just a tool for balancing feature ratios; its fundamental purpose is to learn the optimal architecture dynamically, adapting the degree of feature sharing and task-specific specialization at each layer for each task. This approach ensures that each task has access to the necessary feature capacity without the interference of gradients from other tasks, while at the same time allowing for maximizing the sharing of relevant features.

It’s important to clarify that our contribution is not merely the gating mechanism itself, which, as the reviewer correctly notes, is not a novel concept in MTL. The novelty of our work lies in the application of this mechanism to dynamically learn and adjust the architecture for optimal feature sharing and specialization. This is a significant departure from traditional MTL models, where such parameters are usually fixed and do not adapt to the specific requirements of each task, while being highly efficient at the inference in solving all tasks at once.

Where to gate, encoder or decoder? We thank the reviewer for bringing up this interesting discussion. We absolutely agree that in typical MTL setups, the encoder is responsible for encoding out the shared features across all tasks. As highlighted in the above discussion, task interference is inherently the problem that arises with sharing representations across several tasks and the encoder is a natural choice because it is fully shared among all tasks. We would like to provide three additional major reasons for why we chose to resolve task interference within the encoder rather than the decoder:

1. Computational efficiency and specialization of decoders: Our design is grounded in the conventional MTL architecture, which typically comprises a large shared backbone and multiple task-specific decoders, as depicted in Figure 1 (left). In such architectures, the majority of computational resources are consumed by the encoder (backbone). For instance, in the context of NYUD-v2 with three tasks, the cumulative computational load of the HRNet18 decoders constitutes merely 17% of the total. Given that the decoders are designed to be highly specialized and independent for each task to optimize performance, it is both computationally efficient and logical to address task interference predominantly in the backbone.

2. Diversity in task decoding requirements: In MTL, particularly in complex applications like autonomous driving, tasks can vary greatly in their nature, such as combining dense prediction tasks, object detection tasks, and classification tasks such as weather prediction. In such scenarios, decoders might possess distinctly different architectures or layers. Forcing feature sharing at the decoder level could be detrimental, as it disregards the unique requirements and structures of each task’s decoder.

3. Encoder feature quality and decoder success: The quality of features extracted by the encoder is pivotal for the success of the specialized decoders. Significant task interference at the encoder stage can lead to degraded feature quality, resulting in irreversible information loss and sub-optimal input for the decoders. Ensuring high-quality feature extraction by the encoder is therefore crucial for the effective functioning of the task-specific decoders.

Finally, our empirical evaluations strongly demonstrate the effectiveness of the encoder-based approach to mitigating task-interference. We leave the investigation of decoder-based approaches to future work and discuss that in the manuscript.

Mean Rank Metric: We thank the reviewer for this suggestion. We added the rank mean measure to all tables in the paper and updated the manuscript. This new metric further strengthens the merits of GatedMTL, consistently showing strong performance across both the $\Delta_{MTL}$ and rank mean measures.

Baselines: In response to the reviewers feedback regarding the selection of baselines, we have extensively expanded our experiments to include the suggested state-of-the-art methods, i.e. PCGrad [1], CAGrad [2], RDW [3], and MGDA [4]. The results are presented in the table below. As can be seen, GatedMTL substantially outperforms all MTO baselines across all the metrics. We also integrated the implementation of NashMTL into our MTL baselines, however, despite extensive hyperparameter search, the experiments diverged which are not reasonable to be reported.

Updated performance comparison on NYUD-v2 using HRNet-18 backbone. The new baselines are shown in bold.

---

Updated performance comparison on NYUD-v2 using ResNet-50 backbone.

Please note that the MTL setting we are targeting, solves multiple tasks simultaneously in one forward pass. In contrast, Adashare and existing MoE-based approaches, can only solve one task at a time, requiring one forward pass per task. Therefore, these models are not suitable for the setup we are targeting in this work.

Q) Different Lagrange multipliers for L1 regularization: Yes, we indeed experimented with varying Lagrange multipliers for the L1 sparsity term for each task. However, we observed that the Lagrangian multipliers did not yield similar training dynamics for the gates as compared to the hinge loss. We hypothesize that, without a hinge target rate, the sparsification endlessly continues throughout the training making it difficult for the model to converge and the task-specific and shared representations to co-adapt.

We trust that we have addressed all the comments raised by the reviewer and are happy to further clarify any points as needed.

References:

[1] Gradient Surgery for Multi-Task Learning. NeurIPS, 2020.

[2] Conflict-Averse Gradient Descent for Multi-task Learning. NeurIPS, 2021.

[3] Reasonable Effectiveness of Random Weighting: A Litmus Test for Multi-Task Learning. TMLR, 2022.

[4] Multi-Task Learning as Multi-Objective Optimization. NeurIPS, 2018.

We thank the reviewer for their thorough review and constructive feedback on our manuscript. We appreciate the recognition of the strengths of our work, finding our paper well-written, and the proposed solution to be super intuitive, as well as highlighting our empirical evaluations to be impressive and the empirical results being quite positive. While we understand the concerns raised regarding limitations, we would like to provide additional insights, clarifications, and further baseline experimentations to address the reviewer's feedback.

Training time: The primary focus of GatedMTL is on inference-time efficiency, crucial in real-world applications. An example use case is XR applications requiring edge deployment for solving multiple tasks concurrently. We acknowledge that GatedMTL is not optimized for training time efficiency and we have added a limitation section in the paper, but we do not concur that this poses a significant drawback for practitioners. The training cost of GatedMTL is comparable, if not more efficient, than most existing MTO approaches. During training, many MTO approaches require $T$ backward passes and additional computations such as computing gradient conflicts/storing per task gradients simultaneously. GatedMTL, however, allows for sparsification into both forward and backward passes, meaning parameters that are not chosen do not get gradients and hence enables significant memory and compute saving (depending on the sparsity level). At inference, where GatedMTL is primarily designed for, we outperform all competing MTO baselines offering a significantly better performance-compute trade-off.

Scaling to an extremely high number of tasks: We agree that GatedMTL may not be the optimal choice for scenarios requiring simultaneous training/inference for an exceptionally high number of tasks, such as 40. However, we observe that this limitation is not unique to GatedMTL but is a common challenge across the MTL field. Our framework, therefore, remains relevant and effective within the practical range of multi-task learning scenarios that the vast majority of the related literature is focusing on.

Hyperparameter $\tau$: We believe setting of this parameter to be both intuitive and straightforward. A lower value indicates more reliance on shared features, whereas a higher value suggests a need for more task-specific ones. We would like to ask for clarification of what the reviewer finds unintuitive about $\tau$? In addition, in MTL research and practice, establishing single-task baselines is a standard practice to gauge the relative performance of MTL models. Therefore, having access to single task models is not unreasonable from a practical perspective. In addition, as shown in Appendix C the parameter $\tau$ itself is very robust.

Citations: We acknowledge the issues with our citations and bibliography and have meticulously revised this section, ensuring that all references are accurately and consistently cited.

We thank the reviewer for their response and clarifying these points further.

Training and scaling: In response to the reviewer's concern about training time efficiency, we report a comprehensive benchmarking of our GatedMTL model against various MTL approaches on our most complex dataset, Pascal-Context with 5 tasks. We include the average time for a single iteration for forward pass, backward pass, as well as the training time. As observed in the table below, while GatedMTL is generally cheaper than MTO approaches and Cross-Stitch, it adds non-substantial training overhead compared to standard MTL yielding a much higher $\Delta_\text{MTL}$ measure. For instance, our GatedMTL variant with a training time of 8.4h adds 11% to the training time of standard MTL while improving the $\Delta_\text{MTL}$ from -4.14 to -1.35. More details about the setup of the experiments are provided on Appendix F of the paper.

With regards to experiments on CelebA, we considered this as a toy classification dataset, particularly for some of our ablation studies on the capacity experiments. However, we are fine with moving the CelebA materials to the appendix if the reviewer recommends that.

Hyperparameter $\tau$: The reviewer suggests that if $\tau$ represented "the desired percentage of task-specific parameters for a task", it would offer a more intuitive and practical understanding for tuning. We would like to clarify that this is indeed exactly the intended interpretation of the $\tau$ parameter, as implemented in our method.

As detailed in Equation (4), the 'target rate' is designed to control the proportion of active gates at each specific layer, with each gate governing the activation of a single convolutional kernel. Thus, $\tau$ sets a soft upper limit for active task-specific parameters through these gates.

Accordingly, we have edited the manuscript to further clarify this.

Using STL weights: We respectfully disagree with the assertion that utilizing single-task model (STL) weights should not be practiced in multi-task learning (MTL). Our position is that the inability of existing MTL solutions to effectively incorporate STL weights mainly because of inherent design limitations, represents a notable limitation in the field. This should not, however, be a deterrent to the development and adoption of innovative approaches that can successfully leverage STL weights, as we propose in our work.

We consider our method’s ability to harness STL weights for enhancing the performance of MTL models during inference as a strength and a novel aspect of our research. Far from being a questionable practice, employing STL data in this manner aligns with practical machine learning principles and optimizes the use of available resources.

In addition, in numerous real-world applications, the gains in efficient inference over time, especially when scaled across millions of edge devices, can significantly offset the additional requirement of collecting STL weights. We would like to emphasize that single task pre-trained models for tasks such as semantic segmentation, depth estimation, etc. are abundantly available for a multitude of model architectures that could be plugged-in off-the-shelf.

However, to ensure full transparency with readers, we have acknowledged the necessity of collecting STL weights in the limitations section of our manuscript to clearly inform readers of any potential trade-offs involved in our approach.

Baselines: We are indeed using MGDA-UB algorithm proposed by Sener and now call the algorithm accordingly in our tables.

We are are attaching an updated version of our manuscript by highlighting the changes in blue.

We thank the reviewer for their thorough review of our paper and insightful comments.

Open-source Code Availability: We acknowledge the critical role of open-sourcing in advancing research and reproducibility. We have initiated the process for obtaining legal approval to open-source our code. Additionally, to facilitate understanding and reproducibility in the interim, we reported all detailed hyperparameters and have included detailed pseudo-code for GatedMTL in Appendix E.

Interference impact on gating: This is indeed a very interesting point raised by the reviewer. While our gating mechanism does indeed select between the shared and task-specific features, potentially influenced by task interference, it also serves as a mechanism to mitigate the negative impact of task-interference by avoiding the selection of representations with severe task interference issue. This is also supported by the results in Table 6 (first row) showing that in the absence of a sparsity objective, the gates inherently and strongly favor selecting task-specific over shared features. Additionally, our empirical findings show that in practice even with strong sparsity regularization, the gates can cope with task-interference by selecting the least interfering shared features, suggesting that task-interference at this meta-level might be less of a concern for the gates.

Multiple runs with random seeds: All reported results for the baselines in our study were averaged across three random seeds. We are sorry for the oversight of not mentioning this in the original manuscript and thank the reviewer for notifying us. We have now included this information in Section 4.1. We additionally report the standard deviation for the runs in the appendix. For GatedMTL runs, we inherently had numerous results for the Pareto optimal performance vs. FLOPs curve. But based on the reviewer's suggestion we are providing the GatedMTL results for three random seeds as well. Please note that the experiments are running and in the Table below we provide the mean+std for the completed runs and will update the final manuscript with the full set of results. Note that in the Table below we additionally include results for 4 new baselines suggested by wdHJ and 7nfe.