简单通用：视觉基础网络最高3倍无损训练加速，清华EfficientTrain++入选TPAMI 2024

机器之心

发布于 2024-05-22 16:52:59

1640

发布于 2024-05-22 16:52:59

文章被收录于专栏：机器之心

本论文作者王语霖是清华大学自动化系 2019 级直博生，师从吴澄院士和黄高副教授，主要研究方向为高效深度学习、计算机视觉等。他曾以第一作者在 TPAMI、NeurIPS、ICLR、ICCV、CVPR、ECCV 等期刊、会议上发表论文，曾获百度奖学金、微软学者、CCF-CV 学术新锐奖、字节跳动奖学金等荣誉。个人主页：wyl.cool

本文主要介绍刚刚被 IEEE Transactions on Pattern Analysis and Machine Intelligence （TPAMI）录用的一篇文章：EfficientTrain++: Generalized Curriculum Learning for Efficient Visual Backbone Training。

论文链接：https://arxiv.org/pdf/2405.08768
代码和预训练模型已开源：https://github.com/LeapLabTHU/EfficientTrain
会议版本论文（ICCV 2023）：https://arxiv.org/pdf/2211.09703

近年来，「scaling」是计算机视觉研究的主角之一。随着模型尺寸和训练数据规模的增大、学习算法的进步以及正则化和数据增强等技术的广泛应用，通过大规模训练得到的视觉基础网络（如 ImageNet1K/22K 上训得的 Vision Transformer、MAE、DINOv2 等）已在视觉识别、目标检测、语义分割等诸多重要视觉任务上取得了令人惊艳的性能。

然而，「scaling」往往会带来令人望而却步的高昂模型训练开销，显著阻碍了视觉基础模型的进一步发展和工业应用。

面向解决这一问题，清华大学的研究团队提出了一种广义课程学习（generalized curriculum learning）算法：EfficientTrain++。其核心思想在于，将「筛选和使用由易到难的数据、逐步训练模型」的传统课程学习范式推广至「不进行数据维度的筛选，一直使用全部训练数据，但在训练过程中逐步揭示每个数据样本的由易到难的特征或模式（pattern）」。

EfficientTrain++ 具有几个重要的亮点：

即插即用地实现视觉基础网络 1.5−3.0× 无损训练加速。上游、下游模型性能均不损失。实测速度与理论结果一致。
通用于不同的训练数据规模（例如 ImageNet-1K/22K，22K 效果甚至更为明显）。通用于监督学习、自监督学习（例如 MAE）。通用于不同训练开销（例如对应于 0-300 或更多 epochs）。
通用于 ViT、ConvNet 等多种网络结构（文中测试了二十余种尺寸、种类不同的模型，一致有效）。
对较小模型，训练加速之外，还可显著提升性能（例如在没有额外信息帮助、没有额外训练开销的条件下，在 ImageNet-1K 上得到了 81.3% 的 DeiT-S，可与原版 Swin-Tiny 抗衡）。
对两种有挑战性的常见实际情形开发了专门的实际效率优化技术：1）CPU / 硬盘不够强力，数据预处理效率跟不上 GPU；2）大规模并行训练，例如在 ImageNet-22K 上使用 64 或以上的 GPUs 训练大型模型。

接下来，我们一起来看看该研究的细节。

一．研究动机

近年来，大型基础模型（foundation models）的蓬勃发展极大促进了人工智能和深度学习的进步。在计算机视觉领域，Vision Transformer（ViT）、CLIP、SAM、DINOv2 等代表性工作已经证明，同步增大（scaling up）神经网络尺寸和训练数据规模能够显著拓展识别、检测、分割等大量重要视觉任务的性能边界。

然而，大型基础模型往往具有高昂的训练开销，图 1 给出了两个典型例子。以使用 8 块 NVIDIA V100 或性能更强的 GPU 为例，GPT-3、ViT-G 仅完成一次训练即需要等效为数年甚至数十年的时间。如此高昂的训练成本，无论是对学术界还是工业界而言，都是较难负担的巨大开销，往往只有少数顶尖机构消耗大量资源才能推进深度学习的进展。因此，一个亟待解决的问题是：如何有效提升大型深度学习模型的训练效率？

图 1 示例：大型深度学习基础模型的高昂训练开销

对于计算机视觉模型而言，一个经典的思路是课程学习（curriculum learning），如图 2 所示，即模仿人类渐进式、高度结构化的学习过程，在模型训练过程中，从最「简单」的训练数据开始，逐步引入由易到难的数据。

图 2 经典课程学习范式（图片来源：《A Survey on Curriculum Learning》，TPAMI’22）

然而，尽管动机比较自然，课程学习并没有被大规模应用为训练视觉基础模型的通用方法，其主要原因在于存在两个关键的瓶颈，如图 3 所示。其一，设计有效的训练课程（curriculum）并非易事。区别「简单」、「困难」样本往往需要借助于额外的预训练模型、设计较复杂的 AutoML 算法、引入强化学习等，且通用性较差。其二，课程学习本身的建模存在一定不合理性。自然分布中的视觉数据往往具有高度的多样性，图 3 下方给出了一个例子（从 ImageNet 中随机抽取的鹦鹉图片），模型训练数据中包含大量不同动作的鹦鹉、离镜头不同距离的鹦鹉、不同视角、不同背景的鹦鹉、以及鹦鹉与人或物的多样化的交互等，将如此多样化的数据数据仅仅以「简单」、「困难」的单一维度指标进行区分，事实上是一个比较粗糙和牵强的建模方式。

图 3 阻碍课程学习大规模应用于训练视觉基础模型的两个关键瓶颈

二．方法简介

受到上述挑战的启发，本文提出了一种广义课程学习（generalized curriculum learning）范式，其核心思想在于，将「筛选和使用由易到难的数据、逐步训练模型」的传统课程学习范式推广至「不进行数据维度的筛选，一直使用全部训练数据，但在训练过程中逐步揭示每个数据样本的由易到难的特征或模式」，这样就有效避开了因数据筛选范式引发的局限和次优设计，如图 4 所示。

图 4 传统课程学习（样本维度） v.s. 广义课程学习（特征维度）

这一范式的提出主要基于一个有趣的现象：在一个自然的视觉模型训练过程中，虽然模型总是可以随时获取数据中包含的全部信息，但模型总会自然地先学习识别数据中包含的某些比较简单的判别特征（pattern），而后在此基础上再逐步学习识别更难的判别特征。并且，这一规律是比较普遍的，「比较简单」的判别特征在频域和空域都可以较方便地找到。本文设计了一系列有趣的实验来证明上述发现，如下所述。

从频域的角度来说，「低频特征」对于模型而言「比较简单」。在图 5 中，本文作者使用标准 ImageNet-1K 训练数据训练了一个 DeiT-S 模型，并使用带宽不同的低通滤波器对验证集进行滤波，只保留验证图像的低频分量，在此基础上报告 DeiT-S 在训练过程中、在低通滤波的验证数据上的准确率，所得准确率相对训练过程的曲线显示于图 5 右侧。

我们可以看到一个有趣的现象：在训练早期阶段，仅使用低通滤波的验证数据不会显著降低准确性，且曲线与正常验证集准确率间的分离点随滤波器带宽的增大而逐渐右移。这一现象表明，尽管模型始终可以访问训练数据的低频和高频部分，但其学习过程自然地从仅关注低频信息开始，识别较高频特征的能力则在训练后期逐步习得（这一现象的更多证据可参考原文）。

图 5 频域角度上，模型自然倾向于先学习识别低频特征

这一发现引出了一个有趣的问题：我们是否可以设计一个训练课程（curriculum），从只为模型提供视觉输入的低频信息开始，而后逐渐引入高频信息？

图 6 研究了这个想法，即仅在特定长度的早期训练阶段对训练数据执行低通滤波，训练过程的其余部分保持不变。从结果中可以观察到，尽管最终性能提升有限，但有趣的是，即使在相当长的一段早期训练阶段中仅向模型提供低频分量，模型的最终精度也可以在很大程度上得到保留，这也与图 5 中「模型在训练初期主要关注学习识别低频特征」的观察不谋而合。

这一发现启发了本文作者关于训练效率的思考：既然模型在训练初期只需要数据中的低频分量，且低频分量包含的信息小于原始数据，那么能否使模型以比处理原始输入更少的计算成本高效地仅从低频分量中学习？

图 6 在相当长的一段早期训练阶段中仅向模型提供低频分量并不会显著影响最终性能

事实上，这一思路是完全可行的。如图 7 左侧所示，本文作者在图像的傅里叶频谱中引入了裁切操作，裁切出低频部分，并将其映射回像素空间。这一低频裁切操作准确地保留了所有低频信息，同时减小了图像输入的尺寸，因此模型从输入中学习的计算成本可以呈指数级降低。

如果使用这一低频裁切操作在训练早期阶段处理模型输入，可以显著节省总体训练成本，但由于最大限度地保留了模型学习所必要的信息，仍然可以获得性能几乎不受损失的最终模型，实验结果如图 7 右下方所示。

图 7 低频裁切（low-frequency cropping）：使模型高效地仅从低频信息中学习

在频域操作之外，从空域变换的角度，同样可以找到对于模型而言「比较简单」的特征。例如，没有经过较强数据增强或者扭曲处理的原始视觉输入中所包含的自然图像信息往往对于模型而言「比较简单」、更容易让模型学习，因为它们是从现实世界的分布中得出的，而数据增强等预处理技术所引入的额外信息、不变性等往往对于模型而言较难学习（图 8 左侧给出了一个典型示例）。

事实上，现有研究也已观察到，数据增强主要在训练较晚期阶段发挥作用（如《Improving Auto-Augment via Augmentation-Wise Weight Sharing》, NeurIPS’20）。

在这一维度上，为实现广义课程学习的范式，可以简单地通过改变数据增强的强度方便地实现在训练早期阶段仅向模型提供训练数据中较容易学习的自然图像信息。图 8 右侧使用 RandAugment 作为代表性示例来验证了这个思路，RandAugment 包含了一系列常见的空域数据增强变换（例如随机旋转、更改锐度、仿射变换、更改曝光度等）。

可以观察到，从较弱的数据增强开始训练模型可以有效提高模型最终表现，同时这一技术与低频裁切兼容。

图 8 从空域的角度寻找模型 “较容易学习” 的特征：一个数据增强的视角

到此处为止，本文提出了广义课程学习的核心框架与假设，并通过揭示频域、空域的两个关键现象证明了广义课程学习的合理性和有效性。在此基础上，本文进一步完成了一系列系统性工作，在下面列出。由于篇幅所限，关于更多研究细节，可参考原论文。

融合频域、空域的两个核心发现，提出和改进了专门设计的优化算法，建立了一个统一、整合的 EfficientTrain++ 广义课程学习方案。
探讨了低频裁切操作在实际硬件上高效实现的具体方法，并从理论和实验两个角度比较了两种提取低频信息的可行方法：低频裁切和图像降采样，的区别和联系。
对两种有挑战性的常见实际情形开发了专门的实际效率优化技术：1）CPU / 硬盘不够强力，数据预处理效率跟不上 GPU；2）大规模并行训练，例如在 ImageNet-22K 上使用 64 或以上的 GPUs 训练大型模型。

本文最终得到的 EfficientTrain++ 广义课程学习方案如图 9 所示。EfficientTrain++ 以模型训练总计算开销的消耗百分比为依据，动态调整频域低频裁切的带宽和空域数据增强的强度。

值得注意的是，作为一种即插即用的方法，EfficientTrain++ 无需进一步的超参数调整或搜索即可直接应用于多种视觉基础网络和多样化的模型训练场景，效果比较稳定、显著。