多任务学习新篇章 | EMA-Net利用Cross-Task Affinity实现参数高效的高性能预测

简介: 多任务学习新篇章 | EMA-Net利用Cross-Task Affinity实现参数高效的高性能预测

多任务学习(MTL)因其能够同时预测多个任务,在实现每个任务性能更好且使用比单任务学习更少的每个任务模型参数的优势下而得到了广泛关注。最近,关注解码器的架构在利用其他相关任务特征优化任务预测方面取得了显著的改进,从而显著提高了多任务性能。然而,这些优化方法在参数高效的情况下无法同时捕捉局部、全局和跨任务特定的表示,以及跨任务模式。

在本文中,作者引入了高效的Multitask Affinity Learning Network(EMA-Net),它是一个轻量级框架,可以增强多任务网络的任务细化能力。EMA-Net巧妙地利用作者新提出的Cross-Task Affinity Learning(CTAL)模块来捕获局部、全局和跨任务交互。CTAL的关键创新在于它能够以最佳方式调整任务相似性矩阵,而无需担心信息损失,同时应用参数高效的组卷积。

作者的结果表明,作者使用比使用更少的模型参数实现了最先进的多任务学习性能,适用于基于CNN的解码器关注模型。

代码:https://github.com/Armanfard-Lab/EMA-Net

1 Introduction

现代AI研究正在迅速融入作者的日常生活。然而,大多数最先进模型庞大且依赖远程访问,因为它们无法轻松部署在边缘设备如移动电话、智能配件或可穿戴医疗设备上。对于许多应用,模型在本地运行且无需依赖网络的需求突显了在模型设计中平衡性能与参数效率的重要性。

近年来,多任务学习(MTL)作为一种参数高效的学范式,通常通过同时优化多个损失函数来学习多个任务。因此,使用单个网络意味着多个任务之间共享多个层或特征。在许多密集预测情况下,跨任务共享特征已被证明可以提高每个任务性能,同时使用较少的每个任务模型参数。这是利用相关任务之间的特定知识来提高泛化能力的产物。现代MTL方法的主要研究方向要么是优化策略,要么是深度多任务架构的设计。在深度架构设计方面,Vandenhende等人进一步将设计划分为解码器关注和编码器关注架构。正如其名所示,解码器关注的模型在解码器内部采用特征共享机制。

具体而言,解码器关注的架构引入机制来捕获任务之间的特征模式。任务之间的一种模式是分割边缘与深度值的不连续性对齐。这种细化过程也称为任务预测蒸馏。通过捕获这些任务间的关系,解码器关注的架构在MTL中始终实现最先进性能,并可被视为多任务架构设计的主导研究方向。

解码器关注的基础方法PAP-Net引入了“任务亲和度”的概念,它是一种给定任务特征对之间相似度的度量。具体而言,他们构建了亲和矩阵来存储给定任务中每个可能特征对之间的相似信息。使用亲和矩阵的好处是它们捕获了长程依赖性,并引入了很少的额外模型参数。

为了进行任务预测蒸馏,他们使用加权求和的方法将所有任务中的亲和矩阵组合在一起,并将这种相似信息扩散到每个初始任务预测的特征中。这种方法的一个问题是,学习一个每个亲和矩阵的权重表明特征空间中所有配对相似性模式同等重要。然而,作者认为跨任务模式是细微的,并且在特征空间中具有重要性。这种简单的跨任务机制可能是因为处理这些矩阵可以变得昂贵,尤其是在更大的特征尺度上。基于作者的假设,作者认为亲和表示(i)中存在大量未开发的潜力。

更近的一项工作,MTI-Net,声称在某个特征尺度上亲和度较高的任务不一定会在不同特征尺度上具有较高的亲和度。因此,他们模型多个尺度的任务交互以实现“多尺度多模态蒸馏”,从而显著提高了性能。然而,这种性能改进需要为多个尺度的深度监督增加额外的参数,这使得这种方法对简单数据集容易过拟合。因此,作者认为需要一种更参数高效的框架,仍然可以利用多尺度处理的优点(ii)。

InvPT和TaskPrompter继续了使用更多的模型容量来提高性能的趋势,转向基于视觉 Transformer (ViT)的架构。他们认为基于CNN的当前多任务注意力机制在跨任务模式建模方面有局限性。他们的多任务注意力方法可以捕获局部、全局和跨任务关系,但主要是因为它们在压缩特征上操作,这是由于使用了参数预算显著更高的特征提取器的结果。然而,使用轻量级CNN基础特征提取器来建模局部、全局和跨任务关系还尚未实现(iii)

尽管他们最近取得了成功,但当前解码器关注的方法尚未解决(i)、(ii)和(iii)这些问题。作者通过引入高效的Multitask Affinity Learning Network(EMA-Net)来解决这些问题,该网络使用作者新颖的Cross-Task Affinity Learning(CTAL)模块进行任务预测蒸馏的改进。CTAL通过精心重排和交错连接将亲和矩阵对齐,这使作者能够利用分组卷积实现模型参数的显著减少。

此外,由于作者在任务内部和跨任务中完全建模每对特征关系,作者可以使用这些分组卷积而不会担心信息损失。这有效地解决了问题(i),因为作者能够充分利用亲和矩阵表示的潜力。同时,作者还解决了问题(iii),因为作者使用轻量级CNN基础架构明确地建模所有任务内和跨任务内的每对相互作用。EMA-Net也扩展到多尺度框架,通过在多个尺度上应用深度监督来初始化任务预测。然而,与MTI-Net不同,作者从每个尺度上融合初始任务预测,然后在单一尺度上进行任务预测蒸馏。因此,EMA-Net框架在参数效率方面更优,并解决了问题(ii)。因此,作者可以总结作者的贡献如下:

  1. 在基于CNN架构中,一种处理亲和矩阵的参数高效方法,同时全面建模所有内任务和跨任务关系。这解决了(i)和(iii)问题。
  2. 一个轻量级多尺度框架,可以利用多尺度深度监督的优势,只需要一个尺度进行任务预测蒸馏。这解决了(ii)问题。
  3. 实现了其他解码器关注技术相比学可参数较少时显著的多任务性能改进。

3 EMA-Net

在文献中,有两种主要的将注意力图应用于基于CNN架构生成的图像特征的方法。

  1. 第一种方法涉及使用卷积块处理特征并获取具有相同形状的注意力Mask。注意力Mask经过激活(即Sigmoid)后,所有值设置为0和1之间。通过逐元素乘法将注意力图应用于特征,这在STL和MTL中应用。作者将这种称为逐元素乘法注意力(EM attention)。
  2. 第二种方法旨在通过计算特征的Gram矩阵(所有对列向量的内积)来明确建模特征之间的长程依赖性,然后通过矩阵乘法将相应的矩阵(即亲和矩阵)扩散到原始特征。作者将这种方法称为矩阵乘法注意力(MM attention)。MM attention 也在STL和MTL中使用。

在MTL中,PAD-Net的EM attention和PAP-Net的MM attention 实现了几乎相同的结果。然而,PAP-Net在模型参数方面使用较少,但需要更多的浮点运算(FLOPs)。他们的可比性能很可能归因于亲和矩阵明确捕获每个特征对之间的长程依赖性,并通过MM attention将其扩散到所有特征;而EM attention方法通过训练整个空间维度上的局部块卷积滤波器来隐式学习长程依赖性。

PAP-Net中使用简单的加权求和进行跨任务融合的MM attention能够实现与使用EM attention(通过元素乘法)几乎相同的性能,这是令人惊讶的。作者认为这些亲和表示(i)中存在大量未开发的潜力,因此作者受到启发,开发了一种参数高效的跨任务注意力机制来优化局部和全局交互的建模。

Intra-Task Modelling

如图1所示,对于给定的任务,作者遵循标准的程序生成亲和矩阵,。这涉及提取初始预测的特征,,将空间维度展开,对每个列进行L2归一化,并计算Gram矩阵(所有列向量的内积)。

在此过程中,中的每一行包含特征与中其他特征的余弦相似度。接下来,作者将 Reshape 为。这恢复了特征的原始空间维度,但现在给定二维位置的个通道包含特征与所有其他特征的余弦相似度,因此,在这种情况下,整个第一个通道对应于所有与的相似度,这也是在位置处对齐的。

这是在后续处理过程中保持空间一致性的有用属性。

Inter-Task Modelling

接下来,为了将 Reshape 的亲和矩阵进行融合,作者首先执行一个交错拼接操作。如图2所示,这涉及将每个(即和)的第一通道拼接,然后第二通道,以此类推,直到所有个通道。这给作者得到联合亲和矩阵,,其中是多任务学习(MTL)系统的任务数量。根据数据的空间维度,可以非常大,因此使用标准卷积处理它将非常昂贵。相反,作者旨在利用参数高效的分组卷积[18, 19, 20]。作者战略性地组织强烈地证明了使用分组卷积进行作者的多任务融合而无需担心丢失重要跨任务信息是合理的。这是因为每个分组中的个通道已经包含了给定特征与所有其他任务中的所有特征之间的任务交互。因此,不仅为每个任务融合所需的模型参数数量显著减少,而且还允许作者学习个专注于学习给定特征与所有其他任务中的所有特征之间关系的空间一致核。当作者考虑在上进行传统卷积时,作者需要个大小为的核,而EMA-Net只需要个大小为的核,其中是卷积滤波器的尺寸。这相当于减少了倍参数。例如,在考虑3个任务,的特征大小和的NYUv2数据集上,作者只使用了187K参数,而单标准卷积层需要1.29B参数。

在处理每个任务的共享后,作者得到个矩阵,并将每个矩阵 Reshape 为以获得每个。现在,在给定的中,每一行包含关于单个特征与所有其他任务中的所有特征之间关系的信息。由于作者使用核,其中,作者还可以嵌入跨任务空间交互模式。接下来,作者转置矩阵,以便在扩散过程后,包含所有与相关的信息的行存储在位置,这完全保持了整个注意力过程中的空间一致性。

Task-Specific Diffusion

随后的扩散过程涉及执行矩阵乘法以获得扩散后的特征,。

其中是的重新排列卷积投影。通过这个矩阵乘法,中的每个值都是中包含所有个特征的行向量与中包含特征与其他所有个特征之间的交叉任务亲和模式信息的列向量的点积的结果。当存在跨任务的高亲和模式时,相对于其他亲和模式较低的特征,的值将变得更大,类似于EM注意力中的缩放行为。如方程2所示,扩散后的特征然后与原始特征使用元素加法加标量加权参数进行混合,以获得精炼特征。这种混合确保精炼特征不会偏离原始特征太多。

Multiscale EMA-Net

所提出的EMA-Net架构有两种变体。第一种变体使用单个特征尺度(SS)进行初始预测,另一种使用多个特征尺度(MS)。后者如图3所示,用于与多尺度多模态蒸馏方法(如MTL-Net)竞争。然而,作者的模型更高效和可扩展,因为作者只需要一个单一的蒸馏模块(即CTAL),而不是每个尺度都有一个模块。作者通过在CTAL之前将每个尺度的初始预测特征与跨尺度融合(CSF)块结合来实现这一点。

对于跨尺度融合,作者遵循与[20]中相同的过程,涉及将所有特征上采样到1/4输入尺度,沿通道维度 ConCat 它们,并通过卷积块进行组合。对于在单个特征尺度上操作的其他方法,作者对共享编码器生成的多尺度特征执行相同的跨尺度融合,并只生成一个单一的初始预测集。

4 Experimental Setup

Datasets

作者在NYUv2和Cityscapes数据集上进行实验,这两个数据集都是多任务学习非常流行的数据集。NYUv2包含1449个密集标记的RGB-深度图像,用于室内场景。原始数据集中包含带有不完整深度值的图像,这些值在训练期间被屏蔽。与该数据集相关的任务包括13个标签的语义分割、深度估计和表面法向量预测。该数据集不包含表面法向量标签,因此作者使用了[1]中获得的伪地面法向量数据,其中与相应的深度图位置具有相同的不完整值。训练和验证集分别包含795和654张图像,图像分辨率分别为。

Cityscapes是一个更大的数据集,包含3475个户外城市街道场景,这些场景来自50个城市,并在几个月内进行了精细标注。从这些精细标注中,作者得到了2975个训练图像和500个验证图像。与该数据集相关的任务包括19个标签的语义分割和深度估计。使用的标签来自其官方文档,该文档将多个标签合并为一个空类,并指定了在训练期间应使用19个其他标签。图像分辨率分别为。

Tasks and Performance Metrics

语义分割涉及将图像中的每个像素分配一个类标签。在训练期间,目标是使预测的类标签与目标标签之间的深度交叉熵损失最小,对于所有像素。作者还评估作者的模型在均交点率(mIoU)和像素精度(absolute pixel accuracy)上的性能。然而,mIoU是一个更好的语义理解指标。

深度估计涉及预测每个像素的深度值。在训练期间,作者的目标是使预测的深度值与目标深度值之间的绝对误差(L1范数)最小化。作者还报告相对深度误差。

表面法向量预测涉及估计图像中物体的法线方向,这使得可以获取几何和结构场景信息。作者训练模型以最小化归一化预测值与目标值之间的元素点积。对于评估,作者考虑与之间的平均角度距离。作者还报告了预测值落在11.25、22.5和30.0度误差范围内的比例。

最后,MTL Gain是相对于单个任务学习 Baseline ,对方法的总体多任务改进的聚合度量,如方程3所示。

在这里, 表示指标  的较低值更有利,否则为0。作者将  视为百分比在作者的评估中使用。尽管作者在评估中使用了多个指标,但为了在计算  时确保每个任务都得到公平的权重,作者将每个任务中最好的指标选择为最能展示泛化性能的单个指标。

因此,在计算  时,作者将使用分割的 mIoU、深度的相对误差和表面法向量的平均角度距离。作者还展示了在附录中使用其他指标组合仍然可以实现优越的 MTL Gain。

在作者的结果中,表示更 Favorable 的指标用()表示,表示更不 Favorable 的指标用()表示。每个任务特定的性能指标的公式也可以在附录中找到。

Baselines

在多任务学习中,将作者的模型与传统的单任务学习(STL)和多任务学习(MTL) Baseline 进行比较是一种标准做法。STL Baseline 涉及为每个任务使用单个网络,其中每个网络使用与提出的模型相当的可比 Backbone 网络和输出头进行公平比较。

MTL Baseline 使用一个共享参数的硬多任务网络,其中所有任务共享 Backbone 层,然后将共享的特征表示传递到每个任务特定的输出头。为了使这些 Baseline 更具竞争力,并保持与最新状态的一致性,作者为其配备了高分辨率网络 Backbone (HRNet18),以生成多尺度特征,这些特征由特定尺度的输出头处理并汇总为最终输出。

因为作者提出了一种新的跨任务注意力学习机制,用于基于CNN架构,所以作者必须与该领域的当前最佳方法进行比较。因此,作者将与PAD-Net和PAP-Net进行评估,它们分别作为EM注意力和MM注意力的 Baseline 。此外,它们也作为作者的单尺度 Baseline。接下来,作者还将与MTI-Net进行比较,因为它是目前基于CNN解码器关注的模型的最新状态。这将作为作者的多尺度 Baseline 。

对于作者的模型和 Baseline 的所有实验,作者进行了3次,每次使用不同的种子。所有模型使用相同的3种子以保持一致性。所有表格中的结果包含每个模型所有3次实验中收敛的平均值。

Implementation Details

网络: EMA-Net和所有 Baseline 都配备了预训练的HRNet18多尺度特征提取器 Backbone 。EMA-Net和 Baseline 的单尺度变体将使用上述CSF过程融合输入特征。初始预测的输出头包括两个残差块[1] followed by an output convolution layer。用于任务预测蒸馏的初始预测是第二个残差块的输出。单尺度模型的最终输出头使用与初始预测的输出头相同的架构。

对于多尺度 Baseline ,作者将所有输出上采样到1/4输入尺度,按通道维度 ConCat ,并通过一个两层卷积块处理,以获得最终预测。除了PAP-Net,作者自行实现了除PAP-Net以外的所有 Baseline 网络的实现代码。

超参数: 作者使用Adam优化器训练作者的模型,权重衰减为。NYUv2和Cityscapes上的学习率分别为和。作者为每个模型进行了小的学习率搜索,以确保此配置对所有 Baseline 都有利。作者还使用余弦退火学习率调度器以实现平滑收敛。多尺度模型在Cityscapes上倾向于收敛较早,因此作者为它们使用余弦退火学习率调度器,以促进探索并避免局部最小值。

对于这两个数据集,作者使用批量大小8,混合因子(与PAP-Net相同),滤波器大小(对于所有模型)。和的值没有针对每个数据集进行调整,并且作者的模型对这些参数不太敏感(见附录)。作者在单个NVIDIA RTX A5000 GPU上分别训练200和75个epoch的NYUv2和Cityscapes数据集。

5 Results

Ablation Study

为了观察CTAL在作者的单尺度(SS)和多尺度(MS)配置中的效果,作者首先评估了没有CTAL或CSF块的EMA-Net基础架构的性能。因此,它只包含一个用于初始预测特征和相应任务特定解码器的投影层。接下来,作者在单尺度变体中添加CTAL,并在多尺度变体中进行跨尺度融合。

如表2所示,没有CTAL的EMA-Net在NYUv2上未超过STL Baseline ,在Cityscapes上实现了+6.47%的MTL Gain。添加CTAL后,作者可以看到与STL Baseline 相比,每个数据集的MTL Gain分别提高了+2.64%和+12.67%。最后,具有CSF和多个尺度深度监督的CTAL,作者在NYUv2和Cityscapes上分别实现了+4.76%和+14.85%的MTL Gain。总体而言,作者发现在作者的SS和MS配置中,在两个数据集上实现了巨大的MTL Gain。

Comparison to State-of-the-Art

表1显示了EMA-Net在单尺度(SS)和多尺度(MS)配置中与所有 Baseline 的性能。表格分为3部分,分别分离传统STL和MTL Baseline 、SS模型和MS模型。如作者所见,对于两个数据集,作者在SS和MS配置中实现了所有任务指标的显著提高。作者可以看到,作者的EMA-Net(SS)在NYUv2上即使没有从多个尺度获得深层监督,也能与MTI-Net竞争。

在Cityscapes上,作者也可以看到,MTI-Net在简单的2任务设置中,在较小的输入图像尺度上挣扎。MTI-Net在Cityscapes上的结果可能以前从未报告过,因为它们容易过拟合。尽管作者努力减轻过拟合(例如,空间dropout、warm restart调度器、数据增强、架构修改、超参数调优),但作者在表1中看到的性能无法实现更好的表现。然而,EMA-Net(MS),也采用多尺度初始预测的深层监督,并没有表现出相同的过拟合行为。事实上,作者在EMA-Net(SS)的基础上进一步提高了性能。

由于本研究的目标是轻量级基于CNN的方法,因此作者没有与该参数范围内的Transformer模型进行比较。然而,作者想指出,当配备可比的Transformer Backbone 时,当前的CNN Baseline 表现与InvPT和TaskPrompter [20]非常相似。

Resource Analysis

正如前面提到的,使用显著较少的模型参数的代价是引入额外的FLOPs。作者认为这种权衡是有利的,因为在各个 Level 上都有更多的优化机会来减少额外FLOPs的影响,而与模型参数相比。例如,作者可以算法地优化,提高硬件利用率,利用稀疏矩阵操作(如果适用),并调整特征尺度。

表3总结了每个模型在NYUv2上的资源使用情况,作者可以看到,通过仅将任务预测蒸馏的尺度降低,作者可以大大减少FLOPs的数量。作者还看到,作者可以在与其他方法相当的时间内将墙时钟时间降低或相匹配。作者观察到,在SS设置中,作者使用不到一半的参数,并且FLOPs更少,超过了PAD-Net。在MS设置中,作者使用更少的模型参数,并仍然实现更好的多任务性能,即使特征尺度较小。

总的来说,EMA-Net在效率方面取得了几个优势。这可以归因于作者如何组织特征以便更有效地使用参数。这还减少了过拟合的风险,如Cityscapes数据集的结果所示。总的来说,随着使用更多的模型容量来获得更好的性能的趋势持续下去,作者必须更加注意如何优化效率,以便这些模型可以在具有严格内存限制的实时环境中部署。

Qualitative Analysis

图4是作者在NYUv2验证集上的EMA-Net(MS)和STL Baseline 的预测可视化。如作者所见,作者的模型在分割图上产生了显著较少的伪影。

此外,作者的模型生成的深度和表面法向量图比STL更平滑和更细致。例如,作者可以看到作者的模型与STL相比,对橱柜等结构的扭曲较小,如第二张图像中的橱柜。此外,作者的模型在深度和法向量任务的未定义区域内 generalize 得更好,如第一张图像中的窗户。

6 Conclusion

作者提出了作者的EMA-Net多任务架构,其中包含新的CTAL模块用于任务预测蒸馏。作者的网络是第一个基于CNN架构,可以明确并充分地建模所有内任务和跨任务特征对关系的模型。令人惊讶的是,作者使用比竞争对手STL和MTL模型更少的模型参数实现了这一点,同时在高特征尺度上实现了更好的多任务性能。

作者还证明了,通过调整作者的特征尺度,作者可以显著减少引入的FLOPs数量,同时仍然在更高的特征尺度上优于竞争对手。正如作者所提到的,作者认为与减少模型容量相比,性能改进 justification了额外的FLOPs,尤其是考虑到与模型容量相比,优化机会更多。向前看,探索进一步减小FLOPs对墙时钟时间的影响的技术,例如稀疏化亲和矩阵,将是非常有趣的。

参考

[1].EMA-Net: Efficient Multitask Affinity Learning for Dense Scene Predictions.

相关文章
|
6月前
|
API C# 图形学
DotNetGuide新增C#/.NET/.NET Core充电站(让你学习不迷路)
DotNetGuide新增C#/.NET/.NET Core充电站(让你学习不迷路)
|
7月前
|
移动开发 网络协议 NoSQL
.NET Core WebSocket实现简易、高性能、集群即时通讯组件
.NET Core WebSocket实现简易、高性能、集群即时通讯组件
122 0
|
3月前
|
NoSQL 关系型数据库 MongoDB
【DotNetGuide】C#/.NET/.NET Core学习、工作、面试指南
【DotNetGuide】C#/.NET/.NET Core学习、工作、面试指南
107 0
|
3月前
|
开发框架 .NET API
C#/.NET/.NET Core推荐学习书籍(23年12月更新)
C#/.NET/.NET Core推荐学习书籍(23年12月更新)
234 0
|
4月前
|
程序员 数据库
VB.NET—Bug调试(参数话查询、附近语法错误)
VB.NET—Bug调试(参数话查询、附近语法错误)
25 0
|
19天前
|
开发框架 .NET API
C#/.NET/.NET Core推荐学习书籍(已分类)
C#/.NET/.NET Core推荐学习书籍(已分类)
175 0
|
5月前
|
算法 程序员 Linux
MKL.NET:为.NET开发者提供高性能数学计算支持的开源库
MKL.NET:为.NET开发者提供高性能数学计算支持的开源库
62 0
|
5月前
|
开发框架 安全 前端开发
一个高性能类型安全的.NET枚举实用开源库
一个高性能类型安全的.NET枚举实用开源库
32 0
|
5月前
|
程序员 Linux 网络安全
一个基于.Net高性能跨平台内网穿透工具
一个基于.Net高性能跨平台内网穿透工具
49 0