用于高效RGBT追踪的单级Transformer

简介

论文来源: Unified Single-Stage Transformer Network for Efficient RGB-T Tracking

现有的网络通常以单独的方式提取模态特征, 缺乏模态之间的交互与指导, 三阶段融合策略极大地限制了跟踪速度. 本文作者将提出一个统一的单级 Transformer RGB-T 追踪方法, 即 USTrack, 它将三阶段统一为一个 ViT 主干, 通过双层嵌入自注意力机制, 让模型能够在模态相互交互的前提下进行特征融合. 此外, 作者还引入了一种新颖的特征选择机制, 以减轻无效模态对预测的影响.

方法

网络的输入包括 4 张图: RGB 模板 $z_{image}^{rgb} \in \R^{H_z \times W_z \times 3}$, RGB 搜索区域 $x_{image}^{rgb} \in \R^{H_x \times W_x \times 3}$, 热红外模板 $z_{image}^{t} \in \R^{H_z \times W_z \times 3}$, 热红外搜索区域 $x_{image}^{t} \in \R^{H_x \times W_x \times 3}$.

与 ViT 一致, 它们将被切割为 patch, 并使用投影矩阵 $E_{rgb}, E_{t}$ 来映射到 $D$ 维潜在空间中. 不同点在于, ViT 只添加了位置嵌入 $\bm{P}$, 而本文还添加了模态嵌入 $\bm{M}$:

$$\begin{array}{rl} \tilde{\bm z}_{rgb} &= [\bm{z}_{rgb}^{1} \bm{E}_{rgb}; \bm{z}_{rgb}^{2} \bm{E}_{rgb}; \cdots; \bm{z}_{rgb}^{N_z} \bm{E}_{rgb}] + \bm{P}_z + \bm{M}_{rgb} \\ \tilde{\bm z}_t &= [\bm{z}_t^{1} \bm{E}_{t}; \bm{z}_t^{2} \bm{E}_{t}; \cdots; \bm{z}_t^{N_z} \bm{E}_{t}] + \bm{P}_z + \bm{M}_{t} \\ \tilde{\bm x}_{rgb} &= [\bm{x}_{rgb}^{1} \bm{E}_{rgb}; \bm{x}_{rgb}^{2} \bm{E}_{rgb}; \cdots; \bm{x}_{rgb}^{N_x} \bm{E}_{rgb}] + \bm{P}_x + \bm{M}_{rgb} \\ \tilde{\bm x}_t &= [\bm{x}_t^{1} \bm{E}_{t}; \bm{x}_t^{2} \bm{E}_{t}; \cdots; \bm{x}_t^{N_x} \bm{E}_{t}] + \bm{P}_x + \bm{M}_{t} \end{array}$$

这些序列将被拼接为新的标记序列 $\bm H = [\tilde{x}_{rgb}; \tilde{x}_t; \tilde{z}_{rgb}; \tilde{z}_t] \in \R ^{(2N_x + 2N_t) \times D}$, 对其使用多头注意力:

$$\bm{M} = \bm{A} \cdot \bm{V} = Softmax\left(\frac{\bm{Q} \bm{K}^\top}{\sqrt{d_k}}\right) \cdot \bm{V} \\ \bm{Q}\bm{K}^\top = [\bm{Q}_{rgb}^x; \bm{Q}_{t}^x; \bm{Q}_{rgb}^z; \bm{Q}_{t}^z] [\bm{K}_{rgb}^x; \bm{K}_{t}^x; \bm{K}_{rgb}^z; \bm{K}_{t}^z]^\top \\ \bm{V} = [\bm{V}_{rgb}^x; \bm{V}_{t}^x; \bm{V}_{rgb}^z; \bm{V}_{t}^z]$$

将 $\bm{Q}\bm{K}^\top$ 展开得到:

$$\begin{array}{ll} &\bm{Q}\bm{K}^\top \\ =& [\bm{Q}_{rgb}^x {\bm{K}_{rgb}^x}^\top, \bm{Q}_{rgb}^x {\bm{K}_{t}^x}^\top, \bm{Q}_{rgb}^x {\bm{K}_{rgb}^z}^\top, \bm{Q}_{rgb}^x {\bm{K}_{t}^z}^\top; \cdots] \\ =& [\bm{W}_{x_{rgb}}^{x_{rgb}}, \bm{W}_{x_{t}}^{x_{rgb}}, \bm{W}_{z_{rgb}}^{x_{rgb}}, \bm{W}_{z_{t}}^{x_{rgb}}; \cdots] \end{array}$$

最终的注意力表达式可以写作:

$$\bm{M} = [\bm{W}_{x_{rgb}}^{x_{rgb}} \bm{V}_{rgb}^x + \bm{W}_{x_{t}}^{x_{rgb}} \bm{V}_{t}^x + \bm{W}_{z_{rgb}}^{x_{rgb}} \bm{V}_{t}^z + \bm{W}_{z_{t}}^{x_{rgb}} \bm{V}_{t}^z; \cdots]$$

其中 $\bm{W}_{x_{rgb}}^{x_{rgb}} \bm{V}_{rgb}^x$ 表示 RGB 搜索区域图像特征提取, $\bm{W}_{x_{t}}^{x_{rgb}} \bm{V}_{t}^x$ 根据两个模态之间的语义相似性来聚合热模态特定信息, $\bm{W}_{z_{rgb}}^{x_{rgb}} \bm{V}_{t}^z$ 负责聚合 RGB 模板图像特征, 以获取搜索区域与模板间的关系信息, $\bm{W}_{z_{t}}^{x_{rgb}} \bm{V}_{t}^z$ 负责聚合热红外模板图像特征, 获取热红外模板与 RGB 搜索区域之间的关系.

仅通过自注意力的全局感知能力, 作者就已将特征提取, 特征融合与关系建模统一到单个 ViT 主干中, 使得该网络能够在模态交互的作用下直接提取模板和搜索区域的融合特征.

ViT 骨干的输出是两个融合特征, 它们分别以 RGB 和热红外作为主导, 并辅以另一模态的信息. 由于两个融合特征都包括了模态融合信息, 以及模板与搜索区域的关系信息, 因此两者都可以直接用于目标位置预测. 上图 (b) 展示了训练过程中的损失函数计算方式, 两个融合特征具有相同的损失函数, 可靠性评估模块输出可靠性权重, 用以组合成最终损失函数. 该方案允许不可靠模态产生较差的结果, 然后由可靠性评估模块给出较低权重, 而更可靠的模态将得到更大的权重.

作者采用 OSTrack 的预测头作为本文预测头, 两个预测头对应的损失函数如下:

$$\begin{array}{} \mathcal{L}_{RGB} = \mathcal{L}_{cls_{RGB}} + \lambda_{giou} \mathcal{L}_{giou_{RGB}} + \lambda_{L_1} \mathcal{L}_{1_{RGB}} \\ \mathcal{L}_{T} = \mathcal{L}_{cls_{T}} + \lambda_{giou} \mathcal{L}_{giou_{T}} + \lambda_{L_1} \mathcal{L}_{1_{T}} \end{array}$$

其中 $\mathcal{L}_{RGB}, \mathcal{L}_{T}$ 是每个预测头的整体损失函数, $\mathcal{L}_{cls_{RGB}}, \mathcal{L}_{cls_{T}}$ 是 $focal\ loss$, $\mathcal{L}_{giou_{RGB}}, \mathcal{L}_{giou_{T}}$ 是广义 IoU 损失. 评估模块由多个 $Conv-BN-ReLU$ 组成, 为了防止权重被归零, 使用 $Softmax$ 处理后的结果作为新的权重:

$$\lambda_{RGB}, \lambda_{T} = Softmax(R_{RGB}, R_{T}) \\ \mathcal{L}_{total} = \lambda_{RGB} \mathcal{L}_{RGB} + \lambda_{T} \mathcal{L}_{T}$$