[论文解读] Learning to Branch for Multi-Task Learning

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

Domain Adaptation and Few-Shot Learning参考文献 50被引用 42

一句话总结

本文提出 LearnToBranch，一种端到端可训练的方法，能够在多任务学习中自动学习在网络中何处共享或分支，使用由 gumbel-softmax 采样引导的可微分树状拓扑。它在人造数据、CelebA 和 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.

研究动机与目标

• 在没有人工设计的任务相关性假设的前提下，自动搜索多任务的最优网络共享与分支结构。
• 通过可微分分支将树状拓扑拼接成一个以最小化多任务损失为目标的结构。
• 提供一个端到端的训练框架，联合优化结构与权重。
• 在合成数据、CelebA 和 Taskonomy 数据集上证明其有效性。

提出的方法

• 将网络表示为带有分支块的有向无环图（DAG），其中每个子节点通过一个可学习的分类分布 p_theta 采样父连接。
• 使用 gumbel-softmax 使离散分支决策在训练过程中可微，并逐渐退火温度以收敛到硬树。
• 定义分支操作 x_j^{l+1} = E_{d_j ~ p_theta_j}[d_j · Y^l]，实现对拓扑和权重的端到端优化。
• 通过交替从设计空间抽样网络配置并通过反向传播同时更新架构概率和网络权重进行训练。
• 训练结束后，通过对 theta 取最大值（无噪声）来选择最终架构，并从头开始重新训练以获得最终性能。
• 堆叠分支块以构建更深的树状多任务网络，同时使任务数量与叶节点对齐。

实验结果

研究问题

• RQ1 一个可微分的树状分支机制是否能够自动确定多任务应在何层共享或分支？
• RQ2 端到端优化架构与权重是否能带来比人工设计或静态拓扑更好的多任务性能？
• RQ3 在没有先前任务相关性信息的情况下，任务分组是否能自然从反向传播信号中出现？
• RQ4 学习到的拓扑在合成数据、CelebA 和 Taskonomy 数据集上的效果如何？

主要发现

• 该方法学习到将相关任务聚类的任务分组结构，并在任务出现分歧时分支扩展，且无需人类先验。
• 在 CelebA 上，LearnToBranch 使用更少的参数达到具有竞争力或更高的准确性，相比若干基线。
• 在 Taskonomy 上，LearnToBranch 在五个任务（分割、法线、深度、关键点、边缘）上优于 AdaShare 等基线，且参数量更少。
• 学习到的架构在多次运行中显示出一致的共享模式，表明自动任务分组稳定。
• 训练包含拓扑搜索阶段（耗时：hours）和从头重新训练最终架构，结合端到端优化实现强大性能。