[Paper Review] Learning to Branch for Multi-Task Learning

Pengsheng Guo|arXiv (Cornell University)|Jun 2, 2020

Domain Adaptation and Few-Shot Learning50 references42 citations

TL;DR

The paper introduces LearnToBranch, an end-to-end trainable method that automatically learns where to share or branch within a network for multi-task learning, using a differentiable tree-structured topology guided by gumbel-softmax sampling. It demonstrates improved task grouping and performance on synthetic data, CelebA, and Taskonomy.

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.

Motivation & Objective

• Automatically search for optimal network sharing and branching structures for multiple tasks without hand-crafted task relatedness assumptions.
• Stitch together a tree-structured topology that minimizes the multi-task loss via differentiable branching.
• Provide an end-to-end training framework that jointly optimizes architecture and weights.
• Demonstrate effectiveness on synthetic data, CelebA, and Taskonomy datasets.

Proposed method

• Represent the network as a DAG with branching blocks where each child node samples a parent connection via a learnable categorical distribution p_theta.
• Use gumbel-softmax to make discrete branching decisions differentiable during training and gradually anneal temperature to converge to a hard tree.
• Define the branching operation x_j^{l+1} = E_{d_j ~ p_theta_j}[d_j · Y^l], enabling end-to-end optimization over topology and weights.
• Train by alternating sampling network configurations from the design space and updating both architecture probabilities and network weights via backpropagation.
• After training, select the final architecture by argmax over theta without noise and retrain from scratch for final performance.
• Stack branching blocks to build deeper tree-structured multi-task networks while keeping the number of tasks aligned with leaves.

Experimental results

Research questions

• RQ1 Can a differentiable, tree-structured branching mechanism automatically determine which layers to share or split for multiple tasks?
• RQ2 Does end-to-end optimization of architecture and weights yield better multi-task performance than hand-crafted or static topologies?
• RQ3 Can task grouping emerge naturally from back-propagation signals without prior task relatedness information?
• RQ4 How effective is the learned topology on synthetic, CelebA, and Taskonomy datasets?

Key findings

• The method learns task-grouping structures that cluster related tasks and branch out when tasks diverge, without human priors.
• LearnToBranch achieves competitive or superior accuracy with fewer parameters compared to several baselines on CelebA.
• On Taskonomy, LearnToBranch outperforms AdaShare and other baselines across five tasks (segmentation, normal, depth, keypoint, edge) with a smaller parameter count.
• The learned architectures show consistent sharing patterns across runs, indicating stable automatic task grouping.
• Training includes a topology search phase (time: hours) followed by retraining the final architecture from scratch, achieving strong performance with end-to-end optimization.