Understanding and Improving Information Transfer in Multi-Task Learning

TL;DR

This paper analyzes how task data alignment affects multi-task learning performance, providing theoretical insights and practical methods to improve transfer and robustness across NLP and image tasks.

Contribution

It offers a theoretical framework for understanding task alignment effects and proposes alignment and reweighting techniques that enhance multi-task learning performance.

Findings

Alignment improves transfer learning outcomes.

Misalignment can cause negative transfer.

Proposed methods yield significant performance gains.

Abstract

We investigate multi-task learning approaches that use a shared feature representation for all tasks. To better understand the transfer of task information, we study an architecture with a shared module for all tasks and a separate output module for each task. We study the theory of this setting on linear and ReLU-activated models. Our key observation is that whether or not tasks' data are well-aligned can significantly affect the performance of multi-task learning. We show that misalignment between task data can cause negative transfer (or hurt performance) and provide sufficient conditions for positive transfer. Inspired by the theoretical insights, we show that aligning tasks' embedding layers leads to performance gains for multi-task training and transfer learning on the GLUE benchmark and sentiment analysis tasks; for example, we obtain a 2.35% GLUE score average improvement on 5 GLUE tasks over BERT-LARGE using our alignment method. We also design an SVD-based task reweighting scheme and show that it improves the robustness of multi-task training on a multi-label image dataset.