引言
多任务学习（MTL）旨在让模型经过一个训练过程中，让模型具备处理多种任务的能力。简单来说，MTL能够在不同任务之间共享信息，有效提高模型的泛化能力和数据效率。多任务学习的关键主要体现在参数共享、联合损失函数、权重调整等方面。
当前将大模型作为骨干模型，进行多任务学习，是高效利用大模型能力一种方法。但是现有的MTL策略在LLMs微调过程中，「会存在两个问题：1）计算资源要求高；2）无法保证多任务的同时收敛」。为此，今天给大家分享的这篇文章，为了解决这两个问题，「提出了一种新型MTL方法：CoBa」，即在训练过程CoBa可以动态地调整任务权重，促进各任务收敛平衡，降低了计算资源要求；结果表明：该方法可以让LLMs的性能最高提升13%。
代码：https://github.com/codefuse-ai/MFTCoder
本文源于蚂蚁集团AI native团队自研项目，现已被「EMNLP2024」接受为Oral Presentation。
背景介绍
近年来，大语言模型（LLMs）因其强大的性能成为研究热点。LLMs通过预训练积累大量通用知识，并在微调阶段针对具体任务进一步优化。但由于每个任务需要单独微调，部署变得复杂，且模型规模大，资源消耗高。
为了解决这些问题，多任务学习（MTL）应运而生。MTL通过一个模型支持多个任务，显著节省计算和存储资源，同时还能提升任务的整体性能和泛化能力。LLMs大规模的参数空间使其具有很强的适应性，能够在支持多任务的同时实现较好的效果，比如OpenAI的GPT-3.5和GPT-4就展示了这种能力。
为了在LLMs上实现高效的MTL，需要满足两个关键要求：「减少额外计算成本」，以及「确保所有任务可以同时收敛」，并优化至共享的最优点。虽然在小模型上有效，但在处理LLMs时计算成本高且整合复杂。例如，GradNorm、FAMO和MetaWeighting等方法的计算成本随着任务数量显著上升（如下表所示），此外，像Muppet和ExT5这样的多任务模型未能有效解决任务收敛不均衡，模型的整体效果未能达到最佳。
表1：现有MTL方法的时间复杂度对比，F和B分别表示一个批数据前向和后向传播的计算开销，K是任务个数，是模型参数量，是常量，其中 ，表示权重是基于验证损失计算得来。
CoBa介绍
CoBa的核心策略是基于验证数据中的收敛趋势，动态调整每个任务的训练损失权重。
作者首先假设模型参数量为，任务数量为，定义任务在第个迭代中的损失函数为，训练损失的权重为，那么MTL优化的整体训练损失可以构造为：
「下图1：CoBa在真实训练场景下的任务权重计算过程」
「准则一：平衡收敛速度」：当所有任务的验证损失在下降时，减小收敛较快任务的权重，增大收敛较慢任务的权重。
「准则二：避免任务发散」：对接近发散（可能过拟合）的任务减小权重，对仍然稳定收敛的任务增大权重。
为实现这两项准则，CoBa引入了相对收敛分数（RCS）和绝对收敛分数（ACS），其中RCS旨在评估任务之间的相对收敛速度，ACS旨在衡量每个任务当前的收敛速率与其历史速率的对比。同时作者还引入一个发散因子（DF），确保在所有任务都在收敛时，RCS对权重的影响占主导地位，而当某一特定任务开始发散时，ACS将起主导地位。
收敛斜率
不同任务的收敛速度可以通过观察验证损失曲线的斜率来直观地衡量。在图1(a)展示的某种情况中，绿色曲线任务B在5900步前的收敛速度都快于红色曲线任务C。如图1(b)所示，这一现象体现在了任务B的曲线斜率更陡，从而表示着任务B的收敛斜率绝对值高于任务C。作者在选定的迭代范围内拟合一个由定义的线性模型，其中的就是该段迭代内的收敛斜率的估计。为应对初始收敛斜率的不准确性，作者方法引入了warm-up参数。warm-up期间各任务权重被统一设置为以确保平衡的起点。warm-up期结束后，将根据滑动窗口内的收敛斜率更新权重。
相对收敛分数（RCS）
RCS的目标是根据任务的收敛速度动态分配权重，即对收敛较快的任务分配较小的权重，对收敛较慢的任务分配较大的权重：
其中，表示对任务的维度进行softmax操作。首先归一化收敛斜率，即 ，避免受到任务损失本身尺度的影响。考虑到当增加时，该归一化值趋于零，我们通过乘以来进行补偿。最后采用softmax计算出各个任务的RCS值。
在实际训练场景中，如图1(a)所示，可以观察到任务B（蓝色曲线）的收敛速度最慢，其RCS值是最高的（图1(c)）。此外，任务C（红色曲线）在第8000步前的收敛速度慢于任务A（绿色曲线），任务C的RCS高于任务A。
绝对收敛分数（ACS）
仅凭RCS无法同时满足前文提到的两个多任务学习的需求，因为任务B虽然已经开始发散，但仍然获得了最高的RCS（如图1(a)和1(c)所示）这可能使整体模型性能下降。因此ACS是必要的，通过为即将发散任务分配较小的权重，而对仍在收敛的任务赋予更大的权重，降低此类风险。ACS对数学表达式为：
ACS的计算方式与RCS不同，先沿着迭代维度进行归一化，随后才在任务维度上进行softmax操作，且在归一化过程中仅考虑任务自身的历史表现。此外，ACS式子中的反向归一化能让持续收敛的任务获得一个比开始发散任务更高的值。最后在任务间应用softmax分配ACS权重，实现增强收敛任务和抑制发散任务效果。
如图1(a)，1(b)和1(d)所示，任务B由于最早开始发散而获得最小的ACS。此外，由于任务A和任务C在5900步前后的收敛斜率相对大小发生变化，其ACS的相对大小也发生了转变。
发散因子（DF）和最终权重
在训练初期，任务通常表现为收敛模式，此时RCS应在权重分配占主导。随着训练进行，部分任务可能发散，此时ACS应优先影响权重分配。为实现从RCS主导到ACS主导的无缝过渡，我们引入发散因子（DF）来监控发散趋势。则最终的权重向量可以表示为：
提取每次迭代时所有任务中（考虑符号后）的最大收敛斜率，再通过以下公式进行量化：
在softmax函数中，我们将分子乘以当前步数，以防止随着增加而急剧下降，温度参数确保了softmax输出分量的区分度。
CoBa与现有方法的差别
当前的收敛平衡方法主要通过减缓快任务和加速慢任务来实现，作者提出的RCS同样有效地实现了这一目标。然而，当某些任务开始发散时，这种方式会产生反作用，影响收敛平衡。为此，CoBa引入了ACS和DF，有助于抑制某些任务过早的发散趋势，从而确保多任务训练的整体稳定性。
实验结果分析
作者在「四个不同数据集」上评估了CoBa的表现：代码补全（CC）数据集、代码相关任务（CRT）数据集、XTREME-UP数据集以及多领域问答数据集；除此之外，作者将「CoBa与八个最先进的基准模型（SOTA）进行对比」：单任务学习（STL），Uniform，GradNorm，LBTW和FAMO。值得注意的是，后三种方法最初是基于训练损失设计的。为了提高泛化能力，作者已将这些方法调整为关注验证损失，分别记作LBTW、GradNorm和FAMO*。
「下表2为：Phi-1.5-1.3B模型在CC数据集上的实验结果」
「下表4为，CRT数据集上Code Completion任务的实验结果」
「下图6为XTREME-UP数据集的3任务设定实验结果」
