预训练视觉 Transformer （ViT）的核心：特征表示还是注意力模式？

传统智慧认为，预训练视觉 Transformer （ViT）通过学习有用的表示可以提高下游性能。这是否真的如此？作者对此进行了调查，发现预训练期间学习的特征和表示并非至关重要。
令人惊讶的是，仅使用预训练的注意力模式（即指导信息如何在 Token 之间 Stream ）就足以使模型从零开始学习高质量特征并实现可比的下游性能。作者通过引入一个简单的称为注意力转移的方法来证明这一点，其中注意力模式只从预训练的教师 ViT 中传递给学生，通过复制或提炼注意力图。
注意力转移使学生学习自己的特征，将其与微调的教师相结合进一步提高了在 ImageNet 上的准确性。
作者对注意力图的充分性进行了系统研究，包括在分布转移设置下它们表现不佳的情况。
作者希望作者的探索能提供更好的理解预训练所实现的内容，并导致一种有用的替代方案，即标准微调的替代方案。
可重复作者的结果的代码已在 https://github.com/alexlioralexli/attention-transfer。

1 Introduction

预训练作为一种主导的机器学习范式已经显著提高了各种任务上的性能。在计算机视觉领域，自监督表示学习方法和弱监督方法使得从大量图像中学习成为可能。普遍认为这些方法有效的原因在于，它们教导模型具有对下游任务有用的有用特征。但是，这个故事实际上是真实的吗？也许在预训练过程中，还学习到了另一种能力，足以解释其益处。

在本文中，作者提出了一种替代解释：预训练使模型了解信息如何在 Token 之间传递。作者特别关注视觉 Transformer （ViT）[12]，不仅因为它们是扩展的最受欢迎架构，而且因为 Transformer 明确地将这种信息流解耦。 Token 间的通信仅由注意力完成，而其余的大部分计算都是针对每个 Token 独立进行的 Token 内操作。相比之下，其他架构如卷积神经网络（ConvNets）[33; 20]同时扩展感受野并提取特征，使得隔离信息流的效果难以实现。作者假设，由 Token 内操作计算的特征对于解释预训练的优势并不重要，而预训练的注意力图通常对于下游任务是足够的。

作者通过引入一组名为注意力迁移的新方法来测试作者的假设。具体来说，作者将预训练的 ViT 模型作为老师，在下游任务上训练学生模型，同时只将老师的注意力模式进行迁移。与将所有权重（混合特征和注意力图的效果）迁移的常见微调范式不同，作者只迁移注意力图之间的信息。这样，学生必须从头学习特征，同时隔离出在预训练过程中学习的注意力图的优点。

作者研究了两种类型的注意力迁移。

第一种是注意力复制，它直接“ Copy-Paste ”注意力图。学习过程完全解耦，因为跨 Token 计算完全由教师完成，学生只学习由教师注意力图指导的跨 Token 模式。这非常适合作为科学 Prob ，但由于推理过程中需要传递两个网络，因此实际应用较少。

第二种是注意力蒸馏，学生在训练后仅从教师中提炼注意力模式，训练后不再使用注意力图。这种方法实际，同时有助于识别教师跨 Token 信息流的的重要性。

虽然两种注意力迁移变体都很简单，但作者发现它们 _高度有效_。图 1 说明了这一点，使用了 Masked Autoencoding (MAE) [22]预训练的 ViT-L [12]。与无转移（从头训练）和全转移（微调所有 MAE 权重）相比，注意力复制可以关闭大部分性能差距，而注意力蒸馏可以在 ImageNet-1K 分类[10]上达到微调的准确性。这是通过只传输相同模型的跨词流实现的。此外，由于注意力迁移需要学生从头学习特征，这些特征与教师（图 5）的特征显著不同，并在与教师（图 6）集成时提高 ImageNet-1K 准确性得分至 86.3（+0.6）。

总结一下，作者做出了以下贡献：

针对注意力图的充分性进行详细分析。作者发现，仅使用预训练的注意力模式通常足以实现与在 ImageNet-1K 上微调相同的下游精度。此外，作者观察到实际效益，因为与注意力迁移结合显著提高了 ImageNet 性能。这质疑了普遍认为预训练仅关于特征学习的故事。尽管作者的主要观察对不同的模型和预训练方法都具有鲁棒性，但作者确实发现有些设置，在这些设置中，预训练特征实际上是必要的，才能实现预训练的全部收益。作者的注意力迁移的最小解决方案比权重调优更受数据分布转移的影响。

注意力迁移方法。作者提出了注意力复制（Attention Copy）和注意力蒸馏（Attention Distillation）方法，这些方法在利用预训练教师 ViT 的注意力图训练 ViT 在下游任务上，只使用预训练教师 ViT 的注意力图。这些方法帮助作者理解预训练期间特征与学习到的注意力模式的角色的差异。通过进一步研究，注意力迁移可能提供了一个潜在的替代方法，以结束为期十年的微调预训练视觉模型的实践。微调流水线的几乎所有方面都已进行全面检查，这表明可能的配方已经饱和。权重共享也可能面临安全风险。作者希望作者的注意力迁移系统检查能为如何利用预训练 ViTs 提供新的洞察，并在权重转移不适用时，帮助确立这种方法的有效替代。

Attention Preliminaries

在多头自注意力（MSA）块中，每个 Token 使用多个注意力分布。注意力图确定其他 Token 的值如何聚合，每个 Token 使用多个注意力分布。对于每个 Query ，软 max 函数在_attention map_上计算。注意力图是确定如何聚合其他 Token 值的，在多头中，每个 Token 使用同一多头自注意力（MSA）块中的多个注意力分布。

对于L层 Transformer，MSA（Multi-Head Self-Attention）块与 MLP（Multi-Layer Perceptron）块交错排列，每个 Transformer 层包含每种块类型的一种。

然而，每种操作的相对重要性尚不明确，Transformer 通常通过联合微调所有权重进行训练。

偏离常见的联合权重调优做法，作者提出两种注意力转移方法，旨在探索解耦训练用于 ViTs，以下进行描述。

Attention Copy

在这个设置中，作者使用两个单独的网络：一个预训练的教师网络仅执行一次前向传播以计算其注意力图，一个学生网络直接从教师网络复制注意力图，但计算所有其他激活。学生的权重是随机初始化的，并通过反向传播进行训练，而教师网络的权重保持冻结。这种设置完全将注意力图与应用于它们的特征隔离，因此非常适合在学生网络学会执行其他任务（例如图像分类）时测量预训练注意力模式的效用。

作者称这种方法为 Attention Copy,因为作者“复制并粘贴”教师网络中的注意力图到学生网络中。图 2（左）显示了这种方法的基本原理。请注意，在推理时间内需要通过教师和学生的网络进行前向传播。由于额外的计算，Attention Copy 并非完全实用的方法。作者将在下一节中解决这个问题。

Attention Distillation

与 Attention Copy 相比，Attention Distillation 更具实用性。在训练后，老师无需再参与，学生可以作为独立模型使用。与从头训练 ViTs 相比，唯一的添加是 distillation 损失，这意味着大多数优化（例如，学习率，动量）和正则化（例如，权重衰减，dropout 率[50]）超参数可以遵循从零开始的配方，只需进行微调。这确实引入了一个新的超参数 λ，该超参数加权并平衡了 distillation 损失和任务损失。

注意力蒸馏可以被视为一种泛化的知识蒸馏形式，但它与 Hinton 等人提出的架构设计存在几个关键差异。注意力蒸馏训练学生匹配教师的注意力图层，而不是最终的教师输出。这赋予了从任何任务训练的模型中蒸馏的能力，而不仅仅是相同最终任务训练的模型。这一特性非常适合今天的“预训练和迁移”范式，其中预训练任务（例如，重建）和下游任务（例如，分类）通常不同。然而，注意力蒸馏确实增加了架构需要计算注意力图层的约束。作者将其作为未来的工作，继续探索这一想法在其他架构上的应用。

总的来说，虽然更复杂的设计可以用于注意复制和注意浓缩，但作者选择保持它们简单，以便更清晰地评估其有效性。

Connection to Transformer Training Dynamics

作者的研究还与最近试图从理论上理解 Transformer 的训练动态的努力有关。具体而言，预训练注意力图中的跨 Token 流可以被视为数据集中的一个发现的潜在层次结构。自注意力可以快速捕获经常同时出现的 Token 对[31, 52]。然而，更多的偶然同时出现需要由高层次结构来解释，而不是直接在低层次上学习[53]。这是因为存在许多潜在的虚假相关性[30]，特别是在超参数化的设置中。从训练好的教师模型中传输注意力图可以减少这些虚假的跨 Token 相关性，这样学生就可以专注于 Token 内的学习（即，计算有用的特征）。

2 Main Results

如图 1 所示，尽管其简单，但注意力转移非常有效。具体来说，作者使用预训练的 ViT-L 模型[12]进行 Masked Autoencoding（MAE）[22]进行图像 Net-1K 分类[10]。请注意，这是确立 MAE 预训练有效性的标志性结果：与从零开始训练的 ViT-L（准确率为 83.0）相比，使用相同数据集预训练的 MAE 模型进行微调，结果显著提高至 85.7。2 对于注意力转移，作者使用与教师相同的预训练 MAE 模型，由于从零开始训练可以看作没有转移，而微调权重转移所有权重，上述两个结果可以作为作者注意力转移方法有效性的自然下限和上限。作者从表 1 和图 1 中得到两个观察（见表 1 和图 1）。

Attention Copy can largely close the gap between scratch training and full weight fine-tuning

作者报告了 Attention Copy 的准确率为 85.1，这基本上关闭了从 scratch 到 full weight tuning（准确率提高 2.7 个百分点，即 87.8%）之间的差距。这令人惊讶，因为在预训练另一个任务（图像重构）后，老师的注意力图像是冻结的，学生必须从头开始学习 everything else（包括 intra-token 操作）。

作为另一个上界，作者还尝试了从_fine-tuned_模型中获取注意力复制。这达到了 85.6 的准确度，几乎与教师的性能相匹配（85.7），表明将注意力图适应特定任务仍然是有益的，但并非至关重要，尤其是在进行相同的数据分布上的 MAE 预训练。

Attention Distillation can match fine-tuning performance

更为惊奇的是，作者发现注意力蒸馏可以达到 85.7 的性能，与微调 ViT-L 权重在 MAE 上相媲美。由于注意力蒸馏和权重调优都导致了相同大小的模型，这需要相同的计算预算进行推理，因此这个结果表明，当后者不太适用时（例如，如果权重共享存在安全风险，作者可以选择发送老师的注意力图），注意力蒸馏可以成为权重调优的有效替换方案。

作者假设蒸馏比复制更好，因为学生可以选择它与教师注意力图的匹配程度，以便更好地适应任务目标。这也得到了从微调的 MAE 中复制的 85.6%的准确性的支持，因为它具有正确的任务特定注意力图。

3 Analysis

接下来，作者进行广泛的分析以更好地理解注意力迁移的有效性。总的来说，这些探索主要受到以下两个问题的驱动：

不同激活、层和头的重要性有多大？(第 4.1 节)

注意力转移是否需要重新学习一切？(第 4.2 节)

Variants of Attention Transfer

作者研究了四种注意转移变体。作者使用本节内的注意复制，因为它是一种完全解耦的设置，非常适合科学研究。

作者建议复制 Q 为模型提供偏离教师注意力图的灵活性，并使用更适合下游任务的注意力模式。这一观点得到了支持，因为复制 Q 和 Attention Copy 从微调模型中取得的准确率相同，均为 85.6。附录 B.3 对此假设进行了深入研究，发现复制 Q 的注意力图类似于教师注意力图，但在其他迁移方法中更为宽松。虽然可以在未来的工作中进行更多调查，但作者的发现表明 Q 比 K 更重要，这与在文本序列建模中先前的发现一致，即每个层的 Key 和 Value 的数量可以显著减少 [49]。

最后，作者测试了是否可以将 Q 蒸馏到比全注意力蒸馏更好的程度。然而，表 3 显示 Q 蒸馏显著表现更差。作者假设这可能是由于学生在 Q 仍在学习的过程中，学习有用的键 K 更具挑战性，而且注意力蒸馏已经具有调整其注意力图以适应下游任务的灵活性。

部分迁移：层。接下来，作者改变 Transformer 层数迁移的数量，以确定哪些层从教师中更有价值。 Baseline 将所有层都进行了迁移。在图 3 中，作者尝试仅从第一层或最后层迁移注意力图。对于其余层，学生学习计算自己的注意力图。

作者得出以下观察：

i）作者发现转移更多层通常更有帮助。这有点令人惊讶，因为人们可能会预期针对 MAE 的 Patch 式重建任务的注意力模式在像分类这样的高级识别任务上是不够理想的。

ii）作者发现转移后期的注意力图通常更好。事实上，当将最后 15 个注意力图（总共 24 个）转移到一个总性能大致饱和的状态时，性能大致饱和。这表明 ViTs 可以在被告知如何将这些特征组合成更高层次的表示时，学习良好的低层次表示；但反之则不行。这加强了 Tian 等人[53]的理论，即关于高层次层次指导的重要性，因为有很多可能性，注意力转移可以减少低层次的潜在虚假相关性。

部分迁移：头。最后，作者回到所有层都进行迁移，但改变每个 MSA 块复制的头数。具体来说，作者不仅可以从每个头的注意力图进行复制，还可以选择性地使用教师头的一个子集。学生可以计算未使用头上的自己的注意力模式。图 4 显示了每个层使用较少头的影响。随着头数的增加，性能提高，尽管在 16 个头时性能几乎饱和。请注意，作者简单地遵循一个基于多样性的选择策略，并使用第一组头；更先进的基于多样性的策略可能会在减少头数时提高鲁棒性。

Are We Re-Learning the Teacher?

由于迁移学习为学生提供了大量来自教师的信息（ViT-L 注意力图总共约有 10M 个激活值），一个自然的问题是学生表现良好是否仅仅是因为它简单地重新学习了教师的表示。作者在学生模型的不同方面彻底测试了这一假设。

表示相似性。 学生可以通过学习相同的中间特征来复制教师。作者使用中心核对齐（CKA）[32]来衡量表示的相似性，这是一种对使用不同初始化或架构的神经网络有效的相似性度量。CKA 是一种层级的比较，范围从 0（完全不同）到 1（相同），对正交线性变换和特征的 isotropic 缩放不变。图 5 显示了作者的微调模型和作者的注意力迁移方法之间的 CKA。作者还展示了预训练模型和从头训练的 ViT-L 模型供参考。作者计算 CKA 相对于微调模型，尽管作者从预训练的 MAE 模型中复制或蒸馏，因为微调期间特征变化显著，使其更具任务特定性。总体而言，注意力复制和注意力蒸馏都与微调的 MAE 模型非常不同，与原始模型的相似性趋势相同。作者的验证通过，因为 CKA 显示预训练和微调的 MAE 在早期层具有非常相似的表示。这是预期的，因为使用分步学习率衰减[8]微调意味着微调过程中最早的层变化很小。

预测相似性和集成学习。 作者的 CKA 分析可能无法捕获中间表示的某些相似性，因为 CKA 无法检测到相同信息的全部形式（例如，可逆非线性变换）[32]。由于中间表示可能无法完整地讲述全部信息，作者还研究了网络输出的相似性。

作者使用网络集成来衡量这一点：给定来自微调模型的软 max 预测和来自另一个模型的，作者测试它们的平均值的准确性。模型预测越独立，其集成精度越高。图 6 比较了在集成微调模型之前和之后的准确性。注意转移的相似性足够高以实现高集成准确性，并将注意力蒸馏与微调 MAE 集成，实现了 86.3，比微调 MAE 模型高 0.6。

最后，附录 B.4 展示了由注意分散学习获得的注意力图，表明它们在经过提炼的注意力块中相匹配，但对于未经过提炼的层来说则有极大的差异。

4 Generalization and Limitations

在本节中，作者将测试作者的注意力迁移研究成果在不同预训练和微调数据集、预训练方法、模型大小和任务方面的适用性。

Pre-training and fine-tuning datasets

迄今为止，作者主要关注在 ImageNet-1K 上预训练和评估的 MAE 模型。如果作者在不同的数据集上进行预训练或评估，又会发生什么呢？作者通过在两个新的数据集上预训练 MAE ViT-L 模型进行测试，这两个数据集分别是 ImageNet-22K [10]和 COCO [37]。这些数据集与 ImageNet 的偏差在多个轴上都有所不同，如外观、类平衡和多样性。表 4 显示，在 ImageNet-1K 上微调得到的 MAE 模型在性能上仍然保持相对良好，最多下降 0.5。然而，注意力复制准确率下降更多，甚至达到 2.1。在表 5 中，作者使用在 ImageNet 上预训练的 MAE 模型并将其转移到 iNaturalist 数据集[56]上，同样可以看到，当预训练数据集与转移数据集不匹配时，注意力复制准确率会显著下降。作者认为，冻结的教师注意力图对于微调数据集来说并不适合，这限制了性能。

Out-of-distribution robustness

显著的是，标准的微调 MAE 模型表现出轻微的“有效鲁棒性”，_i.e._，其异分布（OOD）准确性略好于基于其内分布（ID）准确性预期[14]。作者测试了是否注意力蒸馏，实现相同的 ID 准确性，也有同样的益处于 OOD。表 6 显示，注意力蒸馏仍然表现良好，但在作者尝试的所有 4 种分布转移上的准确率都低于微调的 MAE。这些结果表明，注意力图并未完全解释鲁棒性的全部优势，而 MAE 在预训练期间学习的特征在非 ID 情况下也有助于 OOD。

Pre-training methods

作者迄今为止主要关注了基于重建的预训练方法（MAE）。现在作者来检查，如果老师使用不同的算法进行预训练，注意力迁移是否仍然有效。具体来说，作者测试了 MoCo-v3 [7]，一种自监督对比学习方法，以及 FLIP [36]，一种图像文本对比学习方法。表 7 显示，对于每种预训练方法，注意力复制仍实现了大部分性能改进。令人印象深刻的是，ViT-L 仅通过从 FLIP 转移注意力图，就能达到 86.6。这证实了在学习过程中，学习适当的注意力模式确实是一个重要的 Bottleneck 。请注意，作者使用的 FLIP 模型是在 LAION-2B [48]上预训练的，但其对 ImageNet-1K 的分布转移的影响较小。

Model size

作者测试了注意力迁移是否适用于模型大小的变化。迄今为止，作者一直在使用 ViT-L；在这里，作者尝试了来自较小（ViT-B）和较大（ViT-H）的注意力复制模型，这两种模型都使用 MAE 进行预训练。表 8 显示，随着规模的增长，注意力复制继续提高，甚至达到了 86.1%的准确率（在 ViT-H 上）。尽管在 ViT-L 大小的 scratch 模型性能已经饱和，但它仍然可以做到这一点。这再次表明，模型需要适当的跨 Token 路由来学习好的泛化特征。否则，它们无法充分利用增加的模型容量。

Object Detection

最后，作者研究了在标准 ViTDet Pipeline [35]中注意力迁移在 COCO 目标检测中的性能。作者将从头开始训练与微调与从 COCO 上的 MAE ViT-B 预训练进行注意力迁移进行比较，以减轻分布漂移的影响。为了进行公平的比较，作者使用了一个 448✖️448 的输入来消除窗口注意力和位置嵌入插值的影响，并消除了相对位置嵌入的影响。表 9 表明，注意力蒸馏在稠密预测设置中恢复了预训练的大部分收益。根据表 8，作者预计 ViT-L 将使微调与注意力迁移之间的差距减小，但受计算资源限制。

5 Related Work

注意力模式的结构。许多先前的研究已经研究了预训练视觉 Transformer 的注意力模式[59, 63, 43]。这些工作仅将这些差异视为定性观察，而作者能将注意力模式独立出来，并显示它们是微调性能差异的主要原因。其他方法，如 Lora [28]或 Prompt-to-Prompt [25]，依赖于预训练网络中高质量注意力模式的重要性，但它们也利用预训练特征，并未能像作者一样得出这些特征通常对作者所研究任务不必要的结论。Trockman 和 Kolter [55]观察到训练有素的监督网络中注意力层权重乘积的斜对角结构。他们表明，使用这种结构的权重初始化可以在训练早期对小模型提高准确性。与作者工作最相似的是 Zhang 等人[68]在语言领域的研究，他们发现预训练 BERT 模型在几个特定的合成任务上提高了长度泛化。他们归因于几个特定 Head 的注意力模式，并表明将这些模式硬编码到网络中可以实现同样的益处。作者的工作是互补的，并强调注意力图的重要性超过了特征。

解耦跨词和跨词内部操作。 GLoMo [64] 也试图解耦特征应该如何组合。他们使用无监督预训练来训练一个网络以输出一个图，该图后来用于组合特定任务的特征。作者发现没有必要开发一个专门的架构来实现这一点 - 视觉 Transformer 已经自然地做到这一点了。

知识蒸馏 知识蒸馏是一种流行的框架，用于训练小型、高性能的学生网络[26]。知识蒸馏方法通常在学生网络中添加一个损失，以鼓励学生网络匹配教师网络的 logits，但变体通常使用其他特征统计数据作为目标，例如最终表示[41, 13]，中间特征相似度[19]或中间特征幅度[66, 34]。这种最后的方法也 previously 被称为“注意力转移”，但他们的方法确实不同，实际上指的是在卷积神经网络（ConvNets）中提取空间激活幅度。这些知识蒸馏方法都假设学生需要明确地学习正确的特征。相比之下，作者的注意力转移分析表明，注意力图足以恢复预训练的所有收益。一些论文将注意力蒸馏作为一种辅助损失，以帮助较小的模型更有效地学习教师输出[62, 61]。然而，这些研究仅考虑在模型大小之间传递相同的功能，而实际上是指从预训练模型到不同的下游任务传递知识。

论文中关于彩票票假设的联系** 彩票票假设[15]提出，大的密集神经网络包含小而稀疏的子网络（"中奖票"），当从零开始训练时，可以匹配甚至超越原密集网络的表现。这尤其有趣，因为这些稀疏子网络仅通过其原始初始化值保持性能；它们之间的神经元连接在某种程度上具有特殊性。作者的发现形成了一种类比，表明在预训练的 ViTs 中，由注意力模式控制的"patches"之间的连接也具有类似的特殊性。Frankle 和 Carbin [15]进一步推测，过拟合可以提高性能，因为更大的模型包含指数更多的稀疏子网络叠加在一起，因此更有可能包含一个"中奖票" - 这个假设得到了后续的实证和理论工作[46, 44, 42, 3]的支持。然而，在作者的设置中，由于 ViT 注意力模式只有每层少数几个注意力图（而不是数千个神经元），这种现象并未出现。因此，良好的注意力模式不太可能通过偶然出现，而必须在学习预训练过程中进行学习。作者希望一种新的模型架构能够有效地结合每层更多的注意力图，以解决这种限制，并比现有的 ViTs 更好地从零开始学习。

6 Conclusion

尽管 Transformer 在流行度上不断攀升，但作者的使用方式仍然保持不变：预训练，然后微调权重。在这项工作中，作者提出了注意力转移，这是 ViT 微调的一个简单替代方案，它将每个 Token 内的操作（如何为每个 Token 提取更多有用的特征）与 Token 间的操作（如何组合这些特征）解耦。

作者的关键发现是，注意力模式（ Token 间的操作）是预训练有效性的关键因素 - 作者的注意力蒸馏方法在 ImageNet-1K 上完全匹配微调。

作者确实发现了一些局限性：如果预训练和转移数据集不同，注意力转移效果不佳，而且它失去了一定程度的泛化鲁棒性。然而，作者的发现提供了关于预训练 ViTs 中注意力的洞察，作者希望未来的工作修复注意力转移的缺陷，并探索这种新转移方法的优势。

一些未来的研究方向特别有趣。首先，对 Query 进行更深入的研究可以帮助作者更好地理解其重要性，并可能产生更好的迁移策略。其次，注意力转移消除了微调所需的一些技巧，如层逐步学习率衰减。层逐步学习率衰减增加了早期层应该比后期层变化较少的先验。

然而，这种先验对于下一代模型可能过于限制，因为它阻止了早期特征的变化，删除它可能为新一代模型开辟新的机会。最后，由于注意力图是，其中是序列长度，注意力图可以在模型大小之间更容易地转移。

相比之下，权重调优在模型具有不同维度时更为困难。预训练一个较小的模型并将其注意力模式转移到较大的下游模型可能比当前的微调实践更实际。

参考文献

[0]. On the Surprising Effectiveness of Attention Transfer for Vision Transformers.

END

作者：小书童
来源：集智书童

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。

预训练视觉 Transformer （ViT） 的核心：特征表示还是注意力模式？