Scaling Language-Image Pre-training via Masking (2023)

[YoojLee]

Summary

CLIP training 시에 이미지 50~75% 정도 마스킹해서 encoding하게 되면 (sparse computation), 동일한 정확도로 더 빠르게 training 가능함 (동일 training 시에 more epochs achieved). 또한, 동일한 메모리 사용량을 가정했을 때 배치를 늘림으로써 (larger batch size) 같은 시간이라고 가정하면 더 높은 정확도 achieve (favorable accuracy/time trade-off) -> 그 중간쯤에서, 더 빨리 더 정확하게가 가능함. 다양한 downstream task에서 FLIP의 우수성 입증.

Method Highlights

MAE에서 영감을 받긴 했지만, autoencoding은 진행하지 않고 높은 비율의 masking ratio를 적용하여 image encoder에 넣어줌. text token에 masking하는 것은 그렇게 큰 이점은 없어서 final architecture는 image masking으로만 진행.

contrastive loss에 들어가는 negative pair는 in-batch negative로 구성 (SimCLR, MoCo discovered that large batch일 수록 유리 -> more prominent in language-supervised learning; 이건 왜지?)

MAE처럼 그냥 바로 pre-trained encoder 써줘도 되지만, 실제 inference 시에 만날 unmasked image와의 distribution gap을 없애기 위해서 pre-training 마지막에 적은 step 동안만 masking ratio를 0으로 세팅해서 학습시켜줌 (called unmasking tuning strategy)

Implementation

• ViT를 따라서 patch embedding 후에는 extra LayerNorm 적용 없이 바로 transformer layer에 feed
• image encoder 마지막에는 GAP
• WordPiece tokenizer를 활용하여 sequence length를 32로 조정 (기존 clip 레시피와 비교했을 때, 크게 달라지는 건 없었음)

Experimental Highlights

Ablations

trained on LAION-400M and evaluate zero-shot acc on ImageNet-1K

• Image masking: FLOPs 감소(연산량 줄어듦)하면서 masking 적용하면 accuracy도 competitive (or better) 하게 유지하면서 speed-up은 ⬆️ accuracy 증가하는 방향은 배치 사이즈 키워서 그런 듯.
• Batch size: masking하기 때문에 배치 사이즈 늘릴 수 있고, 늘리면서 성능 향상. 배치 사이즈 늘리면서 발생하는 추가적인 memory cost도 거의 없는 편 (FLOPS 보면 거의 배치 사이즈에 리니어하게 감소하는 경향). (a)의 row1이랑 비교하면, masking해서 동일 batch size로 고정해도 성능이 거의 비슷함. 이는 masking이 어느 정도 overfitting을 방지하는 regularization으로 적용하기 때문으로 보임 (masking으로 인한 정보 손실을 어느 정도 보완함). masking ratio를 더 키우면 정보 손실의 영향이 커지는 걸 확인 가능.
• Text masking: 마스킹 50 정도 적용했을 때 성능 소폭 감소 (information density 차이 때문). encoder가 작아서 masking해도 speed up에는 큰 영향 X. ➡️ 굳이 text에는 masking하지 X (참고로 prioritized 같은 경우에는 패딩이 적용된 토큰부터 먼저 masking하는 옵션)
• Inference unmasking: training 시에는 masked image만 보고, inference 시에는 masking을 적용하지 않으면 training과 test time 시에 distribution shift가 발생할 것임. 그럼에도 불구하고, inference 시에 masking을 적용하는 것보다 온전한 이미지를 feed하는 게 더 성능이 좋았음. 해당 성능 하락은 inference 시 정보 손실에 기인할 수도 있기에 여러 masking된 이미지 (전체 패치를 합치면 결국에는 온전하도록)를 ensemble해서 feed해주었을 때, 해당 perf gap이 조금은 나아졌지만 그래도 여전히 full view 보는 것보단 좋지 X.
• Unmasked Tuning: distribution shift가 있어도 어느 정도 잘 되는 걸 볼 수 있었지만, 조금만(0.32epoch)이라도 unmasked image로 tuning 시켜주면 성능이 향상됨. 이러한 경향은 high masking ratio일 수록 두드러짐.

→ 같은 accuracy를 더 적은 training-time에 achieve 가능. 반대로, 동일한 시간을 들였을 때 더 높은 성능 achieved (favorable accuracy/time trade-off. pre-training에 있어서 acc과 time 간에는 trade-off가 존재할 수밖에 없는데 이미지 하나 당 더 적은 patch를 process함으로써 배치 사이즈를 늘리고 시간도 줄인다..)

• Reconstruction: MAE 따라서 recon objective도 추가해서 디자인한 결과, 그닥 성능 면에서 이득이 없음. 오히려 살짝 하락. 굳이 decoder 추가로 붙이느니 (그에 따른 디자인적인 복잡도나 accuracy/time trade-off 면에서 불리) 안하는 게 낫다.

Compare with CLIP

pre-training data를 통일해서 비교하기 위해 OpenCLIP과 저자들이 reproduce한 CLIP 결과도 추가해서 비교. 결과는 요약하자면, 전반적으로 동일한 pre-training data에서 학습된 clip 비교군 중에는 flip이 우수함. 또한, pre-training dataset에 따른 large systematic gap을 발견.

• 전체적으로 그냥 CLIP보다 약간 뒤처지지만, pre-training data의 영향을 배제하지 않은 결과라서 완전히 fair한 비교는 아님 (이런 pre-trainign data의 영향은 Table 4에서 데이터셋 별로 WIT vs. LAION을 비교했을 때 각 pre-training dataset 별로 성능이 잘나오는 데이터셋이 뚜렷이 구분된다는 점에서 두드러짐). pre-training data를 통일하고 보았을 때 다양한 시나리오 (zero-shot, linear probing, fine-tuning) & downstream task에서 CLIP training보다 전반적으로 더 좋은 성능을 보여줌.
• 참고로, retrieval의 경우에는 flip이 가장 좋은 성능 (이는 다른 말로 하면 WIT 데이터셋이 LAION 대비 retrieval에 큰 도움이 안된다고 볼 수도).
• robustness 테스트에 있어서는 LAION에서 사전학습된 모델들이 정말 확연히 안 좋음 (특히, IN-A). 개중에 FLIP이 선전하지만 그래도 OpenAI의 clip보단 확연히 떨어짐. 그래도 FLIP이 선전한다는 점에서 어느 정도 masking이 regularization으로 작용하여서 모델의 robustness함을 제고했다는 걸 확인.
• captioning과 VQA는 fine-tuning 진행. captioning의 경우에는 FLIP이 전반적으로 우수 (clip이랑 비슷하거나 혹은 좀 더 나음), VQA 같은 경우에는 clip보다 살짝 뒤떨어지지만 LAION-based clip 에서는 우수 (근데 큰 차이 X. 또이또이라고 보면 될 듯..)

Scaling behavior

scaling은 3가지 측면에서 study. 1) Model size 2) Dataset size 3) Training schedule length

각 scaling axis에 따른 IN-1K의 zero-shot acc를 기록. 모델 size의 경우에는 전체적으로 모든 training step에서 특정 gap 이상의 차이가 있었음. 데이터셋 사이즈 같은 경우에는 중간 정도까진 비슷하지만 끝으로 갈 수록 2B가 더 성능이 좋음. 동일한 training step을 유지했다고 했을 때, 다양한 데이터를 볼 수 있는 2B가 더 유리했을 것으로 판단 (2B라고 해서 더 많이 학습시키고 그런 게 아니라는 뜻), 또 training schedule 같은 경우에는 더 오래 step을 밟는다고 해서 의미 있었던 것은 아님.

data scaling의 경우에는 zero-shot에 유리하고, model scaling 같은 경우에는 transfer learning에 유리함 (하지만, model scaling이 transfer learning에 유리하다는 점은 downstream dataset의 크기에 dependent함. 즉, 단정지을 수는 없다).

data scaling을 할 때에 추가로 training을 더 하거나 computation을 더 하지 않았다는 점에서, data scaling을 통한 성능 향상의 요인은 데이터셋 사이즈가 커지면서 이미지 및 텍스트가 다양해짐에 따라 보다 풍부한 정보를 학습할 수 있다고 해석이 가능함.

scaling 시 두 개 이상의 측면에서 scaling을 수행했을 때 실제로 하나의 측면만 scaling을 수행하는 것보다 더 좋은 효과 (증가폭이 단순 합보다 더 큼). 즉 서로서로 시너지가 있다. scheduling scaling 같은 경우에는 단독으로만 두고 봤을 때는 성능 향상 등에 미치는 영향은 거의 없었다고 볼 수 있는데, 모두 다 scale up한 결과 (joint)를 보았을 때 captioning 제외하고는 best. scheduling scaling은 나머지 두 개를 condition으로 했을 때 성능에 향상을 가져올 수 있다. joint scaling 결과를 보았을 때 zero-shot IN-1K가 기존의 OpenCLIP-ViT/H (LAION-2B) 점수보다 더 높음. sampled data의 수가 25.6B로 32B의 openclip보다 더 적음 -> 역시 여기서도 accuracy/time trade-off가 좀 더 favorable하게 작용한다는 점을 확인할 수 있다.

좀 더 detail하게 보면..

Strengths

• 매우 efficient함과 동시에 높은 성능
• 굉장히 단순한 구조
• scaling behavior에 대한 실험
• 실험 설계가 체계적임 (dataset에서 기인한 성능 차이; so called systematic gap와, 모델 training 방식에서의 성능 차이를 제대로 구분해낼 수 있는 설계)
• 논문을 잘 쓴다..contribution을 되게 설득력 있게 잘 전달함 + simple yet effective를 항상 근본으로 가져가는 듯.

Weaknesses

• 아무래도 technical novelty는 떨어짐 (까딱하면 tech report처럼 되기 쉬운 거 같은데 그걸 잘 피해가는 것 같음.)
• weakness라기보다는 captioning에서 model scaling이 잘 나오는 이유가 궁금함. dataset size를 늘려서 다양한 concept에 대한 학습을 한 후에 transfer하는 게 captioning에서 훨씬 유리하지 않을까 싶었는데...