[论文解读] Fine-tuning Image Transformers using Learnable Memory

Mark Sandler, Andrey Zhmoginov+2 more|arXiv (Cornell University)|Mar 29, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文提出通过在视觉变换器（Vision Transformers）中引入可学习的记忆标记——在每一层注入额外的可学习嵌入——以实现高效、参数高效的微调，适用于下游任务。通过使用注意力掩码机制，该方法使模型在扩展至新任务时，能保留对先前任务的性能，同时仅造成极小的准确率下降，其性能接近全量微调，且每层仅增加少量额外参数。

ABSTRACT

In this paper we propose augmenting Vision Transformer models with learnable memory tokens. Our approach allows the model to adapt to new tasks, using few parameters, while optionally preserving its capabilities on previously learned tasks. At each layer we introduce a set of learnable embedding vectors that provide contextual information useful for specific datasets. We call these "memory tokens". We show that augmenting a model with just a handful of such tokens per layer significantly improves accuracy when compared to conventional head-only fine-tuning, and performs only slightly below the significantly more expensive full fine-tuning. We then propose an attention-masking approach that enables extension to new downstream tasks, with a computation reuse. In this setup in addition to being parameters efficient, models can execute both old and new tasks as a part of single inference at a small incremental cost.

研究动机与目标

• 为解决视觉变换器在持续学习场景中全量微调的高成本与不稳定性问题。
• 在仅进行少量参数更新的前提下，实现对新下游任务的高效适应。
• 在将模型扩展至新任务时，保持对先前学习任务的高性能。
• 实现在单个模型中多个任务的计算复用与联合推理。
• 探索适用于视觉变换器的参数高效、可扩展且稳定的终身学习方法。

提出的方法

• 在每个ViT编码器层中增加一组少量可学习的记忆标记，将其附加到输入序列中，以在注意力计算前使用。
• 采用改进的注意力机制，其中记忆标记不关注其他标记，且仅类别标记和图像块标记在各层之间更新。
• 引入一种注意力掩码策略，限制旧任务的标记无法关注新添加的记忆标记和类别标记，从而在推理过程中防止干扰。
• 通过在输入序列中增加新任务特定的记忆标记和类别标记，顺序训练模型以应对新任务，同时保持注意力约束以维持旧任务的性能。
• 通过拼接独立训练的模型并配合注意力掩码，实现模型扩展，支持联合推理，且计算成本仅小幅递增。
• 为所有任务使用共享主干网络，配合任务特定的记忆标记和分类头，实现参数共享与高效推理。

实验结果

研究问题

• RQ1 可学习的记忆标记是否能在仅进行少量参数更新的情况下，显著提升视觉变换器微调的准确率？
• RQ2 注意力掩码是否能防止在将模型扩展至新任务时，对先前学习任务性能的下降？
• RQ3 与仅微调分类头和全量微调相比，记忆增强微调的性能如何？
• RQ4 能否将多个在不同任务上训练的模型组合成一个单一模型，实现极小的准确率损失与增量计算成本？
• RQ5 每层记忆标记的数量是否影响模型性能，是否存在性能饱和点？

主要发现

• 每层仅增加少量记忆标记，即可显著提升准确率，其性能接近全量微调，远超仅微调分类头的方法。
• 通过注意力掩码，顺序在多个任务上训练的模型在所有先前任务上均保持高准确率，避免了朴素拼接方法中常见的20%-40%准确率下降。
• 通过顺序添加任务（如先SUN-397后Places-365，或反之）扩展的模型表现出稳定性能，无性能退化，优于独立训练后拼接的模型。
• 即使仅使用全量微调所需参数的一小部分，加入记忆标记的模型也显著优于基线模型。
• 当每层记忆标记超过10个后，性能增益趋于饱和，表明增加记忆大小的边际收益递减。
• 注意力掩码方法实现了高效的计算复用，使得在所有任务上进行联合推理的计算成本仅随新任务增加而小幅提升。