注意力可视化

Attention Rollout

论文来源：Quantifying Attention Flow in Transformers

问题提出

一般来说，我们直接用的是raw attention map，也就是取最后一层中 CLS token对各个patch的Attention权重，将它们画成热力图。但是最后一层中的Patch token已经和前面的Transformer层相互交换过信息。因此最后一层的Attention map只能告诉你最终的CLS 从导数第二层的哪些token中读取了信息，理论上不能告诉你最终的CLS中有多少信息最初来源于输入图像的每一个Patch

一层Attention表示什么

关于注意力的简述可以参考前文 注意力机制简述。我们这里假设一张图像被切成两个Patch，ViT的输入序列中共有三个token：

$$[\text{CLS}],\quad P_1,\quad P_2$$

某一层的Attention矩阵记为

$$A= \begin{bmatrix} A_{C,C} & A_{C,P_1} & A_{C,P_2}\\ A_{P_1,C} & A_{P_1,P_1} & A_{P_1,P_2}\\ A_{P_2,C} & A_{P_2,P_1} & A_{P_2,P_2} \end{bmatrix}$$

每一行表示接收信息的token，列表示提供信息的token。例如$A_{C,P_1}=0.4$表示更新CLS时，模型从 $P_1$ 中读取了权重为0.4的信息。

同时对于每一行我们使用Softmax对每一个Query对应的Attention权重进行了归一化

Attention Rollout

假设模型有两层。第一层中P1读取了一部分P2的信息

$$P_1^{(1)} = 0.2C^{(0)} + 0.6P_1^{(0)} + 0.2P_2^{(0)}$$

第二层中，最终的CLS又读取了P1‘

$$C^{(2)} = 0.3P_1^{(1)} +\cdots$$

那么第一层的内容间接影响了最终分类，存在一条路径

$$P_2^{(0)} \rightarrow P_1^{(1)} \rightarrow C^{(2)}$$

这条路径带来的贡献为：0.2 * 0.3 = 0.06

残差连接

一层Attention的输出并不是纯粹的：AX；而是X + AX。其中X是上一层token，A是attention矩阵。所以某一个token在更新时，除了读取其他token，也会直接保留自身信息。因此不能直接使用A，而要加入单位矩阵I。加入单位矩阵后，每一行的和变成了2，为了重新将每一行归一化为1，可以除以2

$$\widetilde{A} = \frac{A+I}{2}$$

多层 Attention 累计

假设模型有两层，加入残差后的Attention矩阵分别为$\widetilde{A}^{(1)}, \widetilde{A}^{(2)}$，那么两层之后的信息传播矩阵为：

$$R^{(2)} = \widetilde{A}^{(2)} \widetilde{A}^{(1)}$$

扩展到L层，公式如下

$$R^{(L)} = \widetilde{A}^{(L)} \widetilde{A}^{(L-1)} \cdots \widetilde{A}^{(1)}$$

最终$R^{(L)}_{i,j}$表示高层 token i中，有多少信息可以沿着 Attention路径追溯到输入token j。

Grad-CAM

Grad-CAM关注的是特征图中的哪些空间位置真正推动了模型的分类结果。它并不试图完整还原模型内部的信息流，而更像是一种局部归因方法，通过梯度判断每一类特征对目标类别的重要程度，再将这些特征重新投影回空间位置。

这个方法最初面向CNN提出，核心思想是利用目标类别分数回传到卷积特征图的梯度，生成一张粗粒度的类别判别热力图。

流程

前向传播获得特征图

对于目标卷积层，保存每一个通道的特征图$A^1,\ A^2,\ \cdots,\ A^K$，每一张特征图都有空间结构$A^k \in \mathbb{R}^{H \times W}$。这个特征图从直觉上表示图像中存在哪些高级视觉模式

举个例子，我们前向传播得到两个特征图

$$A^1 = \begin{bmatrix} 1 & 2\\ 0 & 1 \end{bmatrix} \quad A^2 = \begin{bmatrix} 0 & 1\\ 3 & 2 \end{bmatrix}$$

反向传播计算每个通道的重要性

选择一个目标类别，将对应的logit $y^c$向前面的特征图反向传播，计算$\frac{\partial y^c}{\partial A^k_{ij}}$，然后对第 k 个通道中的所有空间位置取平均

$$\alpha_k^c = \frac{1}{HW} \sum_{i=1}^{H} \sum_{j=1}^{W} \frac{\partial y^c}{\partial A^k_{ij}}$$

其中，$\alpha_k^c$ 是 第 k 个通道对于类别 c 的权重。直觉上来说 代表第 k 中视觉特征对该类别有多重要

对于上述例子来说，我们将logit回传到两个特征图后，得到

$$\frac{\partial y^{\text{rose}}}{\partial A^1} = \begin{bmatrix} 2 & 2\\ 0 & 0 \end{bmatrix} \quad \frac{\partial y^{\text{rose}}}{\partial A^2} = \begin{bmatrix} -1 & -1\\ 1 & 1 \end{bmatrix}$$

然后分别求出两个通道的权重

$$\alpha_1^{\text{rose}} = \frac{1}{4} \sum_{i,j} \frac{\partial y^{\text{rose}}}{\partial A^1_{ij}} \\ \alpha_2^{\text{rose}} = \frac{1}{4} \sum_{i,j} \frac{\partial y^{\text{rose}}}{\partial A^2_{ij}}$$

将重要通道重新叠加为空间热力图

将所有特征图按照权重加权求和。直觉上来说代表这些重要模式出现在图像中的什么位置

$$M^c = \sum_{k=1}^{K} \alpha_k^c A^k$$

最后经过ReLU，得到Grad-CAM热力图

$$L_{\text{Grad-CAM}}^c = \operatorname{ReLU} \left( \sum_{k=1}^{K} \alpha_k^c A^k \right)$$

需要ReLU是因为某些位置可能是负数