SMALLCAP: Lightweight Image Captioning Prompted with Retrieval Augmentation

1. Introduction

ClipCap와 I-Tuning의 경우 이미 제작되어있는 pre-trained vision encoder과 language decoder을 사용하는데 이때 pretrained된 요소들이 아닌 mapping하는 부분을 학습하는 방식으로 학습되는 parameter수를 줄일 수 있었다. 하지만 이 역시도 매번 새로운 use-case에 대해서 학습을 진행해야하고 매번 파라미터를 학습을 통해 수정해야한다.

• SMALLCAP은 input image를 기반으로한 외부의 text data들로부터 caption을 찾아서 생성하는 방식이다. 이러한 방식은 (1)lightweight training,(2)training-free domain transfer,(3)exploitation of large data in a training-free fashion의 특징이 있다.
• SMALLCAP은 pre-trained CLIP vision encoder와 GPT-2 language model을 사용하고 cross-attention layers로 2개를 연결한다.
• SMALLCAP은 COCO benchmark로 학습을 한 결과 다른 lightweight-training models과 비슷한 성능을 보이지만 학습할 파라미터수는 83%줄어들었다.
• SMALLCAP can also leverage data in a training-free manner

2. Related Work

2.1. Image Captioning Model

• 최근 image captioning은 encoder-decoder methods을 사용한다.

2.2. Freezing Image Captioning Models

• good generalization를 유지하기 위해 pre-trained weights의 일부 혹은 전체를 frozen한다.
• SMALLCAP은 CLIP과 GPT-2를 cross-attention layers을 통해서 연결한다. retrieval augmentation을 이용해서 성능을 유지하고 학습해야될 파라미터 수를 줄였다.

2.3. Retrieval-Augmented Generation

• Retrieval-augmented language generation은 외부 datastore로부터 검색된 추가적인 정보로부터 생성이 된다.
• SMALLCAP에서는 simple prompt-based
  conditioning method를 적용하였다.

2.4. Prompting Text Generation

3. Proposed Approach

3.1. Model

SMALLCAP은 lightweight-training의 image captioning model로 prompt를 이용해captions검색을 확대한다.

• 인코더에서 CLIP을 이용하여 임베딩된 sequence를 생성한다.
• 디코더에서 GPT-2를 이용한다
• 두 모델이 다른 vector spaces에서 동작하기 때문에 multi-head cross attention으로 인코더의 출력과 디코더의 입력을 연결시킨다.
• 인코더와 디코더를 얼리고 초기값이 랜덤값인 cross-attention layers만 학습하는데 이는 계산량을 줄이고 기존 모델들의 generalization capabilities를 보존하기 위함이다.

3.2. Prompting with Retrieved Captions

• CLIP model을 이용해서 input image와 datastore의 contents를 인코딩한다.
• 코사인 유사도를 이용한 nearest neighbor search를 통해서 주어진 이미지와 유사한 데이터 저장소의 k개의 text item들을 찾는다.
• 찾은 text들은 다음 아래의 prompt template의 빈슬롯에 들어가 prompt를 완성한다.

"Similar images show {caption1}...{captionk}. This image shows ___"

• 이때 template의 빈 슬롯에 들어간 text는 모델에게 어떤 task인지를 알려주는 역활을 해주는 task demonstraion이 된다.
  

• 최종적으로 디코더는 해당 prompt를 input token으로 받고 image features V와 task demonstration를 통해 caption을 생성한다.

• 새로운 task에서도 해당 task에 대해서 새롭게 학습할 필요 없이 데이터 저장소를 이용해서 input image로부터 task demonstratoin를 찾아서 인코딩된 이미지와 task demonstration이 포함된 prompt를 language model에 넣어서 출력을 얻어낸다.

• retraining없이 domain transfer이 가능하다.
  

• retrieved captions을 구분하는 seperator로 콤마, 점, 빈칸등 다양한 방법으로 실행해보았고 방법간에 성능차이는 없었다. 아래는 최종 prompt template이다.
  

• cross attenton을 학습할때 cross-entropy loss를 줄이는 방향으로 학습을 진행한다.

4. Main Experiments

4.1. Experimental Setup

• SMALLCAP의 인코더와 디코더로CLIP-ViT-B/32 and GPT-2Base를 사용한다.
• 12-head cross-attention layer이 GPT-2의 12 layer에 연결되었고 학습을 통해 cross-attention layers만 학습하였다.
• 파라미터 수를 줄이기 위해 cross-attention layers의 차원 d를 기존 64에서 16,8,4로 실험하였다.
• SMALLCAP은 7M의 trainable 파라미터, 총 218M의 파라미터로 구성되어있다.
• cross-attention layers는 coco dataset을 이용해 학습하였다.
• Retrieval은 CLIP-ResNet-50x64을 기반으로 하였다.
• training에서 모델은 이미지당 k = 4 caption을 coco에서 Retrieve하였다.
• inference에서 모델은 beam size 3의 beam search로 디코딩을 진행한다.

4.2. Benchmark Results

nocaps: coco로 학습된 모델의 일반화 성능을 평가하기 위한 dataset이다.

• 모델들의 파라미터 수와 성능 비교.
• SMALLCAPMedium이랑 SMALLCAPLarge는Medium, Large GPT-2 decoders를 사용한 모델이다.

• 데이터 저장소를 이용하여 SMALLCAP의 성능을 높일 수 있다.SMALLCAP+W+H는 추가적인 데이터를 이용했을때의 성능이다.

5. Training-Free Use of Data

Flick30k/VizWiz/MSR-VTT 데이터들

5.1. In-domain Data

In-domain: 데이터 저장소(datastore)가 각각 데이터셋에 연관된 학습 데이터로 채워진 경우를 이야기한다.

5.2. Augmenting the Datastore

In-domain 데이터에 일반화 성능을 높이기 위해 데이터를 추가한다.

+Web Data

• large-scale이지만 automatically labeled되었다
• Flickr30K와 MSR-VTT에서 큰 성능 향상이 있었고 VizWiz는 향상이 미미했는데 이는 VizWiz가 다른 데이터들과 잘 매치되지 않는 데이터이기 때문이라 판단한다.

+ Human-labeled Data

• smaller-scale이지만 clean한 데이터이다.
• datastore에 이미지를 잘 설명하는 어느 형태의 text도 담을 수 있고 image-caption pairs에 어떤 제약도 없다.
• 그래서 image captions(COCO, Flickr30k, VizWiz)뿐만아니라 video captions (MSRVTT , VATEX, TGIF), audio captions (Clotho
  ),d localized narratives (LN ADE20k, LN COCO, LN Flickr30k, LN OpenImages)를 모두 고려하였다.
• Human-labeled Data를 추가했을때도 VizWiz에서 성능 향상이 미미했다.
• MSR-VTT에서 In-domain + Web에서 보다 성능 향상이 있었는데 이는 video captioning을 포함한 다양한 task들이 데이터 저장소에 포함되었기 때문이라한다.

5.3. Domain-agnostic Datastore

Web datastore을 in-domain의 경우와 비교해봤을때 성능이 유사하거나 높았다. 이를 통해 SMALLCAP의 image captioning 성능을 web data를 통해서 transfer할 수 있다는 가능성을 확인할 수 있다.

5.4. Results with the Best Configuration

• 각 모델에서의 최고 성능을 3가지 데이터를 이용하여 비교하였다, SOTA는 각 데이터셋에서large-scale 의 pre-trained V&L models을 respective datasets으로 finetune하여 얻은 결과이다.

• 다른 lightweight-training models(clipcap, camel)이 out-of-domain의 일반화가 부족하지만 smallcap은 추가적인 학습없이 단순히 데이터 저장소의 데이터를 변경하는것 만으로도 좋은 성능을 얻을 수 있다.

6. Discussion

6.1. Qualitative Examples

• smallcap은 datastore의 text로 힌트와 새로운 개념을 배우고 caption을 생성한다.
• 2번째 사진은 SMALLCAP이 어떻게 datastore을 이용해 새로운 domain에 적용되는지를 보여준다.
• retrieved captions을 통해서 tutu, Swanson(brand name), Pokemon을 정확히 맞췄는데 coco 데이터에서 tutu랑 Swanson는 전혀 나온적이 없었고 Pokemon는 6번밖에 나온적이 없었다.

6.2. Analysis of the Retrieved Captions

• datastore의 setting을 In-domain + Web + Human-labeled로 하여 각 데이터셋의 관계와 영향력을 실험하였다.
• MSR-VTT는 비디오 domain의 데이터셋인만큼 비디오 데이터셋에서 많이 retrieval하였다.
• VizWiz가 독특한 데이터셋인만큼 in-domain retrieval의 비중이 컸다.
  
• Human-labeled data는 다양한 domain에서 text를 자유롭게 구성할 수 있었다.
• 위 표는 다양한 Human-labeled data들이 성능에 얼마나 연관/영향을 미치는지를 측정하였다.

6.3. The Impact of Retrieval

• without retrieval의 경우는 retrieval caption없이 prompt를 'This image shows'만 넣은 경우이다.
• smallcap이 단순히 retrieval caption을 되새김해서 caption을 만드는것은 아닌지를 확인하기 위해 visual modality를 없애서 visual features을 0으로 인코딩하여 실험하였더니 성능이 현저히 떨어졌다.

6.4. Alternative Decoders

Inference time

• SMALLCAP은 ligh weight-training captioning model로 추가적인 학습이 필요없다는 장점이 있지만 Inference time또한 매우 짧다.
• NVIDIA A100 GPU환경에서 1,000개의 랜덤한 coco이미지로부터 결과를 내는데 한 이미지당 0.22 ~ 0.58가 걸린다.