[논문 리뷰] CoCa : Contrastive Captioners are Image-Text Foundation Models

CoCa : Contrastive Captioners are Image-Text Foundation Models

https://arxiv.org/abs/2205.01917

CoCa: Contrastive Captioners are Image-Text Foundation Models

*YAI 9기 박준영님이 "나의 야이 아카데미아"팀에서 작성한 글입니다.

---

들어가며…

여러 task에서 대용량의 데이터셋으로 large-scale pretrained network를 학습하는 이유는 다양한 downstream tasks에 학습된 representation을 빠르게 적용하기 위함이고, 이를 딥러닝에서는 representation transfer 혹은 knowledge transfer 관점에서 접근한다. 리뷰할 CoCa 논문도 결론적으로 말하자면 image/text pair를 활용하여 보다 다양한 task에 robust하게 활용될 수 있는 large network를 학습하고자 하는 방법론을 제시한다. CoCa는 CLIP이나 ALIGN에서 하는 encoder 단위의 contrastive learning(understanding task)와 SimVLM과 같은 encoder-decoder 단위에서의 generative learning(generation task) 구조를 통합하고자 한다.


기존의 트랜스포머 기반 encoder-decoder 구조에서는 encoder의 output에 대해 decoder의 attention layer가 cross-attention 구조를 가지지만, CoCa에서는 이를 분리(절반은 unimodal encoder로 사용하고, 나머지 절반을 multimodal decoder로 사용함)하면서 embedding에 대한 contrastive 학습과 captioning을 동시에 학습할 수 있는 구조를 제시하였다. 이 논문의 제목인 CoCa(Contrastive + Captioner)가 바로 이러한 맥락에서 나왔다고 보면 된다.


학습 구조는 뒤에서 보다 디테일하게 설명할 예정이지만, 단순히 디코더 구조를 분리함으로써 loss term 두 개를 사용할 수 있는 것이 중요한게 아니라 하나의 computational graph로 end-to-end 학습이 가능하다는 점이다. 사실 contrastive learning과 generative learning을 같이 활용할 수 있는 방법에 대한 연구는 CoCa가 유일하진 않지만, 보다 간단한 구조로 scratch부터 학습하면서 VLP의 가장 큰 두 기둥이라고 볼 수 있는 contraative learning과 generative learning을 융합시켰다는 점이 중요한 포인트가 된다.

---

Why new architecture is needed?

사실 VLP task에 대한 논문을 읽다보면 introduction이 다 비슷하게 시작되는 것을 볼 수 있다. 그럼에도 불구하고 이 논문의 연구 과정에서 문제로 삼고 넘어간 부분에 대해 제대로 이해하기 위해서 정리하도록 하겠다.


Vision 문제들은 딥러닝 연구가 시작되며 모두 CNN(Convolutional Neural Network)와 같은 single-encoder 구조를 가지고, 각 task에 맞게 학습된다. 이때 주로 활용되는 네트워크가 우리가 많이 알고 있는 ResNet, GoogleNet, EfficientNet 등등의 pretrained backbone이고 image understanding(classification)이나 video understanding(action recognition, grounding 등등)을 해결하는 과정에서 특정 head에 대한 fine-tuning을 진행하는 등의 representation transfer를 사용한다. 그러다가 최근에 나온 CLIP과 ALIGN 계열의 연구로부터 VLP(Vision-Language Pretraining)이 제안되었고, 학습법은 web에서 획득할 수 있는 대용량의 image-text pair를 기반으로 배치 단위의 contrastive learning을 진행하는 것이었다. Image와 Text pair가 서로 매칭된다면 positive sample로서 학습이 되고, 서로 매칭되지 않는다면 negative sample로서 학습이 되는 구조이다. Contrastive learning을 목적으로 하는 VL task에서는 image encoder와 text encoder가 각각 unimodal로 학습되는 dual-encoder 구조가 되며, 각각 담당하는 modality를 embedding space로 매핑하면 contrastive loss로 하여금 cross-modal alignment가 진행되는 구조이다.

CLIP과 ALIGN이 contrastive 관점에서의 image와 text를 align하는 문제로 접근했다면, SimVLM과 같이 generative 문제의 관점으로 접근한 연구들도 있다. 차이점이라면 서로 다른 modality인 image와 text 간의 alignment를 어떻게 진행하는가에 대한 부분인데, dual-encoder로 처리한 contrastive learning은 직접 각 encoder가 implicit하게 modality를 이해하는 중간 과정에 explicit supervision을 주지 않는 입장이라면 이와는 반대로 generative learning은 이미지 혹은 이미지 + prompt에 대한 정보와의 attention을 통해 이미지를 묘사하는 text를 생성하는 학습 구조를 가지기 때문에 decoder 중간에서 alignment가 진행된다고 볼 수 있다. 물론 encoder-decoder 구조의 encoder의 output 또한 query에 대한 key/value가 되기 때문에 간접적으로 alignment가 이루어진다고 볼 수 있다.

결국 두 방법 모두 학습의 주된 목적은 이미지와 텍스트를 잘 연관짓는 것인데, 각각의 task는 구조 특성상 task agnostic하지 않다는 것이 드러났다. 이를테면 contrastive learning의 경우에는 zero-shot classification과 같은 understanding 기반 downstream task의 성능이 매우 좋으나, 이미지의 각 부분에 대한 captioning/reasoning 성능이 떨어지는 문제가 발생하기도 하고, generative learning의 경우에는 encoder-decoder 구조를 가지기 때문에 image 없이는 text-only representation을 생성할 수 없는 문제(text modality에 대한 unimodal encoder가 부재하기 때문)가 있거나, image encoder의 output을 사용한 understanding task의 probing 과정에서도 representation의 성능이 떨어진다는 점 등의 문제가 발생한다.


그렇기에 이 논문에서 single-encoder/dual-encoder/encoder-encoder 연구들의 모든 패러다임을 융합할 수 있는 형태를 제안하고자 한 것이고, 학습 과정도 단일화하여 image/text 관계성을 기반으로 end-to-end 최적화가 가능한 알고리즘을 만들고자 한 것이다.

위의 그림을 보면 우측에 나와있는 세 가지의 큰 틀의 패러다임이 존재한다. 가장 좌측에 보이는 것이 multimodal을 제외하고 볼 수 있는 vision recognition(understanding) task가 되고, 중간이 CLIP/ALIGN과 같은 VL contrastive learning 접근법이고, 마지막으로 우측이 generative learning에서 사용하는 encoder-decoder 구조가 된다. 개인적으로는 우측 그림에서 multimodal decoder 하나만 있는 것이 논문에서 말하고자 하는 흐름과 더 맞을 것 같다는 생각을 했다.

---

Is it first time?

..라고 한다면 그렇지는 않다. ALBEF와 같은 논문들이 이미 image/text와 관련된 여러 task를 하나의 large network로 해결하고자 학습법을 제시하였다.

그러나 그림에서 볼 수 있듯이 학습 과정에서 단일 batch에 대해 single process로 학습되지 못한다는 점, encoder 구조를 사용했기 때문에 generative task에서는 추가로 decoder head가 필요하다는 점 그리고 BERT와 같이 pre-trained model의 representation에 의존하여 학습되기 때문에 scratch부터 학습에 사용할 수 있는 방법론은 아니라는 점이 차이가 있다(+momentum module을 통해 distillation하는 과정도 필요).

---

Approach

앞에서도 언급했지만 논문에서 목적으로 삼은 것은 single encoder/dual encoder/encoder-decoder에 대한 모든 supervision을 natural language로부터 획득하는 것이다. 보통 단일 encoder로 classification하는 경우 cross entropy loss를 통해 human annotated class vector와의 KL divergence를 최소화하는 방향으로 score network가 학습된다.

$$
\mathcal{L}_\text{Cls} = -p(y) \log q_\theta (x)
$$

$p(y)$ 는 one-hot encoding된 class vector이며 $q_\theta(x)$는 parameter set $\theta$로 구성된 score network의 prediction을 의미한다. 하지만 단순히 class vector에 대해 학습하는 것은 label 정보를 natural language로 간주하지 않고 discretized label로 간주하기 때문에 language와의 유기적인 학습이 불가능하게 된다. 이를 해결하기 위해 제시된 방법 중 contrastive learning은 image encoder와 text encoder $\mathcal{I}$와 $\mathcal{T}$에 대해 인코딩된 image embedding $x$, text embedding $y$를 contrastive하게 모델링한다.

$$
\mathcal{L}_\text{Con} = -\frac{1}{N} \left( \sum_i^N \log \frac{\exp (x_i^\top y_i / \sigma)}{\sum_{j=1}^N \exp(x_i^\top y_j/\sigma)} + \sum_i^N \log \frac{\exp (y_i^\top x_i / \sigma)}{\sum_{j=1}^N \exp(y_i^\top x_j/\sigma)} \right)
$$

Captioning은 비교적 간단한데, 보통 autoregressive한 language decoder를 학습할 때 특정 시점을 기준으로 causality를 만족하는($t$ 이전의 prompt만 참고하는 형태) conditional likelihood에 대한 negative log likelihood를 모두 더해서 사용한다. 따라서 captioning loss는 각 generation이 독립적이라는 가정 하에 모든 caption이 제대로 생성될 수 있게끔 조건부 확률을 최적화한다.

$$
\mathcal{L}_\text{Cap} = -\sum_{t=1}^T \log P_\theta (y_t \vert y_{< t}, x)
$$

sequence to sequence 구조에서는 학습 효율성을 위해 teacher-forcing을 사용한다(prediction이 조건부 기준이 되지 않고 ground truth를 조건부의 기준이 되는 것).

소개한 개념 중에서 contrastive loss와 captioning loss를 더해주기만 하면 masked decoder가 추가된 encoder-decoder 구조(decoder가 물론 두 파트로 나뉨)에서 end-to-end 학습의 objective function으로 사용할 수 있다.

$$
\mathcal{L}_\text{CoCa} = \lambda_\text{Con} \cdot \mathcal{L}_\text{Con} + \lambda_\text{Cap} \cdot \mathcal{L}_\text{Cap}
$$

---

Decoupled text decoder and CoCa architecture

앞서 언급한 loss를 기반으로 학습하기 위해서는 captioning과 contrastive learning이 동시에 진행되어야한다. 식을 보면 알 수 있지만 captioning은 conditional likelihood를 최적화하게 되고, contrastive는 unconditional likelihood를 최적화한다. 보다 정확하게 말하자면 각각의 modality에 대한 encoder에 대한 output가 서로 독립이라고 가정했을 때의 joint distribution에 대한 negative log likelihood를 최적화하는 것과 같다.


만약 단일 구조의 decoder를 사용하게 되면 하나의 모델에서 두 가지의 objective를 학습할 수 없으므로(conditional likelihood와 unconditional likelihood), 저자들이 제시한 방법은 앞서 설명했던 것과 같이 decoder를 두 부분으로 decouple하여, unimodal과 multimodal을 담당하도록 분리하는 것이다. Unimodal decoder layer는 input text를 latent vector로 인코딩하고, 이때 input은 일반적인 transformer decoder에서 연산하는 것과 같이 causality에 따라 masking한다(하단의 $n_\text{uni}$ layers). 그리고 이렇게 연산된 text encoder에서 나온 output을 기반으로 image encoding 결과와 cross-attention을 진행하는 상단의 $n_\text{multi}$ 레이어가 존재하게 된다. Decoder layer는 정확히 절반으로 나누었다고 한다($n_\text{uni} = n_\text{multi}$).

Transformer 기반 vision task encoder에서는 ‘cls’ 토큰과 같이 이미지 전체에 대한 정보를 인코딩하는 부분이 존재하고, 이때 encoder로부터 나오는 토큰은 각각 $n$차원의 벡터로 생각하면 된다. 마찬가지로 contrastive learning을 진행하기 위해서 저자는 image encoder 외에 unimodal text decoder의 문장 끝에도 ‘cls’ 토큰을 붙여서 문장 전체에 대한 representation을 통해 contrastive loss를 계산한다.

---

Attention pooling

Contrastive loss와 같이 understanding을 베이스로 하는 경우에는 단일 token에 각 modality 전체에 대해 요약하는 형태가 보통 일반적이고 잘 이용되지만(single pooled image embedding), 이와는 다른 downstream tasks(논문에서 말하고자 하는 것은 generation task)의 경우 단순히 이미지 전체/텍스트 전체를 요약하고자 하는 것이 아니기 때문에 보다 많은 visual token이 사용되는 것이 좋은 성능을 보장할 수 있다.


따라서 위의 그림에서 보는 바와 같이 contrastive learning에 사용될 attention pooler와 multimodal text decoder에서 cross attention에 사용될 attention pooler를 정의하였고, 여기서 attention pooler는 image encoder의 output을 key/value로 가지면서 $n$개의 query를 파라미터로 최적화한다고 볼 수 있다. Generative loss를 위한 query는 $256$, contrastive loss를 위한 query는 $1$을 사용하였다.

---

CoCa for downstream tasks

학습된 CoCa는 image/text에 대해 다양한 zero-shot transfer가 가능하다. CLIP/ALIGN과 같은 기존 방식들이 가능한 classification부터, image and text retrieval, video and text retrieval 등등이 있다.


이외에도 zero-shot이 아닌 task에 대해 frozen encoder를 사용할 수도 있다. 앞서 설명했던 attention pooler가 각 task마다 적절한 query를 학습하게 되므로, task 마다의 attention pooler를 학습하는 과정이 결국 encoder가 학습한 representation을 다양하게 쓸 수 있는 방법이 된다.

예를 들어 video action recognition의 경우, 위와 같이 각 프레임을 encoder에 통과한 embedding을 토대로 attention pooling을 진행, 비디오에서 흔히 얻고자 하는 spatial/temporal information에 대한 query를 정의한 뒤에 해당 query에 attention된 정보를 통해 action recognition task에 맞는 softmax cross-entropy를 계산하게 된다. Pooler는 query를 하나만 가지므로(spatial 혹은 temporal) 연산이 크게 어렵지 않다.


Video-text retrieval에서는 더 간단하게 $16$ 프레임의 embedding 평균을 사용하게 되고, 이때 target embedding으로는 각 video의 캡션을 인코딩해서 사용한다.

---

Experiments

CoCa의 장점이라고 할 수 있는 것은 모든 task에 대해 쉽게 적용 가능한 학습법 및 네트워크 구조를 고안했다는 것이다. 실제로 각 task의 기존 SOTA 성능을 노란색으로 표현했을 때, CoCa가 전반적으로 성능들을 뛰어넘는 것을 확인할 수 있다.

위의 표를 보면 알 수 있듯이 CoCa는 encoder가 frozen 된 상태에서도 충분히 기존 네트워크에 필적하는 성능을 보여주며, task에 따라 fine-tuning하게 되면 SOTA 성능을 모두 뛰어넘는 것을 볼 수 있다.

CoCa는 image encoder에 대해 ViT와 동일한 구조를 가지기 때문에 표에서 볼 수 있는 것과 같이 표현된다. 가장 parameter 수를 적게 가지는 네트워크부터 시작하여 가장 parameter 수를 크게 가지는 네트워크까지 classification에 대해 fine-tuned 및 zero-shot 성능을 확인해볼 수 있다.

가장 큰 네트워크 기준 Flickr와 MSCOCO에 대한 image-text retrieval 성능은 위와 같다. Image to text retrieval이란 이미지를 기반으로 매칭되는 텍스트를 찾는 task이고 text to image retrieval은 텍스트를 기반으로 매칭되는 이미지를 찾는 task가 된다.

또한 domain generalization에 관련된 dataset인 ImageNet-어쩌구 종류들에 대해서도 zero-shot 성능이 매우 좋은 것을 볼 수 있다. SOTA인 BASIC 기준으로 모든 dataset에 대해 성능이 좋고, 상대적으로 domain shift가 빡센 dataset에 대한 경향성도 유사하게 나온다. 이외에도 video/text retrieval, captioning 등 여러 task에 대한 결과가 있지만 결국 성능이 잘나왔다는 것을 보여주는 것이기 때문에 넘어가도록 하겠다.

---

Ablations

저자가 한 모든 ablation은 위와 같다. 참고로 표에 나와있는 LE는 linear evaluation, ZS는 zero-shot, AE는 attention evaluation 그리고 FT는 fine-tuning을 의미한다. Ablation은 사실 리뷰어들 방어용이라 논문 이해에는 큰 무리가 없으니 넘어가도록 하겠다.

---

결론

논문을 보다보니, 정말 깔끔하다는 생각을 했다. 네트워크를 구성한 형태나 논문 디벨롭한 구조, 성능 그리고 마지막으로 ablation을 진행한 것까지 크게 나무랄 것 없는 논문이었다. 조금 아쉬운 점이 있다면 CLIP, ALIGN 등등 large scale VLP들에 모두 해당되는 이야기겠지만, 파라미터 수를 지나치게 많이 사용하려는 경향 때문인지 google research와 같이 hardware concern이 없다면 상관없지만 구조를 단순화했다고 해도 여전히 네트워크의 학습에는 큰 load가 있을 수 밖에 없다는 점이다. 사실 이건 논문 자체의 limitation으로 볼 수는 없고 VLP task가 가지는 한계라고 할 수 있겠다. 추가로 논문에 대한 제안인데, 만약 attention module이 task에 비해 오버로딩되는 상황(too much하게 사용되는 상황)을 시각화에서 문제로 삼을 수 있다면 VLP task에 사용되는 encoder/decoder 구조를 parameter-free layer를 통해 단순화할 수 있는 연구를 진행해볼 수 있지 않을까 생각해본다.