[ICML 2021] ALIGN: Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision

안녕하세요, 허재연입니다. 오늘 리뷰할 논문은 google research에서 2021년 ICLR에 게재한

ALIGN: Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision 이라는 논문입니다. CLIP과 비슷한 시기에 나온 비슷한 컨셉의 논문으로, 학습 방법론도 크게 다르지 않지만 역시 학습 데이터셋에 차이가 있습니다. 한번 살펴보도록 하겠습니다.

Introduction

딥러닝 모델의 크기를 키울수록, 학습에 대량의 데이터를 사용할수록 일반적으로 모델의 성능을 끌어올릴 수 있음이 알려지게 되면서, 점점 대규모 학습을 활용한 방법론들이 많이 제안되었습니다. 특히, NLP에서 대규모 사전학습의 성공이 비전 분야에도 조금씩 영향을 미치며, 대규모 언어 모델의 표현력을 활용하거나 대규모 Vision-Language pretraining을 통한 멀티모달 표현력을 활용하고자 하는 시도가 시작되었고, self-supervised learning에 기반한 방법론들이 좋은 결과를 보였습니다.

기존에는 representation learning에서 visual 및 vision-language은 서로 다른 데이터셋 소스를 활용해 개별적으로 연구되었습니다. visual represntation learning에서는 ImageNet, JFT300M, OpenImage와 같은 대규모 지도학습 데이터셋을 통한 사전학습이 transfer learning을 통한 downstream task 성능 향상에 중요한 역할을 한다는 점이 입증되었지만 당연하게도 labeling의 어려움으로 인해 대규모 데이터셋을 구축하는데 한계가 있었고, 좋은 representation을 학습하기 위한 self-supervised 사전학습 방법의 연구가 진행되었습니다. vision-language modeling에서도 사전학습을 활용하는 것이 일반적인 접근법이었는데, Conceptual Captions, Visual Genome Dense Captions, ImageBERT와 같은 visual-language 사전학습 데이터셋들은 데이터셋의 클래스 간 균형을 맞추고, 정제하고, 어노테이션을 하는데 특히 더 많은 자원을 필요로 했기에 그 크기를 늘리는데 한계가 있어 1000만개 이상의 크기로 확장하기 쉽지 않았다고 합니다. 이는 단순한 vision 데이터셋과 비교해도 굉장히 작은 규모이고, NLP 학습을 위해 웹에서 수집된 대규모 텍스트 corpus와 비교하면 너무 작은 규모입니다.

이에 본 논문에서 저자들은 10억개 이상의 image-alt text 쌍으로 구성된 데이터셋을 활용하여 visual / visual-language representation learning을 활용하였습니다. 데이터셋 구축 방법의 전반은 Conceptual Captions dataset에서 제안된 방법론을 따랐지만, 복잡한 필터링 / 후처리 과정을 거치지 않고 단순히 빈도 기반 필터링만을 적용하였습니다. 이렇게 데이터셋을 구축하형 비록 노이즈가 데이터셋에 포함되지만, 기존 conceptual captions 데이터셋보다 두 자리 수 이상 더 큰 규모의 데이터셋 확보할 수 있게 되었다고 합니다. 실험 결과, 저자들이 구축한 초대규모(exascale) 데이터셋을 활용한 사전 학습이 다양한 task에서 좋은 결과를 보이는 것을 확인하였다고 합니다. 제가 얼마 전에 읽은 BLIP에서는 noise없이 깨끗한 데이터의 필요성을 강조하며 이를 위한 방법론을 제안하였었는데, 이번 논문에서는 노이즈가 좀 껴도 데이터의 크기를 늘릴 수 있다면 괜찮다는 느낌이어서 데이터 규모 및 노이즈 비율에 대해 어느 정도 trade-off가 있는지 생각해 볼 필요성이 있겠네요. 어느정도 노이즈가 포함되더라도 압도적인 scale로 찍어눌러 성능 향상을 시킨 느낌이긴 합니다(당시까지는 아직 Vision-Language Pretraining에서 web scale의 데이터를 활용하는게 보편적이진 않았던 것 같습니다).

저자들은 visual-language pretrained model인 ALIGN: A Large-scale ImaGe and Noisy-text embedding을 제안합니다. 모델 학습을 위해서는 간단한 dual-encoder 구조를 사용해 동일한 latent embedding space에 시각적 표현(visual representation)과 언어적 표현(lauguage representation)을 정렬(align)하도록 하였다고 합니다. ALIGN에서는 이미지와 텍스트 인코더를 학습할 때, 정규화된 소프트맥스(normalized softmax)로 정의된 contrastive loss을 적용하여 매칭된 이미지-텍스트 쌍의 임베딩을 가깝게 하고, 매칭되지 않은 쌍의 임베딩을 멀어지도록 하는 contrastive learning을 활용합니다. 당시 contrastive learning은 SimCLR, MoCo, BYOL등의 성공을 통해 자기 지도 학습에 상당히 좋은 학습법임이 받아들여지고 있었습니다. 이를 영상 간 매칭이 아닌 영상 및 텍스트 간 매칭으로 수정한 것이죠. 이렇게 영상 및 텍스트 representation 간 정렬로 학습된 모델은 cross-modality matching이나 retireval task에 적합합니다. ALIGN은 Flickr30K 및 MS COCO 데이터셋에서 zero-shot 및 fine-tuning에서 기존보다 좋은 성능을 달성하는데 성공했습니다. 또한, 이런 cross-modal matching은 class 이름을 text encoder에 입력하는 방식으로 zero-shot image classification이 가능하게 하며, ImageNet의 학습 샘플을 전혀 사용하지 않고 76.4%의 top-1 accuracy를 달성하였습니다. downstream task에 적용하였을때는 ImageNet에서 88.64%의 top-1 accuracy를 달성하였습니다.

A Large-Scale Noisy Image-Text Dataset

저자들은 visual / vision-language representation learning의 규모를 키우는 것을 목표로 하였습니다. 이를 위해 기존보다 훨씬 더 큰 규모의 데이터셋을 활용했죠. 영어로 된 image-alt text pair를 얻기 위해 기존의 Conceptual Captions 데이터셋의 구축 방법론을 따라 데이터를 수집하였습니다.

conceptual captions는 2018년에 구글에서 공개된 image-caption 데이터셋인데요, MS-COCO 데이터셋의 문제점을 개선하기 위해 330만 개의 다양한 이미지-캡션으로 구축한 데이터셋입니다. 인터넷에서 데이터를 수집했기 때문에 정제된 데이터셋을 구축하기 위해 이미지와 텍스트의 필터링을 수행하는 파이프라인을 적용하였습니다. 이미지의 경우 문제가 될 만한 자극적인 사진을 거르고(음란물,비속어 탐지기에 감지되지 않는 이미지), 이미지의 가로 세로 모두 400픽셀 이상이며 그 비율이 2:1 ~ 1:2 사이에 있는 이미지를 필터링하여 전체 이미지의 65%를 거르고, 텍스트의 경우 webpage HTML로부터 명사 혹은 전치사가 없거나 너무 많지 않은 텍스트, 토큰이 많이 반복되지 않는 텍스트, 대소문자 구별이 잘 된 텍스트, alt-text의 모든 토큰이 위키피디아에 5번 이상 등장하는 텍스트 .. 등등의 필터링을 거쳐 정제한 데이터라고 합니다. 이렇게 필터링을 거치면 대략 3%의 후보 쌍만 남게 된다고 하네요. conceptual captions의 구축 과정이 더 궁금하신 분은 해당 논문에 구축 방법이 자세히 기술되어 있으니 참고하시면 좋을 것 같습니다.

conceptual captions의 경우 위와 같은 강력한 필터링 및 후처리 과정을 거쳐 정제되었지만, 본 논문의 저자들은 데이터의 규모를 확장하기 위해 이러한 정제 과정을 대부분 생략하는 전략을 취했습니다. 원본 연구에서 적용한 복잡한 정제 과정을 완화해 어느 정도 데이터 품질을 일정 부분 희생하는 대신, 데이터의 규모를 극대화한 것이죠. 그래도 필터링을 아예 안 할 수는 없으니, 다음과 같이 최소한의 빈도 기반 필터링만 적용하였습니다. 그 결과, 18억개의 image-text 쌍으로 구성된, 기존보다 훨씬 더 큰 대규모 데이터셋을 구축할 수 있었습니다. Figure 2에서 그 예시를 볼 수 있는데, jpg가 걸러지지 않는 등 caption이 정리되지는 않은 듯한 모습을 보입니다.

Image-based filtering

Conceptual Caption의 방법론에 따라, 음란물에 해당하는 이미지를 제거하고, 각 변의 길이가 200픽셀 이상이며 aspect ratio가 3 미만인 이미지만을 남겼습니다. 또한, 1,000개 이상의 alt-text와 연관된 이미지는 제외하였다고 합니다. 또한 test 이미지로 학습하는것을 방지하기 위해 downstream 평가 데이터셋(ImageNet, Flicker30K, MS COCO)와 중복되거나 유사한 이미지를 제거해 모델이 테스트 셋과 겹치는 데이터로 학습되지 않도록 하였다고 합니다. near-duplicate image를 찾아내기 위해 먼저 별도의 고품질 이미지 임베딩 모델을 학습시키고 4000개의 k-means cluster를 만들어 ALIGN 데이터셋의 각 쿼리 이미지와 downstream task set의 index 이미지에 임베딩 거리 기반으로 top-10 nearest cluster를 찾은 다음 유사도가 높은 것들은 제외하였다고 합니다.

Text-based filtering

먼저 10개 이상의 이미지가 공유하는 alt-text들을 제외하였습니다. 보통 이러한 alt-text들은 “1920×1080”, “alt img” .. 와 같이 영상 내부 컨텐츠와 관련이 없는 캡션이었다고 합니다. 또한 너무 희귀한 토큰을 포함한 alt-text도 제외하고, 너무 짧거나(<3 unigrams) 긴(>20 unigrams) 캡션도 제외하였다고 합니다. 이를 통해 “image tid 25&id mggqpuweqdpd& cache 0&lan code 0”와 같은 noisy한 텍스트나, 사용하기에 너무 일반적인 텍스트를 제외할 수 있었다고 합니다.

Pre-training and Task Transfer

Pre-training on Noisy Image-Text Pairs

저자들은 dual-encoder 구조를 활용해 ALIGN을 사전학습하였습니다. 모델은 image encoder와 text encoder로 구성되며, 마지막에 cosine-similarity를 비교할 수 있도록 하였습니다. 이미지 인코더로는 global pooling이 적용된 efficientNet을 사용하였고, 텍스트 인코더로는 [CLS] token embedding이 있는 BERT를 사용하였습니다(training dataset에서 100k의 wordpiece vocab을 생성하였다고 합니다). 이미지 인코더 출력과의 차원을 맞추기 위해 BERT의 끝단에 fc layer와 linear activation을 추가하였으며, 두 인코더 모두 scratch로부터 학습하였다고 합니다.

이후에는 익숙한 contrastive loss로 학습을 하게 됩니다. normalized softmax 로 두 인코더를 학습하는데, 학습 과정에서 매칭되는 image-text pair를 postiive로, 배치 내부의 매칭되지 않는 다른 random image-text들을 nagetive pair로 간주하여 학습을 진행합니다. loss는 두개의 term을 합친 형태가 되는데, 하나는 아래 수식 (1)의 image-to-text classification을 위한 loss고, 다른 하나는 (2)번의 text-to-image classification을 위한 수식입니다. 순서만 바뀌었다고 생각하면 되겠네요.

{x}_{i}와 {y}_{j}는 각각 i번째 pair의 영상과 j번째 pair의 텍스트의 normalized embedding 벡터입니다. N은 배치 사이즈이고, sigma는 contrastive loss에서 항상 함께 등장하는 temperature loss입니다. logit의 smoothness를 조정해주는 역할이라고 생각하시면 됩니다. temperature value의 경우 학습으로 적절한 값을 찾았다고 하는데, BLIP도 그렇고 VLM 모델들은 temperature을 learnable parameter로 사용하는데 SimCLR와 같은 이전의 image contrastive learning 방법들과 다른 점 같아 눈에 띄네요.

Transferring to Image-Text Matching & Retrieval

저자들은 ALIGN을 image-to-text 및 text-to-image retrieval task에서 finetuning 없이 평가하였습니다. 데이터셋으로는 Flicker30K와 MSCOCO를 사용했고, 추가적으로 MS COCO에 추가적인 주석을 달아 확장한 Crisscrossed Caption(CxC) 데이터셋을 사용하였습니다. 학습 셋은 본래의 MS COCO와 동일하므로, MSCOCO에서 파인튜닝된 ALIGN모델을 CxC의 annotation을 통해 평가할 수 있다고 합니다

Transferring to Visual Classification

먼저, ALIGN 모델을 ImageNet 벤치마크 및 ImageNet-R(예술, 만화, 스케치와 같은 비자연적 이미지), ImageNet-A(모델이 처리하기 어려운 이미지), ImageNet-V2와 같은 변형된 데이터셋에 대해 zero-shot transfer 방식으로 적용하여 visual classification task에 대한 성능을 평가하였습니다. 변형 데이터셋들은 모두 ImageNet의 class 에 속해있지만, ImageNet-R 및 ImageNet-A와 같은 경우 원본의 이미지넷과 매우 다른 분포에서 샘플링한 이미지들로 구성되어있다고 합니다.

또한, image encoder를 다양한 다른 데이터셋에서 평가했다고 합니다. Oxford Flowers-102,Oxford-IIIT Pets, Stanford Cars와 같은 소규모의 fine-grained 분류 데이터셋에서 평가가 수행되었습니다. ImageNet의 경우 두 가지 측면에서 평가가 이루어졌는데, 이미지 인코더를 고정(freeze)하고 맨 마지막 classificatio layer만 학습하는 linear evaluation과 전체 네트워크를 파인튜닝하는 세팅에서 평가되었습니다. 이와 달리 finegrained classification benchmark에서는 linear evaluation 없이 전체 모델을 finetuning한 실험만을 수행하였습니다.

Experiments and Results

저자들은 ALIGN모델을 scratch부터 학습했으며, 이미지 인코더로는 오픈소스로 공개되 EfficientNet을, 텍스트 인코더로는 BERT로 사용했습니다. 이미지 인코더에 영상을 입력할 때 289×289 픽셀로 입력하는데, 이 때 입력 이미지를 346×346으로 resize하고 random crop하였다고 합니다(이외에도 flipping 등의 augmentation이 적용되었습니다). 텍스트 인코더인 BERT의 경우 입력 텍스트가 20개 이사의 unigram으로 이루어져 있어 최대 64개의 토큰을 갖는 wordpiece 시퀀스를 사용했다고 합니다. optimizer로는 LAMB를 사용했다고 하고, 배치 사이즈는 16384를 사용했다고 합니다. (보통 contrative learning에서 배치가 너무 크면 기존의 optimizer들로는 제대로 학습이 되지 않아 LARS와 같이 기존과는 다른 optimizer를 사용합니다)

Image-Text Matching & Retrieval

cross-modal retrieval benchmark 데이터셋으로는 Flickr30K와 MSCOCO를 사용하였습니다. Flickr30K에서 표준 1K test set에 대해서 평가하였고, 30k training set으로 finetuning하였다고 합니다. MS COCO의 경우, 5K의 test set에서 평가를 진행했고, 82K training set 및 5K validation이나 5K test set에 속하지 않은 추가적인 30K validation image로 finetuning했습니다.

Table 1에서 기존 방법론들과의 성능 비교를 나타내었는데, Flickr30k와 MS COCO 벤치마크에서 ALIGN이 SOTA를 달성하였습니다. zero-shot setting에서는 ALIGN이 기존 SOTA였던 CLIP보다 R@1에서 부려 7%나 개선된 결과를 보여주었습니다. fine-tuning을 했을 때에도 기존 방법론들보다 ALIGN이 큰 차이로 가장 좋은 결과를 보여주었습니다.

Table 2에서는 Crisscrossed Captions(CxC) retrievla task에서 ALIGN의 성능을 나타낸 것입니다. 여기서도 ALIGN이 모든 방면에서 SOTA를 달성하였는데, 특히 image-to-text 와 text-to-image에서 20%의 R@1이 넘는 큰 차이를 보였습니다.

Zero-shot Visual Classification

클래스 이름(classname)의 텍스트를 직접 텍스트 인코더에 입력하면, ALIGN은 이미지-텍스트 검색(image-text retrieval)을 통해 이미지를 후보 클래스 중 하나로 분류할 수 있습니다. Table 4에서는 ImageNet 및 그 변형 데이터셋에서 ALIGN과 CLIP을 비교한 결과를 나타내었습니다. CLIP과 마찬가지로, ALIGN도 다양한 이미지 분포에서 강력한 분류 성능 및 높은 robustness을 보였습니다. 공정한 비교를 위해, 저자들은 해당 실험에서 CLIP과 동일한 프롬프트 앙상블링 방법을 사용하였다고 합니다. 각 클래스 이름은 CLIP에서 정의한 프롬프트 템플릿(ex: “A photo of a [classname]”)을 활용하여 확장하였습니다. 클래스 임베딩(class embedding)은 모든 프롬프트 템플릿을 사용하여 계산된 임베딩의 평균을 구한 후, L2 정규화를 수행하여 사용하였습니다. 이러한 앙상블링 기법을 적용했을 때 ImageNet에서 top-1 accuracy가 2.9% 향상되는 것을 확인할 수 있었다고 합니다.

Visual Classification w/ Image Encoder Only

ImageNet 벤치마크에서, 먼저 학습된 visual feature를 freeze한 상태에서 classification head만을 학습시킨 다음 모든 계층을 fine-tuning하였습니다. 학습 과정에서는 배치 사이즈 1024를 사용했으며, optimizer는 SGD를 사용하였습니다(사전학습과 학습 세팅이 다르므로 다른 배치와 optimizer를 사용한 듯 합니다) Table 5에 ImageNet Benchmark에서의 기존 방법론들과 ALIGN 성능 결과를 비교하였는데, frozen feature에서 ALIGN은 CLIP보다 큰소하게 나은 결과를 보여 SOTA를 달성하였습니다. fine-tuning 이후에는 ViT 및 BiT보다 높은 정확도 성능을 보이며, imagenet 학습과 large-scale unlabeled data 간 더 많은 연관관계를 필요로 한 Pseudo Label를 제외한 방법론들과의 비교에서 가장 좋은 성능을 달성하였습니다.

fine-grained classification benchmark에서도 간단하게 fine-tuning 한 뒤 평가를 진행하였다고 합니다. 평가 결과를 Table 7에 나타내었는데, ALIGN은 복잡한 튜닝을 시도하지 않았음에도 기존의 SOTA 성능에 견주는 결과를 보일 수 있었다고 합니다.

google에서 발표한 VLM 모델인 ALIGN을 살펴보았습니다. 전반적으로 학습 방법론은 CLIP과 매우 유사했고, CLIP이 그러했듯이 핵심은 사전학습에 활용할 대규모 데이터셋을 구축하는 것이었습니다. 특히 눈에 띄는게 데이터의 정제를 어느정도 포기함으로써 학습 규모를 대폭 키울 수 있었다는 점에서 최근 읽었던 VLM 논문들과 그 철학이 약간 달라 흥미로웠는데요, 대량으로 학습 데이터를 때려박으니 노이즈가 있음에도 불구하고 어쨌든 성능은 좋아져서 신기하긴 했습니다. 정제 방법은 이후에도 많이 나왔으니 이후의 기법들을 함께 적용하면 보다 좋은 학습 set을 구축할 수 있겠네요.

리뷰 마무리하도록 하겠습니다. 감사합니다.