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

📄 arXiv: 2504.00037v1 📥 PDF

作者: Guoyizhe Wei, Rama Chellappa

分类: cs.CV, cs.AI

发布日期: 2025-03-30

💡 一句话要点

ViT-Linearizer：通过知识蒸馏将二次复杂度ViT模型转化为线性复杂度视觉模型

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

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

📋 核心要点

  1. Vision Transformer计算复杂度高，在高分辨率图像处理中面临挑战。
  2. ViT-Linearizer通过知识蒸馏，将ViT的知识迁移到线性复杂度的循环模型。
  3. 实验表明，该方法显著提升了高分辨率任务的推理速度，并在ImageNet上取得了优异的性能。

📝 摘要（中文）

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

🔬 方法详解

问题定义：Vision Transformer (ViT) 虽然性能优异，但其全局自注意力机制导致计算复杂度为输入尺寸的二次方，这在高分辨率图像处理中成为瓶颈，限制了其在资源受限设备上的应用。现有方法通常通过近似自注意力或使用其他线性复杂度的架构来降低计算成本，但往往会牺牲模型性能。

核心思路：ViT-Linearizer 的核心思路是通过知识蒸馏，将 ViT 模型中蕴含的丰富知识迁移到线性复杂度的循环神经网络 (RNN) 风格的模型中。通过让学生模型学习教师模型的 token 依赖关系和上下文信息，使其在保持高效计算的同时，尽可能地逼近 ViT 的性能。这种方法旨在弥合 ViT 的高性能和 RNN 的高效率之间的差距。

技术框架：ViT-Linearizer 采用跨架构的知识蒸馏框架，包含一个预训练的 ViT 教师模型和一个线性复杂度的学生模型（例如，基于 Mamba 的架构）。训练过程主要包括两个阶段：激活匹配和掩码预测。激活匹配通过最小化教师和学生模型中间层激活之间的差异，来约束学生模型学习教师模型的 token 依赖关系。掩码预测则要求学生模型根据上下文信息，预测教师模型中被掩码的 token 的表示，从而学习更丰富的上下文信息。

关键创新：该方法最重要的创新点在于其知识蒸馏策略，它不仅仅是简单地模仿教师模型的输出，而是通过激活匹配和掩码预测，更深入地学习教师模型的内部表示和 token 依赖关系。这种方法能够更有效地将 ViT 的知识迁移到线性复杂度的模型中，从而在保持高效计算的同时，获得更好的性能。与现有方法的本质区别在于，ViT-Linearizer 关注的是知识的迁移，而不是简单地降低计算复杂度。

关键设计：在激活匹配阶段，使用均方误差 (MSE) 作为损失函数，最小化教师和学生模型对应层激活之间的差异。在掩码预测阶段，随机掩码一定比例的 token，并使用 MSE 作为损失函数，最小化学生模型预测的被掩码 token 的表示与教师模型对应表示之间的差异。学生模型的架构可以根据具体任务进行选择，例如，可以使用基于 Mamba 的架构来实现线性复杂度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ViT-Linearizer 在 ImageNet 分类任务上取得了显著的成果，使用基础尺寸的 Mamba 模型达到了 84.3% 的 top-1 准确率，与同等规模的 ViT 模型相比具有竞争力。此外，该方法在高分辨率图像处理任务中表现出显著的加速效果，验证了其在实际应用中的价值。这些实验结果表明，ViT-Linearizer 能够有效地将 ViT 的知识迁移到线性复杂度的模型中，从而在保持高效计算的同时，获得良好的性能。

🎯 应用场景

ViT-Linearizer 具有广泛的应用前景，尤其是在需要处理高分辨率图像或视频的场景中，例如自动驾驶、医学图像分析、遥感图像处理等。该方法可以降低模型的计算成本和内存占用，使其能够在资源受限的设备上运行，从而加速这些领域的应用落地。此外，该方法还可以促进 RNN 风格模型在视觉任务中的应用，为未来的视觉模型设计提供新的思路。

📄 摘要（原文）

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.