Capturing the Attention Shift with MLLM-LAWS

对于具有 $L$ 层的MLLM，我们定义Layer-level Attention Weights（MLLM-LAWS）为：

$${\text{LAWS}}_{{\mathcal{V}}_{\ast }}=\left[a_{{\mathcal{V}}_{\ast }}^1,a_{{\mathcal{V}}_{\ast }}^2,\cdots ,a_{{\mathcal{V}}_{\ast }}^L\right],a_{{\mathcal{V}}_{\ast }}^l=\sum _{i=0}^s\sum _{j\in {\mathcal{V}}_{\ast }}{a}_{ij}^l\left(h^{l-1}\right)$$

$a_{ij}^l$ 表示MLLM的第 $l$ 层中第 $i$ 个token对第 $j$ 个token的注意力权重，$h^l$ 表示第 $l$ 层的隐藏状态，$\forall i,{\sum }_{j=0}^s{a}_{ij}^l\left({\mathbf{h}}^{l-1}\right)=1$，${\mathcal{V}}_{\ast }$ 表示一个token序列，可以是 ${\mathcal{V}}_{\text{itself}}$、${\mathcal{V}}_{\text{before}}$ 或 ${\mathcal{V}}_{\text{after}}$，分别表示视觉序列、视觉序列前的文本序列和视觉序列后的文本序列。

${\text{LAWS}}_{{\mathcal{V}}_{\ast }}$ 可以表示MLLM对当前序列 ${\mathcal{V}}_{\ast }$ 在所有MLLM层上的注意力动态曲线。

• (a): 视觉编码器和投影器的视觉特征token被插入到文本特征序列中。
• (b): 文本层中在文本token的前图像、图像本身和后图像上的注意力权重比例。红色曲线来自在纯文本任务表现更好的MLLM，而蓝色曲线来自在纯文本任务表现更差的MLLM。
• ©: 实验表明，在100多个MLLM中，视觉token前后的皮尔逊相关系数与MLLM在纯文本上的表现存在正相关关系。

因此，注意力转移可通过视觉序列前后 $\text{LAWS}$ 之间的皮尔逊相关系数量化为：

$$\text{Attention Shift}={\mathbb{E}}_{\mathbf{x}}\left[-\rho \left({\text{LAWS}}_{{\mathcal{V}}_{\text{before}}},{\text{LAWS}}_{{\mathcal{V}}_{\text{after}}}\right)\right]+1$$

WINGS

我们引入了与主要注意力平行的额外模块，作为增强Learner来补偿注意力转移。我们在一侧训练Visual Learner，减轻一些转移的注意力。然后，我们基于路由转移的注意力权重协同训练Visual Learner和Textual Learner。

WINGS由Low-Rank Residual Attention（LoRRA）模块构建，其中前一个隐藏状态作为Query，而视觉/文本特征作为Key和Value。训练从Visual Learner和Projector开始，接着是动态注意力Router。

Learner:

Learner ∗ ( Q = h l , K , V = x ∗ ) ∗ ∈ { V , T } = Softmax ( h l ( 1 + W Q ) ⋅ ( x ∗ ( 1 + W K ) ) ⊤ d head ) x ∗ ( 1 + W V ) W O \text{Learner}^* \left( \rm{Q} = \mathbf{h}^l, \rm{K}, \rm{V} = \mathbf{x}_* \right)_{* \in \{\mathbf{V}, \mathbf{T}\}} = \text{Softmax} \left( \frac{\mathbf{h}^l (1 + \mathbf{W}^{\rm{Q}}) \cdot (\mathbf{x}_* (1 + \mathbf{W}^{\rm{K}}))^\top}{\sqrt{d_{\text{head}}}} \right) \mathbf{x}_* (1 + \mathbf{W}^{\rm{V}}) \mathbf{W}^{\rm{O}} Learner∗(Q=hl,K,V=x∗​)∗∈{V,T}​=Softmax(dhead​ ​hl(1+WQ)⋅(x∗​(1+WK))⊤​)x∗​(1+WV)WO

Router接受一个注意力权重作为输入，通过单层MLP和Softmax处理，然后将Learner的输出叠加到主注意力上。

Att WINGS = Att main + ∑ ∗ ∈ { V , T } Router ( a ) ⋅ Learner ∗ ( h l , x ∗ ) \text{Att}^{\text{WINGS}} = \text{Att}^{\text{main}} + \sum_{* \in \{V, T\}} \text{Router}(\mathbf{a}) \cdot \text{Learner}^* \left( \mathbf{h}^l, \mathbf{x}_* \right) AttWINGS=Attmain+∗∈{V,T}∑​Router(a)⋅Learner∗(hl,x∗​)

训练

WINGS的架构包含四个元素：Vision Encoder、Projector、LLM以及带有Router的Learner。

• Stage 1
  • ❄️ Vision Encoder、LLM
  • ???? Projector、Visual Learner
  • Visual Learner的输出直接加到主分支上
• Stage 2
  • ❄️ Vision Encoder
  • ???? LLM、Projector、Router、Visual Learner、Textual Learner