End-to-End Multi-Task Learning with Attention

Abstract

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.

arXiv commentAccepted at Computer Vision and Pattern Recognition (CVPR), 2019

References in corpus (11)

• Adam: A Method for Stochastic Optimization
• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Distilling the Knowledge in a Neural Network
• Wide Residual Networks
• Theoretical Models of Learning to Learn
• The Cityscapes Dataset for Semantic Urban Scene Understanding
• Learning multiple visual domains with residual adapters
• Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics
• Residual Attention Network for Image Classification
• R-CNNs for Pose Estimation and Action Detection
• Piggyback: Adapting a Single Network to Multiple Tasks by Learning to Mask Weights

Cited by in corpus (6)

• Gradient Surgery for Multi-Task Learning
• Adaptive Scheduling for Multi-Task Learning
• AuxNet: Auxiliary tasks enhanced Semantic Segmentation for Automated Driving
• Feature Partitioning for Efficient Multi-Task Architectures
• JEPOO: Highly Accurate Joint Estimation of Pitch, Onset and Offset for Music Information Retrieval
• Exceeding the Limits of Visual-Linguistic Multi-Task Learning