TransNeXt: Robust Foveal Visual Perception for Vision Transformers
Dai Shi 独立研究人
TransNeXt的命名和其模型本身关联我尚未发现，或许只是为了蹭一蹭ConvNeXt的热度。作者也是借鉴了非常多的设计技巧, 这篇论文很利于了解最新的各种Tricks。
Aggregated Pixel-focused Attention
像素聚焦注意力，为了贴合生物视觉中的聚焦机制，能同时捕捉图像中一个点局部的信息和全局的信息。
人在观察一个物体时，是“打量”式的扫视，而非一眼望到底。见下图。
作者提出的 Pixel-focused Attention (PFA) 的计算图设计为两路分支. 我们先假设输入和输出的特征图尺度相同, 那么对于输出图的每一个点 (i, j) 的像素值该如何计算呢? 两路分支有着不同的计算方式: 第一路关注局部, 采用 swin 模式, 输入为一个
k
×
k
k \times k
k×k 的小窗; 第二路关注全局, 采用池化, 输入为整个特征图. 然后两路分支的 affinity matrix 需要进行拼接, 与value相乘, 最后再过一层 Linear 即可得到像素值.
到这里整个 PFA 的计算图设计思路就讲解完毕了, 而文章在网络结构的细节上下了不少功夫.
Activate and Pool
作者认为, 在关注全局的分支中直接使用无参的全局平均池化有些太蠢了, 可以设计得更加高级一些, 何不添加一层 Linear 呢?
σ
(
X
)
=
L
a
y
e
r
N
o
r
m
(
A
v
g
P
o
o
l
(
G
E
L
U
(
L
i
n
e
a
r
(
X
)
)
)
)
\sigma(X) = LayerNorm(AvgPool(GELU(Linear(X))))
σ(X)=LayerNorm(AvgPool(GELU(Linear(X))))
看似复杂, 实则普通的 Linear 添加了平均池化, 而作者表示这一层是能提升信息压缩率的. 不可置否, 但这里也警示了我一点: 都2024年了, 激活函数要用 GELU, 归一化要用 LayerNorm!
Query embedding
近年来 Multi-Modal 任务研究如火如荼, 例如 Vision-language 模型, 就是来同时处理视觉和文本的模型. 为了能够更好完成结合文本的图像理解等任务, 需要解决跨模态的信息交互问题. 例如 QKV 注意力机制即可设计为两种模态的输入, Query 接受文本向量, Key 接受图像展平向量, 从而可以完成视觉问答任务 (Visual Question Answering). 不难发现, 文本向量长度远远小于图像向量长度, 这样会让 Query 的信息量很少. 为了缓解这个问题, 我们可以直接在 Query 之前拼接一定长度的可学习向量.
这种思路可以推广到一般的注意力层中, 我们都可以无脑地为 Query 加一个前缀, 从而补充信息.
但是在单纯的视觉问题中, 一般采用自注意力, Query 并非文本. 既然没有文本, 我们就造个"文本"出来, 视觉的下游任务是明确的, 那么这些下游任务的名称即可作为文本嵌入, 然后再补充可学习前缀. 最终作者把这一"人造"向量嵌入到原始 Query 中, 即为本文的 Query embedding, 这类新机制又叫 Learnable-Key-Value (LKV) .
Positional attention
既然 Query 能够添加可学习的向量, Key 能不能如法炮制呢? 作者直接造了一个完全不依赖输入的可学习的 Key, 然后与Query进行结合.
设计理念也有的说, 如此设计可以打破原始 QKV 中 KV 一一对应的映射, 而是多个 value 对应一个 Key, 每个 value
对应着各自不同的 Query, 能够在训练过程中使得 Key 接收到多个 Query 的信息. 从反向传播的角度来看, 如此设计能使训练完毕的 Key 关注到各个 Query, 最主要的是 Query 的位置信息: 如果两个 Query 对应的 Value 相似, 则暗示其位置可能相距不远, 这种相对位置关系能够被 Key 隐式地学习到, 所以这种自定义的可学习 Key 有点充当位置编码的角色.
注意, 位置注意力并非代替了所有的注意力, 而是另起一路, 直接与 Query "相乘"了. 其余路仍然是基于自注意力, 并添加了 Query embedding 的.
Length-scaled cosine attention
作者还感觉, Query 和 Key 直接使用 scaled dot product attention (即两者点积) 也很蠢, 采用余弦计算 scaled cosine attention 更好. 而余弦计算中通常会乘以一个可学习的因子
λ
\lambda
λ. 这种因子的设计也该有讲究, 于是作者从熵不变性的角度出发, 解释说点积注意力的熵大约是:
H
i
=
l
o
g
n
−
0.24
λ
b
+
σ
(
1
)
\mathcal{H}_i = logn -0.24\lambda b + \sigma(1)
Hi=logn−0.24λb+σ(1)
所以可以设计
λ
\lambda
λ 为
λ
=
τ
log
n
\lambda = \tau \log{n}
λ=τlogn ,设计
τ
\tau
τ 为可学习的. n为序列长度.
那么定义的 length-scaled cosine attention 为:
A
t
t
e
n
t
i
o
n
(
Q
,
K
,
V
)
=
s
o
f
t
m
a
x
(
τ
log
N
∗
Q
^
K
^
T
)
V
Attention(Q,K, V) = softmax(\tau\log{N}*\hat{Q}\hat{K}^T)V
Attention(Q,K,V)=softmax(τlogN∗Q^K^T)V
其中
N
=
ρ
+
σ
−
μ
N=\rho +\sigma-\mu
N=ρ+σ−μ,
ρ
\rho
ρ 为一路分支中局部窗口的大小,
σ
\sigma
σ 为另一路分支中全局窗口的大小, 而
μ
\mu
μ 为填充的数量.
Position bias
对于两路分支, 都设计一个位置偏置, 尤其是在全局分支中. 全局分支中, 作者采用了 log-CPB 来计算偏置, 这是一个二层神经网络. 而局部分支中, 作者直接设计了一个可学习的向量, 因为局部的位置信息不太缺乏.
Aggregated attention
最终整合上述, 得到了 aggregated pixel-focused attention, 简称为 Aggregated Attention (AA).
Convolutional GLU
在 ViT 出世之前, 有个CNN算法一度刷新了Image-Net的SoTA, 即 Squeeze-and-Excitation (SE) 机制.
SE强在两方面, 一方面设计了通道注意力来门控各个通道的通量, 另一方面使用了全局平均池化来获取全局感受野. 当然由于 ViT 在全局感受野方面已经无人能敌, 但是 ViT 没有通道注意力的机制. 我们何不试着将其缝合入模型中以查漏补缺呢?
作者基于一个原始的门控通道 Gated Linear Unit (GLU) 方法进行了改进, 原始的门控见下图:
而改进的如下:
只是添加了一层卷积, 那么称为 Convolutional GLU(ConvGLU).
TransNeXt总体架构
总体的架构也没啥好说的, 直接看图吧.
思考
这篇文章是一个大杂烩, 基本将能加的小 trick 都加了一遍, 对于快速了解近些年的发展还是有很大帮助的.
另外值得注意的点是, 本文的启发灵感来源于生物视觉, 而文章提出的方案实际上和生物视觉有些不同. 我们的聚焦系统是一个深度系统, 而非固定窗口大小的局部感知. 如何利用RGB-D图像或者双目图像或者3D感知来进一步提升模型性能是一个值得思考的问题.
扫一扫
522
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
