1简介
Vision Transformers吸引了大量的研究，并成为各种图像识别任务的Backbone之一，如分类、分割和检测。
但MSA引入的ViT的两个局限性已经被认识到。首先，MSA层的一个众所周知的问题是二次时间和内存复杂性，这阻碍了vit的大规模开发和部署，特别是对于长序列的建模。为此，人们已经做出了许多努力来开发高效的Transformer，例如采用修剪方法来修剪不太重要的ViT组件。
另一个基本问题是MSA层缺乏用于编码局部信息的归纳偏差，这对于处理low-level特征(如边缘和形状)至关重要。为了解决这个问题，现有技术在ViT中插入卷积层来引入归纳偏差，例如，在ViT编码器之前或MSA层之前的内部引入前馈网络层(FFNs)。
解决这两个问题的一个统一解决方案是使用NAS剪枝方法搜索出卷积替换部分MSA层。因此，可以消除一些昂贵的MSA层，以显著提高ViT模型的效率，同时添加适当的局部性，而不需要启发式。
最近，一些单次NAS方法建议将候选MSA和卷积操作分别包含到搜索空间中，其中每个操作都被表述为一个单独的可训练路径(参见图1(a))。然而，由于在搜索过程中，每一层的所有候选操作都需要独立维护和更新，因此这些多路径方法会带来昂贵的搜索成本和具有挑战性的优化。
为了解决这一局限性，本文提出了一种新的Single-Path Vision Transformer剪枝方法，称为SPViT；SPViT以低搜索成本和高准确率将预训练的 ViT 剪枝到紧凑的模型中。
具体来说，由于观察到MSA层具有建模局部特征的能力，并且 MSA Head 有时会关注局部特征，作者开发了 MSA 和卷积操作之间的权重共享方案。因此，可以利用预训练的MSA参数获得较好的卷积核权值初始化，并在搜索过程中通过索引MSA中间结果轻松地得到卷积运算的输出。
通过进一步将所有候选操作封装到每个层的单个MSA中形成了如图1(b)所示的单路径搜索方案。不需要像在多路径方法中那样选择不同的操作，相反，这里将搜索问题定义为在每个MSA层中寻找权重子集，从而在选择卷积层时将MSA剪枝到卷积层。
此外，考虑到FFN层消耗了很大比例的计算，例如当使用DeiT处理224×224图像时，17.5G的Mult-Adds中FFN占有11.1G，作者提出了一种简单而有效的方法来搜索细粒度的MLP扩展比率，为每个FFN层定制隐藏维度的数量。具体来说，SPViT使用可学习门来确定每个FFN隐藏维度的重要性，然后去掉不太重要的维度。
本文的主要贡献：
- 提出了一种新的 MSA 和卷积操作之间的权重共享方案，该方案允许在单路径框架中将所有候选NAS操作编码到MSA层中。然后，将搜索问题转换为寻找MSA参数的子集，从而显著降低搜索成本;
- 在单路径方案的基础上，提出了SPViT，它可以自动将昂贵的全局MSA操作精简为轻量化的局部卷积操作，并在期望的效率约束下搜索细粒度的MLP扩展比；
- 在CIFAR-100和ImageNet-1k实验表明SPViT能够显著降低搜索成本，同时性能优于Baseline方法。作者还对搜索的体系结构进行了一些有趣的观察，揭示了ViT的architecture preferences。
2相关工作
2.1 Transformer剪枝
Transformer通常有昂贵的计算成本。为了减少这种成本，模型剪枝是产生高效Transformer的主要方法之一。现有的Transformer剪枝技术大致可以分为module剪枝和token剪枝。
module剪枝对不显著的Transformer modules进行修剪，如MSA层中的heads、linear projections通道和权重神经元。
最近，有学者提出动态识别和修剪不太重要的token特征维度。Chen等人在lottery ticket假设下提出了Transformer参数的prune-and-grow方法。
token剪枝侧重于修剪不太重要的token，如DynamicViT分层地修剪冗余token，从而在分类任务中实现FLOPs减少。然而，随着一些token的消除，将token剪枝方法应用于密集预测任务是一个挑战，如分割。
本文工作旨在对modules进行剪枝，但与之前的工作不同的是，本文着重于对MSA层中多余的全局关联进行剪枝。并没有以多路径的方式将卷积和MSA独立地添加到搜索空间中，而是考虑了这两种操作之间的内在关系，并将MSA层剪成卷积层，这允许以低搜索成本快速部署高效的Transformer模型。
2.2 卷积 vs 自注意力
本文工作是由最近探索卷积和自注意力层的性质的研究推动的。例如，卷积层在提取带有卷积归纳偏差的局部纹理线索方面表现出很强的能力。另一方面，自注意力层倾向于通过建模全局相关性来强调物体形状。
为了使用这两种操作的优点，有一种方法独立对待卷积层和自注意力层，并形成混合模型；例如，在ViTt中插入卷积层或将自注意力叠加在CNN之上。另一项工作探讨了两种操作之间的内在关系。其中，有研究表明，具有可学习的二次位置编码和丰富的Head数量的自注意力层可以表示任何卷积层。ConViT将soft inductive bias作为自注意力层的部分注意力得分，Transformed CNN学习将预处理卷积层转换为自注意力层。
本文的方法也属于后一种工作的范围。与之前的工作相比，针对MSA层强大的表达能力，作者提出了一种新的权重共享方案，表明MSA操作参数子集可以表达卷积操作。权重共享方案能够在单路径框架中将所有候选NAS操作编码到MSA层中，以降低搜索成本。作者还进一步设计SPViT来自动有效地将MSA操作精简为轻量级的卷积操作，同时享受局部性的好处。
3本文方法
3.1 MSA和卷积运算之间的权重共享
1、回顾卷积和自注意力机制
权重共享方案是指在不同的操作之间共享一个参数子集。该方案的有效性已经被广泛的NAS方法所证明，例如，在候选卷积操作之间共享权重大大降低了搜索成本。在本文中提出了一种新的MSA和卷积运算之间的权重共享方案。
卷积层是CNNs的基本构建模块。设和为输入特征，卷积的为kernel-size为k，其中、、和分别为空间宽度、高度、输入维数和输出维数。标准卷积层将局部感受野中的特性进行聚合，这些聚合的特征被定义为：
形式上，设为特征宽度和高度的索引集，处的标准卷积层的输出可以表示为:
Transformers以MSA层为主要构件。然而，MSA层的感受野更大，覆盖了整个序列，而不是只聚合邻近的特征。将Eq.(1)中的输入X作为维度嵌入的集合，并让作为任何关键特征嵌入的索引，可以定义处MSA层的输出为:
式中，，，，，为MSA层的Head的数量。给定第h个head 的输出维度为，定义、和作为相应的可学习value、query和key线性投影。因此，表示第h个head的注attention map，在和的索引下，成为一个标量。为了简单起见，省略了MSA层中的位置编码，以及卷积和MSA层中伴随可学习投影的所有可学习偏置项。
MSA层的计算复杂度为。在这种情况下，当建模高分辨率特征()，新二次项建模全局相关性主导计算,计算复杂度也会随之剧增。同时，在标准卷积层输出具有相同宽度和高度的特征映射时，卷积层的计算复杂度仍然为。
2、权重共享方案
首先，由于卷积运算只处理局部区域内的特征，因此在用MSA参数表示卷积运算时，每次嵌入时，将非局部相关的注意力分数固定为0，将局部相关的注意力分数固定为1。
代回式(2)，得到：
在这里，使用来表示Eq.(3)中具有局部性限制的输出。在Eq.(4)中，强制MSA操作的head只关注以为中心的大小为的局部窗口内的位置。
其次，可以进一步将MSA操作参数配置为卷积内核。一种方法是定义一个bijective映射，它将选定的head子集中的head分配到局部窗口中的特定位置。
然后可以用带入h，即
在Eq.(5)中，每个head都关注局部窗口中的某个位置。然而，许多ViT变体都有，如DeiT-Ti和DeiT-S，这将限制权重分配方案仅适用于大型架构。此外，在MSA层中，不同的head倾向于关注不同的区域，因此，定义子集和映射来识别最适合局部位置的head是很重要的。为了使权重分配方案适用于一般ViT架构，并减轻定义和的困难，作者使用可学习参数和softmax函数将MSA集合到卷积kernel位置:
在这种情况下，具有缩放因子的MSA head集合学会关注局部窗口内预定义的位置。通过在z上使用softmax函数来确保head集合保持原来的比例。
然后，通过将V替换为并应用结合律，可以将剖分成Eq.(6)中的卷积核:
相反，出于复杂性考虑，选择bottleneck卷积运算并进行修改:
注意，在这种形式下，卷积运算的输出维数，通过将投影回。可以很容易地发现，Eq.(7)中的和Eq.(10)中的都不大于，即head的尺寸。但前者的计算复杂度为，后者的计算复杂度为。由于，因此选择使用bottleneck卷积来获得更好的效率。
进一步添加BN和ReLU，得到由和表示的kernel size为k的bottleneck卷积运算:
3、Remark
提出的权重共享方案鼓励MSA操作的head不仅建模全局区域，同时学习集合来处理局部区域特征。对这种方法的直观认识主要来自两个方面。
- 首先，局部区域确实是全局区域的一部分，因此MSA head集合自然具有处理局部区域的能力。
- 其次，由于ViT编码器中的一些head经常关注query像素周围的局部区域，用这些head表示卷积操作的行为可能会得到合理的输出。然而，考虑到不同的ViT模型添加locality的最佳位置可能有所不同，因此引入SPViT，它自动学习选择最优操作。
3.2 Single-path ViT剪枝
在权重共享方案的基础上进一步提出SPViT，通过自动将部分MSA层裁剪为卷积层，得到最优的精度-效率权衡的架构。此外，还研究了控制FFN层中隐藏维度数量的细粒度MLP扩展比率。
1、Search formulation
- 搜索MSA和卷积层
第一步：在ViT中寻找最优的位置将MSA层剪成卷积层。具体来说，在每个ViT块中，将每个原始的MSA层替换为一个Unified MSA(UMSA)层，该层通过3.1节描述的权重共享方案同时获得MSA的输出和卷积操作。
然后，定义一系列可学习的二进制门对应的操作，以编码体系结构配置。
为此，在UMSA层中，当进行MSA操作时，可以根据Eq.(8)对V进行索引和缩放，直接得到预定义核的卷积输出，而不需要将每个候选卷积操作保持为单独的路径。
图2(a)
这个过程如图2(a)所示，然后通过Eq.(12)得到卷积运算输出BConvk(X)。
通过这种方法可以以较低的计算成本得到了BConv卷积和MSA的输出。
根据计算复杂度由低到高的顺序可以将BConv卷积和MSA表示为一个有序集P，其中|P| = |K|+1。
第二步：定义一系列二进制门来编码操作选择。对于任何具有指数的操作，从具有可学习参数的伯努利分布中采样，该分布可以通过梯度自动优化。
具体来说，使用确定选择第p个操作的likehood:
出于效率考虑，在每个层中只选择一个操作。通过这种方式，强制每个UMSA层最多有一个打开的gate，将替换为:
式(14)本质上说明，一旦选择了一个更复杂的门，所有不那么复杂的门都将关闭。因此，UMSA层的输出可以表示为:
其中是第p个操作的输出。然后，每个UMSA层后面跟着残差连接和层归一化，就像ViT做的那样。所有门关闭时采用shortcut操作。因此，UMSA搜索空间包括3种类型的候选操作:ShortCut、不同kernel size的BConv卷积以及MSA。
- 搜索细粒度的MLP扩展率
MLP扩展率控制FFN层的隐藏尺寸数量。之前的工作将MLP扩展率的搜索空间定义为一个粗粒度值列表，例如[3,3.5,4]。相反，给定预训练模型定义的MLP扩展率，首次尝试通过一个简单但有效的Unified FFN(UFFN)层搜索每个FFN层的细粒度MLP扩展率()。给定输入特征X, FFN层的输出可以表示为:
这里用和来表示预定义的FFN层的MLP扩展率和隐层维数。和是可学习的全连接层，前者将X的通道维数映射到，后者将通道维数映射回。GeLU激活被添加在2个可学习的全连接预测之间。
与式(13)一样，通过编码具有二进制门的FFN隐层维配置来寻找细粒度的MLP扩展率。通过在搜索过程中对每个隐藏维度应用二进位门，定义 UFFN层输出为：
这可以删除不重要的隐藏维度。
class UnifiedMlp(nn.Module): def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0., theta=0., ffn_indicators=None): super().__init__() self.in_features = in_features out_features = out_features or in_features self.hidden_features = hidden_features = hidden_features or in_features self.act = act_layer() self.drop = nn.Dropout(drop) if ffn_indicators is None: self.fc1 = nn.Linear(in_features, hidden_features) self.fc2 = nn.Linear(hidden_features, out_features) # Threshold parameters self.register_parameter('ffn_thresholds', nn.Parameter(torch.tensor([theta] * hidden_features))) # The indicators self.register_buffer('assigned_indicator_index', nn.Parameter(torch.zeros(hidden_features))) self.fine_tuning = False else: self.fc1 = nn.Linear(in_features, ffn_indicators.nonzero().shape[0]) self.fc2 = nn.Linear(ffn_indicators.nonzero().shape[0], out_features) self.fine_tuning = True def forward(self, x): if not self.fine_tuning: return self.search_forward(x) else: return self.finetune_forward(x) def search_forward(self, x): ffn_probs = F.sigmoid(self.ffn_thresholds) if self.training: ffn_indicators = bernoulli_sample(ffn_probs) else: ffn_indicators = (ffn_probs > 0.5).float() x = self.fc1(x) x = self.act(x) x = ffn_indicators.unsqueeze(0).unsqueeze(0) * x x = self.drop(x) x = self.fc2(x) x = self.drop(x) # We derive the FFN indicators by expectation, and # ffn_indicators are kept to calculate complexity loss self.ffn_indicators = (ffn_probs > 0.5).float() - torch.sigmoid( ffn_probs - 0.5).detach() + torch.sigmoid(ffn_probs - 0.5) self.register_buffer('assigned_indicator_index', self.ffn_indicators) return x, torch.sum(self.ffn_indicators) def finetune_forward(self, x): x = self.fc1(x) x = self.act(x) x = self.drop(x) x = self.fc2(x) x = self.drop(x) return x, None
- 搜索目标
为了获得具有理想效率约束的架构，利用辅助计算复杂度损失来优化网络。具体地说，定义了一个查找表，其中包含候选操作和模块的计算复杂性。计算代价定义为:
其中F(X)为当前网络计算复杂度，为目标复杂度。最后，将总体搜索目标定义为， 为超参数，用来平衡和交叉熵损失。
2、Fine-tuning
在微调过程中，使之前3.2.1节中的随机二进制门具有确定性，即:
通过应用Eq.(14)，得到，并在微调过程中在UMSA层中最多选择一个操作。输出可以表示为:
图2(b)
如果选择的操作是bottleneck convolution，可以将训练良好的配置为一个卷积kernel，以寻求高效的硬件实现，如Eq.(10)所示，如图2(b)所示。结合Eq.(12)，可以将微调时的bottleneck convolution定义为：
类似地，对于UFFN层，可以自动找到隐藏维度T的子集，该子集表示选定的隐藏维度。因此，为，微调时的UFFN层输出可由。在微调期间，将目标损失设置为。
