[论文解读] End-to-End Multi-Task Learning with Attention

Shikun Liu, Edward Johns+1 more|arXiv (Cornell University)|Mar 28, 2018

Advanced Neural Network Applications参考文献 30被引用 20

一句话总结

本文提出多任务注意力网络（MTAN），一种参数高效、端到端的多任务学习架构，通过任务特定的软注意力模块从共享全局特征池中动态选择特征。MTAN 在图像分割、深度估计和图像分类任务中均达到最先进性能，同时对损失加权方案具有鲁棒性，并且相比先前方法参数量更少。

ABSTRACT

We propose a novel multi-task learning architecture, which allows learning of task-specific feature-level attention. Our design, the Multi-Task Attention Network (MTAN), consists of a single shared network containing a global feature pool, together with a soft-attention module for each task. These modules allow for learning of task-specific features from the global features, whilst simultaneously allowing for features to be shared across different tasks. The architecture can be trained end-to-end and can be built upon any feed-forward neural network, is simple to implement, and is parameter efficient. We evaluate our approach on a variety of datasets, across both image-to-image predictions and image classification tasks. We show that our architecture is state-of-the-art in multi-task learning compared to existing methods, and is also less sensitive to various weighting schemes in the multi-task loss function. Code is available at https://github.com/lorenmt/mtan.

研究动机与目标

• 解决多任务学习中有效特征共享与损失平衡的双重挑战。
• 设计一种统一架构，自动学习任务共享与任务特定特征，无需人工干预。
• 提升多任务网络中的参数效率与可扩展性，尤其在任务数量增加时。
• 降低对损失权重超参数调优的敏感性，此类调优常导致多任务训练的不稳定性。
• 在多样化的多任务基准上实现最先进性能，涵盖密集预测与图像分类任务。

提出的方法

• 该架构使用单一共享主干网络从输入数据生成全局特征池。
• 针对每个任务，在每个卷积块上应用软注意力模块，根据任务相关性重新加权共享特征。
• 注意力掩码是可微的，并端到端训练，实现从共享表示中自动选择任务特定特征。
• 该方法兼容任意前馈神经网络（如SegNet或Wide ResNet），支持灵活的主干网络集成。
• 提出一种新颖的动态加权平均（DWA）损失加权方案，根据每个任务损失的变化率自适应调整任务权重。
• 网络使用标准优化方法端到端训练，无需任务特定头分离或复杂正则化。

实验结果

研究问题

• RQ1 多任务学习架构是否能自动学习共享与任务特定特征，而无需显式架构分离？
• RQ2 与固定特征共享相比，基于注意力的特征选择在性能与鲁棒性方面有何提升？
• RQ3 所提方法在多任务训练中在多大程度上降低了对损失权重超参数的敏感性？
• RQ4 该架构在显著提升参数效率的同时，能否保持高性能，相比现有多任务网络？
• RQ5 该方法是否能在包括密集预测与图像分类在内的多样化任务上，在基准数据集上实现良好泛化？

主要发现

• MTAN 在CityScapes数据集上，于语义分割、深度估计与表面法线预测任务中均达到最先进性能，参数量仅2941（比单任务基线小2.9倍）。
• 在Visual Decathlon挑战中，MTAN 每项任务的累积得分为96.88 / 1000（占最大值的96.88%），优于多数基线，且在无需复杂正则化的情况下达到最先进性能。
• 随着任务复杂度增加，该方法表现出更大的性能增益，在多任务设置中优于单任务注意力网络（STAN），尤其在复杂任务上。
• 注意力掩码的可视化结果表明任务特定特征选择有效，其中深度任务显示出更高对比度的掩码，表明对任务特定特征的更强依赖。
• MTAN 对多种损失加权方案具有鲁棒性，包括所提出的动态加权平均（DWA）方案，后者提升了训练稳定性和收敛性。
• 该架构具有极高的参数效率，在Visual Decathlon上10项任务仅需2941个参数，相比具有显式任务特定分支的方法显著减小了模型规模。