[논문 리뷰] 4M: Massively Multimodal Masked Modeling

NIPS 2023 paper code demo page
David Mizrahi, Roman Bachmann, Oğuzhan Fatih Kar, Teresa Yeo, Mingfei Gao, Afshin Dehghan, Amir Zamir
Google
Submitted on 11 Dec 2023

 

 

Abstract

현재 Vision AI 모델들은 극히 하나의 모달리티와 태스크에 특화된 모습을 보이고있다. 반면 최근 유행하는 LLM은 굉장히 넓은 범위의 능력을 갖고있다. 이는 Vision AI 분야에서도 유사하게 가능성이 있음을 암시한다. 본 연구에서는 멀티 모달 학습 전략인 일명 '4M'을 선보인다. 이는 text, images, geometric, semantic 모달리티는 물론 feature map을 포함하는 광범위한 input/output 모달리티 범위를 갖는 masked modeling objective를 사용하는 하나의 Transformer encoder-decoder로 구성된다. 4M은 모든 모달리티를 개별 토큰으로 매핑하고 토큰의 작은 무작위 집합에서 멀티모달 마스크 모델링을 수행하여 모든 모달리티의 표현 공간을통합함으로써 확장성을 달성한다.

4M은 몇 가지 주요 능력들을 갖고 있다.

1. 다양한 vision 작업을 즉시 수행할 수 있다.
2. 새로운 다운스트림 태스크에서 새로운 입력 양식을 위해 fine-tuning 할 때 좋은 성능을 보인다.
3. 임의의 모달리티를 조건화할 수 있는 생성 모델로 작동하여 뛰어난 유연성으로 다양한 표현의 멀티모달 편집 기능을 구현할 수 있다.

 

Main Contributions

1. Method : 멀티모달 마스크 모델링 objective를 사용하여 Vision 작업을 위한 다용도 및 확장 가능한 Foundation Model을 학습하는 프레임워크인 4M을 제안하였다. 이러한 접근 방식을 통해 풍부한 표현을 학습하고 태스크별로 조정할 필요 없이 다양한 태스크에서 우수한 성능을 발휘하는 모델을 만들 수 있다.
   
   
2. Performance : 광범위한 실험과 벤치마크를 통해 저자의 접근 방식의 효율성을 입증하고, 이러한 모델이 많은 주요 Vision 태스크를 즉시 수행할 수 있을 뿐만 아니라 다루지 않았던 다운스트림 태스크를 fine-tuning 할 때 매우 경쟁력 있는 성능을 달성할 수 있다.
   
   
3. Generative Capabilities : 4M을 사용하여 학습된 모델의 유연하고 조정 가능한 생성 기능을 선보이며, 임의의 모달리티에 대한 컨디셔닝을 활용하여 다양한 멀티모달 편집 작업을 수행할 수 있다.
   
   
4. Experimental Study : 4M의 성능에 영향을 미치는 요인을 연구하기 위해 광범위한 Ablation Study를 수행하여 이러한 모델의 동작과 설계에 대한 중요한 인사이트를 제공한다.

Fig. 2. Model Overview.

Method Description

먼저, 4M 아키텍쳐와 objective는 심플하면서도 효율적인 컨셉을 유지하되 모달리티의 타입과 수에 대해 호환가능성과 확장가능성에 초점을 맞추어 설계되었다. 저자는 다음 주요 기능들을 결합하여 설계하였다.

1. Tokenizing modalities
   이미지, 텍스트. 희소 데이터, feature map 등 모든 모달리티를 시퀀스 또는 개별 토큰 셋으로 매핑하여 모달리티별 복잡성을 추상화 한다. 이를 통해 모달리티 간의 모든 가능한 매핑을 하나의 시퀀스 또는 토큰셋을 다른 토큰셋에서 예측하는 것으로 볼 수 있다.
   
   
2. Training a single compatible network on all modalities
   vision, NLP 및 기타 영역의 다양한 태스크에는 전통적으로 매우 다른 모델링 선택, 아키텍터 및 손실함수가 필요했기 때문에 여러 모달리티에 대한 공동 학습이 어려웠다. 저자는 모든 모달리티를 통합 표현 공간(unified representation space)으로 토큰화하면 (병렬 또는 직렬화된 autoregressive) 토큰 예측을 통해 서로 다른 모달리티 간에 매핑할 수 있는 단일 Transformer Encoder-Decoder를 학습하였다.
   
   
3. Multimodal masked pre-training objective
   Transformer는 다양한 태스크에서 데이터와 모델 크기에 따라 뛰어난 확장성을 보여주었으며, 특히 masked reconstruction과 같은 확장 가능한 pre-training objective와 결합할 때 더욱 그렇다.

 

Modalities & data

Pre-training modalities

 

저자는 다양한 모달리티(RGB, captions, depth, surface normals, semantic segmentation maps, bounding boxes, tokenized CLIP feature maps) 셋을 이용하여 4M 모델을 학습하였다. 이러한 모달리티들을 선정한 근거는 다음과 같다.

1. 이들은 Semantic inofrmation(captions, semantic segmentation, bounding boxes, CLIP)과 geometric information(depth, surface normals), RGB를 갖고있다. 이 모달리티들을 입력으로 사용하는 경우, scene geometry와 semantic content에 대한 유익한 사전정보(prior)가 될 수 있다. 출력으로 사용할 경우, 어떤 종류의 표현으로 학습할지 조종할 수 있다.
   
   
2. 정보를 인코딩하는 측면에서 다양한 형식을 갖고있다. : 조밀한 시각적 모달리티(RGB, depth, surface normals, semantic segmentation), 희소 및 시퀀스 기반 모달리티(captions, bounding boxes), 신경망 feature maps(CLIP)
   
   
3. 생성을 목적으로 하는 모델로써 다양하고 충분한 상호작용을 할 수 있다. 예를들어, caption과 segmentation map, bounding box는 semantic 정보를 조건으로 주어 생성을 하도록 하고,  geometric 모달리티들은 3D 정보를 가지고 생성할 수 있도록 한다.
   
   

 

Pseudo labeled multimodal training dataset

 

본 모델을 학습하기 위해서는 위에서 언급한 모든 모달리티 및 태스크를 포함하는 충분히 다양하고 큰 규모의 멀티모달 및 멀티태스크 데이터셋이 필요하다. 하지만 대부분의 멀티모달 데이터셋은 양이 너무 적거나 충분히 다양하지 않다. 이러한 이유로 저자는 공개적으로 사용 가능한 Conceptual Captions 12M(CC12M)을, 강력한 기성용 모델을 사용한 binding dataset으로 사용하는 pseudo labeling 방법을 선택할 수 밖에 없었다고 한다. 이러한 방식은 오직 RGB 이미지에만 접근하면 되므로 거대한 웹 스케일로의 데이터셋 확장이 가능하다.

 

 

Tokenization

모든 모달리티들은 모달리티 특화 토크나이저(modality-specific tokenizers)를 이용하여 개별 토큰들의 시퀀스 혹은 셋으로 구성된다. caption과 bounding box는 모두 텍스트로 처리되며 WordPiece를 이용하여 인코딩 된다. bounding box 모델링을 위해 object detection 태스크를 sequence prediction 태스크로 변환하는 Pix2Seq 방식을 사용했다. RGB, depth, normals, semantic segmentation maps, CLIP feature maps은 이미 학습된 vector quantized autoencoders(VQ-VAE)를 이용하여 토큰화가 이루어진다. RGB로 pre-trained VAE-GAN을 사용하는 모든 이미지 유사 모달리티들을 표현하는 Unified-IO과 달리, 저자는 모달리티 특화 토크나이저를 사용한다. 이를 통해, 기존에 이미지 토크나이저로 표현하기 어려웠던 신경망 feature map을 포함할 수 있게되었다고 한다. 모달리티를 토큰에 매핑하거나 그 반대로 매핑하는 경우(토큰 -> 모달리티) 추론 과정에서 약간의 계산 오버헤드가 발생하지만, 멀티모달 데이터셋을 조립하는 동안 토큰을 미리 계산하여 사전 학습 중에 이러한 오버헤드를 방지할 수 있다.

 

Fig. 3. Chained multimodal generation. partial RGB와 bbox로 부터 MaskGIT를 이용하여 full RGB를 생성하는 예시

Multimodal transformer

저자의 모델은 효율성, 확장성, 단순성을 염두하며 설계되었으며 기존 Transformer의 Encoder-Decoder를 기반으로 하지만 다양한 모달 간의 결합을 위해  몇 가지 수정사항이 있다.

 

Multimodal encoder

 

모달리티 별 각 토큰에 학습가능한 모달리티 임베딩과 1D(for sequences) 혹은 2D(for dense modalities) 사인-코사인 위치 임베딩을 추가한다. 전이학습이 가능하게 하기 위해, Encoder 부분은 추가적으로 학습 가능한 patch-wise linear projection을 사용하여 RGB를 수용 가능하도록 추가로 설계되어 Vision Transformer backbone으로도 두 배가 될 수 있다고 한다.

 

Multimodal decoder

 

decoder는 각 모달리티별 토큰을 다루기 위해, 각기 다른 접근 방식을 다룬다. 이 중 다음 두 가지는 모든 토큰에 반영된다.

1. 모든 모달리티 토큰들은 cross-attention layer에서 자유롭게 아무 encoder 토큰과 연산될 수 있다. 한마디로 모든 encoded information에 접근이 가능하다.
2. 다른 모달리티들의 decoder 토큰과 분리하기 위해 attention mask를 사용한다. 이는 decoder가 동시에 생성되는 다른 출력에 관계없이 각 특별한 모달리티를 위한 일관된 출력을 수행한다는 것을 의미한다. decoder의 역할은 이 mask를 예측하는 것이다. decoder는 다음 토큰을 예측하는 태스크이다. 각 토큰은 앞선 토큰들에서만 영향을 받도록 해야하기에 저자는 표준 autoregressive model로써 self-attention에 causal mask를 적용하였다. 모든 타겟 태스크들은 개별 토큰으로 구성되어있기에 모든 토큰들을 통해 cross-entropy loss를 계산할 수 있으며, 일르 통해 태스크별 loss balancing이 필요없어지고 학습 안정성을 향상시킨다.
   
   

Multimodal masking strategy

멀티모달 사전학습의 경우, 모든 모달리티에서 보이는 토큰/패치의 작은 셋을 샘플링 및 인코딩하고 모델을 학습시켜 cross-modal 예측 코딩을 수행한다는 점에서 MultiMAE와 유사한 사전학습 전략을 사용하였다. 

Input & target masking

 

마스크 이미지 모델링을 수행할 때 마스킹된 토큰을 삭제하고 보이는 토큰의 작은 셋만 인코딩하면 학습 효율이 크게 향상된다고 하며, 여러 모달리티에 대해 학습할 때 매우 중요하다고 저자는 언급한다. 일반적으로 적은 수의 입력 토큰과 훨씬 많은 수의 목표 토큰의 불균형은 적은 수라도 디코더에서 상당한 계산 비용을 유발할 수 있다. 저자는 마스킹된 모든 토큰을 디코딩하지 않고 샘플링된 서브셋만 디코딩하는 타겟 마스킹을 사용할 것을 제안하였다. 랜덤하게 샘플된 입력 및 타겟 토큰들의 수를 고정했을 때, 4M은 낮은 학습 비용을 유지하며 많은 모달리티들의 사전학습을 가능케 한다. MultiMAE와 유사하게 집중 매개변수 $\alpha$가 있는 symmetric Dirichlet distribution을 사용하여 모달리티 당 입력 토큰 수를 샘플링한다. 타겟 토큰에도 마찬가지로 위 방법을 적용한다. 모달리티별 입력 및 타겟 토큰 샘플링 이후, 시퀀스 모달리티에 대해 span masking을 수행한다. 

 

Transfer Experiments

Table 1. Transfer learning study. ImageNet-1K로 transfer leaarning을 하기 위해 우선 ImageNet-21K로 학습을 하였다. 4M은 ImageNet-1K를 제외한 모든 태스크에서 가장 높은 성능을 유지했는데, 다른 모델과 달리 4M은 Data Aug.를 진행하지 않았다는 차이점이 있다.

 

저자는 4M 모델에 대해 4M-B(86M)과 4M-L(303M) 모델을 선보였다(파라미터는 인코더 기준). 이후, 이 모델들을 이용하여 각 다운스트림 태스크에 transfer learning을 수행하였다. 이때, 다운스트림 성능에 큰 영향을 미칠 수 있는 Dataset, Augmentation, Model Architecture, Compute를 더 잘 제어하기 위해 MAE(Masked RGB $\rightarrow$ RGB)와 BEiT-v2(Masked RGB $\rightarrow$ CLIP)와 개념적으로 유사한 자체 baseline을 추가로 보여주었다.

 

Transfer tasks and dataset

• ImageNet-1K classification
• COCO detection and instance segmentation
• ADE20K semantic segmentation
• NYUv2 depth estimation

다른 ViT backbone과의 경쟁력을 위해 저자는 4M 모델에 transfer task-specific heads(e.g. Cascade MAsk R-CNN)을 encoder에 붙히고 디코더를 제거하는 방법을 선택했다. 저자는 비교를 위해 다른 모델들의 논문 속 설정값들을 따르며 4M을 학습하였다고 한다.

 

위 Table 1.에서 4M 모델이 ImageNet-1K를 제외한 모든 태스크에서 가장 높은 성능을 달성했는데, 이는 다양한 Vision  태스크에 대해 4M 모델이 유연하게 사용가능하다는 것을 시사한다. 이외에도 저자는 4M-XL(2.7B)에 대해 실험을 진행하였으나 pseudo labeled CC12M 데이터셋에 과적합이 나타나며 추가적인 실험이 제한되었다고 한다. 따라서, 더 큰 규모의 모델로 학습하기 위해서는 데이터셋 또한 규모가 더 커야한다는 의견을 제시하였으며, 생성 부분에서는 품질이 많이 올라갔으므로 이후 섹션에서는 4M-XL에 대한 실험 결과를 공유한다.

 

Fig. 4. Multimodal generation and editing

Fig. 5. Probing the learned representation through manipulation

 

Generative Capabilities & Probing the Learned Representation

4M은 Fig. 3.에 보이는 바와 같이 반복적인 디코딩을 통해 모든 pre-trained 모달리티의 생성을 위해 사용될 수 있다. 또한 다음 두 특징적인 수용력을 갖고 있다.

1. 4M은 학습에 사용되었던 어떠한 모달리티도 생성할 수 있으며 심지어 cond or uncond 방식으로 어떠한 모달리티든 생성할 수 있다. (Fig. 3. Top)
2. 4M은 Masking 방식으로 학습되었기 때문에 cond 방식으로 인페인팅과 아웃페인팅을 수행할 수 있다. (Fig. 1 and 4)

이러한 수용력을 통해 4M은 멀티모달 에디팅과 멀티모달 weighted guidance 또한 가능하다. (Fig. 4.)

 

Ablations

Masking objective를 사용하며 대규모 모달리티 셋에 대한 사전 학습을 하면 다음과 같은 질문이 제기된다고 한다.

• 어떤 모달리티를 사전학습해야 하는지
• 최적의 마스킹 비율은 얼마인지
• 각 모달리티에서 마스킹할 토큰 수를 어떻게 선택해야 하는지

저자는 위 질문에 대한 답을 하기 위해 Ablation study를 진행했으며 이전 섹션들은 이 Ablation study의 결과로 정해진 수치들을 적용한 결과들이다.

 

저자는 기본적으로 reference setting을 따르고 나머지는 고정된 상태로 유지하면서 모델 성능의 편차를 측정하였다. 성능은 대규모 다운스트림 태스크 셋으로 모델을 전송하고 validation set을 통해 성능을 측정한다. 이때 밴치마크의 목적은 4M의 특정 인스턴스화가 새로운 타겟 태스크뿐만 아니라 unseen input 모달리티로 모두 transfer 하는데 얼마나 우수한지 측정하는 것이다. 이를 위해 RGB를 입력으로하여 새로운 타겟(object detection, semantic segmentation 등)으로 변환하는 등을 수행한다.

 

단순성을 위해 사전학습과 유사하게 각 벤치마크 태스크를 다른 토큰셋에서 예측하는 것으로 모델링한다. 이는 사전학습과 동일한 방식으로 이러한 모달리티와 작업에 대해 tokenizer를 학습함으로써 달성된다. 모든 transfer 시 이미지 해상도는 $224 \times 224$로 수행된다. 모달리티와 태스크 간의 비교를 위해 별개의 metric이 아닌 validaiton set에 대한 cross-entroy 간 성능을 측정한다.

 

Reference model

이전 연구들과 마찬가지로 저자의 모델도 12 encoder와 12 decoder layer로 구성되었다. CC12M 학습 시 $224 \times 224$ 해상도에 $14 \times 14$ 패치 크기를 사용하였으며 cropping이나 color와 같은 augmentation을 적용하지 않았다. 총 학습 길이는 100B 토큰으로 고정하였으며 약 400M masked sample을 사용하였다. 무작위로 샘플된 입력 및 타겟 토큰 수를 각 12개로 설정하고 매개변수 $\alpha = 0.2$인 symmetric Dirichlet distribution을 사용하여 각각 샘플링하였다.

 

input modalities and target tasks

Table 2. Pre-training input and target modalities ablation

Importance of target tasks for representation learning

4M은 멀티모달 및 멀티태스크 학습 체계이며 대상 태스크를 선택하는 것은 모델이 학습하는 표현을 조정하는 강력한 방법이다. ablation을 위해  입력 모달리티를 RGB 또는 모든 모달리티("All")로 고정하고 대상 태스크를 변경한다. 위 Table 2.의 결과에서 사전 학습 설정의 최적 선택은 수행되는 transfer 태스크에 따라 크게 달라지며 모든 transfer에서 최상의 성능을 발휘하는 단일 사전 학습 task는 없다는 것을 보인다. 그러나 모든(all) target 모달리티에 대한 사전 학습은, 사전 학습 중에 어떤 입력 모달리티를 사용했든 평균 손실 측면(Avg. Loss)에서  다른 단일 태스크 및 다중 태스크보다 일관되게 성능이 뛰어나다. 따라서 4M은 향후 타겟 태스크가 불분명할 때 선호되는 모델이 되기에 적합하다.

 

Multimodal masking strategy

Fig. 6. Ablations results

멀티모달 마스킹은 4M의 핵심으로, 본 섹션에서 얼마만큼의 토큰을 샘플링해야하는지 나타낸다.

 

우선 Dirichlet 매개변수 $\alpha$의 경우, 낮으면 샘플링 절차는 대부분의 토큰이 하나의 모달리티에서만 샘플링 되는 경우를 선택한다. 반대로 높으면 대부분의 샘플에는 모든 모달리티에서 동일한 비율로 토큰이 포함된다. Fig. 6. (a)의 결과에서 균일한 샘플링이 평균적으로 가장 우수하지만 큰 차이는 나타나지 않는 모습이다.

 

멀티모달 마스킹 모델링 작업의 난이도는 주로 보이는(마스킹 되지 않은) 입력 토큰의 수에 의해 결정된다. 보이는 토큰의 서브셋만 인코딩하면 학습 효율성을 크게 향상시킬 수 있다. Fig. 6. (b)는 고정된 학습 토큰 예산으로 128-256 입력 토큰으로 학습하는 것이 잘 수행됨을 보여준다.

 

나머지 모든 mask-out 토큰의 서브(랜덤)셋만 디코딩하기로 결정할 수 있으며, 이는 대규모 디코더로 효율적인 멀티모달 학습을 가능하게 하는 데 특히 중요하다. Fig. 6 (c)에서 볼 수 있듯이, 모든 대상의 서브(랜덤)셋만 디코딩하면 (고정된 총 학습 토큰 수에 대해) 성능이 우수하면서도 계산 비용도 절감할 수 있다.

 

How well does 4M scale?

저자는 4M 모델의 확장성을 실험하기 위해 다음 세 가지 측면에 대한 실험을 진행하였다. 실험 결과는 위 Fig. 6.에서 확인할 수 있다. 

• 데이터셋 규모 : CC12M의 다양한 서브셋에서 전체 데이터셋의 $\frac{1}{64}$ 까지 4M 모델에 대해 실험하였다.
• 학습 길이 : 표시되는 토큰의 총 수를 변경하여 실험, 모델이 학습한 입력 토큰과 타겟 토큰의 총 수로 간주되는 토큰을 정의하여 실험하였다.
• 모델 크기 : Tiny(4M-Ti, 24M)부터 Large(4M-L, 705M)까지 실험하였다.