[ Paper Review ] ALBEF (Align before Fuse: Vision and Language Representation Learning with Momentum Distillation, 2021)

Align before Fuse: Vision and Language Representation Learning with Momentum Distillation

 

0. Abstract

대규모 vision과 language representation 학습은 다양한 vision-language task에서 유망한 개선을 보여주었다.

대부분의 기존 방법론들은 transformer 기반의 multimodal encoder를 사용하여 visual 토큰과 word 토큰을 함께 모델링하는데, visual 토큰과 word 토큰이 정렬되지 않기 때문에 multimodal encoder가 이미지-텍스트 상호작용을 학습하는 것은 도전적이다.

 

본 논문에서는 이미지와 텍스트 representation이 결합하기 전에 cross-modal attention을 통해 정렬하는 contrastive loss를 도입하여 더 확고한 vision 및 language representation 학습을 가능하게 하는 ALBEF를 소개한다.

노이즈가 많은 웹 데이터로부터 학습을 개선하기 위해, momentum 모델에 의해 만들어진 pseudo-target으로부터 학습하는 self-training 방법론인 momentum distillation을 제안한다.

 

ALBEF는 다수의 downstream vision-language task에서 SOTA 성능을 달성하고, image-text retrireval에서 훨씬 더 큰 데이터셋으로 pre-trained된 방법론들을 능가한다.

 

 

1. Introduction

Vision-and-Language Pre-training(VLP)는 대규모 이미지-텍스트 쌍으로부터 multimodal representation을 학습하는 것을 목표로 한다.대부분 기존 VLP 방법론들은 pre-trained object detector를 사용하여 영역 기반 이미지 feature를 추출하고, multimodal encoder를 사용하여 이미지 feature를 단어 토큰과 결합한다.

 

이러한 VLP 프레임워크는 몇 가지 주요 한계가 있다.

1. 이미지 feature와 단어 토큰 embdding은 각 공간에서 존재하여, multimodal encoder가 그들의 상호작용을 모델링하는 것을 어렵게 한다.

2. object detector는 pre-traing동안 bounding box 주석을 요구하고, inference 할 때 고해상도 이미지를 요구하기 때문에, 주석 작업과 계산 자원에 많은 비용이 든다.

3. 널리 사용되는 이미지-텍스트 데이터셋은 웹으로부터 수집되었고 본질적으로 노이즈가 많으며, MLM과 같은 기존의 pre-training 방법들은 노이즈가 많은 텍스트에 과적합되고 모델의 일반화 성능을 저하시킨다.

 

본 연구에서는 이러한 한계를 다루기 위해 새로운 VLP 프레임워크 ALign BEfore Fuse(ALBEF)를 제안한다.

먼저 detector 없이 이미지 encoder와 텍스트 encoder로 이미지와 텍스트를 독립적으로 encoding한다.

그 후 multimodal encoder를 사용하여 cross-modal attention을 통해 이미지 feature와 텍스트 feature를 결합한다.

unimodal encoder로부터의 representation에 중간 image-text contrastive(ITC) loss를 도입하는데, 이는 세 가지 목적을 가지고 있다.

1. 이미지 feature와 텍스트 feature를 정렬하여 multimodal encoder가 cross-modal 학습을 더 쉽게 수행하도록 한다.

2. unimodal encoder가 이미지와 텍스트의 semantic 의미를 더 잘 이해할 수 있도록 한다.

3. 이미지와 텍스트의 embedding할 저차원 공간을 학습하여, 이미지-텍스트 매칭 목표가 더 유익한 샘플을 찾을 수 있도록 한다.

 

노이즈가 있는 supervision에서 학습을 개선하기 위해, Momentum Distillation(MoD)를 제안하여 선별되지 않은 웹 데이터셋을 모델이 사용할 수 있게 한다. 학습하는 동안, momentum model은 추가적인 supervision으로서 사용할 pseudo-target을 생성한다. MoD를 사용하면 웹 주석과 다른 합리적인 출력을 생성해도 모델이 페널티를 받지 않는다.

 

본 논문에서는 ITC와 MLM이 이미지-텍스트 쌍의 다른 관점에서 생성된 부분 정보를 이용하여 상호 정보의 lower bound를 최대화하는 것을 보여준다. ALBEF는 다양한 downstream V+L task에서 효과성을 보여준다.

 

 

2. Related Work

2.1. Vision-Language Representation Learning

대부분의 기존 visiong-language representation learning은 2가지로 나눌 수 있다.

 

1. transformer 기반의 multimodal encoder로 이미지와 텍스트 feature의 상호작용을 모델링하는데 중점을 둔다.

- 복잡한 reasoning을 요구하는 task에서 우수한 성능을 보이지만, 고해상도 이미지와 pre-trained object detector를 요구한다.

2. 이미지와 텍스트의 분리된 unimodal encoder를 학습하는데 중점을 둔다. ex. CLIP

- retreival task에서 우수한 성능을 얻지만, 복잡한 상호작용을 모델링하기에는 부족하다.

 

ALBEF는 두 카테고리를 통합하여, unimodal과 multimodal representation을 얻고 강력한 retreival과 reasoning 모두에서 우수한 성능을 얻는다.

 

2.2. Knowledge Distillation

지식 증류(Knowledge Distillation)은 student 모델의 예측을 teacher 모델의 예측과 맞추는 방식으로 teacher 모델로부터 지식을 증류하여 student 모델을 향상시키는 것을 목표로 한다.

대부분의 방법이 pre-trained teacher 모델로부터 지식을 증류하는 반면, online distillation은 여러 모델을 동시에 훈련시키고 앙상블을 teacher로 사용한다.

 

본 논문에서 소개하는 momentum distillation은 online self-distillation으로 해석될 수 있으며, student 모델의 시간적 앙상블이 teacher로 사용된다.

 

 

3. ALBEF Pre-training

3.1. Model Architecture

ALBEF는 image encoder , text encoder, multimodal encoder로 구성되어 있다.

이미지 encoder는 12-layer ViT-B/16 구조이며, ImageNet-1k로 pre-trained된 가중치로 초기화한다.

텍스트 encoder와 multimodal encoder는 6-layer transformer이며, 각각 $BERT_{base}$의 처음 6 layer와 마지막 6 layer로 초기화한다.

 

input 이미지 $I$는 $\{v_{cls}, v_1, \cdots, v_N\}$ 으로, input 텍스트 $T$는 $\{w_{cls}, w_1, \cdots, w_N\}$ 으로 encoding된다. 그 후 이미지 feature와 텍스트 feature는 multimodal encoder의 각 layer에서 cross attention을 통해 결합된다.

 

3.2. Pre-training Objectives

ALBEF를 ① image-text contrastive learning(ITC) on unimodal encoder, ② masked language modeling(MLM), ③ image-text matching(ITM) on multimodal encoder 세 가지 목적으로 pre-traing한다.

 

- Image-Text Contrastive Learning

결합 전에 unimodal representation을 더 잘 학습하는 것을 목표로 한다.

유사도 함수 $s = g_v(v_{cls}) \top g_w(w_{cls})$를 학습하여 평핸한 이미지-텍스트 쌍이 더 높은 유사도 점수를 갖도록 한다.

$g_v$, $g_w$: [CLS] embdding을 정규화된 저차원(256C차원) representation으로 매핑하는 선형 변환

 

momentum unimodal encoder로부터의 최근 M개의 이미지-텍스트 representation을 저장할 2개 큐를 유지한다.

$g'_v(v'_{cls})$, $g'_w(w'_{cls})$: momentum encoder로부터 정규화된 feature

 

$S(I, T) = g_v(v_{cls})\top g'_w(w'_{cls})$

$S(I, T) = g_w(w_{cls})\top g'_v(v'_{cls})$

 

각 이미지와 텍스트에 대해, softmax-normalized image-to-text, text-to-image 유사도를 다음과 같이 계산한다.

 

$y^{i2t}(I)$, $y^{t2i}(T)$: ground-truth one-hot 유사도 (negative: 0, positive: 1)

 

image-text contrastive loss는 $p$와 $y$ 사이의 cross-entropy H로 정의된다.

 

다시 말해, [CLS] 토큰에 대한 embdding으로 같은 쌍은 유사도가 높아지고, 다른 쌍은 유사도가 낮아지도록 학습하는 것이다. (CLIP과 유사하며, 해당 논문에서의 M이 배치 사이즈와 같은 역할이다).

 

- Masked Language Modeling

이미지와 문맥적 텍스트 모두를 활용하여 마스킹된 단어를 예측한다.

 

$\hat{T}$: 랜덤으로 마스킹된 텍스트

$p^{msk}(I, \hat{T})$: 마스킹된 토큰에 대해 예측한 확률

$y^{msk}$: one-hot vocabulary 분포로, 실제 ground-truth 토큰(마스킹된 토큰)은 1이다.

 

- Image-Text Matching

이미지와 텍스트 쌍이 positive(matched)인지 negative(not matched)인지를 예측한다.

 

multimodal encoder의 [CLS] token output embdding을 이미지-텍스트 쌍의 공동 representation으로 사용하고, FC layer를 거친 후 softmax로 2개의 class(매칭 여부)에 대한 확률 $p^{itm}$을 예측한다.

 

$y^{itm}$: ground-truth 라벨을 나타내는 2차원 one-hot 벡터

 

negative 이미지-텍스트 쌍은 비슷한 의미를 갖지만 세부적인 다른 특성을 가지도록 하여, fine-grained한 디테일까지 잘 학습할 수 있도록 한다.

이미지를 기준으로, mini-batch 에서 유사도가 높은 텍스트를 negative로 추출하며, 텍스트도 마찬가지로 샘플링한다.

 

ALBEF의 전체 pre-trianing loss는 다음과 같다.

 

3.3. Momentum Distillation

pre-training에 사용되는 이미지-텍스트 쌍은 웹에서 수집되어 노이즈가 많으며,  positive pair가 실제로 크게 관련이 없을 수 있다.

ITC 학습에서 이미지에 대한 negative 텍스트가 이미지의 내용과 매칭될 수도 있으며, MLM에서 이미지를 잘 설명하는 다른 단어가 존재할 수도 있다.

그러나, 두 학습에 대한 one-hot label은 일치성에 관계 없이 모든 negative 예측에 페널티를 부여한다.

 

본 연구에서는 momentum 모델로부터 만들어진 pseudo-target으로부터 학습하는 것을 제안한다.

momentum 모델은 지속적으로 진화하는 teacher로, unimodal과 multimodal encoder의 EMA 버전으로 구성된다.

 

※ EMA(Exponential-Moving-Average)

: 최근 데이터에 더 큰 가중치를 부여하면서 과거 데이터의 영향을 지수적으로 감소시키는 방법

 

기본 모델을 momentum 모델의 예측과 일치하도록 훈련한다.

ITC의 경우, momentum unimodal encoder로부터의 feature를 사용하여 유사도를 계산한다.

$s'(I, T) = g'_v(v_{cls}) \top g'_w(w'_{cls})$, $s'(T, I) = g'_w(w_cls) \top g'_v(v'_{cls})$

 

그 다음, 식 (1)에서 $s$를 $s'$로 대체하여 soft pseudo-target $q^{i2t}$와 $q^{t2i}$를 계산한다.

 

원래 loss에, 모델의 예측과 pseudo-target(momentum 모델의 예측)의 KL-divergence를 더한 형태이다.

 

MLM도 마찬가지로, $q^{msk}(I, \hat{T})$를 마스킹된 토큰에 대한 momentum 모델의 예측 확률로 하면 loss는 다음과 같다.

 

결국, momentum 모델이 지수 이동 평균을 통해 기본 모델보다 더 안정적인 예측을 제공하고, 이는 one-hot이 아닌 확률 분포 형태로 제공되기 때문에 soft pseudo-target으로 사용하는 것이다.

 

pseudo-target에 대한 예시는 다음과 같다.

 

3.4. Pre-training Datasets

2개의 웹 데이터셋(Conecltual Captions, SBU Captions)와 2개의 도메인 내 데이터셋(COCO, Visual Genome)으로 구성되어 있으며, 총 unique 이미지는 4백만 개이고 이미지-텍스트 쌍은 5.1백만개이다.

 

3.5. Implementatinon Details

8개의 A100 GPU에서 배치사이즈 512로 30 epoch동안 pre-train했다.

이미지-텍스트 contrastive learning에 사용된 큐의 사이즈는 65536이며, distillation 가중치 $\alpha$를 0에서 0.4까지 선형적으로 증가시켰다.

 

 

4. A Mutual Information Maximization Perspective

ITC, MLM, MoD는 이미지-텍스트 쌍의 다양한 view를 만드는 방법으로 해석될 수 있다.

 

$a$, $b$: 이미지-텍스트 쌍의 서로 다른 변형을 한 두 변수

a와 b의 MI(Mutual Information)을 최대화함으로써 view의 변경에 불변하는 representation을 학습할 수 있다.

 

다음과 같은 InfoNCE loss를 최소화함으로써 $MI(a, b)$의 lower bound를 최대화할 수 있다.

 

※ InfoNCE loss

: Contrastive Learning에 사용되는 Loss 중 하나로, 해당 연구에서 데이터의 두 view 사이의 유사성을 평가한다.

 

$s(a, b)$: score function (e.g. a dot product)

$\hat{B}$: positive sample b와 나머지 negative sample로 구성되어 있는 집합

 

- ITC

InfoNCE loss를 식(2)의 ITC loss로 재정의할 수 있다.

 

ITC는 이미지와 텍스트 view 사이의 MI를 최대화시키는 것과 같다.

 

- MLM

MLM은 마스킹된 단어 토큰과 맥락(이미지 + 마스킹된 텍스트) 사이의 MI를 최대화하는 것과 같다.

 

식 (3)의 MLM loss를 다음과 같이 쓸 수 있다.

 

$\psi(y): V \rightarrow R^d$, 단어 토큰 y를 벡터로 변 ($V$: 전체 단어 집합)

$f(I, \hat(T)$: 마스킹된 토큰에 일치하는 multiomdal enoder의 마지막 hidden state

 

따라서 MLM은 이미지-텍스트 쌍의 두 가지 관점인 무작위로 선택된 단어 토큰과 해당 단어가 마스킹된 이미지와 문맥 텍스트를 고려하는 것이다.

 

- MoD ITC

KL divergence 수식은 다음과 같다.

 

식 (6)에서, $KL(p^{i2t}(I), q^{i2t}(I))를 최소화하는 것은 다음의 목적식을 최소화하는 것과 같다.

 

momentum 모델은 원본 이미지-텍스트 쌍에는 없는 다양한 view를 생성하고, 기본 모델이 view에 불변한 의미 정보를 포착하는 representation을 학습하도록 하는 것과 같다.

 

 

5. Downstream V+L Tasks

- Image-Text Retreival

Flickr30K와 COCO 벤치마크로 평가하였다.

fine-tuning 동안, ITC loss를 사용하여 unimodal featre의 유사도를 기반으로 이미지-텍스트 scoring function을 학습하였고, ITM loss를 사용하여 이미지와 텍스트 사이의 fine-grained한 상호작용을 학습했다.

 

- Visual Entailment

이미지이와 텍스트의 관계가 함의인지, 중립인지, 모순인지를 예측하는 fine-grained한 reasoning task이다.

multimodal encoder의 [CLS] 토큰 representation에 MLP를 적용하여 각 class에 대한 확률을 예측한다.

 

- Visual Question Answering

VQA를 answer generation 문제로 여기고, 6-layer transformer를 사용하여 답변을 생성한다.

multimodal encoder weight로 초기화하고 conditional language-modeling loss로 fine-tuning한다.

 

- Natural Language for Visual Reasoning

주어진 텍스트가 두 개의 이미지 쌍을 설명하는지를 예측한다.

두 개의 이미지를 추론하기 위해 multimodal encoder의 각 layer는 두 개의 연속되게 위치한 transformer block으로 복제된다.

 

이미지 쌍을 encoding하기 위한 새로운 multimodal encoder를 위해 추가적인 pre-training을 수행한다.

text-assignment task로, 주어진 이미지 쌍과 텍스틍 대해, 모델은 텍스트가 첫번째 이미지에 해당하는지, 두번째 이미지에 해당하는지, 혹은 어느 것도 아닌지를 할당한다. 이를 classification 문제로 여기고, FC layer를 사용한다.

 

- Visual Grounding

이미지에서 특정 텍스트 설명에 해당하는 영역을 localize한다.

bounding box annotation이 없는 weak-supervised setting으로 연구한다.

 

 

6. Experiments

6.1. Evaluation on the Proposed Methods

여러 방법론들의 조합으로 downstream task를 수행한 결과이다.

 

baseline(MLM+ITM)과 비교하여, ITC를 추가하는 것이 상당한 성능 향상이 있었다.

 

6.2. Evaluation on Image-Text Retrieval

 

image-text retreival에서, 대규모로 학습된 CLIP과 ALIGN을 능가하며 ALBEF가 SOTA를 달성한다.

 

6.3. Evaluation on VQA, NLVR, and VE

 

4M의 상대적으로 적은 데이터로 pre-training 했음에도 V+L task에서 SOTA를 달성하며, 14M의 이미지로 pre-training 했을 때는 기존 방법론들에 비해 훨씬 우수한 성능을 보인다.

ALBEF가 detector를 사용하지 않고 더 낮은 해상도의 이미지를 요구하기 때문에, inference 혹도가 훨씬 빠르다.

 

6.4. Weakly-supervised Visual Grounding

image encoder의 마지막 layer에서 self-attention map과 muldimodal encoder의 3번째 layer의 cross-attention map에서 Grad-CAM 시각화 결과를 확인한다.

 

인간이 결정을 내릴 때 보는 곳과 매우 관련이 있으며, 객체 뿐만 아니라 속성과 관계고 파악할 수 있다.

 

6.5. Ablation Study

$s_{itc}$를 사용하여 상위 k의 후보를 필터링하는데, k의 변화에는 크게 민감하지 않았다.

hard negative가 없을 시 성능이 감소한다.

 

다음은 NLVR에 대해 TA(Text-assignment)의 효과와 파라미터 공유에 대한 성능을 나타낸다.

 

 

TA가 없는 경우, 전체 block을 공유하는 것이 더 나은 성능을 보이며, TA를 사용할 경우 CA를 공유하는 것이 가장 좋은 성능을 보였다.

 

 

7. Conclusion and Social Impacts

ALBEF는 이미지와 텍스트의 unimodal representation을 정렬한 후, multimodal encoder를 사용하여 결합한다.

기존 방법들과 비교하여 ALBEF는 여러 V+L task에서 더 나은 성능과 빠른 추론 속도를 제공한다.