Giving each task what it needs -- leveraging structured sparsity for tailored multi-task learning

TL;DR

This work introduces Layer-Optimized Multi-Task (LOMT) models that utilize structured sparsity to refine feature selection for individual tasks and enhance the performance of all tasks in a multi-task scenario.

Abstract

In the Multi-task Learning (MTL) framework, every task demands distinct feature representations, ranging from low-level to high-level attributes. It is vital to address the specific (feature/parameter) needs of each task, especially in computationally constrained environments. This work, therefore, introduces Layer-Optimized Multi-Task (LOMT) models that utilize structured sparsity to refine feature selection for individual tasks and enhance the performance of all tasks in a multi-task scenario. Structured or group sparsity systematically eliminates parameters from trivial channels and, sometimes, eventually, entire layers within a convolution neural network during training. Consequently, the remaining layers provide the most optimal features for a given task. In this two-step approach, we subsequently leverage this sparsity-induced optimal layer information to build the LOMT models by connecting task-specific decoders to these strategically identified layers, deviating from conventional approaches that uniformly connect decoders at the end of the network. This tailored architecture optimizes the network, focusing on essential features while reducing redundancy. We validate the efficacy of the proposed approach on two datasets, i.e., NYU-v2 and CelebAMask-HD datasets, for multiple heterogeneous tasks. A detailed performance analysis of the LOMT models, in contrast to the conventional MTL models, reveals that the LOMT models outperform for most task combinations. The excellent qualitative and quantitative outcomes highlight the effectiveness of employing structured sparsity for optimal layer (or feature) selection.

Figures (5)

• Figure 1: A simple schematic representation of the proposed work.
• Figure 2: Sparsity pattern of the layers of ResNet-50 backbone for (a) NYU-v2 and (b) CelebAMask-HQ dataset. The first row represents the dense network, with all layers having non-zero parameters indicated in blue. The subsequent rows illustrate the sparsity pattern obtained after training a single-task model using group sparsity. The orange square stands for the last layer, after which all the parameters are zero, and this layer is selected for phase 2 to create LOMT models. The names of tasks $T_n$ for both the figures can be found in Tables. \ref{['tab:NYU_results_res']} and \ref{['tab:celebA_results']}.
• Figure 3: For the celebA dataset, the figure shows the RGB(input) image, the segmentation ground truth mask, and the predicted mask, along with the activation maps for the segmentation classes, as well as all the classification tasks for the dense MTL and LOMT models. Here act. map stands for the class activation mapsjacobgilpytorchcam.
• Figure 4: The graph represents the progression of parameter sparsity during training for different training data sizes for edge detection of the NYU-v2 dataset. For the rest of the task, a similar pattern can be observed.
• Figure 5: The plots illustrate the impact of the regularization parameter $\lambda$ on parameter sparsity (in %) and the single-task performance metrics across four tasks within the NYU-v2 dataset.