Fine-tuning Image Transformers using Learnable Memory

• Published in: 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) pp. 12145 - 12154
• Main Authors: Sandler, Mark, Zhmoginov, Andrey, Vladymyrov, Max, Jackson, Andrew
• Format: Conference Proceeding
• Language: English
• Published: IEEE 01-06-2022
• Subjects: Adaptation models; Computational modeling; Computer vision; Costs; Deep learning; Deep learning architectures and techniques; Efficient learning and inferences; Transfer/low-shot/long-tail learning; Memory management; Transformers
• ISSN: 2575-7075
• DOI: 10.1109/CVPR52688.2022.01184

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.