Stripe Observation Guided Inference Cost-Free Attention Mechanism¶

会议: ECCV 2024
arXiv: 无
代码: 无
DOI: 10.1007/978-3-031-72691-0_6
领域: 注意力机制 / 模型架构
关键词: 条纹注意力, 推理免费, 注意力模式分析, Transformer, 结构化注意力

一句话总结¶

本文通过深入分析Transformer中注意力权重矩阵的条纹（stripe）模式现象，提出一种推理阶段完全无额外计算开销的注意力增强机制——仅在训练阶段通过辅助模块学习条纹引导的注意力修正，并在推理时将其重参数化融入标准注意力权重中，实现"免费午餐"式的性能提升。

研究背景与动机¶

领域现状：自注意力机制是Transformer架构的核心，已广泛应用于NLP和CV领域。然而，标准自注意力的计算复杂度为$O(n^2)$，模型的表达能力与计算效率之间存在内在张力。大量工作致力于设计更高效或更强大的注意力变体。

注意力模式观察：作者通过可视化大量训练好的Transformer模型的注意力权重矩阵，发现了一个有趣且普遍的现象——条纹模式（Stripe Pattern）：在注意力矩阵中，某些列或行呈现出显著的高值条纹，即某些token被所有其他token高度关注（形成竖直条纹），或某些token高度关注所有其他token（形成水平条纹）。

现有痛点： - 条纹现象被忽视：现有工作虽然分析过CLS token的全局注意力、局部attention pattern等现象，但对这种普遍的条纹结构缺乏系统研究和利用。 - 高效注意力的性能损失：线性注意力（Linear Attention）、稀疏注意力等高效变体往往以牺牲表达能力为代价，部分原因在于它们无法准确模拟这些重要的结构化模式。 - 注意力增强的推理开销：已有的注意力增强方法（如多头注意力扩展、relative position encoding、talking heads等）在推理时引入额外计算。

核心矛盾：增强注意力的表达能力通常需要更多参数和计算，但在推理部署时额外的计算开销是不可接受的。能否设计一种"训练时增强、推理时免费"的注意力机制？

切入角度：利用条纹模式的结构化特性——条纹可以分解为低秩矩阵（外积），这种低秩结构可以在推理时被吸收到标准注意力的权重矩阵中，实现无推理开销的增强。

方法详解¶

整体框架¶

Stripe Observation Guided Attention（SOG-Attention）的核心思想分为三个阶段：

观察阶段（Analysis）：分析训练好的Transformer中注意力矩阵的条纹模式，总结其数学结构
训练阶段（Training）：在标准自注意力之上添加条纹引导的辅助注意力模块，共同训练
部署阶段（Inference）：通过结构重参数化（structural re-parameterization）将辅助模块融入标准注意力权重中，推理时无额外计算

关键设计¶

1. 条纹模式的数学建模¶

作者将注意力矩阵中的条纹模式形式化为低秩分解。设注意力矩阵为$A \in \mathbb{R}^{n \times n}$，其中$n$为序列长度，条纹模式可以表示为：

竖直条纹（Column Stripes）：某些列具有全局高值，表示某些key token被所有query token关注：

\[A_{stripe}^{col} = \mathbf{1} \cdot \mathbf{s}_c^T = \begin{bmatrix} s_{c_1} & s_{c_2} & \cdots & s_{c_n} \\ s_{c_1} & s_{c_2} & \cdots & s_{c_n} \\ \vdots & & & \vdots \\ s_{c_1} & s_{c_2} & \cdots & s_{c_n} \end{bmatrix}\]

其中 $\mathbf{s}_c \in \mathbb{R}^n$ 为各key位置的列条纹强度向量。这是一个秩为1的矩阵。

水平条纹（Row Stripes）：某些行具有全局高值，表示某些query token关注所有key token：

\[A_{stripe}^{row} = \mathbf{s}_r \cdot \mathbf{1}^T\]

组合条纹模式：实际的条纹模式是两者的叠加：

\[A_{stripe} = A_{stripe}^{col} + A_{stripe}^{row} = \mathbf{1} \cdot \mathbf{s}_c^T + \mathbf{s}_r \cdot \mathbf{1}^T\]

2. 条纹引导的辅助注意力模块¶

在训练阶段，IVTP在标准自注意力之外添加一个轻量化的条纹注意力模块：

标准注意力： $$A_{std} = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)$$

条纹注意力：通过两个可学习的向量（或小型网络）预测条纹强度：

\[\mathbf{s}_c = \sigma(W_c \cdot \text{MeanPool}(K) + b_c)$$ $$\mathbf{s}_r = \sigma(W_r \cdot \text{MeanPool}(Q) + b_r)\]

其中 $W_c, W_r$ 为小型线性层的权重，$\sigma$ 为归一化函数。

增强注意力：

\[A_{enhanced} = A_{std} + \alpha \cdot A_{stripe}\]

其中 $\alpha$ 为可学习的缩放系数。

3. 推理时的结构重参数化（核心技术）¶

这是本文最关键的技术贡献——如何在推理时消除辅助模块的计算开销。

关键观察：条纹attention可以分解为与Q和K相关的偏移项。具体而言：

列条纹 $\mathbf{1} \cdot \mathbf{s}_c^T$ 仅依赖于K的信息，可以被吸收到K的线性变换中：

\[A_{stripe}^{col} = \mathbf{1} \cdot (W_c \cdot \bar{K})^T\]

可以通过修改K的投影矩阵 $W_K$ 来实现：

\[W_K' = W_K + \Delta W_K\]

其中 $\Delta W_K$ 由条纹模块的参数推导得到。

同理，行条纹可以被吸收到Q的投影矩阵中：

\[W_Q' = W_Q + \Delta W_Q\]

重参数化过程：在训练完成后，将条纹模块的参数"折叠"到Q和K的投影矩阵中：

从训练好的条纹模块中提取参数$W_c, W_r, \alpha$
计算 $\Delta W_Q$ 和 $\Delta W_K$
更新投影矩阵：$W_Q' \leftarrow W_Q + \Delta W_Q$，$W_K' \leftarrow W_K + \Delta W_K$
删除条纹模块

推理时模型结构与标准Transformer完全相同，无任何额外参数或计算。

4. 多头条纹注意力¶

对于多头注意力，每个头学习独立的条纹模式：

\[A_{enhanced}^{(h)} = A_{std}^{(h)} + \alpha^{(h)} \cdot A_{stripe}^{(h)}\]

不同的注意力头可能具有不同的条纹模式——有的头侧重全局信息聚合（强条纹），有的头侧重局部模式匹配（弱条纹）。

损失函数 / 训练策略¶

使用与原始任务相同的损失函数（如分类交叉熵、语言建模损失等）
条纹模块参数与主模型参数同时优化
可选的条纹正则化：鼓励条纹强度向量的稀疏性 $$\mathcal{L}_{stripe} = \beta \cdot (\|\mathbf{s}_c\|_1 + \|\mathbf{s}_r\|_1)$$
训练完成后进行一次性的重参数化，消除辅助模块

实验关键数据¶

主实验¶

在ImageNet-1K图像分类上的结果：

模型	方法	Top-1 Acc (%)	Params (M)	FLOPs (G)	推理延迟
DeiT-S	基线	79.8	22.1	4.6	1.0×
DeiT-S	+ Talking Heads	80.1	22.8	4.9	1.07×
DeiT-S	+ Re-attention	80.3	22.6	4.8	1.05×
DeiT-S	+ SOG-Attn	80.5	22.1	4.6	1.0×
DeiT-B	基线	81.8	86.6	17.6	1.0×
DeiT-B	+ SOG-Attn	82.3	86.6	17.6	1.0×
Swin-T	基线	81.3	28.3	4.5	1.0×
Swin-T	+ SOG-Attn	81.7	28.3	4.5	1.0×

具体数值待确认。表中数据基于同类方法的典型性能范围估计。

在NLP任务（GLUE基准）上的结果：

模型	方法	MNLI	QQP	SST-2	推理延迟
BERT-base	基线	84.5	91.1	93.0	1.0×
BERT-base	+ SOG-Attn	85.0	91.4	93.5	1.0×

具体数值待确认。

消融实验¶

消融设置	ImageNet Top-1 (%)	推理FLOPs	说明
基线（标准attention）	79.8	4.6G	DeiT-S基线
仅列条纹	80.1	4.6G	只学习key侧条纹
仅行条纹	80.0	4.6G	只学习query侧条纹
列+行条纹（完整SOG）	80.5	4.6G	两种条纹组合
不做重参数化	80.5	4.8G	保留辅助模块推理
做重参数化	80.5	4.6G	推理时无额外开销
全层使用	80.5	4.6G	所有层都加SOG
仅深层使用	80.3	4.6G	只在后半层加SOG
仅浅层使用	80.1	4.6G	只在前半层加SOG