Learning to Branch for Multi-Task Learning
This work addresses the challenge of optimizing network sharing for multi-task learning, which is incremental as it builds on existing branching methods by introducing automated, differentiable optimization.
The paper tackles the problem of negative knowledge transfer in multi-task learning by automating the design of network branching structures, achieving improved performance over prior methods on datasets like CelebA and Taskonomy.
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.