[Paper Review] Fine-tuning Image Transformers using Learnable Memory

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

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes augmenting Vision Transformers with learnable memory tokens—additional learnable embeddings injected at each layer—to enable efficient, parameter-efficient fine-tuning for downstream tasks. By using attention masking, the method allows models to retain performance on previous tasks while extending to new ones with minimal accuracy drop, achieving performance close to full fine-tuning using only a few extra parameters per layer.

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.

Motivation & Objective

• To address the high cost and instability of full fine-tuning in Vision Transformers, especially for continual learning scenarios.
• To enable efficient adaptation to new downstream tasks with minimal parameter updates.
• To maintain high performance on previously learned tasks when extending models to new tasks.
• To enable computation reuse and joint inference across multiple tasks in a single model.
• To explore parameter-efficient, scalable, and stable methods for lifelong learning in vision transformers.

Proposed method

• Augment each ViT encoder layer with a small set of learnable memory tokens, appended to the input sequence before attention computation.
• Use a modified attention mechanism where memory tokens do not attend to other tokens, and only the class and patch tokens are updated across layers.
• Introduce an attention masking strategy that restricts old task tokens from attending to newly added memory and class tokens, preventing interference during inference.
• Train models sequentially on new tasks by extending the input sequence with new task-specific memory and class tokens, while preserving attention constraints to maintain old task performance.
• Allow model extension by concatenating independently trained models with attention masks, enabling joint inference with incremental compute cost.
• Use a shared backbone for all tasks, with task-specific memory and classification heads, enabling parameter sharing and efficient inference.

Experimental results

Research questions

• RQ1 Can learnable memory tokens significantly improve fine-tuning accuracy in Vision Transformers with minimal parameter updates?
• RQ2 Can attention masking prevent performance degradation on previously learned tasks when extending models to new tasks?
• RQ3 How does the performance of memory-augmented fine-tuning compare to head-only and full fine-tuning?
• RQ4 Can multiple models trained on different tasks be combined into a single model with minimal accuracy loss and incremental compute cost?
• RQ5 Does the number of memory tokens per layer affect model performance, and is there a saturation point?

Key findings

• Adding just a few memory tokens per layer significantly improves accuracy over head-only fine-tuning, approaching the performance of full fine-tuning.
• With attention masking, models trained sequentially on multiple tasks maintain high accuracy on all previous tasks, avoiding the 20%-40% accuracy drop seen in naive concatenation.
• The model extended by sequentially adding tasks (SUN-397 then Places-365, or vice versa) achieved stable performance without degradation, outperforming independently trained and concatenated models.
• Models with memory tokens outperformed baseline models without memory, even when using only a fraction of the parameters required for full fine-tuning.
• Performance gains from memory tokens plateaued beyond 10 tokens per layer, suggesting diminishing returns with increased memory size.
• The attention masking approach enabled efficient computation reuse, allowing joint inference on all tasks with only a small incremental cost per new task.