Press "Enter" to skip to content

【Transformer】ATS: Adaptive Token Sampling For Efficient Vision Transformers

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

文章目录

 

 

一、背景

 

尽管现有的 transformer 模型在分类等任务上取得了较好的效果,但计算量还是很高,需要很多的 GFLOPs,不适用于很多边缘设备,虽然GFLOPs 也可以通过降低网络中 token 数量来降低,DynamicViT 使用网络预测每个 token 的得分,从而判断哪个 token 是冗余的。虽然这个方法能够降低网络的 GFLOPs,但得分预测网络也会引入额外的参数,并且如果想要不同的降参比率需要再次进行训练。

二、动机

 

作者认为,对于分类任务,并非需要图中的所有信息来进行分类,因为图像的信息对分类任务来说是冗余的。所以本文提出了一个降低 token 数量的方法,可以适用于任何 transformer,不受降参比率限制,且更高效。

 

三、方法

 

作者提出了一种名为 ” Adaptive Token Sampler (ATS) ” 的模块,是一种动态的从输入 token 中选择重要的 token 的模块。也是一个 parameter-free 的方法,总体结构如图 2 所示。卷积网络中,一般会使用 pooling 来降低计算量,stage 越深,分辨率越小。但 Transformer 中不能直接使用这样的方法,因为 token 是与空间位置无关的,即改变位置不会影响最后的结果。而且如果使用下采样的话会有两个弊端,其一是会丢失目标的细节信息,其二是可能保留很多背景信息,对分类无实质性作用。所以作者提出了一种动态选择每个stage的token数量的方法。

ATS 的过程:

首先,对 N 个输入 token 分配得分,基于得分来确定哪些留下
然后,设定 K 为保留的 token 最大数量,这个 K 会决定 GFLOPs 的上限
sampled tokens K’ 一般会比 K 小,且和输入图像的关系如图 6 所示

 

对于每个实例,图7展示了作者使用少数或多数的 patches ,就可以得到正确的分类,图 3 展示了不同每个 stage 使用的 token 数量。作者也提出了一个对每个图像选择正确 token 数量的方法。如图 6 所示,不同图像在不同stage的 token 数量是不同的。

 

 

3.1 Token Scoring

 

在标准的 self-attention 层,输入的 Q、K、V 都是从输入 token 得来的,然后会得到 attention matrix A:

由于 softmax 的存在,A 的每行和为 1,输出 token 会和 attention matrix 作用,从而加权。

 

A 的每行包含了输入 token 的 attention weights,这个 weights 其实就表示了所有 token 对输出 token 的作用,因为 A 的第一行是 cls token,表示了输入 token 对输出 classification token 的作用,所以作者使用第一行的元素作为修剪 A 的根据,如图 2 所示。作者也做了归一化,重要程度得分如下,对于多头注意力,分别对每个头进行计算,然后加起来:

3.2 Token Sampling

 

对每个 token 得到 score 之后,就可以根据 attention matrix A 对 tokens 进行修剪了。

 

一个比较基础的做法是直接选择 top-K 个 tokens,但是实验结果说明,这种方法没有动态选择 K’ 个 tokens 的效果好。其表现不好的原因在于,直接丢弃了所有得分低的token,但有些 token 其实在浅层可能会比较有用。

 

作者的抽样方法中,从几个相似的 token 中抽象的概率等于这些 token 的得分之和。而且从图 3 中也能看出,本文的抽样机制从浅层抽样的 token 数量比深层的更多一些。

 

方法:

 

因为 token score 是被归一化的,所以可以看出概率,可以计算累计密度函数(CDF):

对 CDF 取反,就得到了采样函数:

四、效果

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注