ViT-AdaLA: Adapting Vision Transformers with Linear Attention¶
会议: NeurIPS 2026
arXiv: 2603.16063
论文: https://arxiv.org/abs/2603.16063
领域: 模型压缩
关键词: 视觉Transformer, 线性注意力, 基础模型自适应, 特征蒸馏, 高分辨率稠密预测
一句话总结¶
针对视觉基础模型（VFM）在长序列与高分辨率下二次方注意力复杂度过高、且从头训练线性注意力代价高昂的问题，提出渐进式三阶段自适应框架 ViT-AdaLA，通过逐层注意力对齐与整网特征对齐使线性 ViT 高效继承原模型先验表征，在保持与 Softmax 原型相近性能的同时显著降低 58% 峰值显存并实现超 2 倍吞吐加速。
研究背景与动机¶
基于视觉 Transformer（ViT）的视觉基础模型（如 DINOv2、CLIP 和 SigLIP）已在图像分类、语义分割、目标检测以及高分辨率生成等领域展现出极强的通用表征能力。然而，标准自注意力机制的计算与显存复杂度随输入 token 序列长度呈二次方（\(\mathcal{O}(N^2 D)\)）增长，在处理高分辨率图像或密集 token 流时迅速遭遇严重的算力与内存瓶颈。现有缓解注意力开销的方法主要包括窗口注意力、状态空间模型（Mamba）以及线性注意力变体。其中，线性注意力通过核技巧将矩阵乘法顺序重排，将复杂度降至线性（\(\mathcal{O}(ND^2)\)），被视为突破高分辨率处理瓶颈的最具潜力的方向之一。
现有线性注意力工作主要分为两大范式，但均无法直接赋能视觉基础模型的高效迁移。第一类是从头训练范式（training from scratch），专注于设计更精巧的低秩或核化逼近方法，但缺乏大规模预训练时下游迁移性能严重崩溃（如在 ADE20K 上仅有二十几点 mIoU），而从头进行千亿 token 级视觉预训练代价极高。第二类是面向大语言模型（LLM）的线性化自适应范式（如 LoLCATS、Hedgehog），利用 LoRA 或轻量映射层逼近预训练注意力。然而，Decoder 架构的语言模型身兼特征提取器与自回归生成器，误差随时间步累积；而视觉模型以 ViT 为通用特征编码器，由下游独立任务头负责生成，注意力逼近误差在深度方向层层累积、在空间维度发散，导致全局语义流形失真，直接迁移 LLM 的线性化策略在稠密预测任务中性能急剧下滑。
本文认为，解决视觉线性化的关键不在于设计极其复杂的注意力逼近结构，而在于如何让线性模型平滑继承并校准已有的高质量视觉先验。核心 idea：提出三阶段渐进式自适应框架 ViT-AdaLA，首先在冻结主干下利用极简 Vanilla 线性注意力进行块级注意力对齐，随后通过全局特征对齐消除跨层累积的逼近误差，最后联合下游任务头微调，以极小训练代价实现高精度、高分辨率扩展的线性化视觉基础模型。
方法详解¶
整体框架¶
ViT-AdaLA 旨在将离线预训练的 Softmax ViT 转换为线性注意力 ViT，其核心工作流分为三阶段：阶段一（注意力对齐，Attention Alignment）在保持原始模型权重与中间特征不变的前提下，独立拟合每个 Transformer 块内的线性注意力投影；阶段二（特征对齐，Feature Alignment）将全网自注意力替换为已初始化的线性注意力，通过全局顶层表征蒸馏消除跨层残差累积误差；阶段三（下游监督微调，Supervised Fine-tuning）接入特定任务预测头，联合微调模型以迁移下游。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入图像与预训练 Softmax ViT 教师模型"] --> B["阶段 1：逐块 Vanilla 线性注意力局部拟合<br/>冻结主干，仅更新各层 W_Q, W_K, W_V"]
B --> C["阶段 2：全网顶层特征对齐与误差修正<br/>全模型解冻，蒸馏对齐最终层语义特征 f_v 与 f̂_v"]
C --> D["阶段 3：下游任务适配与联合微调<br/>外接分类头或 Mask2Former 分割头端到端优化"]
D --> E["高效线性视觉模型输出与高分辨率推理"]关键设计¶
1. 逐块 Vanilla 线性注意力局部拟合：消除复杂结构强行拟合带来的优化瓶颈
针对注意力近似问题，LoLCATS 等方法引入了额外的非线性投影网络（如 Hedgehog 映射），但在视觉骨干中这种强加的结构先验往往与预训练教师模型的固有动态发生冲突。ViT-AdaLA 摒弃繁琐的额外模块，直接采用基于核技巧的 Vanilla 线性注意力，使用数值稳定的 \(\phi(\cdot) = \text{ELU}(\cdot) + 1\) 激活函数保证注意力权重非负性。在阶段一中，冻结原始 ViT 的所有组件（LayerNorm、MLP 以及前序特征输入），仅针对每个 Transformer 块内新增的线性注意力模块更新其投影矩阵 \(\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V\)。
通过最小化各块自注意力输出与线性注意力输出之间的均方误差（MSE），使线性模块在局部快速拟合 Softmax 的行为模式： $\(\mathcal{L}_{att} = \frac{1}{N \cdot D}\sum_{i=1}^{K}\sum_{n=1}^{N}\sum_{d=1}^{D} \left(\mathbf{O}_i^{nd} - \hat{\mathbf{O}}_i^{nd}\right)^2\)$ 式中 \(K\) 为层数，\(N\) 为 token 序列长度，\(D\) 为隐藏维度。Vanilla 线性注意力不受约束的高延展性使其能够更自然地吸收注意力逼近模式，计算损耗更低，且为阶段二的快速收敛奠定了优质初始化。
2. 全网顶层特征对齐与跨层误差修正：恢复稠密预测所需的全局语义一致性
尽管阶段一在每个块内实现了局部的均方误差极小化，但在逐层前向传递时，各层微小的未对齐残差会沿网络深度呈指数级复合放大，严重破坏深层特征空间的拓扑结构。针对这一视觉编码器的固有痛点，阶段二将全网的所有 Softmax 注意力统一替换为阶段一已对齐的线性注意力模块，并将原始未修改的 Softmax ViT 作为冻结教师模型 \(\mathcal{F}_{\theta_0}\)。
输入相同图像 \(\mathbf{X}_v\)，直接约束线性化 ViT \(\mathcal{F}_{\theta}\) 顶层输出特征 \(\hat{\mathbf{f}}_v\) 与教师特征 \(\mathbf{f}_v\) 保持一致： $\(\mathcal{L}_{fa} = \frac{\lambda}{N \cdot D}\sum_{n=1}^{N}\sum_{d=1}^{D} \left(\mathbf{f}_v^{nd} - \hat{\mathbf{f}}_v^{nd}\right)^2\)$ 在此阶段，全模型（包括注意力和 MLP 模块）共同参与微调。由于更新范围覆盖前馈网络，网络具备足够的表征容量来分散吸收蒸馏信号，有效避免了仅调整 QKV 投影导致的特征空间扭曲与过拟合。
3. 下游任务适配与高分辨率零样本迁移扩展：低成本打破长序列显存墙
在完成通用特征对齐后，线性化 ViT 已经完整承袭了基础模型的全局视觉语义。在阶段三中，模型直接接入目标任务预测头（如图像分类的线性分类器或语义分割的 Mask2Former 密集预测头），并在目标数据集上进行标准端到端监督微调。值得注意的是，该框架具备极强的输入分辨率外推能力：模型在固定小分辨率（\(512 \times 512\)）下完成第一、二阶段对齐，但在部署到 \(1024 \times 1024\) 等高分辨率下游任务时，线性注意力的显存开销优势呈倍数释放，且无需重新在超高分辨率下昂贵蒸馏即可取得优异性能。
损失函数 / 训练策略¶
训练过程分为三个严格解耦的连续步骤： 1. 阶段一（注意力对齐）：在 COCO 数据集上训练 4 个 epoch，优化器采用 AdamW，固定学习率 \(1\times 10^{-2}\)，仅更新各层线性注意力 \(W_Q, W_K, W_V\) 权重。在 8 张 H100 GPU 上单 epoch 耗时仅 4~5 分钟。 2. 阶段二（特征对齐）：在 ImageNet-22K 上对整网进行 10 到 30 个 epoch 的预训练。采用多项式学习率衰减（初始学习率 \(1\times 10^{-3}\)，主干采用 0.1 倍学习率比例即 \(1\times 10^{-4}\)），损失缩放权重 \(\lambda\) 针对不同基础模型校准（DINOv2 为 5，SigLIP 为 1，CLIP 为 4000，监督 ViT 为 1）。 3. 阶段三（下游微调）：ImageNet-1K 分类微调 10 个 epoch，ADE20K 与 Cityscapes 分割微调 20 个 epoch，初始学习率设为 \(1\times 10^{-4}\)。
实验关键数据¶
主实验¶
论文在 ImageNet-1K 图像分类与 ADE20K / Cityscapes 语义分割三大主流基准上，对 DINOv2-L、CLIP-L、SigLIP-L 等代表性视觉基础模型进行了全面评测，并与 Hedgehog、LoLCATS、Linformer、Performer、Cosformer、Nyströmformer 及 Monarch Attention 等先进线性化与近似注意力方法进行对比。
ImageNet-1K 图像分类对比（\(512 \times 512\) 分辨率）：
| 模型 / 骨干架构 | 注意力机制类型 | 参数量 (M) | 计算量 (GFLOPS) | 峰值显存 (GB) | 吞吐量 (imgs/s) | Top-1 准确率 (%) |
|---|---|---|---|---|---|---|
| DINOv2-L | Softmax (原始教师) | 304.20 | 310.60 | 1.3181 | 36.52 | 86.8 |
| DINOv2-L | Hedgehog (LLM 线性化) | 305.77 | 265.41 (-14.5%) | 1.2221 (-7.3%) | 37.44 (+2.5%) | 58.8 |
| DINOv2-L | LoLCATS (LLM 线性化) | 305.77 | 265.41 (-14.5%) | 1.2221 (-7.3%) | 37.44 (+2.5%) | 61.6 |
| DINOv2-L | Linformer (低秩投影) | 305.77 | 262.19 (-15.6%) | 1.2221 (-7.3%) | 45.41 (+25.3%) | 68.9 |
| DINOv2-L | Performer (随机特征) | 304.20 | 265.41 (-14.6%) | 1.2164 (-7.7%) | 36.48 (+0.1%) | 75.2 |
| DINOv2-L | Cosformer (余弦重加权) | 304.20 | 265.41 (-14.6%) | 1.2226 (-7.2%) | 39.55 (+8.2%) | 75.1 |
| DINOv2-L | Nyströmformer (低秩采样) | 304.20 | 265.20 (-14.6%) | 1.2163 (-7.7%) | 18.29 (-50.8%) | 82.4 |
| DINOv2-L | Monarch (结构化矩阵) | 304.20 | 269.44 (-13.3%) | 1.2304 (-6.7%) | 18.31 (-48.9%) | 82.7 |
| DINOv2-L | ViT-AdaLA (Ours) | 304.20 | 262.19 (-15.6%) | 1.2163 (-7.7%) | 41.56 (+16.1%) | 86.0 |
| CLIP-L | Softmax (原始教师) | 304.15 | 310.60 | 1.3179 | 36.82 | 86.4 |
| CLIP-L | ViT-AdaLA (Ours) | 304.15 | 262.19 (-15.6%) | 1.2161 (-7.4%) | 42.14 (+15.7%) | 85.5 |
| SigLIP-L | Softmax (原始教师) | 316.74 | 312.46 | 1.3636 | 40.86 | 86.9 |
| SigLIP-L | ViT-AdaLA (Ours) | 316.74 | 264.04 (-15.5%) | 1.2620 (-7.4%) | 46.91 (+13.9%) | 86.4 |
Cityscapes 语义分割分辨率扩展对比（基于 DINOv2-L + Mask2Former）：
| 模型方案 | 测试分辨率 | GFLOPS (下降比例) | 峰值显存 (下降比例) | 吞吐量 (加速比例) | mIoU (%) |
|---|---|---|---|---|---|
| Softmax (原始基线) | \(512 \times 512\) | 310.60 (-) | 1.3181 GB (-) | 36.52 imgs/s (-) | 74.86 |
| Monarch Attention | \(512 \times 512\) | 269.44 (-13.3%) | 1.2304 GB (-6.7%) | 18.31 imgs/s (-48.9%) | 50.84 |
| ViT-AdaLA (Ours) | \(512 \times 512\) | 262.19 (-15.6%) | 1.2163 GB (-7.7%) | 41.56 imgs/s (+16.1%) | 72.40 |
| Softmax (原始基线) | \(1024 \times 1024\) | 1241.0 (-) | 3.2836 GB (-) | 7.07 imgs/s (-) | 80.98 |
| LoLCATS (LLM 线性化) | \(1024 \times 1024\) | 442.24 (-64.4%) | 1.3823 GB (-57.9%) | 15.38 imgs/s (+117.4%) | 33.66 |
| Monarch Attention | \(1024 \times 1024\) | 500.24 (-59.7%) | 1.5469 GB (-52.9%) | 9.74 imgs/s (+37.6%) | 55.18 |
| ViT-AdaLA (Ours) | \(1024 \times 1024\) | 429.35 (-65.4%) | 1.3764 GB (-58.1%) | 15.95 imgs/s (+125.4%) | 78.73 |
消融实验¶
训练阶段配置消融（基于 DINOv2-L，在 ADE20K 评估 mIoU）：
| 阶段 1 (注意力对齐) | 阶段 2 (特征对齐) | ADE20K mIoU (%) | 说明与影响 |
|---|---|---|---|
| \(\times\) | \(\times\) | 22.92 | 未经对齐直接接入下游微调，性能严重崩溃 |
| \(\checkmark\) | \(\times\) | 19.37 | 仅独立对齐各层注意力，跨层残差累积反致深层表征发散 |
| \(\times\) | \(\checkmark\) | 52.46 | 仅直接对齐顶层特征，收敛较慢但大幅恢复基础模型表征能力 |
| \(\checkmark\) | \(\checkmark\) | 55.55 | 完整三阶段流程，逼近 Softmax 原型（56.73%） |
参数微调范围消融（512 分辨率）： - CLIP-L 全参数微调 vs 仅微调 QKV：IN1K 准确率从 82.5% 跃升至 85.5%，ADE20K mIoU 从 35.41% 提升至 42.85%（提升 +7.44%）。 - DINOv2-L 全参数微调 vs 仅微调 QKV：IN1K 准确率从 76.4% 提升至 86.0%，ADE20K mIoU 从 48.82% 提升至 55.55%（提升 +6.73%）。仅微调 QKV 限制了模型表征调整容量，导致蒸馏时发生过拟合与特征扭曲。
关键发现¶
- 阶段一的核心价值在于加速与平滑阶段二收敛：从训练损失曲线可见，经阶段一极低成本（几分钟）注意力对齐后，阶段二特征对齐的初始损失大幅降低，并在数个 epoch 内快速达到优异的语义提取水平。
- 高分辨率零样本外推优势显著：当测试分辨率从 \(512^2\) 扩展至 \(1024^2\) 时，输入 token 数量增加 4 倍，Softmax 显存暴涨 149%（达 3.28 GB），而 ViT-AdaLA 仅需 1.38 GB 显存（节省 58.1%），吞吐量达到 Softmax 的 2.25 倍（15.95 vs 7.07 imgs/s），且 mIoU 从 72.40% 自然提升至 78.73%，证明其保留的全局拓扑具备极强的位置自适应性。
- Vanilla 线性注意力优于复杂逼近结构：在阶段一对比中，使用 ELU+1 激活的 Vanilla 线性注意力的逼近误差明显低于 Softplus、ReLU 与 Softmax 逼近，且避免了复杂核函数的显存开销。
亮点与洞察¶
- 范式转变（从复杂近似到先验继承）：以往视觉线性化研究纠结于复杂的数学低秩展开或核逼近，不仅从头训练困难，而且实际硬件吞吐低下；ViT-AdaLA 揭示通过两级对齐直接将已有 VFM 的表征先验迁移到线性架构中，能以极高性价比换取极小性能损失。
- 阶段二解冻 MLP 避免注意机制过载：发现仅微调 QKV 会破坏隐空间分布并造成严重过拟合，联合微调 MLP 可以让蒸馏梯度在全网络平滑分散，是弥补线性表达力不足的关键实操 trick。
- 模块与架构正交兼容：ViT-AdaLA 的自适应流程对具体的线性注意力结构完全解耦，未来新型线性核或状态空间块均可无缝嵌入该对齐管线。
局限与展望¶
- 计算复杂度分水岭制约：根据理论复杂度 \(\mathcal{O}(ND^2)\) 与 \(\mathcal{O}(N^2 D)\)，在线性注意力下当特征通道维度 \(D\) 远大于序列长度 \(N\) 时（如极小分辨率 \(N=256, D=1024\)），线性注意力并不比 Softmax 更快，其吞吐与显存优势仅在长序列/大分辨率场景下显著体现。
- 阶段二依赖大规模通用数据集：当前特征对齐需在 ImageNet-22K 上训练 10~30 个 epoch，计算耗时仍在数天量级（约数十卡时），未来可探索更轻量的高效蒸馏数据子集与更快的特征对齐策略。
相关工作与启发¶
- vs LoLCATS / Hedgehog: 针对 LLM 提出的线性化方法仅关注单层注意力拟合或 Decoder 时间步自回归，在 ViT 稠密预测中会因误差累积与语义流形断裂而严重劣化；ViT-AdaLA 针对编码器架构提出了全局特征对齐，彻底解决了稠密预测的语义漂移问题。
- vs Monarch Attention / Nyströmformer: 该类方法属于结构化数学近似，虽然无需大量训练，但近似矩阵计算复杂、硬件吞吐低（如 Monarch 在高分辨率下的吞吐量远落后于 Vanilla 线性注意力）；ViT-AdaLA 采用极简的 Vanilla 线性算子，在硬件吞吐和下游精度两端均实现大幅超越。
评分¶
- 新颖性: ⭐⭐⭐⭐ [打破了传统从头训练线性 ViT 的思维定势，系统性揭示了视觉编码器与语言解码器在注意力自适应上的核心差异]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 DINOv2、SigLIP、CLIP 和有监督 ViT 四类基础模型，涵盖分类与分割多尺度 benchmark，消融全面]
- 写作质量: ⭐⭐⭐⭐⭐ [动机叙述逻辑紧密，阶段划分清晰，理论与实验分析翔实]
- 价值: ⭐⭐⭐⭐⭐ [为高分辨率密集感知、扩散模型以及 3D 视觉生成等长视觉序列场景提供了极具实用价值的高效部署范式]