IA-RED²：视觉变换器的可解释性冗余降低

尽管transformer在视觉任务中取得了惊人的成绩，但仍然存在计算量大和内存成本高的问题。为了解决这个问题，作者提出了Interpretability-Aware REDundancy REDuction framework ( )。

作者认为计算量的冗余主要是因为不相关的input patch，因此引入一个可解释的模块用于动态地丢弃冗余patch。分层的网络架构逐渐丢弃不相关的token，大大减少了计算成本。在图像识别任务中，较DeiT有1.4倍加速；在视频动作识别任务中，较TimeSformer有4倍的提升。经过实验可以看到，该模型以极小的准确性为代价，可以达到效率和解释性的双赢。

论文信息

论文地址： https://arxiv.org/abs/2106.12620

论文链接： http://people.csail.mit.edu/bpan/papers/iared-preprint.pdf

代码压缩包： https://link.zhihu.com/?target=http%3A//people.csail.mit.edu/bpan/ia-red/ia-red_files/interpretation_tool.zip

IA-RED²的主要贡献

（1）第一个可解释性感知冗余减少的vision transformer框架；

（2）利用动态推理的框架自适应地计算和筛选输入序列的patch token，动态丢弃信息量少的patch，减少输入序列的长度，降低计算代价；

（3）模型无关性和任务无关性，可以进行不同模型和不同任务的实现；

（4）获得良好的解释性的结果，能更精确地感知到图像的具体的有信息的区域。

方法

本文方法建立在vision transformer（ViT）上，通过动态地在原始输入序列中去掉信息量较少的patch，降低vision transformer的冗余，以最小的计算量进行正确分类。

多头解释器Multi-head Interpreter

给定一个带有position信息的patch token序列 ，利用多头解释器丢弃没有信息的patch token。首先将transformer层平均分为D组，每组包含1个多头解释器和L个MSA-FFN块。patch tokens经过多头解释器得到informative score （i是输入token的位置，j是第几组），将ptch token和policy token进行线性映射后再内积相乘，最后应用sigmoid函数。此处的policy token初始化是全1的向量，后续不断优化不断传入之后的模块。 小于阈值0.5，patch 在第j组将会被完全丢弃，并且在后续组中不再使用。

的公式如下所示，其中 指的是第j个多头解释器的policy token，H指的是多头解释器中head的个数， 指的是在第h个头分别针对patch tokens和policy tokens的线性层，* 指的是点积， 指的是sigmoid激活函数：

class PerBlockAgent(nn.Module):
def __init__(self, dim, num_heads=2, qkv_bias=True, qk_scale=None):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.norm = nn.LayerNorm(dim, elementwise_affine=False)
self.qk = nn.Linear(dim, dim * 3, bias=qkv_bias)

self.mat = matmul()
self.act = nn.Sigmoid()

def forward(self, x, prev_policy=None):
# Prev_mask: BxN
B, N, C = x.shape
if prev_policy == None:
prev_policy = torch.ones(B, N)#在第一个多头解释器里将policy token初始化为1
x = x * prev_policy.unsqueeze(2)#x和policy token相乘
x = self.norm(x)#LayerNorm

qk = self.qk(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
q, k = qk[0], qk[1] # B x num_head x N x C

attn = self.mat(q[:,:,:1], k.transpose(-2, -1)) * self.scale#q和k的转置相乘后进行scale操作
attn = self.act(attn)#sigmoid
attn = attn.squeeze(2).mean(1)
attn = attn * prev_policy #最后和初始策略相乘，保证之前去掉的 token 在后续也去掉
return attn

分层训练计划Hierarchical Training Scheme

训练过程中作者固定了patch embedding、positional encoding和class token的参数，只关注每组内部的参数，还采用了递归学习方式优化网络组。假如有3个组，首先优化1到3组，再2-3组，最后优化第三组。这是为了让解释器在早期能学习对最终正确预测有用的包含所有重要文本信息的patch，后期再聚焦到part-level的信息，此时每个token的信息已经经过了全局的交互和融合。作者通过REINFORCE方法优化多头解释器，用交叉熵损失函数优化MSA-FFN块。

在训练阶段，给定一个patch token序列 输入到第j个多头解释器，多头解释器会为每个输入token生成策略，使用 Bernoulli分布： =1表示保留token， =0表示丢弃token。

除此之外，作者定义了reward函数，鼓励多头解释器用尽可能少的patch token预测正确的结果， 是保留的patch的百分比， 是错误预测的惩罚值，用于控制网络效率和准确性的权衡。

def forward_agent(self, x, prev_policy=None, sample=False, mask=False):
probs = self.agent(x, prev_policy)#生成policy token
policy_map = probs.data.clone()
policy_map[policy_map<=0.5] = 0.0 #小于0.5就去掉
policy_map[policy_map>0.5] = 1.0 #大于0.5就保留
policy_map = Variable(policy_map) #可以反向传播
policy_map[:, :1] = 1 #此处保证class token为1

if sample:
if prev_policy is None:
prev_policy = 1
self.distr = Bernoulli(probs)#伯努利的01分布
policy = self.distr.sample()
policy = policy * prev_policy #保证之前被drop掉的token也不会被使用
policy[:, :1] = 1 #此处保证class token为1
return policy, (probs, self.distr)
if mask:
policy_map = prev_policy
return policy_map, probs

更新解释器时，用reward函数的期望梯度来更新多头解释器的参数，再预测模型结果用交叉熵损失函数去更新之后层的transformer结构。起初是训练前面层的解释器，然后固定解释器去微调后续的所有MSA-FFN块。以第j组为例，训练解释器再固定参数，优化后续的第j组...第D组的MSA-FFN块，此时后续组的解释器是被mask的并且保留所有的token。

可解释性的出现Emergence of Interpretability

通过将不同组的多头解释器预测的信息分数可视化，可以看到输入patch的冗余在不同层次上逐级减少。对于那些在先前组中被删除的patch，将它们的信息得分视为零。因此，我们可以从每个网络组中获得长度等于原始输入序列长度的信息分数序列，重新排列这个得分序列，并将其插值回输入视觉内容(如图像或视频)的大小，就可以为每个网络组画一个热图来解释冗余的部分。

实验

数据集和指标：作者在 ImageNet-1k分类数据集上进行图像识别实验，用top-1和top-5的正确率进行衡量；用ImageNet-Segmentation dataset进行弱监督图像分割实验，用像素精度、平均精度(mAcc)和平均IoU (mIoU)反映分割性能；用Kinetics-400 dataset进行视频动作识别实验，用clip-1和video-1误差评估实验性能。

实现细节：针对图像识别任务，组数D=3，MSA-FNN块数L=4。整个网络有D×30个epoch，每30个epoch里前10个epoch优化多头解释器，后20个epoch优化所有剩下的MSA-FFN模块。针对视频理解任务，D=1（只在输入层选择信息patch），5个epoch训练多头解释器，1个epoch微调主干网络。

可解释性

定性评估可以看到IA-RED²能更加准确地定位到感兴趣的对象，尤其是part-level的，比如卡车头部、鲨鱼的鳍和头部、动物的面部特征等

图像任务

在弱监督图像分割实验中，IA-RED²明显优于其他方法

在图像识别任务中，IA-RED²的效率和准确性的平衡做得最好

图像识别任务中，可视化在不同的组里分层丢弃patch token，第一组模型考虑了几乎所有的token，第二组模型丢弃了背景模型，关注感兴趣的对象，第三组模型更加聚焦在part-level级上。好的模型解释性能精确定位正确分类预测的信息区域。

视频行为识别的冗余减少策略

可以看到IA-RED²优于注意力基线，而比随机和时间差异基线差。虽然不是最优的结果，但模型学会了在数千个输入patch tokens中识别出信息量大的patch。

其他实验

不同减少率下，IA-RED²和权重剪枝方法互相补充，可以获得更高的精度

计算量少；由于冗余是依赖于输入数据的，所以更高分辨率的输入数据将包含更多的冗余。保持计算成本的缩减率不变，比较模型的精度损失，高分辨率的模型比采用低分辨率的模型性能下降更小，表明高分辨率的输入数据往往包含更多冗余。

结论

作者提出了一种创新的可解释感知的冗余减少框架IA-RED²。研究表明，IA-RED²分层降低了计算成本并有效提高了ViT的速度，同时具有可解释性。作者又分别在图像分类和视频理解任务上进行了实验，证明了IA-RED²是模型无关和任务无关的，同时能在效率和解释性上取得双赢的效果。

来源：知乎

作者：小菜鸟

｜深延科技｜

深延科技成立于2018年，是深兰科技（DeepBlue）旗下的子公司，以“人工智能赋能企业与行业”为使命，助力合作伙伴降低成本、提升效率并挖掘更多商业机会，进一步开拓市场，服务民生。公司推出四款平台产品——深延智能数据标注平台、深延AI开发平台、深延自动化机器学习平台、深延AI开放平台，涵盖从数据标注及处理，到模型构建，再到行业应用和解决方案的全流程服务，一站式助力企业“AI”化。