[논문리뷰] Visual Prompt Tuning

ECCV 2022. [Paper] [Github]
Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, Ser-Nam Lim
Cornell University | Meta AI | University of Copenhagen
23 Mar 2022

Introduction

사전 학습된 foundation model을 다운스트임 task에 적용하는 가장 확실하고 효과적인 전략은 해당 task에 맞게 처음부터 끝까지 완전히 fine-tuning하는 것이다. 하지만 이 전략은 각 task마다 backbone 파라미터의 별도 복사본을 저장하고 배포해야 한다. 이는 특히 최신 Transformer 기반 아키텍처의 경우 비용이 많이 들고 실현 불가능한 경우가 많다.

본 논문에서는 사전 학습된 Transformer 자체를 수정하거나 fine-tuning하는 대신, Transformer의 입력을 수정한다. 저자들은 NLP 분야의 프롬프트 튜닝 연구에서 영감을 받아, 다운스트림 비전 task에 맞게 Transformer 모델을 조정하는 Visual-Prompt Tuning (VPT)을 제안하였다. 본 방법은 다운스트림 학습 과정에서 사전 학습된 Transformer backbone 전체를 고정한 상태로, 입력 공간에 task별 학습 가능한 파라미터를 소량만 도입한다. 실제로 이러한 추가 파라미터는 각 Transformer layer의 입력 시퀀스 앞에 추가되어 fine-tuning 과정에서 linear head와 함께 학습된다.

사전 학습된 ViT backbone을 사용하여 다양한 도메인에 걸쳐 24개의 다운스트림 recognition task를 수행한 결과, VPT는 다른 transfer learning baseline을 능가했으며, 20개 경우에서는 full fine-tuning보다도 우수한 성능을 보였다. 또한, 각 task에 대해 backbone 파라미터의 1% 미만이라는 현저히 적은 파라미터를 저장하는 장점을 유지했다. NLP에서 프롬프팅 튜닝은 특정 상황에서만 full fine-tuning 성능과 동등한 수준을 보일 수 있는 반면, VPT는 특히 데이터 양이 적은 환경에서 효과적이며, 데이터 규모에 관계없이 그 장점을 유지한다.

Method

1. Preliminaries: ViT

$N$개의 layer를 가진 일반적인 ViT의 경우, 입력 이미지는 $m$개의 고정 크기 패치 $\{I_j \in \mathbb{R}^{3 \times h \times w} \mid 1 \le j \le m\}$로 나뉜다. 각 패치는 먼저 위치 인코딩을 사용하여 $d$차원 latent space에 임베딩된다.

\[\begin{equation} \textbf{e}_0^j = \textrm{Embed} (I_j) \in \mathbb{R}^d, \quad j = 1, \ldots, m \end{equation}\]

이미지 패치 임베딩 모음 $\textbf{E}_i = \{\textbf{e}_i^j \in \mathbb{R}^d \mid 1 \le j \le m\}$을 $(i+1)$번째 Transformer layer $L_{i+1}$의 입력으로 사용한다. 추가적인 학습 가능한 CLS 토큰 $[\textrm{CLS}]$와 함께 전체 ViT는 다음과 같이 구성된다.

\[\begin{aligned} \left[ \textbf{x}_i, \textbf{E}_i \right] &= L_i ([\textbf{x}_{i-1}, \textbf{E}_{i-1}]) \in \mathbb{R}^{(1+m) \times d}, \quad i = 1, \ldots, N \\ \textbf{y} &= \textrm{Head}(\textbf{x}_N) \end{aligned}\]

($\textbf{x}_i \in \mathbb{R}^d$는 $L_{i+1}$의 입력 공간에서의 $[\textrm{CLS}]$의 임베딩, $[\cdot, \cdot]$은 시퀀스 길이 차원으로의 concat)

각 layer $L_i$는 multi-head self-attention (MSA)과 feed-forward network (FFN)를 LayerNorm 및 residual connection과 함께 사용한다. Classification head는 최종 layer의 임베딩 $\textbf{x}_N$을 예측 클래스 확률 분포 $\textbf{y}$로 매핑한다.

2. Visual-Prompt Tuning (VPT)

사전 학습된 ViT 모델을 기반으로, $\textrm{Embed}$ layer 이후 입력 공간에 차원 $d$의 임베딩으로 구성된 프롬프트 $p$개를 도입한다. Fine-tuning 과정에서는 task별 프롬프트만 업데이트하고, ViT backbone은 고정된 상태로 유지한다. 사용되는 Transformer layer 수에 따라, 본 접근 방식은 VPT-shallow와 VPT-deep의 두 가지 변형으로 나뉜다.

VPT-Shallow

프롬프트는 첫 번째 Transformer layer $L_1$에만 삽입된다. 각 프롬프트 토큰은 학습 가능한 $d$차원 벡터이다. $p$개의 프롬프트 모음 $\textbf{P} = \{\textbf{p}^k \in \mathbb{R}^d \mid 1 \le k \le p\}$에 대해, VPT-Shallow는 다음과 같다.

\[\begin{aligned} \left[ \textbf{x}_1, \textbf{Z}_1, \textbf{E}_1 \right] &= L_1 ([\textbf{x}_0, \textbf{P}, \textbf{E}_0]) \in \mathbb{R}^{(1+p+m) \times d} \\ \left[ \textbf{x}_i, \textbf{Z}_i, \textbf{E}_i \right] &= L_i ([\textbf{x}_{i-1}, \textbf{Z}_{i-1}, \textbf{E}_{i-1}]) \in \mathbb{R}^{(1+p+m) \times d}, \quad i = 2, \ldots, N \\ \textbf{y} &= \textrm{Head}(\textbf{x}_N) \end{aligned}\]

모든 layer와 $\textbf{x}_0$는 고정된 파라미터이며, $\textbf{P}$와 head만 학습 가능한 파라미터이다. 특히 ViT의 경우, 프롬프트는 위치 인코딩 후에 삽입되기 때문에 $\textbf{x}_N$은 프롬프트의 위치에 관계없이 일정하다. 예를 들어, $[\textbf{x}_0, \textbf{P}, \textbf{E}_0]$와 $[\textbf{x}_0, \textbf{E}_0, \textbf{P}]$는 수학적으로 동일하다. 이는 VPT-Deep에도 적용된다.

VPT-Deep

프롬프트는 각 Transformer layer의 입력 공간에 도입된다. $(i+1)$번째 레이어 $L_{i+1}$의 학습 가능한 입력 프롬프트의 집합을 $\textbf{P}_i = \{\textbf{p}_i^k \in \mathbb{R}^d \mid 1 \le k \le p\}$라 하면, VPT-Deep은 다음과 같다.

\[\begin{equation} [\textbf{x}_i, \_, \textbf{E}_i] = L_i ([\textbf{x}_{i-1}, \textbf{P}_{i-1}, \textbf{E}_{i-1}]) \in \mathbb{R}^{(1+p+m) \times d}, \quad i = 1, \ldots, N \\ \textbf{y} = \textrm{Head}(\textbf{x}_N) \end{equation}\]

Storing Visual Prompts

VPT는 여러 다운스트림 task가 있는 경우에 유용하다. 각 task에 대해 학습된 프롬프트와 classification head만 저장하고 사전 학습된 Transformer 모델의 원본을 재사용하면 되므로 저장 비용을 크게 줄일 수 있다. 예를 들어, 86M 파라미터와 $d = 768$인 ViT-Base가 주어졌을 때, 50개의 shallow prompt와 deep prompt는 각각 0.038M과 0.46M의 추가 파라미터를 생성하며, 이는 전체 ViT-Base 파라미터의 각각 0.04%, 0.53%에 불과하다.