[논문 리뷰] Chameleon: Mixed-Modal Early-Fusion FoundationModels (2024)

이번에 리뷰할 논문은 Chameleon: Mixed-Modal Early-Fusion FoundationModels 입니다.

 

https://arxiv.org/abs/2405.09818

Chameleon: Mixed-Modal Early-Fusion Foundation Models

 

 

 

 

 

 

 

 

 

 

Token-Based Unified Approach

최근 multimodal foundation model은 널리 활용되고 있지만, 여전히 이미지 혹은 텍스트와 같은 서로 다른 modality를 개별적으로 모델링하는 경우가 많습니다. 일반적으로 특정 modality를 위한 전용 encoder 혹은 decoder를 사용하기 때문에, 이러한 방식은 modality 간 정보를 통합하는 능력을 제한시킵니다. 예를 들어, 임의의 순서로 이미지와 텍스트가 혼합된 multimodal documents 를 생성하는 데 한계를 보여줍니다. 

 

이런 한계를 해결하기 위해, 본 논문에서는 Chameleon 을 제안합니다. 

 

 

Chameleon은 기존 모델과 달리 처음부터 혼합 모달(mixed-modal) 모델로 설계되었습니다. 위 figure와 같이 각 modality인 텍스트와 이미지를 따로 처리하는 것이 아닌 균일한 transformer architecture에서 처리할 수 있도록 설계되었으며, 이미지, 텍스트, 코드 등의 교차(interleaved) 데이터를 기반으로 처음부터 end-to-end 방식으로 학습되었습니다. 

 

또한, 위 figure에서 알 수 있듯이 이미지와 텍스트 modality 모두에 대해 완전한 token-based 표현을 사용합니다. 이미지의 경우, discrete token으로 양자화(quantization)하여 텍스트의 단어와 유사한 형태로 변환합니다. 이를 통해 별도의 이미지 및 텍스트 인코더 혹은 도메인 별 디코더 없이 동일한 transformer architecture를 활용할 수 있습니다. 

 

Sample generation form Chameleon

 

이와 같이 초기 단계에서 모든 modality를 공유된 표현 공간으로 mapping하는 Early-Fusion 방식을 채택함으로써 modality 간 원할한 추론 및 생성이 가능해졌습니다. 하지만, 이러한 접근 방식은 optimization stability 와 model scaling 측면에서 큰 도전 과제를 동반하게 됩니다. 

 

 

 

 

 

 

Technical Contributions

서로 다른 modality인 이미지와 텍스트를 처음부터 한번에 균일한 architecture에서 처리할 수 있도록 함으로써 modality 간 정보를 통합하는 측면에서는 유리하지만, 학습을 할 때 optimization stability 와 model scaling 문제가 발생합니다. 이러한 문제를 해결하기 위해 Architecture innovations 와 New training techniques 를 도입합니다. 

 

• Transformer architecture의 새로운 수정 사항 적용
  • Query-Key Normalization (QK-Norm)
  • revised placement of layer norms
  • 이를 통해, mixed-modal setting에서 stable training이 가능해짐.

 

 

 

 

 

 

Performance Evaluation

Chameleon 모델은 다양한 작업에서 우수한 성능을 보여주었으며, 광범위한 능력을 갖춘 모델임을 입증하였습니다. 

 

• Visual Question Answering 및 Image captioning 벤치마크에서 SOTA 성능 달성
  • 기존 모델(Flamingo, IDEFICS, Llava-1.5) 보다 우수한 성능 기록
• Text 전용 벤치마크에서도 경쟁력 유지
  • Mixtral 8x7B 및 Gemini-Pro 와 Common Sense Reasoning 및 Reading Comprehension 에서 비슷한 성능 보임.

 

하지만 기존 벤치마크는 정적인 공용 데이터셋을 사용하기 때문에 모델 성능을 평가하는 데 한계가 있을 수 있으며, 이에 저자는 개방형 프롬프트에 대해 Mixed-modal Long Form 응답을 평가하는 새로운 인간 평가(Human Evaluation) 실험을 설계하였습니다. 

• Gemini-Pro 대비 60.4% 선호
• GPT-4V 대비 51.6% 선호

 

 

 

 

 

 

Main Contributions

1. Chameleon 모델 제안
   • Token 기반의 Early-Fusion 혼합 모델로, 텍스트-이미지가 교차된 문서(interleaved image-text documents)를 생성하고 추론할 수 있는 능력을 갖춤.
   • Open Multimodal Foundation Models 에서 새로운 기준 제시
2. Architectural innovation 와 training techniques 제시
   • Early-Fusion 토큰 기반 모델의 stable과 scalable training을 가능하게 함.
   • Mixed-modal setting에서 발생하는 주요 문제를 해결하는 정규화 및 학습 기법 개발
3. Extensive Evaluation
   • 다양한 Vision-Language 벤치마크에서 SOTA 성능 달성
   • Text-only task에서도 경쟁력을 유지하며, 고품질 이미지 생성까지 가능
4. Large-scale Human Evaluation
   • mixed-modal reasoning 과 generation에 대한 human evaluation 을 최초로 수행
   • Chameleon이 Gemini-Pro 및 GPT-4V보다 높은 선호도 기록

 

 

 

 

지금까지, 논문에 대한 전반적인 내용을 알아보았으며,

이제부터, 본 논문에서 제안한 기술 및 방법론들에 대해 자세하게 알아보겠습니다.

 

 

 

 

Pre-Training

Chameleon 모델은 텍스트뿐만 아니라 이미지를 discrete token의 sequence로 변환하여 표현하며, autoregressive transformer의 scaling한 특성을 활용합니다. 본 논문에서는 training을 위해 이미지와 텍스트의 배치를 다양하게 다음과 같이 구성합니다.

• 텍스트 전용 데이터 (text-only)
• 단일 텍스트-이미지 쌍 (single text/image pairs)
• 완전히 교차된 텍스트-이미지 문서 (full interleaved text-image documents)

 

 

 

Tokenization

1. Image Tokenization

본 논문에서는 Make-A-Scene 을 기반으로, 새로운 Image tokenization을 학습하였습니다. 

• 512 x 512 픽셀 크기의 이미지를 1024개의 discrete token으로 변환하도록 설계
• Codebook 크기 : 8192
• Training data : 라이선스가 있는 이미지만 사용하여 학습 진행
• 사람 얼굴 생성의 중요성 고려하여, 사람 얼굴이 포함된 이미지의 비율을 2배로 증가시켜 학습 (up-sampling)
• 이미지 내 텍스트가 많은 경우 reconstruction 품질이 낮아지는 한계 존재 -> OCR 관련 작업에서는 모델 성능이 제한될 가능성 존재

 

2. Text Tokenization

• BPE(Byte Pair Encoding) 토크나이저를 새롭게 학습
• Vocabulary Size : 65,536
  • 8192 개의 image codebook tokens 도 포함
• 토크나이저 구현 : SentencePiece 라이브러리 사용

 

 

 

Pre-Training Data

본 논문에서는 pre-training을 two separate stages로 나누어서 진행합니다. First Stage에서는 전체 훈련 과정의 80%를 차지하며, Second Stage에서는 마지막 20%를 차지합니다. 또한 First Stage에서 사용한 데이터의 가중치를 50% 낮추고, 더 높은 품질의 데이터셋을 추가로 포함하여 Second Stage를 진행하게 됩니다. 

 

또한, 텍스트와 이미지 간의 관계를 더 잘 학습할 수 있도록 하기 위해, 모든 텍스트-이미지 쌍에서 이미지와 텍스트의 순서를 50% 확률로 랜덤하게 변경합니다. 즉, 50% 의 확률로 Image captioning 방식과 같이 이미지가 먼저 등장하고 텍스트가 등장하도록 조정합니다. 

 

 

1. First Stage of Pre-Training

First Stage 에서는 아래와 같은 대규모 unsupervised datasets를 조합하여 학습을 진행합니다. 

 

• 텍스트 전용 데이터 (Text-Only Data)
  • LLaMA-2 및 CodeLLaMA의 사전 학습 데이터를 조합하여 사용
  • 총 2.9조 (2.9 trillion) 개의 텍스트 전용 토큰 사용
• 텍스트-이미지 데이터 (Text-Image Data)
  • Publicly available data sources 와 licensed data를 조합하여 학습
  • 이미지의 경우 512 x 512 픽셀 크기로 resizing 및 center cropping 적용 후 토큰화
  • 총 14억 (1.4 billion) 개의 텍스트-이미지 쌍 사용하여 총 1.5조 (1.5 trillion) 개의 텍스트-이미지 토큰 생성
• 교차 텍스트-이미지 데이터 (Text / Image Interleaved Data)
  • publicly available web sources 에서 수집 (Meta 제품 혹은 서비스에서 수집한 데이터 포함 X)
  • 총 4000억 (400 billion) 개의 교차(interleaved) 텍스트-이미지 토큰 사용
  • 텍스트-이미지 데이터에서 적용한 동일한 이미지 필터링 기법을 사용하여 학습 데이터 정제

 

 

2. Second Stage of Pre-Training

• First Stage에서 사용한 데이터의 가중치를 50% 낮춤
• 더 높은 품질의 데이터셋을 추가로 혼합하여 학습 진행
• 이미지-텍스트 토큰 비율을 유지하면서, 더 정제된 고품질 데이터를 포함
• 다양한 지도 학습 기반 Instruction tuning 데이터셋에서 필터링된 subset 포함

 

 

 

Stability

Chameleon 모델의 크기를 80억(8B) 개 이상의 파라미터 및 1조(1T) 개 이상의 토큰으로 확장시켰을 때, 학습의 안정성을 유지하는 것이 어려웠다고 합니다. 특히 학습 후반부에 불안정성(instabilities)이 늦게 나타나는 경우가 많아 본 논문에서는 Architecture 및 optimization 기법을 조정하여 학습을 안정화하도록 하였습니다. 

 

 

Architecture

우선 본 논문에서 사용한 architecture의 base setting은 다음과 같습니다. 

• model : LLaMA-2
• Normalization : RMSNorm
• Activation Function : SwiGLU
• Positional Embeddings : RoPE

 

하지만 기존 LLaMA 아키텍처를 사용하였을 때 학습 후반부에서 복잡한 divergence 문제가 발생한다고 합니다. 이에 대해 본 논문에서는 Softmax 함수의 변환 불변성(translation invariant) 특성에 의해 발생한다고 주장합니다.

 

$$ softmax(z) = softmax(z + c) $$

 

Softmax 함수의 경우 상대적인 확률값을 추출합니다. 즉, Chameleon 모델의 입력으로 들어가는 text 혹은 image의 norm이 커지게 되어도 softmax의 상대적인 확률값에는 차이가 없습니다. 이와 같은 이유로 text의 norm이 커지게 되면, 최적화를 위해 image의 norm도 같이 커지게 되며, 서로 경쟁하며 norm의 크기가 서로 커지다가 divergence되는 문제가 발생하게 됩니다. 

 

 

그렇다면 왜 text와 image는 서로 경쟁을 할까요?

 

 

text와 image는 서로 다른 modality 입니다. 즉, 서로 다른 entropy의 분포를 가지지만, 하나의 network로 학습하기 때문에, 각 modality가 학습 도중 경쟁하여 자체적으로 norm을 증가시키려는 경향이 생기게 됩니다. 

 

이는 학습 초반에는 크게 문제가 되지 않지만, 훈련이 진행됨에 따라 bf16(Brain Float 16)의 표현 범위를 넘어가는 시점에서 학습이 발산하는 문제가 발생합니다. 

 

 

실제로 위 figure를 보시면 image generation을 제거하였을 때, divergence가 발생하지 않음을 확인할 수 있습니다. 

 

 

위 figure를 보시면, 기존 연구 중 Logit Drift Problem 을 확인할 수 있습니다. figure의 각 수치는 Transformer의 마지막 layer vector의 norm 값이며, 점차 증가하는 현상을 보여줍니다. 이와 같이 Norm이 통제되지 않고 증가하면, loss divergence로 이어질 가능성이 높아집니다. 

 

위와 같은 최적화 문제를 해결하기 위해 저자는 다음 3가지의 solution을 제안합니다. 

 

 

 

Solution 1 : Query-Key Normalization (QK-Norm)

Softmax 함수는 Transformer에서 두 가지 주요 위치에 사용됩니다. 

1. Attention 연산 내부
2. 출력 logit 에 적용

Softmax 함수는 divergence가 발생하는 가장 주된 이유이며, 이를 해결하기 위해 LLaMA 아키텍처에서 벗어나 "Query-Key Normalization (QK-Norm)" 을 적용합니다. QK-Norm의 작동 방식은 Attention의 Query와 Key 벡터에 Layer Norm을 적용하여 Softmax 입력의 Norm을 직접 제어하는 것입니다. 

 

 

QK-Norm을 적용함으로써 Softmax가 modality간 entropy 차이로 인해 불안정해지는 현상을 방지합니다. 만약, QK-Norm을 적용하지 않으면, 위 figure와 같이, 약 20%의 학습 진행 시점에서 loss divergence 하게 됩니다. 

 

 

 

Solution 2 : Dropout 적용

Chameleon-7B 모델의 경우, Attention 및 Feed-forward layer 뒤에 Dropout을 추가하여 Norm 성장을 통제하는 것이 가능했습니다. 

 

 

하지만, 파라미터의 크기가 더 큰 Chameleon-34B 모델 학습의 안정화에 Dropout 적용만으로는 충분하지 않았으며, 추가적인 정규화 순서 변경이 필요했다고 합니다. 

 

 

 

Solution 3 : Normalization Reordering (정규화 순서 변경)

본 논문에서는 기존 LLaMA-2의 정규화 방식과 달리, Chameleon-34B 에서는 Normalization Reordering을 통해 학습의 안정성을 확보하였습니다. 우선, 사용한 정규화 기법은 Swin Transformer의 정규화 전략을 사용하였다고 합니다. 

2023.07.28 - [논문 리뷰/Image classification] - [논문 리뷰] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows(2021)

[논문 리뷰] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows(2021)

 

Swin Transformer의 정규화 전략은 Feedforward 블록의 norm 성장을 제한하는 것입니다. 하지만 SwiGLU activation function의 곱셈적 특성(Multiplicative Nature) 때문에, Feedforward block에서 Norm 성장이 더 심해질 수 있습니다. 

 

$$ SwiGLU(x) = Swish(W_{1}x) \odot W_{2}x $$

 

위 식과 같이 SwiGLU는 $ Swish(W_{1}x) $ 와 $ W_{2}x $ 를 곱셈으로 연결하여 비선형성을 더 강하게 적용합니다. 하지만 이러한 곱셈 연산으로 인해 Norm이 더 빠르게 증가할 가능성이 높아집니다. 즉, 정규화 없이 그대로 적용하면 모델 학습이 불안정해질 가능성이 높아지게 됩니다. 

 

기존 LLaMA-2의 정규화 방식을 보시면, 

 

$$ h = x + attention(attention_norm(x)) $$

$$ output = h + feed_forward(ffn_norm(h)) $$

 

와 같은 순서로 적용됩니다. 

 

즉, Feedforward block에 SwiGLU 활성화 함수가 적용 후 정규화가 수행되는데, SwiGLU 곱셈적 특성으로 인해 FFN Block의 출력 Norm이 비정상적으로 커질 가능성이 있습니다. 결과적으로 학습 후반부에서 Norm이 폭발하는 Gradient Explosion이 발생하여 divergence가 발생할 위험이 증가하게 되는 것이죠. 

 

이를 해결하고자, Chameleon 은 정규화 순서를 다음과 같이 변경하였습니다. 

 

$$ h = x + attention_norm(attention(x)) $$

$$ output = h + ffn_norm(feed_forward(h)) $$

 

위 식과 같이, attention 연산 후에 정규화를 적용하여 attention 출력 벡터의 norm을 먼저 조정하여 불안정성을 줄였습니다. 또한, feed-forward 연산 후 정규화를 적용하였습니다. SwiGLU가 적용된 후 Norm이 급격히 커지는 현상을 방지하기 위해, feed-forward 출력에 대해 바로 정규화를 수행한 것이죠. 즉, SwiGLU로 인해 Norm이 커지는 문제를 정규화 단계에서 해결한 것입니다. 

 

 

위 figure는 Chameleon-34B 모델에 대해 norm reordering을 사용했을 때와 사용하지 않았을 때를 서로 비교한 것입니다. 그리과 같이 사용하였을 때 divergence 가 방지되는 것을 확인할 수 있습니다. 

 

또한 Chameleon-34B 모델의 경우, 2번째와 3번째 solution인 dropout과 norm reordering을 함께 사용하였을 때 오히려 학습 불안정성을 초래한다고 합니다. 이에 본 논문에서는 Chameleon-34B 모델에 대해서는 Dropout을 제거하였다고 합니다. 이후 실험을 통해, Chameleon-7B 역시 norm reordering을 한 경우 Dropout 없이도 안정적으로 학습이 가능함을 확인하였습니다. 

 

하지만 QK-Norm 은 Chameleon-7B 및 Chameleon-34B 모두에서 필수적인 요소라고 합니다.

 

 

 

 

 

Optimization

훈련 과정에서는 다음과 같이 optimization setting을 하였습니다. 

 

• AdamW optimizer 사용
  • $ \beta 1 = 0.9 $ 
  • $ \beta 2 = 0.95 $ 
  • $ \epsilon = 10^{-5} $
• Learning Rate Schedule
  • 처음 4000 step 동안 Linear Warm-up 적용
  • 이후, learning rate를 0으로 수렴시키는 Exponential Decay 방식 적용
• Additional Optimization method
  • Weight Decay : 0.1
  • Global Gradient Clipping : threshold (1.0)
• Dropout 사용 여부
  • Chameleon-7B : 학습 안정성을 위해 dropout 0.1 적용
  • Chameleon-34B : dropout 사용 X 

 

 

 

QK-Norm & z-loss regularization

QK-Norm은 Transformer 내부 Softmax의 안정성을 높이는 역할을 하지만, 최종 Softmax에서 발생하는 Logit Shift 문제를 해결하진 못했습니다. 이를 해결하기 위해, z-loss 정규화 ( Chowdhery et al. (2022) , Wortsman et al. (2023) ) 를 적용합니다.

 

z-loss 정규화 방식은 아주 간답합니다.

Softmax의 분할 함수(partition Function) 인 $ Z $를 정규화 하는 방식으로, $ Z $는 다음과 같습니다.

 

$$ \sigma(x)_i = \frac{e^{x_i}}{Z}, \quad Z = \sum_i e^{x_i} $$

 

그 후, 다음 식과 같이 손실 함수에 다음과 같은 항을 추가하여 정규화합니다.

 

$$ 10^{-5} \log_2 Z $$

 

이를 통해, Logit 값이 불필요하게 커지는 현상을 방지하여 학습 안정성을 높입니다. Softmax가 계산될 때, 모든 Logit 의 합이 정규화되기 위한 스케일링 요소로 작용하기 때문에 Logit Shift 문제를 방지할 수 있게 되는 것입니다.

 

 

 

Training Details

• Batch Size
  • Chameleon-7B : 전역 배치 크기 $ 2^{23} $ (약 8M 토큰)
  • Chameleon-34B : 전역 배치 크기 $ 3 \times 2^{22} $ (약 12M 토큰)
• 총 학습 데이터 규모
  • 전체 데이터셋을 2.1 epoch 동안 학습
  • 총 9.2조 (9.2 trillion) 개의 토큰을 학습
• 학습 진행률
  • Chameleon-7B : 첫 600,000 step에서 전체 학습의 약 55% 진행됨
  • Chameleon-34B : 첫 600,000 step에서 전체 학습의 약 80% 진행됨