CVPR 2023|两行代码高效缓解Vision Transformer过拟合,美图&国科大联合提出正则化方法DropKey

摘要:美图影像研究院(MT Lab)与中国科学院大学突破性地提出正则化方法DropKey,用于缓解Vision Transformer中的过拟合问题。该方法通过在注意力计算阶段随机drop部分Key以鼓励网络捕获目标对象的全局信息,从而避免了由过于聚焦局部信息所引发的模型偏置问题,继而提升了基于Transformer的视觉类算法的精度。该论文已被计算机视觉三大顶会之一CVPR 2023接收。
近期,基于Transformer的算法被广泛应用于计算机视觉的各类任务中,但该类算法在训练数据量较小时容易产生过拟合问题。现有Vision Transformer通常直接引入CNN中常用的Dropout算法作为正则化器,其在注意力权重图上进行随机Drop并为不同深度的注意力层设置统一的drop概率。尽管Dropout十分简单,但这种drop方式主要面临三个主要问题。首先,在softmax归一化后进行随机Drop会打破注意力权重的概率分布并且无法对权重峰值进行惩罚,从而导致模型仍会过拟合于局部特定信息(如图1)。其次,网络深层中较大的Drop概率会导致高层语义信息缺失,而浅层中较小的drop概率会导致过拟合于底层细节特征,因此恒定的drop概率会导致训练过程的不稳定。最后,CNN中常用的结构化drop方式在Vision Transformer上的有效性并不明朗。

图1 不同正则化器对注意力分布图的影响

美图影像研究院(MT Lab)与中国科学院大学在CVPR 2023上发表了一篇文章,提出一种新颖且即插即用的正则化器DropKey,该正则化器可以有效缓解Vision Transformer中的过拟合问题。文章中对以下三个核心问题进行了研究: 第一,在注意力层应该对什么信息执行Drop操作?与直接Drop注意力权重不同,该方法在计算注意力矩阵之前执行Drop操作,并将Key作为基础Drop单元。该方法在理论上验证了正则化器DropKey可以对高注意力区域进行惩罚并将注意力权值分配到其它感兴趣的区域,从而增强模型对全局信息的捕获能力。 第二,如何设置Drop概率?与所有层共享同一个Drop概率相比,该论文提出了一种新颖的Drop概率设置方法,即随着自注意力层的加深而逐渐衰减Drop概率值。 第三,是否需要像CNN一样进行结构化Drop操作?该方法尝试了基于块窗口和交叉窗口的结构化Drop方式,并发现这种技巧对于Vision Transformer来说并不重要。

论文链接:https://arxiv.org/abs/2208.02646

背景
Vision Transformer(ViT)是近期计算机视觉模型中的新范式,它被广泛地应用于图像识别、图像分割、人体关键点检测和人物互相检测等任务中。具体而言,ViT将图片分割为固定数量的图像块,将每个图像块都视作一个基本单位,同时引入了多头自注意力机制来提取包含相互关系的特征信息。但现有ViT类方法在小数据集上往往会出现过拟合问题,即仅使用目标局部特征来完成指定任务。
为了克服以上问题,该论文提出了一种即插即拔、仅需要两行代码便可实现的正则化器DropKey用以缓解ViT类方法的过拟合问题。不同于已有的Dropout,DropKey将Key设置为drop对象并从理论和实验上验证了该改变可以对高注意力值部分进行惩罚,同时鼓励模型更多关注与目标有关的其他图像块,有助于捕捉全局鲁棒特征。此外,该论文还提出为不断加深的注意力层设置递减的drop概率,这可以避免模型过度拟合低级特征并同时保证有充足的高级特征以进行稳定的训练。此外,该论文还通过实验证明,结构化drop方法对ViT来说不是必要的。
DropKey
为了探究引发过拟合问题的本质原因,该研究首先将注意力机制形式化为一个简单的优化目标并对其拉格朗日展开形式进行分析。发现当模型在不断地优化时,当前迭代中注意力占比越大的图像块,在下次迭代过程中会倾向于被分配更大的注意力权值。为缓解这一问题,DropKey通过随机drop部分Key的方式来隐式地为每个注意力块分配一个自适应算子以约束注意力分布从而使其变得更加平滑。值得注意的是,相对于其他根据特定任务而设计的正则化器,DropKey无需任何手工设计。由于在训练阶段对Key执行随机drop,这将导致训练和测试阶段的输出期望不一致,因此该方法还提出使用蒙特卡洛方法或微调技巧以对齐输出期望。此外,该方法的实现仅需两行代码,具体如图2所示。

图2 DropKey实现方法

一般而言,ViT会叠加多个注意力层以逐步学习高维特征。通常,较浅层会提取低维视觉特征,而深层则旨在提取建模空间上粗糙但复杂的信息。因此,该研究尝试为深层设置较小的drop概率以避免丢失目标对象的重要信息。具体而言,DropKey并不在每一层以固定的概率执行随机drop,而是随着层数的不断加深而逐渐降低drop的概率。此外,该研究还发现这种方法不仅适用于DropKey,还可以显著提高Dropout的性能。
虽然在CNN中对结构化drop方法已有较为详细的研究,但还没有研究该drop方式对ViT的性能影响。为探究该策略会不会进一步提升性能,该论文实现了DropKey的两种结构化形式,即DropKey-Block和DropKey-Cross。其中,DropKey- Block通过对以种子点为中心的正方形窗口内连续区域进行drop,DropKey-Cross则通过对以种子点为中心的十字形连续区域进行drop,如图3所示。然而,该研究发现结构化drop方法并不会带来性能提升。

图3 DropKey的结构化实现方法

实验结果

图4 DropKey和Dropout在CIFAR10/100上的性能比较

图5 DropKey和Dropout在CIFAR100上的注意力图可视化效果比较

图6 不同drop概率设置策略的性能比较

图7 不同输出期望对齐策略的性能比较

图8 不同结构化drop方法的性能比较

图9 DropKey和Dropout在ImageNet上的性能比较

图10 DropKey和Dropout在COCO上的性能比较

图11 DropKey和Dropout在HICO-DET上的性能比较

图12 DropKey和Dropout在HICO-DET上的性能比较

图13 DropKey和Dropout在HICO-DET上的注意力图可视化比较

总结
该论文创新性地提出了一种用于ViT的正则化器,用于缓解ViT的过拟合问题。与已有的正则化器相比,该方法可以通过简单地将Key置为drop对象,从而为注意力层提供平滑的注意力分布。另外,该论文还提出了一种新颖的drop概率设置策略,成功地在有效缓解过拟合的同时稳定训练过程。最后,该论文还探索了结构化drop方式对模型的性能影响。
研究团队
本论文由美图影像研究院(MT Lab)和中国科学院大学的研究者们共同提出。美图影像研究院(MT Lab)是美图公司致力于计算机视觉、深度学习、增强现实等领域的算法研究、工程开发和产品化落地的团队,深耕人脸技术、人体技术、图像识别、图像处理、图像生成、AR等领域的前沿性和前瞻性技术研究,对美图秀秀、美颜相机、Wink等美图旗下全系软硬件产品提供技术支持,同时面向影像行业内多个垂直赛道提供针对性SaaS服务,通过领先技术推动美图的产品发展,曾先后参与CVPR、ICCV、ECCV等计算机视觉国际顶级会议并斩获十余项赛事冠亚军,累计发表超30篇论文。