发布人：Google Research Brain 团队研究工程师 Christopher Fifty
许多机器学习 (ML) 模型通常专注于一次学习一项任务。例如，语言模型根据先前单词的上下文，预测接下来可能会出现哪个单词，又如目标检测模型可以识别图像中存在的对象。然而，在某些情况下，同时从多项相关任务中学习也许能够提升建模性能。这一点在多任务学习 (Multi-task learning) 领域中得到了解决，作为 ML 的一个子领域，多任务学习是在同一个模型中同时训练多个目标。
举一个现实生活中的例子：乒乓球运动。在打乒乓球的时候，通常比较有利的做法是，同时判断乒乓球的距离、旋转和即将出现的轨迹，以此来调整姿势并计划挥拍动作。虽然这些任务中的每一项都是独特的（预测乒乓球的旋转与预测其位置有着本质的区别），但提高对球的位置和旋转的推理能力，可能会帮助您更好地预测其轨迹，反之亦然。以此类推，在深度学习领域，训练一个模型来预测三个相关任务（即，乒乓球的位置、旋转和轨迹），会比只预测单一目标的模型性能更加出众。
左图：三个单一任务网络，每个都使用相同的输入，分别预测乒乓球的旋转、距离或轨迹。右图：一个多任务网络，同时预测旋转、距离和轨迹
在 NeurIPS 2021 上，一场关于 “高效识别多任务学习中的任务分组 (Efficiently Identifying Task Groupings in Multi-Task Learning)”的焦点演讲中，我们描述了一种叫做任务相似性分组 (TAG) 的方法，该方法用于确定应该在多任务神经网络中一起训练哪些任务。我们的方案试图将一组任务分成更小的子集，从而让所有任务的性能达到最大化。为了实现这一目标，该方案在单一的多任务模型中一起训练所有的任务，并衡量一个任务对模型参数的 梯度更新 会给网络中其他任务造成何种程度的损失。我们用 任务间相似性 来表示此数值。实验结果表明，选择任务间相似性更高的任务分组能很大程度影响模型的整体性能。
哪些任务应该一起训练？
理想情况下，多任务学习模型会利用它在其中一个任务训练中学习到的信息，来减少在训练网络其他任务中的损失。这种信息转移的结果是，单一的模型不仅可以进行多种预测，而且用为每个任务训练不同的模型相比，这些预测的准确性也可能有所提高。另一方面，为多项任务训练单一的模型可能会导致模型容量竞争，并严重降低性能。当任务不相关时，后一种情况经常发生。回到乒乓球的例子，想象一下，试图预测乒乓球的位置、旋转和轨迹，同时背诵斐波那契数列。这并不是一个有趣的设想，而且很可能不利于乒乓球运动员的进步。
可以用一个直接的方法来选择一个模型应该训练的任务子集，即对一组任务的所有可能的多任务网络组合 (all possible combinations) 进行穷举式搜索。然而，这种搜索的相关费用可能会令人望而却步，在任务繁多的情况下更是如此，因为可能的组合数量会随着任务集的数量呈指数级增长。模型所适用的任务集在其整个生命周期中可能会发生变化，因此情况会更加复杂。随着所有任务集合中任务的添加或删除，需要重复进行这种昂贵的分析以确定新的分组。此外，随着模型的规模和复杂性不断增加，即使是只评估可能的多任务网络子集的近似任务分组算法，也可能会造成过高的成本，耗费大量的评估时间。
构建任务相似性分组
在研究这一挑战时，我们从元学习 (Meta learning) 中获得了灵感，元学习是机器学习的一个领域，它训练的神经网络可以快速适应一个先前从未见过的新任务。MAML 是经典的元学习算法之一，该算法对一个任务集合的模型参数进行梯度更新，然后更新其原始参数集，可充分降低该集合中按照更新的参数值进行计算的任务子集的损失。使用这种方法时，MAML 训练模型来学习表征，这些表征不会将其当前的权重集损失最小化，而是让经过一个或多个步骤训练后的权重降至最低。因此，MAML 训练模型的参数，让其有能力快速适应先前从未见过的任务，因为它的优化针对未来而不是现在。
TAG 采用了一个类似的机制来深入了解多任务神经网络的训练动态。特别之处在于，它只针对单一任务更新模型的参数，并查看这一变化会如何影响多任务神经网络中的其他任务，然后撤消这一更新。之后对其他各个任务重复这一过程，以收集网络中每个任务与其他任务的互动信息。然后继续正常的训练，更新网络中每个任务的模型共享参数。
收集这些统计数据，并在整个训练过程中观察它们的动态，可以发现某些任务始终表现出互益的关系，而有些任务则是相互对立的。网络选择算法可以利用这些数据，以便将任务组合在一起，最大化任务间相似性，而这取决于实践者选择在推理过程中使用的多任务网络数量。
TAG 概览
首先，在计算任务间相似性时，在同一网络中一起训练这些任务。然后，网络选择算法找到任务分组，充分提升任务间相似性。最后，对产生的多任务网络进行训练和部署
结果
我们的实验结果表明，TAG 可以选择非常强的任务分组。在 CelebA 和 Taskonomy 数据集上，TAG 比之前最先进的技术更具竞争力，其运行速度分别提升了 32 倍和 11.5 倍。在 Taskonomy 数据集上，这种速度能够帮助 Tesla V100 GPU 更快地寻找任务分组，让其花费的时间减少 2,008 个小时。
结论
TAG 方法行之有效，可以确定应该在运行单一训练的过程中一起训练哪些任务。该方法着眼于任务在训练中的相互作用，特别是在对一项任务进行训练时更新模型参数会对网络中其他任务的损失值产生影响。我们发现，选择任务组以最大化该分数与模型性能密切相关。
致谢
感谢 Ehsan Amid、Zhe Zhao、Tianhe Yu、Rohan Anil 和 Chelsea Finn 对本研究的重要贡献。还要感谢 Tom Small 设计的动画，以及整个 Google Research 团队打造的充满合作氛围且积极向上的研究环境。
原文：Deciding Which Tasks Should Train Together in Multi-Task Neural Networks
中文：TensorFlow 公众号