【读点论文】RIFormer: Keep Your Vision Backbone Effective But Removing Token Mixer去掉费时的组件,用训练技巧提升点

RIFormer: Keep Your Vision Backbone Effective But Removing Token Mixer

Abstract

  • 本文研究了如何在去除基本构建块中的令牌混合器的同时保持视觉主干的有效性。令牌混合器作为(vit)的自关注器,主要用于在不同空间令牌之间进行信息通信,但存在较大的计算成本和延迟。然而,直接去除它们会导致先验模型结构不完整,从而导致精度显著下降

  • 为此,我们首先开发了一个基于重参数化思想的repidentityformer,研究了令牌混合器的自由模型架构。然后,我们探索了改进的学习范式,以打破简单令牌混合器自由骨干的限制,并将经验实践总结为5条准则。利用所提出的优化策略,我们能够构建一个非常简单的视觉主干,并且具有令人鼓舞的性能,同时在推理过程中具有很高的效率。大量的实验和分析也表明,通过适当的优化策略,可以将网络结构的归纳偏差纳入简单的网络结构中。我们希望这项工作可以作为探索优化驱动的高效网络设计的起点。项目页面: https://techmonsterwang.github.io/RIFormer/

  • Token Mixer是ViT骨干非常重要的组成成分,它用于对不同空域位置信息进行自适应聚合,但常规的自注意力往往存在高计算复杂度与高延迟问题。而直接移除Token Mixer又会导致不完备的结构先验,进而导致严重的性能下降。

  • 本文基于重参数机制提出了RepIdentityFormer方案以研究无Token Mixer的架构体系。紧接着,作者改进了学习架构以打破无Token Mixer架构的局限性并总结了5条指导方针。搭配上所提优化策略后,本文构建了一种极致简单且具有优异性能的视觉骨干,此外它还具有高推理效率优势。实验结果表明:通过合适的优化策略,网络结构的归纳偏置可以被集成进简单架构体系中。本文为后续优化驱动的高效网络设计提供了新的起点和思路。

Introduction

  • 过去几年计算机视觉的巨大进步部分是由视觉主干的革命带来的,包括卷积神经网络(ConvNets)和(vit)。它们的基本构建块中都有特定的模块,用于聚合不同空间位置之间的信息,这些模块称为令牌混合器,例如vit的自关注。尽管令牌混合器的有效性已在许多视觉任务中得到证明,但其计算复杂度通常占用网络的很大一部分。在实践中,由于速度和计算成本的问题,重型令牌混合器使视觉骨干受到限制,特别是在边缘设备上

  • 在文献中已经有几次尝试研究有效的令牌混合器来减少视觉骨干。尽管这些作品已经通过轻量级设计实现了具有竞争力的性能,但它们确实保留了令牌混频器,这带来了不可忽略的延迟增加,如下图所示。

    • 在这里插入图片描述

    • ViTBase中不同组件的延迟分析。(a)令牌混频器(自关注)延时约占骨干的46.3%。(b)我们的动机是移除令牌混频器,同时努力保持性能。

    • 基于修正版12层ViT-B架构,作者进行了系统的延迟分析(可参考上图)。从仅包含Input Embedding的模块出发,逐步添加不同的操作单元,如LN、Attention、MLP等,最终构成了无GAP与分类头的ViT-Base架构。从图示可以看到:Token Mixer部分耗时约为1433.6ms,约占整个架构耗时的46.3%

    • Token Mixer对于模型延迟有非常大的影响,这无疑限制了其实际应用。不有而然升起一个疑问:是否可以移除Token Mixer同时保持其高性能呢?MetaFormer提出了无Token Mixer的架构,但发现会导致不可忽视的性能退化。基于此,本文旨在通过先进的优化策略(如知识蒸馏、结构重参数等)来挖掘极简ViT架构的潜力

  • 最近的工作[Metaformer baselines for vision]发现,删除令牌混频器是可能的,但会导致性能下降。这些对高效令牌混合器的探索启发我们思考,我们是否可以在移除令牌混合器的同时保持视觉主干的有效性?所得到的令牌混合器无视觉主干在实际应用中是高效有效的。

  • 在这项工作中,我们首先回顾了当前的模型架构和学习范式。以往的工作大多集中在结构的改进上,而采用传统的监督学习从头开始优化模型。不同的是,我们建议采用简化的模型架构,并探索学习范式设计,以充分发挥简单模型的潜力。我们的目标是同时保持令牌混合器无视觉主干(即IdentityFormer,如上图-(b)所示)的效率和功效。为此,我们将在以下章节中深入研究简单而有效的学习策略——知识蒸馏(KD)。

  • 我们的主要思想是将知识从强大的教师模型(带令牌混合器)提炼到学生模型(无令牌混合器)。我们实例化了重新参数化的思想,以扩大学生网络的建模能力,同时保持其效率,如下图所示。具体来说,在学生模型中引入简单仿射变换,取代令牌混合器进行训练。在推理过程中将仿射变换参数合并到LayerNorm中,最终使学生令牌混合器自由

    • 在这里插入图片描述

    • RIFormer块的结构重新参数化。

    • 作者在训练阶段采用仿射变换以替代恒等变换,它仅在通道维度进行缩放与移位操作,见上图。在推理阶段,该仿射变换参数可以合并进LN层,故训练时的LN+仿射变换可以等价为LN+Identity

  • 我们通过经验总结我们的学习策略为以下指导方针,希望能对如何学习这个极其简单的模型有所启发。具体而言,

    • 1)不使用基础真值标签的软蒸馏更有效;

    • 2)采用仿射变换而不进行精馏难以调整性能退化;

    • 3)提出的分块知识精馏(称为模块模仿)有助于利用仿射算子的建模能力;

    • 4)接受野大的教师有利于培养接受野有限的学生;

    • 5)将教师模型(除令牌混合器外)的预训练权重加载到学生模型中,提高收敛性和性能。

  • 基于上述准则,我们最终获得了一个性能具有竞争力且效率高的无令牌混频器视觉模型,称为RepIdentityF former (RIF former)。RIFormer与MetaFormer共享几乎相同的宏和微设计,但安全地删除了所有令牌混合器。定量结果表明,我们的网络在ImageNet-1K上以更快的推理速度优于许多主流骨干网。对特征分布和有效接受场(ERFs)的分析也表明,通过适当的优化策略,显式令牌混合器带来的归纳偏差可以隐式地融入简单的网络结构中。综上所述,我们的工作主要贡献如下:

    • 我们提出通过开发简单模型架构的高级学习范式来探索视觉主干,以满足现实应用的需求。

    • 我们实例化了重新参数化的思想,建立了一个无令牌混频器的视觉模型RIFormer,该模型在提高归纳偏置建模能力的同时,还具有较高的推理效率。

    • 我们提出的蒸馏策略的实用指南已被证明是有效的,在保持视觉骨干竞争力的同时消除了令牌混合器。

Related Work

Vision Transformer Acceleration
  • 视觉 transformer 是一种推理速度慢、能量密集的骨干,由于其自关注的二次计算成本,因此不适合部署在资源有限的边缘设备上,因此需要压缩技术。为视觉 transformer 瘦身设计了多种技术路线,如提炼出深度和嵌入维度更小的高效 transformer,修剪或合并不重要的标记,采用节能操作,或设计高效的注意力替代方案等。与此不同的是,我们的工作提出了一个全新的角度,在保持令人满意的性能的同时,完全去除骨干中复杂的令牌混合器。
Structual Re-parameterization
  • 结构重参数化是一种将训练时间和推理时间结构解耦的训练技术。例如,RepVGG 是一个普通的vgg风格的架构,在推理过程中具有吸引人的性能和速度,在训练过程中手动添加了1 × 1分支和身份映射分支,是一个强大的架构。同样,这种技术可以进一步扩展到超大型核卷积神经网络、类mlp模型、网络剪枝和特殊优化器设计。本文采用该技术引入参数,训练后等效吸收到LN层中。适当优化后的额外权重可以帮助模型学习有用的表示。

Preliminary and Motivation

  • 在本节中,我们首先简要回顾令牌混频器的概念。然后,我们通过经验延迟分析重新审视了它们对推理速度不可避免的副作用,从而介绍了本文的动机。
Preliminary: The Concept of Token Mixer
  • 概念令牌混合器是在给定的视觉主干中执行令牌混合功能的结构。它允许来自不同空间位置的信息聚合。例如,自注意模块作为ViT中的令牌混合器,通过在查询、键和值矩阵中的组件之间并行执行注意函数,这些组件从输入特征线性投影。此外,ResMLP通过将空间MLP作为令牌混频器来应用交叉补丁线性子层。上述令牌混合器的计算和内存成本是图像规模的二次元。
Motivation
  • 在本节中,我们将通过对ViT模型的定量延迟分析来关注令牌混频器的副作用。我们从一个修改后的12层ViT-Base模型开始,该模型只包含输入嵌入,在其每个基本构建块中没有任何操作。然后我们逐渐在每个基本块中加入操作组件(如LN、Attention、Channel MLP等),最终得到没有全局平均池化层和分类器头部的ViT-Base模型。对于每个模型,我们使用一个A100 GPU在 22 4 2 224^2 2242分辨率下使用2048个批处理大小,并计算超过30次运行的平均时间来推断该批处理。整个过程重复三次,我们取中间的数字作为统计延迟。如上图所示,在堆叠12层的规则数量后,令牌混频器会带来1433.6ms的额外延迟,约占骨干的46.3%。

  • 根据上述分析,令牌混频器对模型的延迟带来了明显的副作用,限制了其实际应用。观察结果自然提出了一个问题:我们可以在移除令牌混合器的同时保持主干的有效性吗?具体来说,最近的工作介绍了在其基本构建块中没有任何令牌混合器的MetaFormer模型,并发现它会引起不可忽视的性能下降。基于这些发现,我们建议将归纳偏差与先进的优化策略(如知识蒸馏、结构重参数化等)结合起来,以充分利用极简单模型的潜力。我们将在余下的工作中介绍所有的勘探细节。

Exploring RIFormer: A Roadmap

  • 在本节中,我们将介绍从RIFormer的完全监督方法到更高级的训练范例的轨迹。在此过程中,我们研究并开发了类似transformer模型的不同优化方案,同时保持推理时间模型不变。我们使用的基准RIFormer与最近开发的MetaFormer具有完全相同的宏架构和模型大小,不同之处在于在推理期间,其基本构建块中没有使用令牌混合器。我们将RIFormer-S12模型的计算复杂度控制在与poolformers12相当的水平,参数约为12M, MAC为1.8G。本节中所有的RIFormer-S12模型都在ImageNet-1K上进行了120 epoch的训练和评估。超参数的详细信息见附录的第1节。我们的探索路线图如下。
Vision Backbone Without Token Mixer
  • 我们的探索是针对在推理时间模型视觉主干的每个基本块中去除token mixer,以在保持性能的同时获得更高的推理速度。因此,我们从一个RIFormerS12模型开始,该模型具有使用CE损失的完全监督训练方案。作为性能参考,我们将结果与PoolFormer-S12进行比较,因为它仅使用基本池化操作作为令牌混合器,因此性能差距可归因于缺乏基本令牌混合功能。如下表所示。与PoolFormer-S12相比,经过少量监督训练的RIFormer-S12可能导致不可接受的性能下降(2.7%的top-1准确率)。结果表明,如果在每个构建块中没有令牌混合器,则常规监督学习在帮助模型从图像中学习有用信息方面受到限制,需要更高级的训练过程

    • 在这里插入图片描述

    • MetaFormer上不同令牌混频器使用交叉熵损失的结果。

  • 然后,我们研究和修改了一系列的训练范式,以改善较差的基线性能,可以概括为1)知识蒸馏,2)教师类型影响,3)结构重参数化,4)提出的模块模仿技术,5)从教师那里加载部分参数。由于我们的目标是探索不同高级训练配置的影响,而不是网络架构,因此在中间步骤中,推理时间模型架构总是保持不变。接下来,我们分享5条训练RIFormer的有用指南。

Distillation Paradigm Design
  • 我们现在研究了一个通用视觉骨干教师使用token mixer对一个RIFormer学生的知识蒸馏,并总结了如何有效地利用来自强教师网络的“软”标签。准则1:不使用真值标签的软蒸馏对没有令牌混合器的学生有效。基本上,大多数现有的KD方法都是针对带有令牌混合器的模型设计的。例如,通常的做法是帮助学生学习基础真理标签和教师预测的软标签。此外,DeiT的一些观察结果表明,使用硬标签代替软标签作为监督目标,可以显著改善transformer。相比之下,令牌混合器自由主干在其基本块中没有显式的补丁聚合模块。因此,它的蒸馏应该与传统的骨干不同。具体来说,尽管RIFormer与transformer共享相同的宏结构,但它仍然不能被视为学生transformer,因为我们已经故意从每个构建块中删除了令牌混合器。然而,我们也不喜欢将其视为纯粹的connet,因为RIFormer在宏/微级架构设计方面与transformer有相似之处。因此,我们有动力为RIFormer探索一种合适的KD方法,并具有良好的性能。

  • 通常,交叉熵的目标是帮助学生网络重现硬准确的标签,我们认为这个过程可能不适合RIFormer。首先,通过标签平滑正则化将基真硬标签转化为软分布,真标签权值为1−ε,每个类共享ε。非习得的均匀分布在负类中信息量较小,并且可能干扰教师给出的习得的软分布。其次,1×1卷积实际上主导了RIFormer的基本构建块,只“混合”了每个位置的特征,而不是空间信息。这种简化的设计可能需要在监督标签中提供更丰富的信息。为了证明这一点,下表比较了四种不同设置的性能。默认的教师是一个GFNet-H-B (54M参数)。带真标签的硬蒸馏将准确度从72.31%提高到73.51%。结果表明,有令牌混合器的教师对没有令牌混合器的学生有积极的影响。实际上,使用不带真标签的软蒸馏的组合表现最好,将网络性能提高了74.05%。

    • 在这里插入图片描述

    • 以身份映射作为令牌混合器的普通/无标签RIFormer-S12中不同教师类型的结果。

  • 备注1。带真标签的监督学习似乎不是最适合没有token混合器的粗糙模型的方法。使用token mixer的教师可以帮助指导训练,但仍然无法完全恢复取消token mixer后的性能差距,需要其他策略

Re-parameterization for Identity Mapping
  • 准则2:使用仿射变换而不进行定制蒸馏,很难恢复性能退化。在这一部分中,我们采用结构重参数化方法的思想,该方法通常采用强大的模型进行训练,并在推理过程中等效地转换为简单的模型。具体来说,RIFormer中的推理时间令牌混频器模块可以看作是LN层之后的身份映射。因此,训练时间模块至少应满足两个基本要求:

    • 1)允许等效变换的每位置算子;

    • 2)参数算子,允许额外的表示能力。

    • 因此,我们在训练过程中使用仿射变换算子来替换恒等映射,该变换算子只执行按通道的缩放和移位,上如图所示。仿射算子及其前一层LN可以转换为一个修改了权值的LN,从而在推理过程中等价地转换为一个恒等映射。将输入特征记为M∈RN×C×H×W,仿射算子可表示为:

    • A f f i n e ( M , s , t ) : , i , : , : = s i M : , i , : , : + t i − M : , i , : , : , Affine(M, s, t)_{:,i,:,:} = siM{:,i,:,:} + t_i − M_{:,i,:,:}, Affine(M,s,t):,i,:,:=siM:,i,:,:+tiM:,i,:,:,

  • 其中s∈RC和t∈RC是可学习的权向量。由于残差连接,我们遵循在实现过程中添加输入的减法,因此不会合并上等式中的第一项和第三项。然后,我们描述了如何将仿射变换合并到其前面的LN层中,以便训练时间模型可以等效地转换为用于部署的模型,但其块中不再有令牌混合器。我们使用µ,σ, γ, β作为之前LN层的平均值,标准差和学习的比例因子和偏差。设T(a)∈RN×C×H×W, T ’ (a)∈RN×C×H×W分别作为上图-(a)中仿射残差子块的输入和输出。在训练期间,我们有:

    • T ′ ( a ) = A f f i n e ( L N ( T ( a ) , µ , σ , γ , β ) , s , t ) − T ( a ) T^{′(a)} = Affine(LN(T^{(a)}, µ, σ, γ, β), s, t) − T^{(a)} T(a)=Affine(LN(T(a),µ,σ,γ,β),s,t)T(a)

    • 其中LN是LN函数,由PyTorch中的GroupNorm API(将组号设置为1)同[Metaformer is actually what you need for vision]实现。在推理期间,残差子块中只存在一个尾随LN层的恒等映射。因此,我们得到:

    • T ′ ( a ) = L N ( T ( a ) , µ , σ , γ ′ , β ′ ) , T^{′(a)} = LN(T^{(a)}, µ, σ, γ′, β′), T(a)=LN(T(a),µ,σ,γ,β),

    • 其中γ′和β′为合并LN层的权值和偏置参数。基于上等式的等价性,对于∀1≤i≤C,我们有:

    • γ ′ i = γ i ( s i − 1 ) , β ′ i = β i ( s i − 1 ) + t i , γ^{′i} = γ_i^{(si − 1)}, β^{′i} = β_i^{(si − 1) }+ t_i, γi=γi(si1),βi=βi(si1)+ti,

  • 仿射变换和重参数化过程的证明和类pytorch代码分别见附录第2节和第3节。由于LN层在推理时间内没有预先计算的均值和标准差,因此它们的具体值是输入自适应的,不影响变换的等价性。

  • 备注2。对比上表,直接采用结构重参数化方法没有优势。我们将这种现象归因于这样一个事实,即LN层中的仿射变换是一个线性变换,可以直接与我们引入的额外仿射算子合并(如果不添加任何非线性函数)。因此,如果两者都仅由模型的输出来监督,则可能无法充分利用附加参数的潜力。同时,师生的同构设计启发我们探索适合各层模块知识转移的方法。

Module Imitation
  • 指南3:提出的分块知识蒸馏,称为模块模仿,有助于利用仿射算子的建模能力。我们之前尝试的KD方法只关注师生之间网络的输出。我们提出了模块模仿(MI)方法,以利用教师令牌混合器中的有用信息。具体来说,使用预训练的poolformer - s12作为教师网络。如下图所示,我们期望简单仿射算子(其前面的LN层)在训练期间近似于基本令牌混频器的行为。

    • 在这里插入图片描述

    • (a) RIFormer的全面培训程序。RIFormer删除了每个块中的令牌混合器。(b)模块模仿技术旨在通过简单的仿射变换来模仿令牌混合器的行为。

    • 我们将呈现出从全监督到更先进训练体系的探索之路,调研并指定了不同的优化机制。RIFormer采用了与MetaFormer相同的宏观与微观模块,区别仅在于RIFormer未使用Token Mixer

    • 有鉴于上述分析,作者进一步提出了Module Imitation以充分利用老师模型TokenMixer后的有用信息。如上图所示,作者希望:在训练过程中仿射操作能近似模拟TokenMixer的行为。此时,两者之间的MSE可以计算。

  • 设f(·), T ( a ) , m ∈ R N × C × H × W T^{(a),m}∈\R^{N×C×H×W} T(a)mRN×C×H×W,m∈m为仿射算子,其中m为我们使用的中间层集的RIFormer的第m层输入,g(·), T ( T ) , m ∈ R N × C × H × W T^{(T),m}∈\R^{N×C×H×W} T(T)mRN×C×H×W,m∈m分别为教师网络的输入。为简便起见,我们将LN(·,µ,σ, γ, β)缩写为LN(·)。仿射算子LN层与令牌混频器之间输入的均方误差(MSE)可计算为:

    • L i n = α 1 ∣ ∣ L N ( T ( a ) , m ) − L N ( T ( t ) , m ) ∣ ∣ F 2 L_{in}=\alpha_1||LN(T^{(a),m})-LN(T^{(t),m})||^2_F Lin=α1∣∣LN(T(a),m)LN(T(t),m)F2

    • 式中α1 = 1/N CHW。注意,当前层的输入特征是前一层的输出特征。因此,我们建议在实践中匹配该块的输出特征(即下一个后续块的输入特征),这可以看作是transformer中的隐藏状态蒸馏

    • L i n ‘ = α 1 ∣ ∣ T ( a ) , m + 1 − T ( t ) , m + 1 ∣ ∣ F 2 L^`_{in}=\alpha_1||T^{(a),m+1}-T^{(t),m+1}||^2_F Lin=α1∣∣T(a),m+1T(t),m+1F2

    • 然后将基于关系矩阵的隐藏状态蒸馏应用于输出特征:

    • L r e l = α 2 ∣ ∣ R T ( a ) , m + 1 − R T ( t ) , m + 1 ∣ ∣ F 2 L_{rel}=\alpha_2||RT^{^(a),m+1}-RT^{(t),m+1}||^2_F Lrel=α2∣∣RT(a),m+1RT(t),m+1F2

    • 式中 α 2 = 1 / N H 2 W 2 , R ( T ) =   T   T α_2 = 1/N H^2W^2, R(T) = ~ T ~ T α2=1/NH2W2,R(T)= T T, ~ T表示最后一个维的归一化T。考虑仿射算子与令牌混频器之间输出的MSE:

    • L o u t = α 1 ∣ ∣ f ( L N ( T ( a ) , m ) ) − g ( L N ( T ( t ) , m ) ) ∣ ∣ F 2 L_{out}=\alpha_1||f(LN(T^{(a),m}))-g(LN(T^{(t),m}))||^2_F Lout=α1∣∣f(LN(T(a),m))g(LN(T(t),m))F2

    • 结合上3式子,定义具有模态模仿的最终损失函数为

    • L = L s o f t + λ 1 L i n ′ + λ 2 L o u t + λ 3 L r e l L=L_{soft}+\lambda_1L'_{in}+\lambda_2L_{out}+\lambda_3L_{rel} L=Lsoft+λ1Lin+λ2Lout+λ3Lrel

  • 其中,Lsoft为软logit蒸馏目标,λ1、λ2、λ3为寻求损失函数之间平衡的超参数。在选项卡。Feat和Rel表示使用的时代数(L 'in, Lout), Lrel, Layer表示我们使用的中间层数。结果表明,在不同的环境下,模块模仿对学生的RIFormer有积极的影响。在4层设置和仿射算子的使用下,我们得到了75.13%的最佳结果,已经超过了上表中PoolFormer-S12的75.01%的结果。从现在开始,我们将使用这个设置。

  • 备注3。我们认为这种现象的原因可能是模块模仿帮助仿射算子隐式地受益于教师令牌混合器的监督,同时不会失去显式合并前LN层的便利性。此外,我们发现模块模仿可以有效地将特征分布更接近教师网络,并显示出更大的有效接受场(erf)。

  • 准则4:接受野大的教师有利于训练接受野小的学生。下表比较了不同教师架构下学生的表现。虽然GFNet-H-B并没有在教师中显示出最高的ImageNet top-1的表现,但无论是否使用模块模仿,它仍然可以作为更好的选择。

    • 在这里插入图片描述

    • 不同教师对RIFormer-S12 w/o模块模仿(MI)的结果。*表示ImageNet-22K预训练。

    • Hard标签蒸馏可以将模型性能从72.31%提升至73.51%。这说明带TokenMixer的老师模型对于无TokenMixer的学生模型可以起到正面促进作用了;Soft标签蒸馏可以取得最佳性能,学生模型性能可提升至74.05%。

    • 总体来看:监督式训练范式看起来并非无TokenMixer架构的最佳训练方式;带TokenMixer的老师模型有助于引导训练,但仍无法弥补因移除TokenMixer导致的性能损失,即还需要引入其他策略。

  • 备注4。这一事实可能是由于师生之间的接受野差距造成的。归纳偏置可以通过蒸馏从一个模型转移到另一个模型。根据本研究,具有大接受场的模型(例如,在频域具有可学习全局滤波器的GFNet)可以更好地指导具有有限接受场的学生RIFormer。

  • 准则5:将预训练好的教师模型权值(token mixer除外)加载到学生模型中,提高收敛性和性能。我们的方法可以归类为一种模型压缩技术,旨在消除基本块中的令牌混合器以实现加速。在知识蒸馏、量化和模型加速等方法的启发下,我们探索了一种合适的初始化方法,这些方法使用(或部分使用)预训练的重网络的相应权重来初始化轻量级网络的权重。由于我们的目标是只删除令牌混合器,因此剩余部分的权重仍然存在,并且在之前的旅程中没有得到足够的关注。我们观察到,用相应的教师网络初始化RIFormer的权重(除了仿射算子)进一步将性能从75.13%提高到75.36%。这给我们带来了训练RIFormer的最后一个范例。

  • 到目前为止,我们已经完成了探索并发现了训练RIFormer的合适范例。它具有与MetaFormer大致相同的宏设计,但不需要任何令牌混频器。利用所提出的优化方法,RIFormer可以在ImageNet-1K分类中优于带有令牌混合器的复杂模型。这些令人鼓舞的发现激励我们在下一节回答以下问题。1)这种极其简单的架构与我们的训练范式的缩放行为。2)范式对不同教师的普遍性。

Experiments

Image classification

  • 设置。对于具有1.2M张训练图像和50000张验证图像的ImageNet-1K,我们通常在[meta former]中应用训练方案,同时遵循第4节的指导方针。数据扩充包括MixUp、CutMix、CutOut和RandAugment。作为一个移除令牌混频器的模型压缩工作,消除因移除令牌混频器而导致的性能差距绝对是我们的首要任务,而不是提出一个强大的基线。因此,我们延长了600次的训练周期。我们还微调了30个epoch的预训练模型,输入分辨率为 38 4 2 384^2 3842。更多细节见附录。

  • 主要的结果。下表显示了RIFormer对ImageNet分类的结果。我们主要关注吞吐量指标,因为我们的主要考虑是满足边缘设备的延迟要求。正如预期的那样,由于与其他类型的主干相比,RIFormer在其构建块中不包含任何令牌混合器,因此实现了有利的速度优势。

    • 在这里插入图片描述

    • 在ImageNet-1K上使用不同类型令牌混合器的模型结果。表示长达600次的训练。‡表示ImageNet预训练模型对30个epoch进行微调。

  • 令人惊讶的是,通过如此快速的推理,RIFormers成功地使用我们的训练方法删除了所有令牌混频器,而不会影响性能。例如,RIFormer-M36可以以每秒 22 4 2 224^2 2242分辨率处理超过1185张图像,最高精度为82.6%。相比之下,最近的基线PoolFormer-M36带有池化令牌混合器,可以处理1009张相同大小的图像,准确率更低,为82.1%。我们还比较了一个高效的主干,GFNet。它通过一个全局过滤器进行令牌混合,该过滤器由FFT、元素智能乘法和IFFT组成,总计算复杂度为O(N log N)。在吞吐量为939的情况下,GFNet-H-B的准确率为82.9%,而我们的RIFormer-M48仍然可以达到82.8%的准确率,同等吞吐量为897。同时,推理时间RIFormer主体仅以1 × 1的LN变换为主,无需复杂的2D FFT或关注,便于硬件专门化。

  • 值得注意的是,如果没有令牌混合器,RIFormer甚至不能在其构建块中执行基本的令牌混合操作。然而,ImageNet实验表明,使用所提出的训练范式,RIFormer仍然显示出有希望的结果。我们只能认为这背后的原因可能是优化策略起着关键作用。RIFormer很容易成为探索优化驱动的高效网络设计的开始配方,并通过先进的训练方案确保性能。

  • 以PoolFormer-S12为起点,将其中用于TokenMixer的Pooling操作替换为Identity,同时采用常规监督方式进行训练,可以看到:RIFormer-S12会导致不可接受的性能下降(下降约2.7%)。也就是说,当没有TokenMixer操作时,常规的监督式训练在辅助模型学习有用信息时存在一定局限性,需要更先进的训练机制

Ablation studies

  • 模块模仿的有效性。基于蒸馏的模模仿是额外仿射算子学习合适权值的重要方法。因此,我们将其与隐藏状态特征蒸馏方法(带关系)进行了比较。采用无CE损失的软蒸馏,我们得到下表所示的结果。更多细节可以在附录的第4节中找到。特征精馏的精度比模组模仿的精度低0.46%,说明模组模仿对额外权值的优化有积极的作用。

    • 在这里插入图片描述

    • Ablation study of the effectiveness of module imitation.

  • 不同加速策略的比较。接下来,我们讨论删除令牌是否比其他稀疏化策略更好。基于PoolFormer基线,我们首先通过将深度减小到9并保持其宽度(即嵌入维数)的 5 6 \frac 56 65 来构建一个精简的PoolFormer- s9和PoolFormer- xs12,以获得与我们的RIFormerS12相当的推理速度。我们还遵循第4.2节中的软蒸馏范例。下表显示了结果。如果没有耗费延迟的令牌混频器,直接修剪深度或宽度无法提供比我们更好的性能。

    • 在这里插入图片描述

    • Results of comparison with depth or width slimming.

  • 归纳到不同的老师。为了验证所提出的训练范式是一种通用的压缩技术,我们对学生采用 metaformer中的架构修改,并将教师改为其他4个MetaFormer基线,其中教师令牌混合器分别作为rand矩阵、池化、可分离深度卷积和注意力。从下表可以看出,我们的方法在不同的深度设置和教师情况下都有积极的效果。

    • 在这里插入图片描述

    • Results of generalization to other teachers.

Analysis of Module Imitation.
  • 模块模仿(MI)使RIFormer模型的特征分布更接近教师。模块模仿的效果如下图所示。可以看出,poolformer - s12和RIFormer-S12在阶段1和阶段4的特征分布上存在明显差异。应用本文提出的模块模仿后,RIFormer-S12的分布基本向poolformer - s12的分布转移,说明其对帮助学生从老师那里学习有用知识的效果。

    • 在这里插入图片描述

    • PoolFormer-S12和RIFormer-S12第一阶段和最后阶段特征分布的可视化。

  • 模块模仿有助于显示更大的有效感受野(ERF)。ERF反映了经过训练的模型可以响应的图像区域的大小或捕获的对象的大小信息。通过测量输入的每个像素对输出特征中心点的累计贡献来可视化ERF。由于RIFormer删除了所有令牌混合器,因此它显示的ERF预期比PoolFormer小得多,如下图所示。整个区域只有一个正方形的像素出现红色,比PoolFormer小得多。然而,令人惊讶的是,我们可以观察到,经过模块模仿训练后,红色广泛分布在各个位置。似乎虽然没有明显的结构变化,但模块模仿仍然有助于改变学习权值,并显示出更大的ERF

    • 在这里插入图片描述

    • 使用/不使用模块模仿时PoolFormer-S12和RIFormer-S12的有效接受野(ERF)。

Limitations and Conclusion

  • 本文研究了去除视觉主干网中基本构建块的令牌混合器的方法,因为它们具有相当大的延迟成本。为了保持剩余的结构仍然有效,我们彻底地重新审视了训练范式。我们发现,适当的优化策略可以有效地帮助无令牌混合模型从另一个模型中学习有用的知识,提高其性能,并弥补结构不完全造成的差距。限制是更多的视觉任务,包括检测,去模糊等没有讨论,我们将在未来的工作。

A. Detailed hyper-parameters of Sec.4

  • 本文主要提供了第4节探索路线图的一些实验设置。通常,我们在本节中使用RIFormer-S12模型,该模型在ImageNet-1K上进行了120 epoch的训练和评估。在所有情况下,我们采用批大小为512的AdamW优化器。在Sec.4.2和Sec.4.3的蒸馏实验中,GFNet-H-B作为老师,之后进行logit蒸馏。

B. Proof of Eq.4

  • 设T(a)∈RN×C×H×W, T ’ (a)∈RN×C×H×W分别作为主论文图2-(a)中一个仿射残差子块的输入和输出。在训练期间,我们有:

    • T ′ ( a ) = A f f i n e ( L N ( T ( a ) , µ , σ , γ , β ) , s , t ) − T ( a ) ,( 10 ) T′(a) = Affine(LN(T(a), µ, σ, γ, β), s, t) − T(a),(10) T(a)=Affine(LN(T(a),µ,σ,γ,β),s,t)T(a),(10

    • 式中LN为LN层,µ、σ、γ、β为LN层的均值、标准差和可学习的尺度因子和偏置,Affine为仿射变换,s∈RC, t∈RC为其可学习的尺度和移位参数。在推理期间,我们有:

    • T ′ ( a ) = L N ( T ( a ) , µ , σ , γ ′ , β ′ ) ,( 11 ) T′(a) = LN(T(a), µ, σ, γ′, β′),(11) T(a)=LN(T(a),µ,σ,γ,β),(11

    • 根据训练(Eq. 10)与推理(Eq. 11)时仿射残差子块的结构重参数化的等价性,对于∀1≤n≤n、∀1≤i≤C、∀1≤h≤h、∀1≤w≤w,我们有:

    • 在这里插入图片描述

    • 式12可重新表述为 γ ′ i = γ i ( s i − 1 ) , β ′ i = β i ( s i − 1 ) + t i , γ′i = γi(si − 1), β′i = βi(si − 1) + ti, γi=γi(si1),βi=βi(si1)+ti,即等式4得证

Code in PyTorch

PyTorch-like code of our affine operator
  • 我们在Alg. 1中提供了类似pytorch的仿射变换代码,该代码附属于RIFormer块中的训练时间模型。通过将核大小指定为1,并将组号指定为输入通道,可以将仿射变换实现为深度卷积。在实现期间由于残差连接而对输入进行减法。

    • 在这里插入图片描述
PyTorch-like code of the RIFormer block
  • 我们提供了类似pytorch的RIFormer块代码,并在算法2中进行了结构重参数化。

    • 在这里插入图片描述

Visualization of the learned coefficients

  • 为了进一步评估所提出的模块模仿算法的效果,我们可视化了具有(黑色虚线上方)或不具有(黑色虚线下方)模块模仿技术的仿射算子的权值(表示为s)的学习系数。具体来说,我们提供了学习到的浅块(阶段1,block 1)、中间块(阶段3,block 6)和深块(阶段4,block 1)的仿射权值。如下图所示,使用模块模仿训练的仿射权值与未使用该技术训练的仿射权值存在差异。以下图-©为例。未使用模数模拟的仿射权相对更一致,呈现出更多的正数值。相比之下,模块模仿可以帮助仿射算子学习更多的负权重,这可能有助于我们的RIFormer的表达性。类似地,在下图-(b)中,与未使用该方法的仿射权值相比,使用模块模拟的仿射权值具有更适中的振幅。

    • 在这里插入图片描述

    • (a)阶段1,Block 1, (b)阶段3,Block 6, ©阶段4,Block 1中仿射变换的学习系数热图。学习系数的值用不同的颜色表示正数和负数。

Visualization of the activation parts

  • 在[Metaformer is actually what you need for vision]之后,我们提供了由GradCAM获得的四种不同预训练主干的定性结果,分别是RSB-ResNet50、DeiTS、PoolFormer-S24和我们的RIFormer-S24。从[Metaformer is actually what you need for vision中可以看出,transformer模型的map中的激活部分是分散的,而convnet的map中的激活部分是聚集的。有趣的是,还可以做两个额外的观察。首先,用所提出的模块模仿算法训练的RIFormer似乎结合了convnet和transformer的特性。我们认为原因可能是RIFormer具有与transformer相同的通用架构,但没有任何注意(即令牌混合器),因此它本质上是一个convnet。其次,RIFormer映射中的激活部分显示出与PoolFormer相似的特征,这可能是由于通过知识蒸馏过程隐含地从教师模型中引入了归纳偏差。

    • 在这里插入图片描述

    • 在ImageNet-1K上四个不同的预训练主干的Grad-CAM激活图。我们从验证集中抽取4张图像进行可视化。

猜你喜欢

转载自blog.csdn.net/weixin_43424450/article/details/134474122