Locality-Attending Vision Transformer¶
会议: ICLR 2026
arXiv: 2603.04892
代码: GitHub
领域: 分割
关键词: Vision Transformer, 局部注意力, 高斯核, Patch 表示精炼, 密集预测, 分割改进
一句话总结¶
提出 LocAt 模块化插件（GAug + PRR），通过可学习高斯核偏置注意力向局部邻域聚焦并精炼 patch 表示，在不修改训练目标的前提下使 ViT 在 ADE20K 分割上提升超 6%，同时分类精度不降反升。
背景与动机¶
-
ViT 的全局注意力利于分类但不利于分割：ViT 的全局自注意力机制擅长捕获长程依赖，在分类任务上表现优异，但对需要精确定位和细粒度空间细节的密集预测任务（如语义分割）表现不佳。全局注意力会稀释局部线索。
-
分类训练忽视 patch 级表示质量：标准 ViT 分类仅使用 [CLS] token 的输出计算损失，patch 位置的输出不受直接监督，导致最后一层的空间 token 表示质量退化——patch token 逐渐与 [CLS] token 对齐，失去独特的局部结构信息。
-
现有改进方案破坏 ViT 原始架构：层级化 ViT（如 Swin）引入窗口注意力和多阶段设计、卷积混合方案添加额外卷积模块。这些方法虽改善密集预测但改变了 ViT 架构，降低与基础模型（如 CLIP）的兼容性。
-
GAP 作为分类头的局限：全局平均池化（GAP）对所有 patch 施加均匀梯度，迫使背景区域的表示也向分类目标对齐，对分割反而有害（Base 模型上降低分割性能）。
方法详解¶
整体框架：LocAt = GAug + PRR¶
- 功能：为标准 ViT 添加轻量级模块化插件，增强分割能力同时保持分类性能。
- 为什么：ViT 在基础模型中广泛使用，且因架构简洁性被偏好。与其替换 ViT，不如以最小改动增强其密集预测能力。
- 怎么做：(1) GAug 在自注意力 logits 上加可学习高斯核，引导 token 关注局部邻域；(2) PRR 在分类头前添加无参数自注意力，确保 patch 位置获得有效梯度。两者互补——GAug 在 backbone 内部改善特征交互，PRR 在输出端确保梯度回传到 GAug 参数。
关键设计 1：高斯增强注意力 (GAug)¶
在每层自注意力中，将标准 attention 公式修改为：
其中补充矩阵 \(\mathbf{S}\) 由高斯核构造，为每个 patch 提供局部注意力偏置：
- 自适应方差预测：从空间 query 矩阵预测每个 patch 的 2D 高斯方差 \(\mathbf{\Sigma} = f(\mathbf{q}_{sp}\mathbf{W}^\sigma) \in \mathbb{R}_+^{hw \times 2}\)，小方差产生锐利局部聚焦，大方差趋近均匀（全局注意力）
- 高斯核计算：\(\mathbf{G}_{pt} = \exp\left(-\frac{1}{2}\sum_{m=1}^{2}\frac{\mathbf{D}_{ptm}}{\mathbf{\Sigma}_{pm}}\right)\)，其中 \(\mathbf{D}\) 是 patch 间的逐坐标平方距离
- 自适应缩放：\(\bm{\alpha} = \text{softplus}(\mathbf{q}_{sp}\mathbf{W}^\alpha)\) 为每个 query 预测缩放系数，平衡原始 logits 与高斯偏置
GAug 是软性的、数据依赖的局部性机制——网络可学习何时/何地需要局部聚焦，何时保持全局注意。[CLS] token 不参与高斯偏置（无空间坐标），其行列填零。
关键设计 2：Patch 表示精炼 (PRR)¶
标准 ViT 分类仅用 [CLS] 输出计算损失。patch 位置的输出不受监督，导致空间表示退化。PRR 在分类头前添加无参数的多头自注意力：
然后取 \(\mathbf{x}_0^+\)（[CLS] 位置）送入分类头。这一操作使 [CLS] 对不同 patch 产生非均匀的注意力权重，从而将分类梯度不均匀地传播到各 patch 位置，鼓励每个 patch 维持独特的、有区分力的表示。
PRR 可视为 GAP 的替代方案：GAP 对所有 patch 施加均匀梯度，PRR 则根据内容自适应分配梯度。重要的是，PRR 还将梯度路由到最后一个 block 的 GAug 参数，使其能有效学习。
额外参数开销¶
仅引入每层 2 个小权重矩阵 \(\mathbf{W}^\sigma \in \mathbb{R}^{d \times 2}\) 和 \(\mathbf{W}^\alpha \in \mathbb{R}^{d \times 1}\)，PRR 完全无参数。以 Base 模型为例，新增参数仅 2,340 个（0.003% 增长），FLOPs 几乎不增加（17.58G → 17.64G）。
实验¶
实验设置¶
- 预训练：ImageNet-1K 分类训练 300 epochs，AdamW 优化器，batch size 1024
- 分割评估：冻结 backbone 仅训练 1 层 MLP 解码器（20K iterations），评估 ADE20K、PASCAL Context、COCO Stuff
- Backbone 尺寸：Tiny（6M 参数）、Small、Base（86M 参数）
- 基线：ViT、Swin、RegViT（ViT+registers）、RoPEViT、Jumbo
主实验：分割与分类性能¶
| 方法 | ADE mIoU | P-Context mIoU | C-Stuff mIoU | ImageNet Top-1 | 参数(M) |
|---|---|---|---|---|---|
| ViT-Tiny | 17.30 | 33.71 | 20.29 | 72.39 | 6 |
| LocAtViT-Tiny | 23.47 (+6.17) | 38.57 (+4.86) | 26.15 (+5.86) | 73.94 (+1.55) | 6 |
| Swin-Tiny | 25.58 | 36.78 | 28.34 | 81.18 | 28 |
| + LocAt | 26.52 (+0.94) | 37.65 (+0.87) | 29.09 (+0.75) | 81.43 (+0.25) | 28 |
| RegViT-Tiny | 15.98 | 33.45 | 19.58 | 72.90 | 6 |
| + LocAt | 24.39 (+8.41) | 39.90 (+6.45) | 27.38 (+7.80) | 74.08 (+1.18) | 6 |
| ViT-Base | 28.40 | 43.10 | 30.43 | 80.99 | 86 |
| LocAtViT-Base | 32.64 (+4.24) | 45.35 (+2.25) | 33.62 (+3.19) | 82.31 (+1.32) | 86 |
| RegViT-Base | 27.93 | 41.81 | 28.99 | 80.71 | 86 |
| + LocAt | 32.71 (+4.78) | 46.14 (+4.33) | 34.12 (+5.13) | 82.19 (+1.18) | 86 |
消融实验：组件贡献分析¶
| 配置 | ADE (Tiny) | ADE (Base) | ImageNet (Tiny) | ImageNet (Base) |
|---|---|---|---|---|
| ViT 基线 | 17.30 | 28.40 | 72.39 | 80.99 |
| + GAug | 18.98 (+1.68) | 30.26 (+1.87) | 73.16 (+0.77) | 82.00 (+1.01) |
| + PRR | 21.60 (+4.30) | 29.89 (+1.49) | 73.71 (+1.32) | 82.19 (+1.20) |
| + GAug + PRR (LocAt) | 23.47 (+6.17) | 32.64 (+4.24) | 73.94 (+1.55) | 82.31 (+1.32) |
| ViT + GAP | 19.65 | 27.99 | 72.50 | 81.84 |
| ViT - 位置编码 | 15.13 | 24.59 | 69.36 | 79.39 |
| LocAtViT - 位置编码 | 22.69 | 29.73 | 73.10 | 82.17 |
关键发现¶
- 分割提升显著且普适：LocAt 在所有 5 种 baseline（ViT、Swin、RegViT、RoPEViT、Jumbo）和 3 种尺寸上均带来分割提升，最大增幅 +8.41%（RegViT-Tiny on ADE）。
- 分类性能不降反升：所有 LocAt 增强模型的 ImageNet Top-1 均持平或提升（最高 +1.55%），证明局部性偏置不与全局建模冲突。
- GAug 和 PRR 互补且协同：单独使用各自有效（GAug +1.68、PRR +4.30），组合后进一步提升至 +6.17，说明梯度路由（PRR）对 GAug 参数学习至关重要。
- LocAt 编码了位置信息及以上：去除位置编码后 LocAtViT 仍超过有位置编码的 ViT（ADE 22.69 vs 17.30），说明高斯核不仅捕获位置信息还学到更丰富的空间结构。
- 自监督场景同样有效：在 DINO 框架中替换 ViT-S 为 LocAtViT-S，线性分类提升 +2.13%，k-NN 提升 +2.27%。
亮点¶
- 设计极简但效果显著：每层仅增加 3 个参数量极小的权重矩阵（\(\mathbf{W}^\sigma\)、\(\mathbf{W}^\alpha\)），PRR 完全无参数，FLOPs 几乎不增加。
- 模块化即插即用：可直接加到任何 ViT 及其变体上，不改变训练目标或数据增强策略。
- 从梯度流角度揭示 ViT 分割性能差的根因：patch 输出缺乏监督导致表示退化，这一分析具有启发性。
- 高斯核的方差由 query 预测，实现数据自适应的局部/全局平衡，比固定窗口注意力更灵活。
局限¶
- 分割评估仅使用冻结 backbone + 1 层 MLP 解码器，未测试与完整分割 head（如 UPerNet）配合的效果。
- 仅在自然图像上验证，医学影像、遥感等领域的适用性未探讨。
- 未在大规模基础模型（如 CLIP）上验证，计算资源限制了这一重要方向的探索。
- 各向异性 2D 高斯核仅沿行列轴独立，未考虑旋转等更灵活的空间结构。
评分¶
| 维度 | 评分 |
|---|---|
| 新颖性 | ⭐⭐⭐⭐ |
| 有效性 | ⭐⭐⭐⭐⭐ |
| 可复现性 | ⭐⭐⭐⭐⭐ |
| 实用性 | ⭐⭐⭐⭐ |