Learning to Branch for Multi-Task Learning

TL;DR

This paper introduces an automated, end-to-end trainable method for learning optimal branching structures in multi-task neural networks, improving task-specific performance and reducing negative transfer.

Contribution

It proposes a novel differentiable tree-structured design space using gumbel-softmax sampling to learn where to share or branch within a network for multiple tasks.

Findings

Effective network topologies learned on synthetic, CelebA, and Taskonomy datasets.

Outperforms prior methods relying on manual or heuristic branching strategies.

Demonstrates improved multi-task learning performance and reduced negative transfer.

Abstract

Training multiple tasks jointly in one deep network yields reduced latency during inference and better performance over the single-task counterpart by sharing certain layers of a network. However, over-sharing a network could erroneously enforce over-generalization, causing negative knowledge transfer across tasks. Prior works rely on human intuition or pre-computed task relatedness scores for ad hoc branching structures. They provide sub-optimal end results and often require huge efforts for the trial-and-error process. In this work, we present an automated multi-task learning algorithm that learns where to share or branch within a network, designing an effective network topology that is directly optimized for multiple objectives across tasks. Specifically, we propose a novel tree-structured design space that casts a tree branching operation as a gumbel-softmax sampling procedure. This enables differentiable network splitting that is end-to-end trainable. We validate the proposed method on controlled synthetic data, CelebA, and Taskonomy.