ViT-AdaLA：线性注意力适应视觉Transformer

三阶段对齐框架高效地将预训练ViT转换为线性注意力模型

前置知识

Softmax Self-Attention

Transformer核心模块，通过计算查询和键的点积得到注意力分数，再用softmax归一化后与值加权求和。对于输入序列X属于R的N乘D维矩阵，输出Oi等于对所有j的exp乘以QiKj的转置再乘以Vj，除以对所有k的exp乘以QiKk的转置，其中Q等于X乘以WQ，K等于X乘以WK，V等于X乘以WV。计算复杂度为O(N的平方D)，因为需要计算所有N个token之间的两两交互。

ViT的所有性能都建立在self-attention基础上，理解它是理解为什么需要线性注意力以及如何近似它的基础

Linear Attention

利用kernel trick将softmax分解，exp乘以QK的转积等于phi乘以Q乘以phi乘以K的转置，通过重新排列乘法顺序将复杂度降为O(ND的平方)。具体实现是Oi等于phi乘以Qi乘以对所有j的phi乘以Kj的转置乘以Vj，再除以对所有k的phi乘以Kk的转置，常用激活函数phi(x)等于ELU(x)加1确保非负性。关键是将先算注意力矩阵再加权，改为先计算上下文聚合再加权。

本文的核心就是设计有效方法让linear attention继承pretrained softmax ViT的知识，必须理解其原理才能理解为什么简单的vanilla linear attention反而效果更好

Vision Foundation Models

在大规模数据上预训练的通用视觉模型，如DINOv2自监督方法、CLIP图文对比方法、SigLIP等。它们通过自监督学习或对比学习获得了强大的视觉表征能力，可以作为特征提取器用于下游任务如分割、检测、视觉问答等。这些模型通常使用softmax-based ViT作为backbone骨干网络。

ViT-AdaLA的目标就是继承这些pretrained VFMs的知识，避免从零训练linear ViT需要的巨大计算资源。理解VFMs的性质和优势是理解动机的基础

Knowledge Distillation

将大型教师模型的知识迁移到小型学生模型的技术。通过设计蒸馏损失函数，让student学习teacher的输出、中间特征或注意力权重。本文中，softmax ViT作为teacher，linear ViT作为student，通过MSE损失对齐attention输出和最终特征。经典的Hinton KD使用KL散度对齐logits，但本文用MSE更适合对齐特征。

ViT-AdaLA本质上是一个蒸馏框架，理解蒸馏的原理和常见变体有助于理解为什么三阶段对齐策略有效，以及为什么只调QKV矩阵会导致过拟合

Residual Error Accumulation

在每个Transformer block中用linear attention近似softmax attention时会有误差，这些误差会随着深度逐层累积。对于N层ViT，每层的微小近似误差在经过多次MLP变换和残差连接后会放大，导致最终输出特征与原始模型偏差较大。特别是在分割等密集预测任务中，错误在空间和层次上累积会影响全局语义流形。

这是本文设计Stage 2特征对齐的关键原因。理解误差累积机制才能明白为什么只做attention对齐是不够的，必须进一步对齐最终层特征

研究动机

基于ViT的视觉基础模型如DINOv2和CLIP在计算机视觉任务中表现出色，但标准的softmax self-attention具有二次复杂度O(N的平方D)，当处理长序列即高分辨率图像时面临严重的计算和内存瓶颈。随着现代视觉应用需要处理更长的视觉token序列，如高分辨率分割、医学影像、卫星图像，这一问题日益尖锐。现有线性注意力方法分为两类，从零训练的方法需要大规模预训练才能获得强大的先验知识，在资源和数据受限的现实场景下实用性不足，而基于线性化的方法如LoLCATS主要针对decoder-only的LLM设计，直接迁移到视觉模型会导致显著性能下降。实验数据显示，在DINOv2-L上使用Hedgehog和LoLCATS等线性化方法后，ADE20K分割mIoU从56.73%暴跌至22.52%和17.42%，性能损失超过60%。

本文的目标是设计一个有效的框架，能够将预训练的基于softmax attention的视觉基础模型中的先验知识转移到线性注意力ViT中，同时避免大规模从零预训练的计算开销。目标是使线性化后的模型在下游任务上保持接近原始softmax模型的性能，同时获得线性复杂度的推理效率。具体来说，在ImageNet-1K分类任务上希望top-1准确率保持在原始模型1%以内，在ADE20K分割任务上mIoU损失控制在2%以内，内存消耗降低15%以上，吞吐量提升15%以上。

与已有工作不同的是，本文的独特切入角度在于识别了LLM和视觉模型在架构上的根本差异，LLM是decoder-only架构，模型同时作为特征提取器和目标生成器，而视觉模型是encoder-decoder架构，ViT仅作为特征提取器，由独立的任务头负责生成。这导致误差传播机制不同，LLM的误差在时间维度累积，而ViT的误差在空间和层次维度累积。直接移植LLM的线性化方法会破坏视觉任务所需的空间一致性。此外，现有方法倾向于设计复杂的attention近似机制如Hedgehog的双映射模块，但这些刚性结构先验会在蒸馏过程中与教师模型对抗，形成优化瓶颈。本文采用最简单的vanilla linear attention配合渐进特征对齐，首次系统性地解决了视觉模型的线性化问题。

核心方法

ViT-AdaLA是一个三阶段渐进对齐框架，核心思想是通过注意力对齐、特征对齐、监督微调三个步骤，逐步将预训练softmax ViT的知识转移到线性ViT中。方法的直觉是先在局部每个attention block让linear attention学会模仿softmax的行为，再在全局最终层特征纠正累积误差，最后在下游任务上微调。与直接训练linear ViT从零开始相比，这种蒸馏方式充分利用了现有VFMs的强大先验，大幅减少训练开销。与针对LLM设计的线性化方法相比，本文强调特征对齐的必要性，因为视觉任务需要保持空间一致性。整个框架架构无关，可以与各种linear attention变体结合。

核心创新点在于采用最简单的vanilla linear attention只用ELU加1激活而非复杂的近似机制，认为unconstrained的特性避免了优化瓶颈，能灵活学习近似模式。设计三阶段对齐策略，首次明确区分attention对齐和feature对齐，指出后者对视觉模型至关重要。渐进式知识转移而非一次性蒸馏，Stage 1的attention对齐能加速Stage 2的收敛。图5显示vanilla linear attention的training loss低于Hedgehog等方法，同时计算效率更高，证明简单性反而带来更好的近似质量和优化性能。与LoLCATS只调两个额外映射模块不同，本文直接调QKV三个投影矩阵，更充分地利用了linear attention的可塑性。

方法步骤详情

Stage 1注意力对齐，在每个Transformer block中添加一个linear attention模块，冻结原始模型所有组件，只更新新模块的WQ、WK、WV三个投影矩阵。使用MSE损失对齐softmax和linear attention的输出，其中K是层数，N是token数，D是特征维度。在COCO数据集上训练4个epoch，batch size 32，学习率固定为1e减2。Stage 2特征对齐，用Stage 1得到的linear attention替换所有softmax attention，得到linearized ViT，冻结原始softmax教师模型，fine-tune整个linear ViT使其最终层特征与教师对齐。损失函数lambda乘以对所有n和d的教师输出减去linear ViT输出的平方，fv是教师输出，fv hat是linear ViT输出，lambda控制输出损失尺度，DINOv2设为5，CLIP设为1，SigLIP设为4000。在ImageNet-22K上训练10到30个epoch，batch size 16，初始学习率1e减4用polynomial衰减。Stage 3监督微调，给linear ViT添加任务特定头，分类用linear层，分割用Mask2Former，在下游任务数据集上微调，更新backbone和head所有参数，学习率1e减4，权重衰减0.05。

技术新颖性

技术新颖性体现在首次提出针对视觉模型的线性化完整框架，系统分析了decoder-only LLM和encoder-decoder ViT的差异，指出视觉模型需要特征对齐。发现vanilla linear attention比复杂近似方法更适合蒸馏，通过实验对比了softmax、softplus、ReLU、ELU加1四种激活函数，证明ELU加1在convergence和数值稳定性上最优。展示了低分辨率预训练的linear ViT可以scaling到高分辨率而保持性能，如512平方预训练后在1024平方上mIoU从72.40%提升到78.73%，解决了训练成本和推理质量的张力。架构无关设计，可与任何linear attention方法结合，为未来研究提供灵活平台。图6的PCA可视化显示ViT-AdaLA比Monarch attention更好地保持了原始VFM的特征分布，证明其知识提取能力更强。

**Figure 4**: ViT-AdaLA consists of three stages attention alignment feature alignment and supervised fine-tuning

**Figure 5**: Linear attention architecture and Stage 1 training loss comparison with LoLCATS

**Figure 8**: Training loss comparison of different linear attention variants in Stage 1 on DINOv2-L

**Figure 12**: Stage 1 loss comparison across linear-attention activation functions

实验结果

在ImageNet-1K分类任务上，ViT-AdaLA在DINOv2-L上达到86.0% top-1准确率，仅比原始softmax的86.8%低0.8%，显著优于所有baseline，Hedgehog 58.8%、LoLCATS 61.6%、Linformer 68.9%、Performer 75.2%、Cosformer 75.1%、Nystromformer 82.4%、Monarch 82.7%。在CLIP-L上达到85.5%，softmax 86.4%，在SigLIP-L上达到86.4%，softmax 86.9%。在ADE20K分割任务上，ViT-AdaLA在DINOv2-L上达到55.55% mIoU，原始softmax为56.73%，差距仅1.18%，而Hedgehog仅22.52%、LoLCATS仅17.42%。在Cityscapes上，512平方分辨率达到72.40% mIoU，softmax 74.86%，1024平方分辨率达到78.73%，softmax 80.98%，展示了强大的分辨率扩展性。效率方面，内存消耗降低15.6%，从310.60GB到262.19GB，吞吐量提升16.1%，从36.52 imgs每秒到41.56 imgs每秒，GFLOPs保持不变。训练效率方面，Stage 1每轮仅需4到6分钟，Stage 2每轮约10到11小时。Table 3的消融实验显示仅有Stage 2时mIoU为52.46%，加上Stage 1提升到55.55%，证明Stage 1加速了Stage 2的收敛。Figure 7的训练曲线显示有Stage 1初始化时Stage 2 loss下降更快更低。Figure 8显示vanilla linear attention在Stage 1的训练loss低于其他所有variant，证明其优越的近似能力。Table 8显示从零训练linear ViT 200 epochs后ADE20K仅25.80% mIoU，而ViT-AdaLA仅用20 epochs就达到49.75%，证明了继承先验知识的巨大优势。

**Table 1**: Top-1 fine-tuning accuracy comparison on ImageNet-1K under different vision foundation models with multiple linear attention baselines

**Table 2**: mIoU fine-tuning comparison on ADE20K under different vision foundation models with multiple linear attention baselines

**Table 3**: Ablation study of Stages 1 and 2 training using DINOv2-L evaluated on the ADE20K dataset

**Table 4**: mIoU fine-tuning comparison on Cityscapes under different input resolutions 512 vs 1024 based on DINOv2-L

**Table 5**: Hyperparameters of different training stages and VFMs

**Table 6**: Comparison of training efficiency across various Visual Foundation Models during Stages 1 and 2

**Table 7**: Performance comparison using DINOv2-B on ADE20K Cityscapes and ImageNet-1K

**Table 8**: Performance comparison of training-from-scratch and on ADE20K Cityscapes and IN1K

**Table 9**: Comparison of tuning strategies on CLIP-L and DINOv2-L at 512乘以512 resolution

**Table 10**: Top-1 fine-tuning accuracy comparison on ImageNet-1K under CLIP-L with multiple linear attention baselines

**Table 11**: mIoU fine-tuning comparison on Cityscapes under different input resolutions 512 vs 1024 based on IN1K ViT-L and SigLIP-L

**Table 12**: mIoU fine-tuning comparison on ADE20K under ViT-IN1K and SigLIP with multiple linear attention baselines

**Figure 6**: Visualization of PCA-projected features from the final layer of DINOv2-L

**Figure 7**: Stage 2 training loss comparison between with and without Stage 1 pretraining

**Figure 9**: Resolution scalability analysis on the Cityscapes dataset in terms of a mIoU b throughput and c peak memory

**Figure 10**: ADE20K performance across different Stage 2 training epochs for DINOv2-L

**Figure 11**: ADE20K performance of adapting the task model or the VFM for DINOv2-L

**Figure 13**: Visualization of PCA-projected features from the final layer of DINOv2-L comparison with Monarch

查看结构化数据

任务	指标	本文	基线	提升
ImageNet-1K分类	Top-1准确率	86.0%	Softmax 86.8%, Monarch 82.7%, Nystromformer 82.4%	比最佳baseline提升3.3个百分点，比softmax仅降0.8%
ADE20K语义分割	mIoU	55.55%	Softmax 56.73%, Monarch 44.95%, Nystromformer 41.16%	比最佳baseline提升10.6个百分点，比softmax仅降1.18%
Cityscapes语义分割512平方	mIoU	72.40%	Softmax 74.86%, Monarch 50.84%, Nystromformer 48.97%	比最佳baseline提升21.56个百分点，比softmax降2.46%
Cityscapes语义分割1024平方	mIoU	78.73%	Softmax 80.98%, Monarch 55.18%	比最佳baseline提升23.55个百分点，比softmax降2.25%
推理效率	内存降低和吞吐提升	15.6%和16.1%	Nystromformer 14.6%和7.7%, Monarch 13.3%和6.7%	效率提升优于其他linear attention方法

局限与改进

作者承认的局限性包括仅在分类和分割任务上验证，对目标检测、图像生成等其他下游任务的泛化性需要进一步研究。在分割任务上与full softmax attention仍有边际性能gap约1到2个百分点mIoU。Stage 2的训练效率可以进一步提升，引入更先进的蒸馏策略如masked image modeling可能加速先验知识转移。ViT-AdaLA在低分辨率图像时计算开销比softmax attention更大，这是linear attention的共性问题，当序列长度N小于头维度D如N等于256、D等于1024时O(N的平方D)比O(ND的平方)更高效。本文观察到的额外问题是Table 9显示只调QKV矩阵会导致性能显著下降，CLIP-L从85.5%到82.5%，DINOv2-L从86.0%到76.4%，说明full fine-tune对避免过拟合很重要。在DINOv2-B等小模型上性能gap相对更大，ADE20K从53.93%到51.46%，可能需要针对小模型调整超参数。不同VFM的最优lambda值差异很大，从1到4000，需要per-dataset调参。

独立分析的弱点

独立分析的具体弱点包括低分辨率效率悖论，当图像分辨率较低即序列长度N较小时，linear attention的O(ND的平方)复杂度可能高于softmax的O(N的平方D)，限制了在移动设备等低分辨率场景的应用。改进方向可以是adaptive机制，根据序列长度动态切换attention类型。Stage 2训练成本较高，在ImageNet-22K上训练10到30 epochs仍需大量计算，虽然比从零训练好，但对资源受限团队仍有门槛。改进方向包括探索更高效的蒸馏目标，如feature matching的分层版本，引入curriculum learning逐步增加分辨率，或使用知识蒸馏的压缩技术。特征对齐的粗糙性，目前只对齐最终层特征，可能损失中间层的语义信息。改进方向可以设计分层特征对齐，让每个block的输出都与教师对齐，或引入contrastive loss增强特征空间的一致性。激活函数固定，实验只比较了四种激活函数，可能有更优的选择。改进方向可以尝试learnable activation或mixture-of-experts结构，让模型自动学习最优特征映射。

未来方向

作者提出的未来方向包括扩展到视觉大语言模型，这些模型处理长视觉序列时计算开销巨大，替换quadratic attention可显著提升高分辨率图像或长视频输入的可扩展性。结合解释方法提供更透明的洞察，理解visual tokens如何贡献模型预测，在多模态场景下理解跨模态交互仍有挑战。探索更高效和有效的linearization架构设计，作为未来研究的基础平台。基于成果可延伸的方向包括探索与Mamba等state space model的结合，将linear attention替换为SSM模块可能获得更好的长程建模能力。研究跨模态线性化，将框架扩展到音频视觉、文本图像等多模态模型。设计自动化的lambda搜索策略，减少超参数调优成本。探索linear attention在生成任务如diffusion transformer中的应用，可能获得更大的效率提升。研究quantization和pruning与linear attention的协同优化，进一步压缩模型。

复现评估

论文声称reproduce所有baseline以确保公平比较，并在附录A.2详细描述了各baseline的实现细节和超参数设置，如Nystrmformer的landmark点数设为128，Linformer的k等于32。提供了完整的训练配置Table 5，包括三个阶段的epoch数、batch size、学习率、优化器、数据集等。Stage 1在COCO上4 epochs，Stage 2在ImageNet-22K上10到30 epochs，Stage 3在下游任务上微调，使用8乘以H100 GPU。附录A.4提供了训练时间数据，Stage 1每轮4到6分钟，Stage 2每轮10到11小时。但是论文没有提供代码链接或模型权重下载，这是复现的主要障碍。不同VFM的lambda值差异很大，从1到4000，需要per-dataset调参，增加了复现难度。总体而言，复现难度中等，如果后续开源代码则较容易复现，否则需要仔细按照论文描述实现细节。算力要求较高，即8乘以H100，但Stage 2可以用更少epoch或更小数据集简化。

ViT-AdaLA：线性注意力适应视觉Transformer ViT-AdaLA: Adapting Vision Transformers with Linear Attention