ViT-Linearizer: Distilling Quadratic Knowledge into Linear-Time Vision Models

📄 arXiv: 2504.00037 📥 PDF

作者: Guoyizhe Wei, Rama Chellappa

分类: cs.CV, cs.AI

发布日期: 2026-02-28

💡 一句话要点

ViT-Linearizer：通过知识蒸馏将二次复杂度ViT模型转化为线性复杂度模型，提升高分辨率图像处理效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Vision Transformer 知识蒸馏 线性复杂度模型 激活匹配 掩码预测 高分辨率图像处理 Mamba 循环神经网络

📋 核心要点

  1. Vision Transformer虽然性能卓越，但其全局自注意力机制导致计算复杂度呈二次方增长，限制了其在高分辨率输入上的应用。
  2. ViT-Linearizer通过跨架构蒸馏，将ViT的知识迁移到线性复杂度的循环模型，核心在于激活匹配和掩码预测。
  3. 实验表明，该方法显著加速了高分辨率任务的推理速度，并提升了Mamba等RNN架构在ImageNet等基准测试上的性能。

📝 摘要（中文）

本文提出ViT-Linearizer，一种跨架构蒸馏框架，旨在将Vision Transformers (ViTs) 的丰富表征迁移到线性时间复杂度的循环模型中。该方法利用1) 激活匹配，一种中间约束，鼓励学生模型将其token间的依赖关系与教师模型产生的依赖关系对齐；以及2) 掩码预测，一种上下文重建目标，要求学生模型预测教师模型中未见过的（掩码）token的表征。通过这种方式，可以在保持高效复杂度的同时，有效地将二次自注意力知识蒸馏到学生模型中。实验结果表明，该方法在高分辨率任务中提供了显著的加速，显著解决了推理中的硬件挑战。此外，它还提升了基于Mamba的架构在标准视觉基准上的性能，在ImageNet上使用base尺寸的模型实现了具有竞争力的84.3%的top-1准确率。这些结果强调了基于RNN的解决方案在大规模视觉任务中的良好潜力，弥合了理论效率与实际应用之间的差距。

🔬 方法详解

问题定义：Vision Transformer (ViT) 在视觉任务中表现出色，但其全局自注意力机制的计算复杂度为二次方级别，这在高分辨率图像处理时会带来巨大的计算负担和内存需求，限制了其在资源受限设备上的应用。现有方法难以在保持性能的同时，有效降低ViT的计算复杂度。

核心思路：ViT-Linearizer的核心思路是通过知识蒸馏，将ViT模型学习到的丰富表征迁移到计算复杂度更低的线性模型（如基于RNN或Mamba的架构）中。通过让学生模型模仿教师模型的行为，从而在保持甚至提升性能的同时，显著降低计算复杂度。

技术框架：ViT-Linearizer框架包含一个预训练的ViT教师模型和一个线性复杂度的学生模型。训练过程主要包括两个阶段：1) 激活匹配：学生模型学习模仿教师模型中间层的激活值，从而学习token之间的依赖关系。2) 掩码预测：学生模型需要预测教师模型中被掩盖的token的表征，从而学习上下文信息。这两个阶段共同作用，使得学生模型能够有效地学习到教师模型的知识。

关键创新：ViT-Linearizer的关键创新在于其跨架构的蒸馏方法，它不仅关注最终的预测结果，更关注中间层的表征学习和上下文信息的重建。激活匹配和掩码预测的结合，使得学生模型能够更全面地学习到教师模型的知识，从而在降低计算复杂度的同时，保持甚至提升性能。

关键设计：在激活匹配阶段，使用均方误差（MSE）损失函数来衡量学生模型和教师模型中间层激活值的差异。在掩码预测阶段，随机掩盖教师模型的部分token，然后让学生模型预测这些被掩盖的token的表征，同样使用MSE损失函数进行优化。此外，还使用了标准的分类损失函数来保证学生模型的分类性能。具体的网络结构可以根据实际需求选择，例如可以使用Mamba作为学生模型。

📊 实验亮点

ViT-Linearizer在ImageNet数据集上取得了显著的成果，使用base尺寸的模型达到了84.3%的top-1准确率，与同等规模的ViT模型相比具有竞争力。更重要的是，该方法在高分辨率图像处理任务中实现了显著的加速，有效解决了ViT在高分辨率输入上的计算瓶颈。实验结果表明，该方法能够有效地将ViT的知识迁移到线性模型中，并在保持甚至提升性能的同时，显著降低计算复杂度。

🎯 应用场景

ViT-Linearizer具有广泛的应用前景，尤其是在需要处理高分辨率图像或视频的场景中，例如自动驾驶、医学图像分析、遥感图像处理等。通过降低计算复杂度，该方法可以使得ViT模型能够在资源受限的设备上运行，从而扩展了其应用范围。此外，该方法还可以用于提升现有线性模型的性能，例如Mamba等。

📄 摘要（原文）

Vision Transformers (ViTs) have delivered remarkable progress through global self-attention, yet their quadratic complexity can become prohibitive for high-resolution inputs. In this work, we present ViT-Linearizer, a cross-architecture distillation framework that transfers rich ViT representations into a linear-time, recurrent-style model. Our approach leverages 1) activation matching, an intermediate constraint that encourages student to align its token-wise dependencies with those produced by the teacher, and 2) masked prediction, a contextual reconstruction objective that requires the student to predict the teacher's representations for unseen (masked) tokens, to effectively distill the quadratic self-attention knowledge into the student while maintaining efficient complexity. Empirically, our method provides notable speedups particularly for high-resolution tasks, significantly addressing the hardware challenges in inference. Additionally, it also elevates Mamba-based architectures' performance on standard vision benchmarks, achieving a competitive 84.3% top-1 accuracy on ImageNet with a base-sized model. Our results underscore the good potential of RNN-based solutions for large-scale visual tasks, bridging the gap between theoretical efficiency and real-world practice.