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

本讨论文作者王语霖是清华大学自动化系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. 即插即用地实现视觉基础网络 1.5−3.0× 无损训练加速。上游、下游模型性能均不损失。实测速度与理论结果一致。
2. 通用于不同的训练数据规模（例如 ImageNet-1K/22K，22K 效果甚至更为明显）。通用于监督学习、自监督学习（例如 MAE）。通用于不同训练开销（例如对应于 0-300 或更多 epochs）。
3. 通用于 ViT、ConvNet 等多种网络结构（文中测试了二十余种尺寸、种类不同的模型，一致有效）。
4. 对较小模型，训练加速之外，还可显著提升性能（例如在没有额外信息帮助、没有额外训练开销的条件下，在 ImageNet-1K 上得到了 81.3% 的 DeiT-S，可与原版 Swin-Tiny 抗衡）。
5. 对两种有挑战性的常见实际情形开发了专门的实际效率优化技术：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 從空域的角度尋找模型「較容易學習」 的特徵：一個資料增強的視角

到此處為止，本文提出了廣義課程學習的核心框架與假設，並透過揭示頻域、空域的兩個關鍵現象證明了廣義課程學習的合理性和有效性。在此基礎上，本文進一步完成了一系列系統性工作，並在下面列出。由於篇幅所限，關於更多研究細節，可參考原論文。

1. 融合頻域、空域的兩個核心發現，提出並改進了專門設計的最佳化演算法，建立了一個統一、整合的 EfficientTrain++ 廣義課程學習方案。
2. 探討了低頻裁切操作在實際硬體上高效實現的具體方法，並從理論和實驗兩個角度比較了兩種提取低頻資訊的可行方法：低頻裁切和影像降採樣，的區別和聯繫。
3. 對兩種有挑戰性的常見實際情形開發了專門的實際效率優化技術：1）CPU / 硬碟不夠強力，資料預處理效率跟不上GPU；2）大規模平行訓練，例如在ImageNet-22K 上使用64 或以上的GPUs 訓練大型模型。

本文最終得到的 EfficientTrain++ 廣義課程學習方案如圖 9 所示。 EfficientTrain++ 以模型訓練總計算開銷的消耗百分比為依據，動態調整頻域低頻裁切的頻寬和空域資料增強的強度。

值得注意的是，作為一種即插即用的方法，EfficientTrain++ 無需進一步的超參數調整或搜尋即可直接應用於多種視覺基礎網路和多樣化的模型訓練場景，效果較穩定顯著。

圖 9 統一、整合的廣義課程學習方案：EfficientTrain++

三．實驗結果

作為一種即插即用的方法，EfficientTrain++ 在ImageNet-1K 上，在基本上不損失或提升效能的條件下，將多種視覺基礎網路的實際訓練開銷降低了1.5 倍左右。

圖10 ImageNet-1K 實驗結果：EfficientTrain++ 在多種視覺基礎網路上的表現

EfficientTrain++ 的增益通用於不同的訓練開銷預算，嚴格相同表現的情況下，DeiT/Swin 在ImageNet-1K 上的訓加速比在2-3 倍左右。

圖11 ImageNet-1K 實驗結果：EfficientTrain++ 在不同訓練開銷預算下的表現

##EfficientTrain++ 在ImageNet-22k 上可以取得2-3 倍的表現無損預訓練加速。

圖12 ImageNet-22K 實驗結果：EfficientTrain++ 在更大規模訓練資料上的表現

對於較小的模型，EfficientTrain++ 可以實現顯著的效能上界提升。

图 13 ImageNet-1K 实验结果：EfficientTrain++ 可以显著提升较小模型的性能上界

EfficientTrain++ 对于自监督学习算法（如 MAE）同样有效。

图 14 EfficientTrain++ 可以应用于自监督学习（如 MAE）

EfficientTrain++ 训得的模型在目标检测、实例分割、语义分割等下游任务上同样不损失性能。

图 15 COCO 目标检测、COCO 实例分割、ADE20K 语义分割实验结果

以上是简单通用：视觉基础网络3倍无损训练加速，清华EfficientTrain++入选TPAMI 2024的详细内容。更多信息请关注PHP中文网其他相关文章！