Source-linked AI summary
End-to-End Multi-Task Learning with Attention
Shikun Liu, Edward Johns, Andrew J. Davison
TL;DR
Multi-task learning needs to share useful features without losing task-specific representations, while also balancing task losses. MTAN uses a shared global feature pool with task-specific soft attention masks, and experiments report strong performance, parameter efficiency, and robustness to loss weighting. The paper also assumes that an identity attention mask should be no worse than a shared network that separates tasks only at the end.
Problem
Multi-task learning must express both task-shared and task-specific features while balancing task losses without letting easier tasks dominate.
Method
MTAN combines a single shared network and global feature pool with task-specific soft attention masks learned end-to-end.
Results
MTAN outperforms or is competitive with other methods across dense-prediction and image-classification tasks, while showing robustness to loss-weighting schemes and parameter efficiency.
Takeaways & Limitations
Task-specific attention enables MTAN to share weights through attention masks while retaining task-specific features across multiple evaluated tasks.
Takeaways & Limitations
The architecture assumes that an identity attention mask should perform no worse than a shared network splitting into individual tasks only at the end.
Abstract
from arXiv · showhide
We propose a novel multi-task learning architecture, which allows learning of task-specific feature-level attention. Our design, the Multi-Task Attention Network (MTAN), consists of a single shared network containing a global feature pool, together with a soft-attention module for each task. These modules allow for learning of task-specific features from the global features, whilst simultaneously allowing for features to be shared across different tasks. The architecture can be trained end-to-end and can be built upon any feed-forward neural network, is simple to implement, and is parameter efficient. We evaluate our approach on a variety of datasets, across both image-to-image predictions and image classification tasks. We show that our architecture is state-of-the-art in multi-task learning compared to existing methods, and is also less sensitive to various weighting schemes in the multi-task loss function. Code is available at https://github.com/lorenmt/mtan.
1. Introduction
MTAN addresses multi-task learning’s challenges of sharing features appropriately and balancing task losses. It combines a shared feature pool with task-specific attention and performs strongly across dense-prediction and classification evaluations.
- Multi-task networks can improve memory, inference speed, and data efficiency by sharing informative visual features across related tasks.
- MTL must learn both generalisable task-shared features and task-tailored features to avoid over-fitting and under-fitting.
- Loss weighting should give tasks equal importance without allowing easier tasks to dominate, but manual tuning is tedious.
- MTAN cohesively learns task-shared and task-specific features while gaining robustness to the choice of loss weighting scheme.
- MTAN uses a single shared network and task-specific soft attention masks that select relevant shared features end-to-end.
- MTAN outperforms or competes with existing methods across image-to-image and image-classification tasks while using parameters efficiently.
2. Related Work
Prior multi-task vision work explores architectures for sharing features and methods for balancing task contributions. These approaches include cross-task connections, shared representations, progressive transfer, and adaptive loss weighting.
- Multi-task learning in computer vision covers related tasks including classification across domains, pose and action recognition, and dense prediction.
- Cross-Stitch Networks use one feed-forward network per task with cross-stitch units that share features across tasks.
- Other architectures use regularised combinations from one shared network, image pyramids with task-specific layers, or incrementally trained networks for knowledge transfer.
- Prior analyses report that different task combinations may benefit from different amounts of feature sharing and loss weighting.
- Weight uncertainty and GradNorm are examples of methods that adapt task weighting or gradient norms during multi-task training.
3. Multi-Task Attention Network
MTAN combines a shared network with K task-specific attention networks that apply learned soft masks to shared features. Its objective combines task losses, and its attention modules support dense prediction and classification settings.
- Architecture Design: MTAN consists of a shared network and K task-specific attention networks linked through attention modules.
- Task Specific Attention Module: Each attention module applies a soft mask to a shared-network layer, selecting task-specific features.
- Task Specific Attention Module: Task-specific features are computed by element-wise multiplication of attention masks and shared features.
- Task Specific Attention Module: The first encoder attention module uses shared features, while later modules concatenate shared features with task-specific features from the previous layer.
- Task Specific Attention Module: Sigmoid attention masks lie in [0, 1] and are learned with back-propagation.
- The Model Objective: The total objective is a weighted sum of task-specific losses, L_tot(X, Y_1:K) = Σ_i λ_iL_i(X, Y_i).
- Task Formulations: Image-to-image evaluation uses semantic segmentation, depth estimation, and surface-normal prediction, while classification treats each dataset as a task.
4. Experiments
The experiments evaluate MTAN on image-to-image regression and image-classification settings, covering one-to-many and many-to-many task predictions.
- MTAN is evaluated on one-to-many image-to-image regression and many-to-many image-classification tasks.
4.1. Image-to-Image Prediction (One-to-Many)
The image-to-image evaluation tests MTAN across CityScapes and NYUv2, comparing architectures and loss-weighting schemes while varying task complexity. MTAN matches Dense with fewer parameters on CityScapes, outperforms baselines on NYUv2, remains robust to weighting choices, and benefits as task complexity increases.
- Architectures and baselines: The comparison includes single-task, split, dense, and Cross-Stitch baselines, with baselines designed to have at least as many parameters as MTAN.This controls for whether performance differences arise from attention modules rather than simply from larger networks.
- Results: MTAN performs similarly to Dense on CityScapes with less than half the parameters and outperforms all baselines on NYUv2 across weighting methods and tasks.The NYUv2 comparison covers the more challenging three-task setting.
- Results: MTAN maintains high performance across loss-weighting schemes and is more robust to weighting choices than Cross-Stitch.The paper attributes this robustness to attention masks that automatically learn which shared features to use for each task, reducing the need for manual loss-weight tuning.
- Effect of task complexity: As semantic task complexity increases, multi-task methods improve relative to vanilla single-task learning, while MTAN's relative gain increases at a greater rate.STAN performs best in the simplest 2-class setup, whereas multi-task feature sharing becomes more advantageous with greater task complexity.
- Attention masks as feature selectors: The learned attention masks differ between semantics and depth, with higher-contrast depth masks indicating greater use of task-specific features for depth estimation.The masks act as feature selectors over the shared representation.
4.2. Visual Decathlon Challenge (Many-to-Many)
The Visual Decathlon Challenge evaluates MTAN on 10 image-classification tasks using per-task accuracies and a cumulative score. MTAN surpasses most baselines and remains competitive with state-of-the-art methods without several specialized regularization strategies.
- Evaluation setting: The challenge comprises 10 individual image-classification tasks and reports per-task accuracies plus a cumulative score capped at 10,000.The cumulative score assigns up to 1,000 points per task.
- Results: MTAN surpasses most baselines and is competitive with current state-of-the-art performance on the online test set.The evaluation uses MTAN with a Wide Residual Network backbone across all 10 classification tasks.
- Results: MTAN achieves this performance without complicated regularization strategies such as DropOut, dataset regrouping by size, or adaptive dataset-specific weight decay.
5. Conclusions
MTAN combines a global feature pool with task-specific attention modules to learn shared and task-specific features end to end. Across dense prediction and image-classification benchmarks, it matches or exceeds competing methods, remains robust to loss weighting, and is parameter efficient.
- Conclusion: MTAN uses a global feature pool and task-specific attention modules to learn task-shared and task-specific features end to end.
- Conclusion: Experiments on NYUv2, CityScapes, and the Visual Decathlon Challenge show that MTAN outperforms or is competitive with other methods.The evaluation covers multiple dense-prediction and image-classification tasks.
- Conclusion: MTAN is robust to the particular task-weighting schemes used in the multi-task loss function.
- Conclusion: Attention masks allow MTAN to share weights while achieving state-of-the-art performance with high parameter efficiency.