Arxiv 2305 | Refocusing Is Key to Transfer Learning

Arxiv 2305 | Refocusing Is Key to Transfer Learning

• 论文：https://arxiv.org/abs/2305.15542
• 代码：https://github.com/bfshi/TOAST

论文中沿用了作者自己另一篇论文中的设计，即Top-Down Visual Attention from Analysis by Synthesis中设计的Top-Down Visual Attention。

本文提出的TOAST强调的是在预训练模型迁移过程中进一步强化目标任务相关的注意力信息是很重要的。

Transformer 模型通常是自下而上的，即其注意力仅取决于输入，因此，它通常突出输入信号中的所有显着特征。与自底而上的注意力相反，由顶而下的注意力赋予了根据高级目标或任务调整注意力的能力，即它只关注与任务相关的特征而忽略其他特征。

而核心模块的特征选择操作设计比较简单，即基于固定数量的任务相关的可学习token利用Attention交互，来促使模型特征中任务相关的信息被逐步强化，强化后的表征被加到第二次前向传播中的attention里的value，并重新参与计算。从而在每一层中都增强与任务相关的信息，实现了自上而下的注意力。

# https://github.com/bfshi/TOAST/blob/645284c74ef140b2fd6dd7804dc1a17862d3f146/visual_classification/src/models/vit/vit_top_down.py#L283-L287
# init
self.prompt = torch.nn.parameter.Parameter(torch.randn(self.embed_dim), requires_grad=True)
self.top_down_transform = torch.nn.parameter.Parameter(torch.eye(self.embed_dim), requires_grad=True)

# feature selection and feedback
cos_sim = F.normalize(x, dim=-1) @ F.normalize(self.prompt[None, ..., None], dim=1)  # B, N, 1
mask = cos_sim.clamp(0, 1)
x = x * mask
x = x @ self.top_down_transform
td = self.feedback(x)

由于附加了可学习的结构，所以仍然需要对附加结构的参数进行优化。优化过程中，预训练backbone都会被冻结，仅仅调整附加的top-down attention module的相关组件。为此作者设计了一个两阶段调整策略：

1. Pre-tuning Stage：在公开数据集上进行pre-tuning来获得一个更好的初始化，这一阶段中除了常规的的监督损失和无监督损失外，也引入了Top-Down Visual Attention from Analysis by Synthesis中的variational loss，从而鼓励反馈路径可以从输出重构输入，作为对于反馈权重的正则化。这一阶段对于模型最终的性能起到了明显的帮助。
2. Tuning stage：在下游任务重，同样仅仅微调top-down attention module。由于组件的设计，实际将会有大约15%的参数被更新。其中大部分都位于反馈层，且都是大的方阵，所以作者们通过向反馈层中引入LoRA技术，从而有效降低了实际需要调整的参数到1%左右，这一变体为TOAST-Lite，实际性能与TOAST接近。

整体推理流程：

1. 将输入通过前馈路径来获得初始输出。
2. 选择输出中哪些特征对当前任务有用。这里从通道和空间token两个层面上进行了筛选与加权。消融实验中，通道上的操作展现出了重要的价值。
3. 所选特征通过自定义的附加反馈路径，发送回每个自注意力模块。
4. 再次运行前馈传递，但每个自注意力接收自上而下的信号作为附加输入。

前面的描述是针对Transformer类型的架构，对于卷积网络同样可以迁移：

• 需要将所提结构中反馈路径里的线性层更换为转置卷积层从而匹配形状。
• 由于卷积层没有自注意力，所以反馈回来的信号黑直接加到每个自底而上的卷积层的输入上。
• pre-tuning和tuning阶段保持与Transformer的一致。

可以看到TOAST获得的注意力图，对于目标更加专注。

而在文本和视觉、Transformer和ConvNet的多个实验中，获得了相较于Linear、Fine-tune、VPT、LoRA更好的表现。除了在语义分割任务的测试中，比微调稍差一点，但是仍然优于VPT和LoRA。

Visual Prompt Tuning中发现Parameter-Efficient Fine-Tuning (PEFT)方法在类似语义分割这样的密集预测任务中比不过微调策略。这可能是因为预训练在图像分类上的backbone在所需要的视觉表征的层次和语义上和分割任务差异较大。

论文中也提到了当前主要的局限性，即重复的两次前向传播实际上加倍了计算复杂度。论文中通过调整模块的计算位置，从而设计了两种变体：

• early feedback：早期反馈意味着来自中间层而不是最后一层的反馈。这样，在第一个前馈中，只需要运行中间层之前的部分层。
• late feedback：后期反馈意味着从最后一层反馈到中间层，而不是第一层。这样就可以在两个前馈运行中共享输入位置之前的输出。因为它们没有收到任何反馈，因此在两个前馈运行中具有相同的输出。

然而比较特殊的是，对于语义分割，作者发现early feedback实际上比常规feedback设置更好（在ADE20K上45.11 vs 43.59 mIoU），这可能是因为分割需要更多的细粒度和低级特征，中间层比最后一层包含更多的低级信息。