LiT🔥 : Zero-Shot Transfer with Locked-image text Tuning 논문 리뷰

The overview of this paper

 이 논문에서는 각 모델의 pre-training 장점은 유지하며 이미지와 텍스트 모델을 정렬하기 위한 contrastive training을 사용하는 간단한 method를 제안하였다. 논문의 실험에 따르면 locked pre-trained image model & unlocked text model이 가장 좋은 성능을 보여주었다. 이러한 contrastive-tuning을 'Locked-image Tuning' (LiT)라고 부른다. LiT는 새로운 task에 대해 pre-trained image model로부터 좋은 representation을 읽어내기 위한 text model만을 가르친다. LiT 모델은 새로운 vision task에 대해서 zero-shot transfer 능력을 얻을 수 있었다. 제안된 LiT는 광범위하게 활용 가능하다: 다양한 pre-training methods & 3개의 서로 다른 image-text 데이터셋을 사용한 다양한 architecture에도 사용 가능

 

 

Table of Contents

1. Introduction

2. Methods

   2-1. Contrastive pre-training

   2-2. Contrastive-tuning

   2-3. Design choices & Locked-image Tuning

3. Image-text datasets

4. Experiments

   4-1. Comparison to the previous SoTA

   4-2. Evaluation of design choices

   4-3. Technical advantages of locked image models

5. Discussion

 

 

1. Introduction

 이전의 연구들에서 이미지 모델(or image tower)과 텍스트 모델(or text tower)이 동시에 학습되는 contrastive learning 프레임워크를 제안하였다. 이 두 타워는 contrastive loss를 최소화하기 위해 학습된다: 쌍을 이루는 이미지 & 텍스트의 representation은 유사하게 쌍을 이루지 않은 이미지 & 텍스트의 representation은 다르게 함으로써. 테스트 시에 결과 모델은 image embedding과 textual class 설명 임베딩 간의 비교를 통해 zero-shot 이미지 분류에 사용될 수 있다.

 

 논문에서는 contrastive learning 프레임워크를 적용하고 더 많은 데이터와 계산 효율적인 전략인 contrastive-tuning을 제안하였다. 핵심 아이디어는 image-text 데이터를 사용하여 text tower를 조정하는 동시에 pre-train된 강력한 이미지 모델을 image tower로 사용하는 것이다. 학습 중에 두 타워의 가중치는 잠겨질 수도 있고 잠겨지지 않을 수도 있는데, 그림 2에서 서로 다른 디자인 선택에 대해서 설명하고 있다. 논문에서는 그림 1에서 나타나 있는 것처럼 image tower를 잠그는 것이 가장 좋은 결과를 보여준다는 것을 발견했다. 이러한 contrastive-tuning은 "Locked-image Tuning" (LiT)라고 부르고, 이 모델은 text model이 pre-trained image model로부터 적합한 representation을 읽어내도록 가르친다. LiT은 이전의 모델들보다 좋은 성능을 보여줬고, out-of-distribution(OOD) ImageNet test에 대해서도 새로운 SoTA를 달성하였다.

 

그림 1. 이전 SoTA method들과의 비교. 왼쪽: 공공 YFCC100m 서브셋에서의 결과 / 오른쪽: 사적으로 모은 데이터에서의 결과

 

 논문에서는 LiT이 잘 작동하는 이유가 데이터 소스의 분해와 image descriptor와 vision-language 정렬의 학습을 위한 기술에 있다고 생각하였다. image-text 데이터는 자연어와 시각 간의 관련성을 학습하는데 좋지만, 동시에 SoTA image descriptor를 결과로 내놓기에는 부정확하거나 깨끗하지 않을 수 있다.

 

 이 논문의 또다른 contribution은 다음과 같다.

 

• 그리 많지 않은 computation resource와 공공 데이터셋에서의 학습으로도 고성능의 zero-shot 모델을 만들 수 있는 방법을 제안. 
• 이미 pre-train된 모델을 재사용함으로써 image model을 학습시키기 위한 computational resource는 줄어들 수 있음.

 

2. Methods

2-1. Contrastive pre-training

 

 자유형식의 텍스트 설명과 쌍을 이루는 이미지(noisy 할 수도 있음)의 모음은 visual model을 학습시킬 수 있는 강력한 자원으로 떠올랐다. 이 데이터의 주된 장점은 사전에 정의된 카테고리에 한정된 세트에 의해서 제한되지 않고, 대신에 제약을 두지 않는 자연어를 사용해서 이미지를 설명한다는 것이다. 그 결과 이 데이터로부터 학습된 모델은 광범위한 task에 대해 zero-shot learner로 여겨진다.

 

 contrastive pre-training은 image-text 데이터로부터 모델을 학습시키는데 효과적인 하나의 방법이다. 논문에서는 이 방식에 대해서 더 자세하게 살펴보고 image-text 데이터로부터 contrastive pre-training을 향상시키는 간단하지만 매우 효과적인 방법을 제안하였다.

 

 contrastive pre-training 방식의 주된 아이디어는 두 개의 임베딩 모델을 학습한다는 것이다: image model과 text model 둘은 똑같은 차원의 representation을 생성한다. 이러한 모델은 contrastive loss를 사용하여 학습된다. 이 loss는 해당 image-text 쌍은 유사한 임베딩을 갖게 하고, 해당하지 않는 image-text 쌍은 별개의 임베딩을 갖게 한다.

 

 이 손실 함수의 중요한 디테일은 loss가 각각의 accelerator device에서 독립적으로 계산되고 축적되거나 모든 디바이스에 대해 공동으로 계산된다. 실험을 진행한 결과 후자의 경우가 일관적으로 더 나은 성능을 보여주는 것을 알 수 있었다. 논문에서는 그래서 모든 실험과 ablation에서 global loss를 사용하였다. 

 

 image tower와 text tower가 학습되고 나면, 이들은 zero-shot 분류에 바로 사용될 수 있다: 클래스 이름 또는 설명은 텍스트 모델과 함께 임베딩된다. 그런 다음 지정된 이미지에 대해 이미지 임베딩에 가장 가까운 임베딩이 있는 레이블이 선택된다. 

 

2-2. Contrastive-tuning

 

 contrastive pre-trainin은 두 개의 task를 동시에 학습하는 것처럼 보일 수 있다. 이 모든 task를 동시에 해결하는데 image-text 데이터에서의 pre-training은 잘 작동하긴 하나, 최적의 방법은 아닐 수도 있다. 

 

1. 이미지 임베딩을 학습하고
2. 이미지 임베딩 공간에 맞춰 정렬하기 위해 텍스트 임베딩을 학습

 

 이미지 임베딩을 학습하기 위한 기존의 방법은 거대하고 비교적 깨끗한 수동적으로 라벨링된 이미지 데이터셋을 사용하는 것이다. 하지만 이러한 방식은 분명한 약점을 가지고 있다: 사전에 정의된 카테고리 세트, 따라서 이 데이터에서 pre-train 된 모델은 이 카테고리에 대해서만 추론이 가능하다. 이와 반대로 image-text 데이터는 자유 형식 텍스트로부터 학습하기 때문에 이러한 제한을 갖지 않는다. 그렇지만 조금 떨어지는 퀄리티를 가질 수는 있다.

 

 논문에서는 이 두 가지 유형의 데이터의 장점을 합치기 위해 contrastive-tuning을 제안하였다. 이를 하기 위한 구체적인 방법은 수동으로 라벨링된 깨끗한 데이터를 사용해서 이미 pre-train 된 이미지 모델을 사용해 contrastive pre-training을 초기화시키는 것이다. 이러한 방식으로 이미지 임베딩과 독립적으로 image-text 정렬을 학습하여 두 데이터 소스의 이점을 모두 활용할 수 있다. supervised pre-trained 이미지 모델을 사용함으로써, 제안된 contrastive-tuning은 어떠한 모델이든지 사용 가능하다. 

 

2-3. Design choices and Locked-image Tuning

 

 pre-trained image 또는 text model을 contrastive learning 세팅에 소개하는 것은 여러 디자인 선택을 포함한다. 첫 번째로, 각 타워(image & text)는 독립적으로 랜덤 하게 초기화되거나 pre-trained model로부터 초기화될 수 있다. pre-trained model에 대해 여기에는 최소 2 가지의 변형이 있다: lock(freeze) or fine-tuning 허용.

 

 pre-trained image-text modal은 서로 다른 representation 사이즈를 가지나, contrastive loss는 똑같은 사이즈의 representation을 예측한다. 이를 해결하기 위해, 논문에서는 일반 차원으로 representation을 매핑하는 옵션적 linear projection(head)를 각 타워에 추가하였다. 시도된 MLP 기반 헤드에 대한 preliminary investigation에서는 이러한 간단한 선형 헤드에 비해 상당한 개선이 이루어지지 않았다.

 

 논문에서는 그림 2에 언급되어 있는 것처럼 잠재적 디자인 선택을 위한 2-character notation을 소개하였다. 각각의 캐릭터는 선택된 image model과 text model에 대해서 세팅을 인코딩한다. 논문에서는 3개의 가능한 세팅을 소개하였다. 예를 들어 Lu는 locked pre-trained image model & unlocked(trainable) 랜덤 하게 초기화된 text model을 의미한다. 

 

• L: locked 가중치. pre-trained model로부터 초기화
• U: unlocked/trainable 가중치. pre-trained model로부터 초기화
• u: unlocked/trainable 가중치. 랜덤하게 초기화

 

그림 2. image-text 데이터에서 contrastive-tuning에 대한 디자인 선택. 두 개의 문자는 image tower와 text tower의 셋업을 나타내기 위함.

 

 이전의 연구들은 밑바닥부터 학습을 하는 uu 모델이다. 논문의 실험에 의하면 Lu 세팅이 가장 좋은 성능을 보여준다는 것을 알 수 있다. 그리고 이것을 Locked-image Tuning(LiT🔥) 라고 이름 붙였다.

 

3. Image-text datasets

 논문에서는 Conceptual Captions(CC12M)와 Tahoo Flickr Creative Commons(YFCC100m) 데이터셋을 사용하였다. 그리고 논문에서 따로 수집한 데이터셋도 사용하였는데, 이 데이터셋은 ALIGN과 똑같은 프로세스를 따라서 똑같은 image-base filtering을 사용하지만 더욱 간편한 text-based filtering을 사용하여 40억 개의 image & alt-text 쌍을 수집하였다. 오해의 소지가 있는 평가 결과를 피하기 위해, 논문에서는 평가하는 데이터셋으로부터 거의 복제에 가까운 이미지들은 제거하였다. 논문에서는 논문의 데이터셋 생성이 주된 contribution이라고 생각하지 않았다; 논문의 method의 효율성을 규모 측면에서 설명하기 위해 ALIGN의 데이터 수집 프로세스를 간단하게 만들었을 뿐이다.

 

4. Experiments

4-1. Comparison to the previous SoTA

 

 이 섹션에서, 논문에서는 데이터셋에서 LiT의 결과를 보여줬다. image tower는 JFT-3B에서 pre-train 된 ViT-g/14 모델을 사용하여 초기화하였다. 그리고 32k의 배치 사이즈를 사용하고 180억 개의 image-text 쌍에 대해서 조정을 하였다.

 

 논문에서는 LiT method와 CLIP과 ALIGN을 포함하는 이전의 SoTA method들과 비교하였다. 표 1에서는 ImageNet 데이터셋, 5개의 out-of-distribution(OOD) 테스트, 7개의 VTAB-natural task에서 zero-shot 분류 결과를 기록하였다. LiT는 ImageNet zero-shot 분류에서 이전의 SoTA 모델을 상당히 능가하는 성능을 보여줬다. CLIP과 ALIGN에서 각각 9%와 8.8%의 성능 향상을 보여줬다. 이는 zero-shot transfer 결과와 supervised fine-tuned 결과 간의 갭을 절반이나 줄인 것이다.

 

표 1. ImageNet, 5개의 OOD test variants, 7개의 VTAB-natural tasks에서의 zero-shot transfer 정확도

 

Robustness.  모든 OOD 변형에서, LiT는 이전의 모델들보다 일관되게 능가하는 성능을 보여줬다. 

 

Diverse downstream tasks.  논문에서는 19개의 다양한 task로 구성되어 있는 VTAB에서 LiT 모델을 평가하였다. 7개의 VTAB-natural task에 대한 결과가 표 1에 나타나 있다. LiT 모델은 supervised fine-tunes ResNet50 baseline과 비교해서 더 유망한 zero-shot 결과를 보여줬다.

 

Data & compute efficiency.  그림 1은 더 적은 image-text 쌍과 함께 조정했을 때의 더 많은 결과를 보여주고 있다. 오직 300M 개의 image-text 쌍과 LiT을 사용했을 때 모델은 zero-shot ImageNet transfer에서 81.7%의 top-1 accuracy를 달성하였다. LiT는 literature에서 이미 pre-train된 모델을 재사용해서 이미지 모델을 재생성하는 데 사용되는 computational resource를 줄이기 위한 방법을 제안하였다. 

 

Results on public datasets.  LiT의 높은 데이터 효율성을 감안할 때 논문에서는 더 작고 공개적으로 사용 가능한 모델 및 데이터 세트만 사용할 때 LiT가 얼마나 잘 수행되는지 조사하였다. 자세하게 살펴보면, YFCC100m-CLIP과 CC12M의 조합 데이터셋에서 pre-train 된 ViT-L/16 모델을 사용하여 ImageNet-21k를 조정하였다. 결과적으로 논문에서는 ImageNet에서 전례 없는 75.7%의 zero-shot transfer을 달성하였다. 이는 오직 공공 데이터 소스만 사용한 이전의 SoTA 모델보다 30.9% 더 향상된 성능이다. 또한 LiT는 광범위한 robustness 데이터셋과 VTAB-natural task에 대해서 강력한 결과를 얻었다.

 

4-2. Evaluation of design choices

 

Small-scale thorough investigation.  논문에서는 pre-trained 가중치와 locked(L) 또는 unlocked(U), 랜덤 하게 초기화되는 unlocked(u)를 사용하여 초기화된 image와 text tower의 다양한 조합에 대해 연구하였다. 논문에서는 YFCC100m-CLIP 데이터셋에서 각 설정을 여러 번 학습하고, 세팅의 궤적을 이해하기 위해 총 단계 수를 2,500에서 60,000으로 변경하고, 오도를 피하기 위해 learning rate 및 weight-decay를 변경하였다. 그림 3은 각각의 세팅에 대한 최고의 결과를 보여준다. 곡선의 각 점은 그 기간에 대한 각각의 full run이다. image tower를 잠그는 것은 거의 항상 최고의 결과를 보여주고 pre-trained image tower를 사용하는 것은 상당히 도움된다. 그치만 pre-trained text tower를 사용하는 것은 조금의 성능 향상 밖에 보여주지 않고 text tower를 잠그는 것은 잘 작동하지 않는다.

 

그림 3. YFCC100m-CLIP 데이터셋에서 LiT의 가능한 locking과 초기화 세팅에 대한 연구

 

This still holds in the near-infinite data regime.  누군가는 pre-trained image tower를 잠그는 것이 잘 작동되는 이유가 YFCC100m-CLIP 데이터셋이 비교적 작기 때문이고, 랜덤 하게 초기화된 image tower는 더욱 큰 image-text 데이터셋에서 locked model을 능가하는 성능을 보여줄 수 있을 것이라고 생각할 수 있다. 그림 3의 Uu와 UU 세팅의 곡선은 이러한 예측을 뒷받침하는 것처럼 보인다. 

 

 하지만 놀랍게도 실험의 결과는 image tower를 잠그는 것이 매우 큰 image-text 쌍의 데이터셋에서 contrastive 하게 조정될 때 이익을 제공한다는 것을 보여주고 있다. 표 2는 3개의 세팅(Lu, Uu, uu)에서 40억 개의 이미지 데이터셋에 대한 contrastive tuning의 결과를 보여주고 있다. uu는 CC12M과 YFCC100m-CLIP 같은 작은 데이터셋에서 더 나은 성능을 보여주고 있다.

 

표 2. 거대 데이터셋에서 디자인 선택에 대한 평가

 

 pre-trained model로부터 초기화된 image tower는 더 나은 성능을 제공하고 CLIP/ALIGN의 비교적 간단한 연장이다. 놀랍게도 frozen setup Lu는 더 나은 결과를 얻게 된다. 반면에 다른 관점에서 LiT는 강력한 이미지 임베더로부터 지식을 추출하는 text tower을 간단하게 학습한다. 이러한 유연하고 성능이 좋은 셋업은 text-embedding tower를 붙임으로써 기존의 vision backbone을 zero-shot learner로 변환할 수 있다.

 

Why is locked(L) better than unlocked(U)?  image tower를 잠그는 것은 image tower가 contrastive-tuning 중에 적응되도록 허락하는 것보다 더 낫다는 것은 어느 정도 놀랍고 직관에 반대되는 사실이다. 그림 4는 이러한 이유에 대해 힌트를 준다.

 

그림 4. LiT에 사용되는 데이터셋에서의 loss(위 행)와 OOD 데이터셋에서의 loss 간의 비교(중간 행)와 pre-logit에서 linear few-shot 평가에 의해 측정된 representation quality(아래 행)

 

 첫 행은 image tower를 잠그는 것이 LiT에 사용되는 데이터셋에서 상당히 안 좋은 loss를 이끄는 것을 보여준다. 반면에 locked image 변형의 loss는 COCO caption과 같은 out-of-distribution(OOD) 데이터셋에서 더 나은 성능을 보여준다 (중간 행).

 

 논문에서는 pre-logits에서 few-shot linear regression으로 얻어진 성능을 통해 이미지 모델의 representation quality를 측정한다 (아래 행). 이 모든 것을 합쳐서 이 그림은 pre-trained image model의 image representation은 매우 잘 일반화되지만, contrastive 한 fine-tuning은 visual representation의 일반성을 악화시킨다. 이는 contrastive 데이터셋에서는 더 좋지만 다른 모든 곳에서는 더 나쁘게 만든다. 이는 조정 중에 image tower를 잠그면(예: LiT) 잘 정렬되지만 사용된 데이터셋에 특화된 image-text 모델과 반대로 이미 강력하고 일반적인 image representation에 잘 정렬된 텍스트 모델로 이어진다는 것을 나타낸다.

 

4-3. Technical advantages of locked image models

 

 locked image tower를 사용하는 것은 많은 이익을 가진다. 이에 대한 이유는 다음과 같다. 이 모든 것이 합쳐져서, 이러한 구현 feature는 대규모 배치 사이즈에서 거대한 모델의 사용을 가능하게 할 수 있다.

 

1. image tower를 위한 기울기 계산이 필요 없기에 학습은 상당히 빨라지고 메모리 사용량은 감소함.
2. augmentation이 사용되지 않으면, 이미지 모델의 임베딩은 한 번에 즉시 계산될 수 있음. 이는 계산 시간 & 메모리 필요를 줄여준다. 

 

5. Discussion

Limitations.  cross-modal retrieval task에서는 Uu 셋업과 UU 셋업에 비해 Lu 셋업은 유의미한 결과를 보여주지 않는다. 매우 긴 tuning에서는 가끔 Uu 또는 UU가 Lu를 넘어서는 성능을 보여주기도 한다. 논문의 결과는 제안된 Lu 셋업이 고정된 예산에서 계산 비용을 아낄 수 있지만, 충분히 큰 예산에서는 만약 zero-shot 분류가 주된 목표가 아니라면 Uu 셋업 또한 유용할 것이라고 말하였다.

 

 

 

 

출처

https://arxiv.org/abs/2111.07991

LiT: Zero-Shot Transfer with Locked-image text Tuning