DeiT:训练ImageNet仅用4卡不到3天的平民ViT | ICML 2021

论文基于改进训练配置以及一种新颖的蒸馏方式,提出了仅用ImageNet就能训练出来的Transformer网络DeiT。在蒸馏学习时,DeiT以卷积网络作为teacher,能够结合当前主流的数据增强和训练策略来进一步提高性能。从实验结果来看,效果很不错

来源:晓飞的算法工程笔记 公众号

论文: Training data-efficient image transformers & distillation through attention

[图片上传失败...(image-3e441f-1659936459728)]

Introduction


[图片上传失败...(image-bbabd8-1659936459729)]

  Vision Transformer一般要先在大型计算设施上预训练数以亿计的图片才能有较好的性能,这极大地提高其应用门槛。为此,论文基于ViT提出了可在ImageNet上训练的Vision Transformer模型DeiT,仅需要一台电脑(4卡)训练不到三天(53小时的预训练和可选的20小时微调)的时间。在没有外部数据预训练的情况下,在ImageNet上达到了83.1% 的最高精度。

  此外,论文还提出了一种针对Transformer的蒸馏策略,通过一个蒸馏token确保student网络通过注意力从teacher网络那里进行学习。当使用卷积网络作为teacher网络时,ImageNet上可达到85.2%的准确性。

  总体而言,论文主要有以下贡献:

  • 通过实验表明,在没有外部数据的情况下,Vision Transformer也可以在ImageNet上达到SOTA的结果,而且仅需要4卡设备训练三天。
  • 论文提出了一种基于蒸馏token的新蒸馏方法,这种用于Transformer的蒸馏方法大幅优于一般蒸馏方法。蒸馏token与class token的作用相同,都参与注意力计算中,只是蒸馏token的训练目的在于复现teacher网络的标签预测。
  • 有趣的是,论文发现在使用新蒸馏方法时,用卷积网络作为teacher要比用另一个相同准确率的transformer的作为teacher的效果要好。
  • 在Imagenet上预训练的模型可以转移到不同的下游任务(如细粒度分类),得到很不错的性能。

Distillation through attention


Soft distillation

  一般的蒸馏方法都是Soft distillation,其核心目标是最小化teacher网络和student网络的softmax输出之间的Kullback-Leibler散度。

  定义Z_t为teacher网络的logits输出(输入softmax的向量),Z_s为student网络的logits输出。用\tau表示蒸馏温度,\lambda表示平衡Kullback-Leibler散度损失(KL)和交叉熵损失(LCE)的权值,\psi表示softmax函数。定义soft distillation的目标函数为:

[图片上传失败...(image-2a5c71-1659936459729)]

Hard-label distillation

  论文提出了一种蒸馏的变体,将teacher网络的预测标签作为蒸馏的GT标签。假设y_t = argmax_c Z_t(c)是teacher网络的预测标签,与之相关的hard-label distillation目标为:

[图片上传失败...(image-db2c9e-1659936459729)]

  对于同一张图片,teacher网络预测的标签可能随着特定的数据增强而有所变化。从实验结果来看,将预测标签作为蒸馏目标的做法比传统的做法更好,不仅无额外参数,概念上还更简单:teacher网络预测的y_t与真实标签y是相同的作用。
  此外,hard label也可以通过label smoothing转换为软标签,其中GT标签具有1 - \varepsilon的概率,其余类共享\varepsilon概率。在相关的实验中,参数固定为\varepsilon = 0.1

Distillation token

[图片上传失败...(image-cbeade-1659936459729)]

  论文提出的蒸馏方案如如图2所示,在输入的token序列中添加一个蒸馏token。蒸馏token与class token类似,通过self-attention与其它token交互并将最后一层中的对应输出作为网络输出,其训练目标为损失函数中的蒸馏损失部分。蒸馏token使得模型可以像常规蒸馏一样从teacher网络的输出中学习,同时与class token保持互补的关系。

  论文发现,训练后的输入层class token和蒸馏token收敛到了完全不同的向量,平均余弦相似度仅为0.06。但随着在网络的计算,class和蒸馏token在越深层中的对应输出逐渐变得更加相似,最后一层达到了较高的相似度(cos=0.93),但没有完全相同。这是符合预期的,因为两个token的目标就是产生相似但不相同的目标。

  论文也尝试替代实验,用另一个class token代替teacher网络的蒸馏token进行伪蒸馏学习。但无论如何随机且独立地初始化两个class token,训练后都会收敛到相同的向量(cos=0.999),其对应的输出也是准相同的。这表明这个代替的class token不会对分类性能带来任何影响,相比之下蒸馏token则能带来显著的提升。

Fine-tuning with distillation

  在分辨率增加的fine-tuning阶段,同样使用真实标签和teacher网络预测标签进行训练。此时需要一个具有相同目标分辨率的teacher网络,可通过FixRes的做法从之前的低分辨率teacher网络中转换。论文也尝试了只用真实标签进行fine-tuning,但这导致了性能的降低。

Classification with our approach:joint classifiers

  在测试时,网络输出的class token和蒸馏token都用于标签分类。论文的建议做法是将这两个token独立预测后再融合,即将两个分类器的softmax输出相加再进行预测。

Transformer models


[图片上传失败...(image-fa1509-1659936459729)]

  DeiT的架构设计与ViT相同,唯一的区别是训练策略和蒸馏token,训练策略的区别如表9所示。此外,在预训练时不使用MLP,仅使用线性分类器。

  为避免混淆,用ViT来指代先前工作中的结果,用DeiT来指代论文的结果。如果未指定,DeiT指的是DeiT-B,与ViT-B具有相同的架构。当以更大的分辨率fine-tune DeiT时,论文会在名字的最后附加分辨率,例如DeiT-B↑384。最后,当使用论文提出的蒸馏方法时,论文会用一个蒸馏符号将其标识为DeiT⚗.。

[图片上传失败...(image-88e3e-1659936459729)]

  如表1所示,DeiT-B的结构与ViT-B完全一样,参数固定为D = 768h = 12d = D/h = 64。另外,论文设计了两个较小的模型:DeiT-S和DeiT-Ti,减少了head的数量,d保持不变。

Experiment


[图片上传失败...(image-421243-1659936459729)]

  不同类型的teacher网络的蒸馏效果。

[图片上传失败...(image-bc687f-1659936459729)]

  不同蒸馏策略的对比实验。

[图片上传失败...(image-345c1-1659936459729)]

  不同网络以及蒸馏策略之间的结果差异,值越小差异越小。

[图片上传失败...(image-b489dc-1659936459729)]

  蒸馏策略与训练周期的关系。

[图片上传失败...(image-5a408b-1659936459729)]

  整体性能的对比。

[图片上传失败...(image-f38b6e-1659936459729)]

  ImageNet上预训练模型的在其它训练集上的迁移效果。

[图片上传失败...(image-7e2511-1659936459729)]

  不同优化器、数据增强、正则化的对比,寻找最佳的训练策略和配置。

[图片上传失败...(image-89e4a2-1659936459729)]

  224分辨率预训练的DeiT在不同数据集上用不同分辨率fine-tune的效果。

Conclusion


  论文基于改进训练配置以及一种新颖的蒸馏方式,提出了仅用ImageNet就能训练出来的Transformer网络DeiT。在蒸馏学习时,DeiT以卷积网络作为teacher,能够结合当前主流的数据增强和训练策略来进一步提高性能。从实验结果来看,效果很不错。



如果本文对你有帮助,麻烦点个赞或在看呗~
更多内容请关注 微信公众号【晓飞的算法工程笔记】

work-life balance.
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 203,547评论 6 477
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,399评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 150,428评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,599评论 1 274
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,612评论 5 365
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,577评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,941评论 3 395
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,603评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,852评论 1 297
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,605评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,693评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,375评论 4 318
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,955评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,936评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,172评论 1 259
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 43,970评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,414评论 2 342

推荐阅读更多精彩内容