Non-Local Attention
\colorbox{orange}{Non-Local Attention}
Non-LocalAttention 设
Q
\mathbb{Q}
Q为搜索范围,查询项(Query items)为
y
i
y_i
yi,关键项(Key items)为
y
j
y_j
yj,且
ϕ
(
y
i
,
y
j
)
\phi(y_i,y_j)
ϕ(yi,yj)表示二者的相似度,
G
(
⋅
)
G(\cdot)
G(⋅)表示embedding function,则自注意力机制可表示为:
x
i
^
=
1
z
i
∑
j
∈
Q
ϕ
(
y
i
,
y
j
)
G
(
y
j
)
,
∀
i
.
(1)
\hat{x_i} = \frac{1}{z_i}\sum_{j\in \mathbb{Q}} \phi(y_i,y_j) G(y_j), \forall i.\tag{1}
xi^=zi1j∈Q∑ϕ(yi,yj)G(yj),∀i.(1)其中
i
∈
I
i\in\mathbb{I}
i∈I表示图像中所有的pixel/patch序号,且
Q
⊂
I
\mathbb{Q}\subset \mathbb{I}
Q⊂I;
y
i
,
y
j
y_i,y_j
yi,yj都是向量,每个向量都代表着1个patch;
z
i
=
∑
j
∈
Q
ϕ
(
y
i
,
y
j
)
z_i=\sum_{j\in\mathbb{Q}}\phi(y_i,y_j)
zi=∑j∈Qϕ(yi,yj)表示归一化值;
G
(
y
j
)
G(y_j)
G(yj)是
y
j
y_j
yj的另一种表达;
x
i
^
\hat{x_i}
xi^是结合注意力之后关于
y
j
y_j
yj的新的表达,这个表达结合了
y
i
y_i
yi和
y
j
y_j
yj的相关性。
在神经网络中,也就是将它用在Transformer中,
y
i
,
y
j
y_i,y_j
yi,yj分别表示
Q
、
K
Q、K
Q、K,而
G
(
y
j
)
G(y_j)
G(yj)往往是一个卷积层的输出结果(也就是
V
V
V,因此
K
、
V
K、V
K、V是同一个token的不同表达),
ϕ
(
y
i
,
y
j
)
=
y
i
T
y
j
\phi(y_i,y_j)=y_i^T y_j
ϕ(yi,yj)=yiTyj。所以式(1)表达的就是第
i
i
i个patch在搜索范围内找出所有的
y
i
y_i
yi并通过彼此之间的相似度来查询相应的
G
(
y
i
)
G(y_i)
G(yi)值,最后进行加权平均之后的值就是第
i
i
i个patch结合所有patch相关性的输出值。
Note:
Softmax计算的是每一个patch对其他patch的相似度:上图表达的就是式(1),
x
i
^
\hat{x_i}
xi^表示更新之后的值,其是token的带注意力的表达。
在Transformer中,
K
、
V
K、V
K、V是同一个token的不同表达,相当于key-value的关系,
Q
Q
Q往往是另一个token的表达。
Local Attention
\colorbox{yellow}{Local Attention}
LocalAttention 用局部算子产生的注意力机制称为Local Attention,比如RCAN中的通道注意力,EDVR中的时空注意力机制,它们都是用卷积、池化这种局部手段产生的。优点在于可以捕捉局部相关性,缺点在于感受野相对较小,无法捕捉长距离的相关性。
为什么用
Q
K
T
QK^T
QKT来衡量相似度? 我们知道学习到的
Q
、
K
Q、K
Q、K矩阵的每一行向量代表1个patch,因此
Q
K
T
QK^T
QKT就可以看成是向量之间的内积,如果每个向量都是单位向量,那么
Q
K
T
QK^T
QKT就相当于衡量余弦相似度。但Transformer中并没有对
Q
、
K
Q、K
Q、K做归一化,因此严格来说
Q
、
K
Q、K
Q、K并不能衡量相似度,合理的做法应该是:
s
o
f
t
m
a
x
(
Q
K
T
d
k
⋅
(
∏
i
=
0
N
−
1
∏
j
=
0
N
−
1
∣
∣
Q
i
∣
∣
⋅
∣
∣
K
j
T
∣
∣
)
)
softmax(\frac{QK^T}{\sqrt{d_k}\cdot (\prod_{i=0}^{N-1}\prod_{j=0}^{N-1}||Q_i||\cdot ||K^T_j||)})
softmax(dk⋅(∏i=0N−1∏j=0N−1∣∣Qi∣∣⋅∣∣KjT∣∣)QKT) ,但个人认为可能不进行归一化仍然取得了不错的性能,同时为了节省一定的计算量所以干脆就不加了。
浅层特征提取
\colorbox{hotpink}{浅层特征提取}
浅层特征提取 从原图中提取出浅层特征
F
s
F_s
Fs,具体表达式如下:
F
s
=
H
S
F
(
I
L
Q
)
.
(2)
F_s = \mathcal{H}_{SF}(I_{LQ}).\tag{2}
Fs=HSF(ILQ).(2)其中
H
S
F
(
⋅
)
\mathcal{H}_{SF}(\cdot)
HSF(⋅)表示浅层特征提取算子。
CAB级联模块
\colorbox{yellow}{CAB级联模块}
CAB级联模块 如上图所示,作者级联了4个CAB,其输入是浅层特征
F
s
F_s
Fs,输出是
F
C
A
B
4
F_{CAB}^4
FCAB4,其作为重建模块的输入。 数学表达式为:
F
C
A
B
i
=
H
C
A
B
i
(
⋯
H
C
A
B
1
(
F
s
)
⋯
)
.
(3)
F^i_{CAB} = \mathcal{H}_{CAB}^i(\cdots \mathcal{H}_{CAB}^1(F_s)\cdots).\tag{3}
FCABi=HCABi(⋯HCAB1(Fs)⋯).(3)其中
F
C
A
B
i
、
H
C
A
B
i
F^i_{CAB}、 \mathcal{H}_{CAB}^i
FCABi、HCABi分别表示第
i
i
i个CAB块的输出feature map以及第
i
i
i个CAB块算子。
作者设计2种深层特征提取子网络,一种较为轻量型(Basic)的:它由6个包含
c
o
n
v
→
B
N
→
R
e
L
U
conv\to BN\to ReLU
conv→BN→ReLU级联组成,每个卷积层的size为
3
×
3
×
64
3\times 3\times 64
3×3×64;另一种是增强型(Enhanced)的残差网络结构,由5个残差块级联组成,每个卷积层都是
3
×
3
×
64
3\times 3\times 64
3×3×64。 此外,如果是第一种版本的FEM,则COLA-Net记为COLA-B;第二个版本的FEM,则COLA-Net就记为COLA-E。
DFM模块: DFM模块是CAB的核心,它将局部注意里和通道注意力分别并行处理,最后进行融合产生具备注意力机制的feature map,即
F
C
A
B
i
F^i_{CAB}
FCABi;其中融合部分也是一个注意力机制的过程。
重建模块
\colorbox{tomato}{重建模块}
重建模块 这个重建模块可以理解为3层意思:①深层特征提取;②从feature-wise重建成image-wise;③特征校正、调整。 除了进一步提取CAB输出的特征以外,还要和来自低层的特征进行融合,这部分结构也是被常用于图像重建任务,其有4个好处:①将不同层级的特征进行融合,从而让特征之间互为补充来重建出更好的图像;②构成残差结构,让网络更好的去学习残差信息,一定程度上减轻残差部分的学习压力,增加稳定性;③skip-connection可以作为正则化项;④正如DRCN所述:将原图像直接输送到网络尾端,可以弥补卷积过程造成的特征损失。 具体数学表达式如下:
I
R
E
C
=
I
L
Q
+
H
D
F
(
F
C
A
B
4
)
=
H
C
O
L
A
(
I
L
Q
)
.
(4)
I_{REC} = I_{LQ} + \mathcal{H}_{DF}(F_{CAB}^4) = \mathcal{H}_{COLA}(I_{LQ}).\tag{4}
IREC=ILQ+HDF(FCAB4)=HCOLA(ILQ).(4)其中
H
D
F
(
⋅
)
、
H
C
O
L
A
(
⋅
)
\mathcal{H}_{DF}(\cdot)、\mathcal{H}_{COLA}(\cdot)
HDF(⋅)、HCOLA(⋅)分别为最后一层卷积层以及整个COLA网络。
损失函数
\colorbox{dodgerblue}{损失函数}
损失函数 L2损失函数:
L
(
Θ
)
=
1
B
∑
b
=
1
B
∣
∣
I
H
Q
b
−
H
C
O
L
A
(
I
L
Q
b
)
∣
∣
2
2
.
(5)
L(\Theta) = \frac{1}{B}\sum^B_{b=1}||I^b_{HQ} - \mathcal{H}_{COLA}(I^b_{LQ})||_2^2.\tag{5}
L(Θ)=B1b=1∑B∣∣IHQb−HCOLA(ILQb)∣∣22.(5)
具体而言,FEM输出的feature map作为输入,即
F
F
E
M
i
F^i_{FEM}
FFEMi,然后局部注意里和全局注意力两个并行分支分别进行注意力提取,分别输出
F
L
i
、
F
N
L
i
F^i_L、F_{NL}^i
FLi、FNLi,数学表达式如下:
{
F
L
i
=
H
L
(
F
F
E
M
i
)
,
F
N
L
i
=
H
N
L
(
F
F
E
M
i
)
.
(6)
\begin{cases} F^i_L = \mathcal{H}_L(F^i_{FEM}),\\ F_{NL}^i = \mathcal{H}_{NL}(F^i_{FEM}).\tag{6} \end{cases}
{FLi=HL(FFEMi),FNLi=HNL(FFEMi).(6)其中
H
L
(
⋅
)
、
H
N
L
(
⋅
)
\mathcal{H}_L(\cdot)、\mathcal{H}_{NL}(\cdot)
HL(⋅)、HNL(⋅)分别表示Local注意力子模块和Non-Local注意力子模块的函数。并且两者的输出通道是一样的,记为
C
C
C。
接下来就要进入融合子网络,它本质也是个通道注意力网络(channel-wise):首先将
F
N
L
i
、
F
L
i
F_{NL}^i、F^i_L
FNLi、FLi进行相加;然后通过全局平均池化输出
v
∈
R
C
v\in \mathbb{R}^C
v∈RC;接着分别用2个全连接层
H
F
C
1
(
v
)
、
H
F
C
2
(
v
)
\mathcal{H}_{FC}^1(v)、\mathcal{H}^2_{FC}(v)
HFC1(v)、HFC2(v)输出2个各自
C
C
C通道的张量,并通过通道合并成
2
C
2C
2C的张量;最后使用softmax获取
[
0
,
1
]
[0,1]
[0,1]的值,然后将结果分为2份,各自为
C
C
C个通道,即
w
N
,
w
N
L
∈
R
C
w_N,w_{NL}\in\mathbb{R}^C
wN,wNL∈RC,并分别与对应的
F
N
L
i
、
F
L
i
F_{NL}^i、F^i_L
FNLi、FLi进行元素相乘得到最终的结果
F
C
A
B
i
F^i_{CAB}
FCABi。 整个融合过程的数学表达式如下:
{
v
=
G
l
o
b
a
l
P
o
o
l
i
n
g
(
F
N
L
i
+
F
L
i
)
,
w
N
L
,
w
L
=
s
o
f
t
m
a
x
(
[
H
F
C
1
(
v
)
,
H
F
C
2
(
v
)
]
)
,
F
C
A
B
i
=
F
N
L
i
⋅
w
N
L
+
F
L
i
⋅
w
L
.
(7)
\begin{cases} v = {\color{mediumorchid}GlobalPooling}(F_{NL}^i + F^i_L),\\ w_{NL},w_L = softmax([\mathcal{H}_{FC}^1(v),\mathcal{H}^2_{FC}(v)]),\\ F^i_{CAB} = F^i_{NL}\cdot w_{NL} + F^i_L\cdot w_L.\tag{7} \end{cases}
⎩⎪⎨⎪⎧v=GlobalPooling(FNLi+FLi),wNL,wL=softmax([HFC1(v),HFC2(v)]),FCABi=FNLi⋅wNL+FLi⋅wL.(7)Note:
这里
w
N
,
w
N
L
w_N,w_{NL}
wN,wNL就起着选择合适的局部注意力和全局注意力的分配方案的作用。
为了验证融合的效果,作者使用热度图来反映,具体热值
h
h
h定义如下:
h
=
1
M
∑
m
=
1
M
a
m
,
a
m
=
{
1
,
w
L
m
≥
w
N
L
m
,
0
,
o
t
h
e
r
w
i
s
e
.
(8)
h = \frac{1}{M}\sum^M_{m=1} a_m,\\ a_m = \begin{cases} 1, w^m_L\ge w_{NL}^m,\\ 0, otherwise \end{cases}.\tag{8}
h=M1m=1∑Mam,am={1,wLm≥wNLm,0,otherwise.(8)其中
M
M
M为
w
L
w_L
wL的长度;
h
h
h反映了对两种注意力的偏好,就像投票一样。 最终热度图结果如下:
上图就是局部注意力子网络,作者设计了输出了2种不同感受野的分支,仍然是基于channel-wise的注意力机制来选择不同的感受野方案。和融合子网络的channel-wise注意力机制不同的是,两个输出的注意力通道是独立的,各自使用softmax来获取各自的注意力,最后结果进行元素级相加。 该子网络的注意力是通过局部操作产生的,其数学表达式为:
F
L
i
=
H
L
(
F
F
E
M
i
)
.
F^i_L = \mathcal{H}_L(F^i_{FEM}).
FLi=HL(FFEMi).
3.4 Non-local Attention submodule
Comparison and Analysis
\colorbox{violet}{Comparison and Analysis}
ComparisonandAnalysis 接下来我们先对目前的一些全局注意力机制进行比较,具体如下图所示:
蓝色区域表示为待重建区域;红色区域表示为相关长距离依赖区域;黄色区域表示搜索区域。 我们分3个指标来进行对比:①全局注意力是基于像素还是patch;②注意力机制的更新是基于像素还是patch,即公式(1)中的
x
i
^
\hat{x_i}
xi^;③搜索范围是窗口还是全部图像。
Patch-wise Non-local Attention Model
\colorbox{lightseagreen}{Patch-wise Non-local Attention Model}
Patch-wiseNon-localAttentionModel
上图就是本文提出的一种用于Vision的新型自注意力模块,这种结构和ViT是不同的。 ①:设输入
F
F
E
M
i
∈
R
C
×
W
×
H
F^i_{FEM}\in\mathbb{R}^{C\times W\times H}
FFEMi∈RC×W×H,它先是使用3个独立的embedding function——
θ
(
⋅
)
、
φ
(
⋅
)
、
g
(
⋅
)
\theta(\cdot)、\varphi(\cdot)、g(\cdot)
θ(⋅)、φ(⋅)、g(⋅),各自的参数记为
W
θ
、
W
φ
、
W
g
W_\theta、W_\varphi、W_g
Wθ、Wφ、Wg,它们是
1
×
1
1\times 1
1×1的卷积层,用于输出
Q
、
K
、
V
Q、K、V
Q、K、V,具体表达式如下:
{
Q
=
θ
(
F
F
E
M
i
)
,
K
=
φ
(
F
F
E
M
i
)
,
V
=
g
(
F
F
E
M
i
)
.
(9)
\begin{cases} Q = \theta(F^i_{FEM}),\\ K = \varphi(F^i_{FEM}),\\ V = g(F^i_{FEM}).\tag{9} \end{cases}
⎩⎪⎨⎪⎧Q=θ(FFEMi),K=φ(FFEMi),V=g(FFEMi).(9) ②:接下来就是核心操作——unfold,其产生patch,和ViT那种直接讲图像分块的方式不同,本文的结构是利用滑动窗口的形式,以
s
s
s为步长来抽取patch,每个patch的size为
C
×
W
p
×
H
p
C\times W_p\times H_p
C×Wp×Hp;对于每一组都进行unfold,每一组包含
N
N
N个3D-patch(
N
≥
W
/
W
p
N\ge W/W_p
N≥W/Wp),具体如上图坐标小框所示。这种设计的好处在于它可以在全局注意力的基础上增加一定的局部相关性,因此其比ViT这种块与块之间不重叠的方式更加灵活,但是缺陷在于这样会增加计算量。
③接下来我们将
Q
、
K
Q、K
Q、K的3D-patch分别进行reshape,每一个patch都flatten为1D的向量,这样就可以获取
N
×
(
C
⋅
W
p
⋅
H
p
)
N\times (C\cdot W_p\cdot H_p)
N×(C⋅Wp⋅Hp)的2D矩阵
Q
~
、
K
~
\tilde{Q}、\tilde{K}
Q~、K~;
M
=
Q
^
K
^
T
M = \hat{Q} \hat{K}^T
M=Q^K^T就是注意力矩阵,其衡量不同3D-patch之间的相关性,其中
M
∈
R
N
×
N
M\in\mathbb{R}^{N\times N}
M∈RN×N。 接下来使用softmax对
M
M
M的每一行进行运算:
ϕ
(
Q
^
,
K
^
)
=
s
o
f
t
m
a
x
(
Q
^
K
^
T
)
.
(10)
\phi(\hat{Q},\hat{K}) = softmax(\hat{Q} \hat{K}^T).\tag{10}
ϕ(Q^,K^)=softmax(Q^K^T).(10) ④:最后我们将
ϕ
\phi
ϕ看成
V
V
V的权重进行加权求和。具体而言,根据公式(1),
ϕ
\phi
ϕ的每一行都分别和
V
V
V进行元素加权求和得到结果
x
^
i
∈
R
C
×
W
p
×
H
p
\hat{x}_i\in\mathbb{R}^{C\times W_p\times H_p}
x^i∈RC×Wp×Hp。因此最后将多行结果进行fold,即上图右下小框所示,为了获取同样size的输出
C
×
H
×
W
C\times H\times W
C×H×W,直接硬拼会造成溢出,因此重叠部分去平均处理;Fold操作可以看成是unfold的相反操作。
⑤:最后接一个卷积层以及和输入端相加来减轻残差训练的压力以及避免卷积层造成的信息损失。
小结一下:
本文提出的自注意力子网络和ViT提出的自注意力网络是两种类型。
本文提出的自注意力子网络由于unfold的时候采用了
s
≤
W
p
s\leq W_p
s≤Wp的滑动窗口来提取patch,因此其一定程度上增加了局部相关性。
ViT的出现延申了Transformer在计算机视觉任务上的应用,但是其直接在image-wise上的不重叠分块势必会破坏原图上的一些空间结构,而本文提出的COLA-Net结构其自注意力模块在feature-wise上按步长为
s
s
s进行重叠分块,这不仅保留了一些空间结构,而且增加了一定的局部相关性。两者的对比如下图所示: