NL-352, Mixture-of-Transformers: A Sparse and Scalable Architecture for Multi-Modal Foundation Models, TMLR 2025

◼ Comment

  • 이 논문은 멀티모달 처리할때 새로운 구조를 제안한거다. 
  • BAGEL, Cosmos v3가 대표적으로 MoT 구조이다.
  • 요약
    • dense 모델로 멀티모달 처리하면 서로 충돌하는 경우가 있다 (제 연구랑 비슷한 finding)
    • 실제 다른 스페이스에 모달리티가 저장되더라 (제 연구랑 비슷한 finding)
    • learned router 기반 MoE가 expert utilization imbalance, load balancing, routing instability 문제를 가질 수 있고, 멀티모달 setting에서는 데이터 분포 차이 때문에 특히 speech validation에서 불안정할 수 있다.
    • 즉 “모달리티별 router”라기보다는 learned expert routing이 멀티모달 분포에서 불안정함 (네모트론 speechx랑 비슷한 finding)
    • 뭐 위 이유 및 이런저런 이유로 MoT 가는데 discrete token만써서 학습할때도 효율적/성능 좋고, diffusion 붙여서 학습하는 환경에서도 효율적/성능 좋다.
  • 구조는 사실 특별한건 없고 
    • FFN, Norm, expert-like module은 모달리티별로 따로 존재한다고 보면 되고, 모달리티 간 직접적인 token-level 정보 공유는 global self-attention에서 일어난다고 이해하면 된다.
    • 즉 그냥 Transformer가 모달리티별로 존재하는데, attention만 같이 하는 느낌?
  • 실험을 상당히 많이 했는데, MoT에서 텍스트 모달리티 블럭만 dense -> moe로 바꿔도 효율 좋다고 함
    • 즉 하이브리드 형식도 시도
  • 근데 어쨌거나 핵심은 모달리티별로 분리하는것
    • attention 빼고 다 분리할지 일부 분리할지 ablation 했음
    • FFN 분리가 가장 큰 이득을 주고, Q/K/V attention projection 분리를 추가하면 더 좋아진다. 
    • LayerNorm 분리는 추가 효과가 거의 작다. 
    • Full MoT가 전체적으로 가장 완전한 구조지만, 성능 향상의 핵심은 FFN + attention projection 분리다.
  • 아쉬운건 아키텍쳐 비교를 할때 벤치마크 성능이 아닌 loss matching과 validation loss 기반임
    • loss matching: Dense 모델이 어떤 loss에 도달했을 때, MoT나 MoE는 그 loss에 몇 step 만에 도달하는가를 비교하는 방법

Abstract

대규모 언어 모델(LLM)의 발전은 텍스트, 이미지, 음성을 하나의 통합된 프레임워크 안에서 처리할 수 있는 멀티모달 시스템으로 확장되었다. 이러한 모델을 학습하려면 텍스트 전용 LLM에 비해 훨씬 더 큰 데이터셋과 계산 자원이 필요하다. 이러한 스케일링 문제를 해결하기 위해, 본 논문은 사전학습 계산 비용을 크게 줄이는 희소 멀티모달 트랜스포머 아키텍처인 Mixture-of-Transformers (MoT)를 제안한다.

MoT는 피드포워드 네트워크, 어텐션 행렬, 레이어 정규화 등을 포함한 모델의 비임베딩 파라미터를 모달리티별로 분리한다. 이를 통해 전체 입력 시퀀스에 대해서는 전역 self-attention을 유지하면서도, 각 모달리티에 특화된 처리를 가능하게 한다.

저자들은 다양한 설정과 모델 규모에서 MoT를 평가한다. Chameleon 7B 설정에서는 텍스트와 이미지 생성을 autoregressive 방식으로 수행하는데, MoT는 dense baseline 성능과 비슷한 성능을 달성하면서도 FLOPs는 **55.8%**만 사용한다. 여기에 음성을 추가한 설정에서는, MoT가 dense baseline과 비슷한 음성 성능을 달성하면서도 FLOPs는 **37.2%**만 사용한다.

또한 Transfusion 설정에서는 텍스트와 이미지를 서로 다른 목적 함수로 학습한다. 이때 7B MoT 모델은 dense baseline의 약 3분의 1 FLOPs만으로 이미지 모달리티 성능을 맞추며, 760M MoT 모델은 주요 이미지 생성 지표에서 1.4B dense baseline보다 더 좋은 성능을 보인다. 시스템 프로파일링 결과에서도 MoT는 dense baseline 수준의 이미지 품질을 47.2%의 wall-clock time으로, 텍스트 품질을 75.6%의 wall-clock time으로 달성하여 실제 학습 시간 측면에서도 효율적임을 보인다.

1 Introduction

파운데이션 모델의 발전은 텍스트, 이미지, 음성 등 다양한 데이터 유형을 하나의 통합된 프레임워크 안에서 처리할 수 있는 멀티모달 대규모 언어 모델(LLM)로 확장되어 왔다. Chameleon (Chameleon Team, 2024)과 같은 최근의 발전은 early-fusion 방식의 혼합 모달 모델이 하나의 아키텍처 안에서 다양한 미디어 유형을 생성할 수 있는 가능성을 보여준다. 이러한 모델은 콘텐츠 생성이나 교차 모달 번역과 같은 응용 분야를 발전시킬 가능성이 있지만, 여러 모달리티의 표현을 동시에 학습해야 하기 때문에 상당한 계산적 어려움을 야기한다.

  • early / late 개념을 좀 보니까 비전 인코더에서 feature 뽑고 adaptor 붙여서 LLM이랑 같이 학습하는 경우. 즉 continuous featrue 쓴다고 하면 late fusion이라 보는거 같음
  • 근데 discrete token 뽑아서 학습하는 구조면 early fusion이라 보는거 같음

Early-fusion 멀티모달 LLM을 학습하려면 단일 모달리티 모델에 비해 훨씬 더 큰 데이터셋과 계산 자원이 필요하다. 예를 들어, Chameleon (Chameleon Team, 2024)은 텍스트 성능 측면에서 LLaMA2 (Touvron et al., 2023b)에 맞먹기 위해 이미지 토큰을 포함하여 9.2조 개의 학습 토큰으로 학습되었는데, LLaMA2는 2조 개의 텍스트 학습 토큰으로 학습되었다. 

각 모달리티는 고유한 최적화 문제를 도입하며, 이는 하나의 통합 모델 안에서 동시에 해결되어야 한다. 경험적으로, 이러한 모달리티들은 dense transformer 모델 안에서 서로 충돌하는 학습 동역학을 보이는 경우가 많으며(Figure 15), 이는 최적화를 복잡하게 만들고 계산 부담을 증가시킨다. 또한 입력을 모달리티별 사전 지식 없이 동일한 토큰처럼 처리하더라도, 서로 다른 모달리티는 feature space에서 서로 다른 영역을 차지한다(Figure 2, Appendix Figure 23). 이는 모달리티가 처리되는 방식에 본질적인 차이가 있음을 보여준다.

  • 내가 경험적으로 얻은 것과 비슷한 결과이네?
  • 같은 모델에서 학습하면 이게 서로 충동하는 역할이 있고, 다른 스페이스에 다르게 존재하는 성향을 보여준다는 것

이러한 스케일링 문제를 해결하기 위한 유망한 접근법은 Mixture of Experts (MoE)와 같은 모델 sparsity이다. 이는 각 입력에 대해 모델 구성 요소의 일부만 활성화함으로써 전체 계산량을 줄이는 방식으로 스케일링을 가능하게 한다. MoE에서는 각 transformer layer의 learned router가 여러 MLP 중 하나를 희소하게 활성화하여, 서로 다른 expert가 데이터의 서로 다른 측면에 집중할 수 있도록 한다(Jacobs et al., 1991; Eigen et al., 2013; Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022; Jiang et al., 2024; Sukhbaatar et al., 2024). 그러나 MoE는 여러 문제를 도입한다. learned router는 종종 expert 사용의 불균형을 초래하며, 이 때문에 학습 중 추가적인 load-balancing 기법이 필요하다. 또한 MoE의 bi-level optimization 특성은 학습 동역학을 복잡하게 만들며, 모델 크기가 커질수록 학습이 불안정해질 수 있다. 이러한 MoE의 문제를 해결하는 것은 여전히 열린 연구 문제로 남아 있다.

멀티모달 맥락에서, 이전 연구들(Bao et al., 2022b; Wang et al., 2022; Shen et al., 2023; Lin et al., 2024)은 transformer의 MoE layer에 modality-aware sparsity를 도입하거나, post-training 과정에서 LLM backbone 위에 모달리티별 모듈을 추가로 fine-tuning하였다(Wang et al., 2023; He et al., 2024). 이러한 접근법은 유망한 결과를 보였으며, 모달리티에 따른 단순한 rule-based routing이 MoE에서 흔히 사용되는 learned routing보다 더 좋은 성능을 낼 수 있음을 시사한다. 이러한 성공은 학습 동역학이 더 안정적이기 때문일 수 있다. 즉, 학습 초기 단계에서 expert와 router가 모두 충분히 학습되지 않아 발생하는 불안정성을 피할 수 있기 때문이다.

  • 그냥 MoE에 멀티모달 학습하면 불안정한 것도 있나 보군

이러한 통찰에서 영감을 받아, 본 논문은 모든 비임베딩 transformer 파라미터에 대해 modality-aware sparsity를 도입하는 희소 멀티모달 transformer 아키텍처인 Mixture-of-Transformers (MoT)를 제안한다(Figure 2a). 

기존 접근법과 달리, MoT는 특정 layer나 모듈에만 modality-aware sparsity를 적용하는 것이 아니라 transformer 전체에 적용한다. MoT는 텍스트, 이미지, 음성 등이 섞인 interleaved multi-modal sequence를 입력으로 받고, 각 토큰에 대해 feed-forward network, attention projection matrix, layer normalization을 포함한 서로 다른 모달리티별 파라미터를 동적으로 적용한다. 따라서 MoT 설계는 dense transformer counterpart와 정확히 동일한 계산 구조와 FLOP 수를 갖는 sparse model을 만든다.

  • 이해하기론 FFN, Norm, expert-like module은 모달리티별로 따로 존재한다고 보면 되고, 모달리티 간 직접적인 token-level 정보 공유는 global self-attention에서 일어난다고 이해하면 된다.
  • 즉 그냥 Transformer가 모달리티별로 존재하는데, attention만 같이 하는 느낌?

저자들은 다양한 멀티모달 설정에서 7B 모델 3개를 포함해 총 13개의 모델을 scratch부터 사전학습하여 MoT를 평가했다. 이 포괄적인 설정을 통해, 점점 더 복잡한 학습 목적 함수와 모달리티를 도입하는 여러 실험 구성에서 MoT의 성능을 평가할 수 있었다. 구체적으로, MoT의 적응성과 효율성 향상을 평가하기 위해 다음과 같은 멀티모달 시나리오에서 실험을 수행했다.

  1. 텍스트와 이미지 모두에 대한 autoregressive objectives (Chameleon).
    Chameleon 설정(Chameleon Team, 2024)에서, 저자들의 7B MoT는 여러 데이터 분포에서 평가했을 때 7B dense baseline의 성능과 비슷한 성능을 달성하면서도 FLOPs는 **55.8%**만 사용했다(Figure 5). 이러한 결과는 다른 여러 모델 규모(37M, 94M, 443M, 1.5B)에서도 일관되게 나타났다(Figure 6, Appendix Figure 24).

  2. 세 번째 모달리티로 음성 도입 (Chameleon: Text+Image+Speech).
    Chameleon 설정에 세 번째 모달리티로 discrete speech token을 추가했을 때, MoT는 모든 모달리티에서 유사한 성능을 달성했으며, 음성 모달리티의 경우 더 적은 학습 FLOPs, 즉 **37.2%**만 필요했다(Figure 8). 이러한 결과 역시 다른 여러 모델 규모에서도 일관되게 나타났다(Figure 8, Appendix Figure 25).

  3. 텍스트에는 autoregressive objective, 이미지에는 diffusion-based objective 적용 (Transfusion).
    Transfusion 설정에서는 텍스트와 이미지를 서로 다른 목적 함수로 학습한다. 텍스트는 autoregressive 방식으로, 이미지는 diffusion 기반 방식으로 학습된다. 이 설정에서, 학습 및 추론 FLOPs가 1.4B dense baseline인 Transfusion의 절반인 760M MoT 모델은 이미지 생성의 CLIP score와 FID score, 이미지 캡셔닝의 CIDEr score, 이미지 모달리티 학습 loss 등 여러 지표에서 dense model을 능가했다(Figure 11). 또한 7B MoT 모델은 이미지 생성의 diffusion validation loss와 이미지 캡셔닝의 CIDEr score에서 dense baseline의 3분의 1보다 적은 FLOPs로 이미지 성능을 맞추었다(Figure 10). 추가로 Transfusion 설정의 세 가지 모델 규모(163M, 760M, 1.4B) 전반에서 MoT는 이미지 모달리티에서 일관되게 큰 속도 향상을 보였고, dense model을 큰 차이로 능가했다(Figure 12).

MoT에 대해 더 깊고 포괄적인 평가를 제공하기 위해, 저자들은 여러 차원에서 MoT의 장점을 검증하는 추가 실험을 수행했다. 이 실험들은 MoT의 계산 효율성, wall-clock time 감소, 그리고 다른 sparse architecture와 비교한 효과를 평가했다.

  1. Wall-Clock Time Comparison
    또한 시스템 프로파일링(AWS p4de.24xlarge 인스턴스와 NVIDIA A100 Tensor Core GPU 사용)은 MoT의 효율성이 실제 wall-clock training time의 큰 감소로 이어진다는 것을 보여주었다. 저자들의 7B MoT는 7B dense model의 이미지 성능을 단 **47.2%**의 시간 안에 맞추었고, 텍스트 성능은 **75.6%**의 시간 안에 맞추었다(Figure 19).

  2. Mixture-of-Experts와 MoT 비교
    MoT의 관찰된 성능 향상이 단순히 추가적인 sparse parameter 때문이 아니라는 것을 검증하기 위해, 저자들은 실험 전반에 걸쳐 4-expert mixture-of-expert model (MoE-4x)을 추가 baseline으로 포함했다. MoE-4x는 모든 실험 설정에서 MoT보다 더 많은 sparse parameter를 포함하지만, 특히 비텍스트 모달리티인 이미지와 음성에서 MoT보다 일관되게 낮은 성능을 보였다. MoT가 MoE-4x보다 갖는 장점은 wall-clock time 기준으로 측정했을 때 더욱 크게 나타났다(Figure 19).

  3. Combining the Best of Both Worlds—Mixing Heterogeneous Transformers
    초기 proof of concept로, 저자들은 MoT 프레임워크 안에서 sparse transformer를 통합하는 hybrid 접근을 탐색했다. 구체적으로, MoT의 텍스트 transformer에는 MoE-4x 아키텍처를 적용하고, 이미지 task에는 기존 MoT 아키텍처를 유지했다. 예비 결과는 이 조합이 Chameleon과 Transfusion 설정 모두에서 이미지 생성 품질을 해치지 않으면서 텍스트 모달리티 성능을 더 향상시킬 수 있음을 보여준다(Figure 16, Figure 17).

2 Method: Mixture-of-Transformers Architecture

2.1 Background: Foundation Models for Multi-Modal Generation

대규모 언어 모델의 최근 발전은 텍스트를 넘어서는 모달리티로 확장되어 왔다. 핵심적인 접근법 중 하나는 이미지나 음성과 같은 비텍스트 데이터를 discrete token sequence로 토큰화한 뒤, 이를 텍스트 기반 모델과 유사하게 autoregressive sequence modeling에 적용하는 것이다(Figure 2a). 예를 들어, Chameleon (Chameleon Team, 2024)은 사전학습된 이미지 토크나이저(Gafni et al., 2022)를 사용하여 이미지를 1,024개의 discrete token으로 토큰화함으로써, 텍스트와 이미지에 대한 통합 학습을 가능하게 한다. 유사한 방법은 음성에도 적용되어 왔다(Nguyen et al., 2024). Transfusion (Zhou et al., 2024)과 같은 대안적 접근법은 이미지와 같은 연속 모달리티의 생성을 개선하기 위해 continuous image token과 diffusion-based training objective를 사용한다(Section 3.4).

멀티모달 파운데이션 모델의 내부 표현을 조사하기 위해, 저자들은 이 모델들의 feature space를 분석했다. 그 결과, 여러 layer에 걸쳐 모달리티별로, 즉 텍스트, 음성, 이미지별로 클러스터가 형성되는 것을 확인했다(Figure 2b, Appendix Figure 23). Principal Component Analysis (PCA)는 입력이 모달리티별 사전 지식 없이 동일한 discrete token처럼 처리됨에도 불구하고, feature space 안에서 서로 다른 모달리티가 서로 다른 영역을 차지한다는 것을 보여준다. 이러한 자연스러운 클러스터링은 모달리티 처리 방식에 본질적인 차이가 있음을 시사하며, 이후 저자들의 접근법에 중요한 근거를 제공한다.

  • 여기서 한가지 궁금했던것은, 따로 space가 존재하는게 굳이 나쁜거라 볼 수 있나? 라는 생각.
  • 논문에서는 이게 나쁘다고 하는것은 아니고, 이렇게 따로 존재할거면 MoT 방법이 낫다고 말하는 것에 가까운듯

2.2 Mixture-of-Transformers Architecture: Modality-Specific Parameter Decoupling

우리는 멀티모달 사전학습을 가속하면서 계산 비용을 줄이기 위해 설계된 새로운 아키텍처인 Mixture-of-Transformers (MoT)를 제시한다. MoT는 표준 transformer 아키텍처를 확장하여, feed-forward network, attention matrix, layer normalization을 포함한 모든 비임베딩 모델 파라미터에 대해 모달리티별 가중치를 도입한다. 이 접근법은 모델이 서로 다른 모달리티를 더 효율적으로 처리하도록 하면서도, cross-modal interaction을 학습할 수 있는 능력을 유지한다.

입력 토큰 시퀀스를 (x = (x_1, \ldots, x_n))라고 하자. 여기서 각 (x_i)는 어떤 모달리티 (m_i \in {\text{text}, \text{image}, \text{speech}})에 속한다. 일반적인 transformer layer는 다음과 같이 표현할 수 있다.

우리가 제안하는 MoT에서는 global self-attention을 유지하면서도, 파라미터를 모달리티별로 분리한다.

global self-attention 메커니즘은 모달리티별 파라미터 분리에도 불구하고 모든 모달리티에 걸쳐 작동하며, cross-modal relationship을 포착한다.

여기서 (W_Q^{m_i}), (W_K^{m_i}), (W_V^{m_i}), (W_O^{m_i})는 모달리티별 projection matrix이고, (\mathrm{LayerNorm}{\mathrm{attn}}^{m_i})와 (\mathrm{LayerNorm}{\mathrm{ffn}}^{m_i})는 모달리티별 layer normalization이다.

이 접근법은 MoT가 각 모달리티의 구체적인 특성에 맞게 처리를 조정할 수 있게 하면서도, 멀티모달 학습을 위한 통합 아키텍처를 유지하게 한다. MoT의 계산 과정은 먼저 입력 토큰을 모달리티별로 그룹화하는 것에서 시작한다(Algorithm 1, lines 3–5). 이후 attention을 위해 모달리티별 projection이 적용된다(line 6). 그다음 모든 모달리티에 걸쳐 global self-attention이 수행된다(lines 8–9). 이후 모달리티별 output projection(line 11), layer normalization, feed-forward network가 적용된다(lines 12–13). 마지막으로 residual connection과 layer normalization을 포함하여 출력이 결합되면서 과정이 마무리된다(lines 14–16).

즉 MoE 구조는 아니고 그냥 dense transformer가 모달리티별로 쪼개진 것

3 Experiments

3.1 Results Overview

우리는 Mixture-of-Transformers (MoT) 아키텍처를 세 가지 멀티모달 실험 설정에서 평가했다. 각 설정은 점진적으로 더 복잡한 학습 목적 함수와 모달리티를 포함한다. 각 설정에서 MoT를 두 가지 baseline과 비교했다. 하나는 dense transformer model이고, 다른 하나는 4개의 expert를 가진 Mixture-of-Experts model (MoE-4x)이다. 모든 모델 구현은 dense model을 기반으로 하며, 학습과 테스트 모두에서 동일한 FLOPs를 유지한다. 이를 통해 효율성과 성능을 직접 비교할 수 있다.

  1. Autoregressive objective를 사용하는 멀티모달 실험 설정 (Chameleon, Figure 4).
    두 모달리티 모두 autoregressive objective로 학습된다. 이미지는 사전학습된 VQ-VAE 모델(Gafni et al., 2022)을 통해 1,024개의 discrete token으로 표현된다. 저자들은 두 모달리티 모두에 대해 학습 및 평가 지표 전반에서 MoT의 성능을 baseline들과 비교했다.

  2. 음성 모달리티를 포함한 확장 멀티모달 실험 (Chameleon: Text+Image+Speech, Figure 7).
    이전 설정에 세 번째 모달리티로 음성을 추가했다. 음성은 사전학습된 speech tokenizer를 통해 discrete token으로 표현된다. 모든 모달리티는 autoregressive objective로 학습된다. 이 설정은 MoT가 추가 모달리티를 처리하면서도 효율성과 성능을 유지할 수 있는지를 평가한다.

  3. 모달리티별 objective를 사용하는 멀티모달 실험 (Transfusion, Figure 10).
    텍스트에는 autoregressive objective를 사용하고, 이미지에는 diffusion-based objective를 사용하는 multi-objective training을 탐구했다. 이 실험은 MoT가 서로 다른 모달리티에 대해 서로 다른 학습 목적 함수를 관리할 수 있는지를 보여준다. 이는 텍스트 생성 능력을 유지하면서 이미지 생성 품질을 개선할 가능성을 강조한다.

다음 섹션들은 각 설정에 대한 자세한 결과를 제시한다. 구체적으로 Chameleon은 Section 3.2, Chameleon+Speech는 Section 3.3, Transfusion은 Section 3.4에서 다룬다. 각 섹션은 다양한 멀티모달 생성 평가 지표에서 MoT와 baseline들을 포괄적으로 비교한다. Section 3.5에서는 transformer의 서로 다른 구성 요소에 모달리티별 decoupling을 도입했을 때 모델 성능에 어떤 영향을 미치는지 보여주는 ablation study를 보고한다.

3.2 핵심 메시지 정리

이 섹션은 Chameleon 설정, 즉 텍스트와 이미지를 모두 discrete token으로 바꾼 뒤 autoregressive next-token prediction으로 학습하는 설정에서 MoT가 dense transformer와 MoE보다 얼마나 효율적인지를 보여준다.

핵심은 다음이야.

첫째, MoT는 7B 규모에서 dense model과 비슷한 성능을 훨씬 적은 학습량으로 달성한다.
7B Chameleon 실험에서 MoT는 dense model이 120k step에서 도달한 최종 training loss를 약 60k step에서 달성한다. step matching 분석에서는 MoT가 dense model 대비 45.5%의 training step만으로 비슷한 전체 pretraining loss에 도달한다고 보고한다. 검증 성능 기준으로는 MoT가 dense baseline 성능을 맞추는 데 55.8% FLOPs만 필요하다고 해석한다.

둘째, MoT의 이득은 특히 이미지 모달리티에서 크다.
모달리티별로 보면, 이미지 training loss에서 MoT는 dense model의 최종 loss를 맞추는 데 34.8%의 training step만 필요하다. 반면 텍스트에서도 MoT가 dense보다 좋지만, 이미지에서처럼 압도적인 차이는 아니다. 즉 이 논문의 중요한 메시지는 모달리티별 파라미터 분리가 이미지 같은 비텍스트 모달리티 학습을 특히 빠르게 만든다는 것이다.

셋째, MoE-4x보다 MoT가 더 안정적으로 좋다.
MoE-4x도 텍스트에서는 어느 정도 dense보다 좋은 모습을 보이지만, 이미지 모달리티에서는 개선이 제한적이다. 특히 모델 규모가 커질수록 MoE-4x의 이미지 쪽 이득은 줄어들고, 7B에서는 거의 사라진다. 반면 MoT는 이미지와 텍스트 모두에서 일관된 이득을 보인다.

넷째, MoT의 효과는 작은 모델부터 큰 모델까지 유지된다.
37M, 94M, 443M, 1.5B, 7B까지 여러 규모에서 실험했을 때, MoT는 모든 규모에서 이미지 모달리티 학습을 빠르게 만든다. 텍스트에서도 dense보다 대체로 좋거나 비슷하게 좋다. 따라서 저자들은 MoT가 특정 크기에서만 우연히 잘 되는 구조가 아니라, 여러 모델 scale에서 일관되게 효율적인 구조라고 주장한다.

정리하면, 3.2의 결론은:

Chameleon-style text-image autoregressive 모델에서 MoT는 dense transformer와 같은 FLOPs-controlled 조건으로 비교했을 때, 특히 이미지 모달리티 학습을 크게 가속하며, 7B에서도 dense baseline 성능을 약 절반 수준의 학습 비용으로 맞춘다.

 

입력 토큰 100개이라고 하면

  • dense 모델은 각 토큰이 dense FFN / attention projection / LayerNorm을 1번씩 통과

MoT 기준에서는 텍스트 토큰 40개/이미지 토큰 40개/음성 토큰 20개 로 모달리티별로 나누고

  • 텍스트 토큰 40개 -> text FFN / text QKV / text Norm
  • 이미지 토큰 40개 -> image FFN / image QKV / image Norm
  • 음성 토큰 20개 -> speech FFN / speech QKV / speech Norm
  • 40개 토큰이 text branch 사용 + 40개 토큰이 image branch 사용 + 20개 토큰이 speech branch 사용 = 총 100개 토큰 처리
  • 100개 토큰 × text branch + 100개 토큰 × image branch + 100개 토큰 × speech branch 가 아니라는 것
그래서 한 시퀀스 안에서 모든 branch가 활성화될 수는 있지만, 각 토큰이 모든 branch를 통과하는 것은 아니기 때문에 FLOPs가 3배가 되지는 않는다는 게 중요해

MoE도 비슷한 개념이고 이들간의 비교는 activation 파라미터 7B 기준으로 비교

3.3 핵심 메시지 정리

이 섹션은 기존 Chameleon text+image 설정에 speech를 세 번째 모달리티로 추가했을 때도 MoT가 잘 작동하는지를 확인하는 실험이야. 즉, 텍스트·이미지·음성을 모두 discrete token으로 표현하고, 세 모달리티를 모두 autoregressive objective로 학습한다.

핵심은 다음과 같아.

첫째, speech는 discrete semantic token으로 변환해서 사용한다.
논문은 SpiRit-LM의 speech 데이터를 사용하고, speech input을 DinoSR 변형의 in-house tokenizer로 변환한다. 이 tokenizer는 vocab size 500의 semantic token을 만들며, 각 token은 약 40ms 오디오를 나타낸다. 학습 데이터는 Chameleon의 text-image 데이터와 speech 데이터를 1:6 sampling ratio로 섞어 만든다.

둘째, 7B 규모에서 MoT는 speech 학습을 매우 크게 가속한다.
7B Chameleon+Speech 실험에서 MoT는 dense model과 MoE-4x보다 speech training loss를 훨씬 빠르게 줄인다. step matching 기준으로 보면, MoT는 dense model과 같은 speech pretraining loss에 도달하는 데 22.9% training steps만 필요하다. 검증 loss 기준으로는 speech 성능을 맞추는 데 37.2% FLOPs만 필요하다고 보고한다.

셋째, speech를 추가해도 text/image 성능이 무너지지 않는다.
중요한 점은 MoT가 speech에서만 좋아지는 것이 아니라, speech를 추가한 뒤에도 기존 image와 text 모달리티에서 효율성을 유지한다는 것이다. Figure 8에서는 MoT가 55.8% training checkpoint에서도 dense model의 최종 image/text validation loss와 비슷하거나 더 낮은 수준을 달성한다. 즉, speech를 넣어도 기존 모달리티 성능을 크게 해치지 않는다.

넷째, 작은 모델 규모에서도 speech 쪽 이득이 일관된다.
443M, 880M, 1.5B에서도 MoT는 세 모달리티 전반에서 일관된 가속을 보인다. 특히 speech modality에서는 dense model의 training loss를 맞추는 데 15.1%–33.6% steps만 필요하다. 즉 MoT의 효과는 7B에서만 나타나는 것이 아니라 여러 scale에서 유지된다.

다섯째, MoE-4x는 speech에서 불안정하다.
MoE-4x는 training loss에서는 좋아 보일 수 있지만, speech validation loss에서는 dense baseline보다 나빠지는 경우가 있다. 논문은 이를 learned routing의 불안정성, 데이터 분포 차이, expert 사용 불균형, 그리고 speech token 수가 상대적으로 적어 overfitting이 생길 수 있는 점과 연결해 설명한다. 반면 MoT는 모달리티 라벨 기반 deterministic routing을 쓰기 때문에 speech에서도 더 안정적이라고 주장한다.

정리하면, 3.3의 결론은:

MoT는 text+image에 speech를 추가한 3-modal setting에서도 잘 확장되며, 특히 speech modality 학습을 크게 가속한다. MoE처럼 learned router를 쓰지 않고 모달리티별 파라미터를 고정적으로 분리하기 때문에, speech처럼 분포가 다른 모달리티에서도 더 안정적인 성능을 보인다.

3.4 핵심 메시지 정리

이 섹션은 Transfusion 설정에서 MoT를 평가한다. 앞의 Chameleon 설정은 텍스트와 이미지를 모두 discrete token으로 바꾸고 autoregressive objective로 학습했지만, Transfusion은 조금 다르다.

Transfusion에서는:

텍스트 → autoregressive language modeling objective
이미지 → continuous image token + diffusion-based objective

즉, 하나의 멀티모달 모델 안에서 텍스트와 이미지를 서로 다른 학습 objective로 학습한다. 이 섹션의 핵심 질문은 “MoT가 단순히 autoregressive text-image 설정에서만 좋은가, 아니면 objective가 다른 복잡한 설정에서도 좋은가?”이다.

첫째, MoT는 Transfusion에서도 이미지 모달리티 학습을 크게 가속한다.
7B scale에서 MoT는 dense model과 같은 image pretraining loss에 도달하는 데 약 30% training steps만 필요하다. 즉, 이미지가 discrete token이 아니라 diffusion objective로 학습되는 경우에도, 모달리티별 파라미터 분리가 이미지 학습을 빠르게 만든다는 것을 보여준다.

둘째, 760M MoT가 1.4B dense baseline보다 좋은 결과를 낸다.
흥미로운 결과는 작은 MoT 모델이 더 큰 dense model을 이긴다는 점이다. 760M MoT는 1.4B dense baseline의 절반 FLOPs만 사용하면서도 이미지 생성 관련 지표에서 더 좋은 성능을 보인다. 예를 들어 CLIP score는 더 높고, FID score는 더 낮으며, 이미지 캡셔닝의 CIDEr score도 더 높다. 즉, MoT는 단순히 “같은 크기에서 빠르다”를 넘어서, 더 작은 activated model로 더 큰 dense model을 능가할 수 있음을 보여준다.

셋째, MoT의 이득은 이미지 쪽에서 가장 뚜렷하다.
Transfusion 설정에서 MoT는 이미지 생성과 이미지 이해/captioning에서 강한 이득을 보인다. 여러 모델 크기인 163M, 760M, 1.4B 전반에서 MoT는 dense보다 더 좋은 FID, CLIP, CIDEr를 보인다. 반면 텍스트-only 성능에서는 MoT가 dense와 비슷한 수준을 유지하지만, Chameleon 설정처럼 큰 향상은 보이지 않는다.

넷째, 텍스트 성능 향상이 작았던 이유를 논문은 objective 차이로 설명한다.
저자들은 Transfusion에서는 텍스트와 이미지가 이미 서로 다른 objective로 분리되어 학습되기 때문에, dense model도 텍스트 쪽에서는 어느 정도 최적화가 잘 된 상태일 수 있다고 해석한다. 다시 말해, objective 자체가 분리되어 있어서 MoT의 모달리티별 파라미터 분리 효과가 텍스트에서는 덜 크게 나타났을 가능성이 있다. 또한 텍스트는 이미지 diffusion에 비해 상대적으로 계산 부담이 작기 때문에 sparse 구조의 이득이 작을 수 있다고 설명한다.

다섯째, fine-tuning 후에도 MoT의 이득이 유지된다.
저자들은 7B MoT와 dense model을 이미지 생성 및 이미지 편집 데이터로 fine-tuning한다. 그 결과, MoT는 fine-tuning 후에도 dense보다 더 좋은 이미지 품질과 prompt faithfulness를 보이며, MoT의 성능 이득이 pretraining 단계에만 국한되지 않음을 보여준다.

정리하면, 3.4의 결론은:

MoT는 Chameleon처럼 모든 모달리티를 autoregressive로 학습하는 단순 설정뿐 아니라, 텍스트는 autoregressive로, 이미지는 diffusion으로 학습하는 Transfusion 설정에서도 효과적이다. 특히 이미지 모달리티에서 큰 효율성과 성능 향상을 보이며, 작은 MoT 모델이 더 큰 dense model을 능가하는 결과도 보여준다.

3.5 핵심 메시지 정리

이 섹션은 MoT에서 어떤 transformer component를 모달리티별로 분리하는 것이 중요한가를 보는 ablation 실험이야. 즉, MoT의 성능 향상이 단순히 “모달리티별 branch를 둔다”에서 오는지, 아니면 특정 component를 분리했을 때 특히 효과가 큰지를 분석한다.

비교한 구조는 크게 네 가지야.

  1. Dense baseline
    모든 모달리티가 같은 transformer 파라미터를 공유한다.

  2. FFN만 모달리티별로 분리
    텍스트용 FFN, 이미지용 FFN 등을 따로 둔다. Attention과 LayerNorm은 공유한다.

  3. FFN + Q/K/V attention matrix 분리
    FFN뿐 아니라 attention projection matrix도 모달리티별로 둔다. 다만 LayerNorm은 아직 공유한다.

  4. Full MoT
    FFN, attention projection, LayerNorm까지 모두 모달리티별로 분리한다.

핵심 결과는:

첫째, FFN을 모달리티별로 분리하는 것만으로도 성능이 크게 좋아진다.
논문은 FFN untying alone이 특히 이미지 모달리티에서 큰 성능 향상을 만든다고 말한다. 이는 transformer에서 FFN이 일종의 “memory component” 역할을 하기 때문에, 텍스트와 이미지가 같은 FFN memory를 공유하는 것보다 각자 다른 FFN을 쓰는 것이 유리하다는 해석으로 이어진다.

둘째, Q/K/V attention matrix까지 분리하면 추가 개선이 있다.
FFN만 분리한 것보다, attention의 Q/K/V projection matrix도 모달리티별로 분리했을 때 성능이 더 좋아진다. Obelisc held-out set 기준으로, FFN만 분리한 경우에 비해 이미지 모달리티에서 약 33.3% FLOPs saving, 텍스트 모달리티에서 약 10% FLOPs saving이 추가로 나온다고 보고한다.

셋째, LayerNorm 분리는 효과가 거의 작다.
FFN과 attention projection을 이미 분리한 상태에서 LayerNorm까지 추가로 분리해도 evaluation performance에는 큰 영향이 없었다. 그래서 full MoT에는 LayerNorm untying도 포함되지만, 실제 성능 향상의 핵심은 주로 FFN 분리와 attention projection 분리에 있다고 볼 수 있다.

정리하면, 3.5의 결론은:

MoT의 성능 향상은 주로 FFN을 모달리티별로 분리하는 데서 크게 나오고, attention Q/K/V projection을 추가로 분리하면 더 좋아진다. 반면 LayerNorm 분리는 상대적으로 영향이 작다.

즉, MoT 전체 구조가 다 중요하긴 하지만, 가장 중요한 부분은:

모달리티별 FFN + 모달리티별 attention projection

이라고 보면 된다.

4 핵심 메시지 정리

이 섹션은 **“정말 모달리티를 각각 분리해야 하는가?”**를 확인하는 실험이야. 앞에서는 MoT가 텍스트/이미지/음성 각각에 별도 transformer parameter branch를 둔다고 했는데, 여기서는 일부 모달리티를 다시 합쳐 보면서 성능이 어떻게 변하는지 본다.

비교한 구조는 다음과 같아.

Full MoT:
text / image / speech 각각 분리

LOO-image:
image는 따로 두고, text+speech를 하나로 합침

LOO-text:
text는 따로 두고, image+speech를 하나로 합침

LOO-speech:
speech는 따로 두고, text+image를 하나로 합침

Dense:
text+image+speech 모두 하나의 transformer 공유

모든 모델은 FLOPs를 맞춰서 비교한다.

핵심 결과는:

첫째, 모달리티를 합치면 대체로 성능이 나빠진다.
Figure 15에서 두 모달리티를 하나의 tower로 묶은 LOO variants는 full MoT보다 training/validation loss가 전반적으로 높아진다. 즉 텍스트, 이미지, 음성은 서로 다른 특성을 가지기 때문에 같은 parameter branch를 공유하면 최적화가 어려워진다는 것을 보여준다.

둘째, 특정 모달리티를 따로 분리하면 그 모달리티 성능은 좋아진다.
LOO-text는 text loss가 가장 낮고, LOO-image는 image loss가 가장 낮고, LOO-speech는 speech loss가 가장 낮다. 즉 어떤 모달리티든 자기 전용 tower를 가질 때 그 모달리티에 유리하다.

셋째, 모달리티 간 간섭 효과는 대칭적이지 않다.
어떤 모달리티를 누구와 합치느냐에 따라 손해가 다르게 나타난다. 예를 들어 text+speech를 합치면 둘 다 성능이 떨어지고, image+speech를 합치면 speech 쪽이 특히 나빠질 수 있다. 논문은 이를 non-reciprocal modality competition effects, 즉 모달리티 간 경쟁/간섭이 서로 같은 방식으로 작동하지 않는 현상으로 해석한다.

정리하면, 4장의 결론은:

MoT에서 모달리티별 parameter separation은 실제로 중요하다.
두 모달리티를 같은 transformer branch에 넣으면 성능이 떨어지고, 각 모달리티에 전용 parameter branch를 주는 것이 전체적으로 더 안정적이고 좋다.

쉽게 말하면, 이 섹션은 **“왜 굳이 text/image/speech branch를 따로 두는가?”**에 대한 실험적 근거를 제공하는 부분이야.

5 핵심 메시지 정리

이 섹션은 MoT와 MoE를 서로 대체 관계로만 볼 것이 아니라, 조합할 수도 있다는 것을 보여주는 preliminary experiment야. 앞 실험들에서는 MoT가 MoE-4x보다 특히 이미지/음성 모달리티에서 더 안정적이고 효율적이라고 보였지만, 여기서는 MoE의 장점은 텍스트 쪽에 선택적으로 쓰고, MoT의 장점은 비텍스트 모달리티에 유지하는 hybrid 구조를 실험한다.

핵심 아이디어는 이거야.

MoT 구조는 유지하되, text transformer tower의 FFN만 MoE-4x로 바꾼다.
이미지 transformer tower는 기존 MoT 그대로 둔다.

즉 구조적으로는:

Text tower  -> MoE-4x FFN 사용
Image tower -> MoT 방식 그대로 사용

5.1 Chameleon setting에서의 조합

Chameleon 설정에서는 텍스트와 이미지가 모두 autoregressive objective로 학습된다.

결과적으로 MoT + Text MoE-4x는 text training loss를 dense, MoE-4x, MoT 단독보다 더 빠르게 줄인다. 중요한 점은, text tower에 MoE를 넣었는데도 image modality에서 MoT가 갖던 이득은 유지된다는 것이다.

검증 성능에서도 이 hybrid 모델은 text validation performance가 가장 좋고, image validation performance도 MoT와 비슷하거나 약간 더 좋다. 즉, Chameleon 설정에서는 텍스트에는 MoE, 이미지는 MoT라는 조합이 꽤 잘 작동한다.

5.2 Transfusion setting에서의 조합

Transfusion 설정에서도 같은 아이디어를 적용한다. 여기서는 텍스트는 autoregressive objective, 이미지는 diffusion-based objective로 학습된다.

마찬가지로 text transformer의 FFN을 MoE-4x로 바꾸고, image transformer는 기존 MoT 구조를 유지한다. 결과적으로 hybrid 모델은 dense와 MoT 단독보다 text training loss를 더 빠르게 줄인다. 동시에 이미지 쪽에서는 MoT의 장점을 유지해서, image training loss와 speedup이 MoT와 비슷하게 나온다.

검증 결과에서도 MoT + Text MoE-4x는 text 성능을 개선하면서, image generation 성능을 크게 해치지 않는다. 반면 MoE-4x 단독은 training loss는 낮아 보여도 text validation loss에서는 개선이 거의 없거나 불안정한 모습을 보인다. 논문은 이를 Transfusion처럼 discrete text token과 continuous image token을 함께 다루는 설정에서는 MoE router가 inference에서 일반화하기 어렵기 때문일 수 있다고 설명한다.

정리

5장의 결론은:

MoT와 MoE는 경쟁 관계만이 아니라 조합 가능하다.
MoT는 이미지 같은 비텍스트 모달리티에서 강하고, MoE는 텍스트 tower 안에 선택적으로 넣으면 text 성능을 더 끌어올릴 수 있다.

쉽게 말하면:

이미지/음성처럼 모달리티 차이가 큰 부분은 MoT처럼 모달리티별 tower를 분리하고,
텍스트처럼 MoE가 잘 먹히는 부분은 text tower 내부에 MoE expert를 추가하는 방식이다.

그래서 이 섹션은 **“MoT가 MoE를 완전히 대체한다”**라기보다, 앞으로는 MoT + MoE hybrid sparse multimodal transformer도 가능하다는 방향을 보여주는 proof-of-concept라고 보면 된다.

7 Related Work

7.1 Foundation Models for Multi-Modal Generation

대규모 언어 모델(LLM)의 최근 발전은 멀티모달 응용으로 확장되어 왔다. 초기 멀티모달 LLM은 생성보다는 이해에 초점을 맞추었으며, 별도로 인코딩된 이미지와 텍스트를 결합하기 위해 late fusion 기법을 사용했다(Alayrac et al., 2022; Liu et al., 2023; Laurençon et al., 2023; Chen et al., 2022). 이러한 모델들은 가벼운 학습이라는 장점이 있었지만, 멀티모달 생성 능력은 부족했다.

멀티모달 생성을 가능하게 하기 위한 핵심 전략은 비텍스트 모달리티를 discrete sequence로 토큰화하는 것이다(Aghajanyan et al., 2022; Yu et al., 2023; Bao et al., 2021; Ramesh et al., 2021; Liu et al., 2024c)(Figure 2a). 예를 들어, Chameleon (Chameleon Team, 2024)과 관련 접근법들(Aghajanyan et al., 2022)은 VQGAN (Esser et al., 2021)과 같은 사전학습 모델을 사용해 이미지를 1,024개의 discrete token으로 토큰화하고, 텍스트-이미지 토큰 시퀀스를 결합하여 학습한다. 유사한 토큰화 방식은 음성에도 적용되어 왔다(Nguyen et al., 2024). 최근 Transfusion (Zhou et al., 2024)과 같은 모델들은 시각 생성 품질을 향상시키기 위해 continuous image token과 diffusion-based loss function을 탐구했다.

본 논문에서 제안하는 mixture of transformers 방법은 이러한 접근법들과 호환되며, dense transformer architecture를 대체하는 drop-in replacement로 통합될 수 있다. 저자들은 Chameleon (Chameleon Team, 2024)과 Transfusion (Zhou et al., 2024)을 포함한 다양한 멀티모달 설정에서 상당한 성능 향상을 보였다.

7.2 Sparse Architectures for Multi-Modal Generation

Sparse architecture, 특히 Mixture of Experts (MoE)는 텍스트 기반 모델에서 가능성을 보여주었다. MoE는 각 입력에 대해 동적으로 파라미터를 선택할 수 있게 한다(Jacobs et al., 1991; Eigen et al., 2013; Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022; Jiang et al., 2024; Sukhbaatar et al., 2024). 최근 연구들은 모달리티 간 feature space gap으로 인해 발생하는 문제를 다루기 위해 MoE를 멀티모달 task에 적용해 왔다(Wang et al., 2022; Shen et al., 2023; Bao et al., 2022a; Long et al., 2023; Lin et al., 2024). 이러한 접근법들은 모달리티별 파라미터 할당이 서로 다른 데이터 유형, 즉 모달리티의 특성을 다루어 성능을 향상시킬 수 있음을 시사한다(Liang et al., 2022).

이전 연구들과 달리, 본 논문은 transformer architecture 내부의 모든 비임베딩 파라미터를 분리함으로써 MoE 개념을 일반화하는 Mixture-of-Transformers (MoT) 프레임워크를 제안한다. MoT는 전체 파라미터 수를 통제했을 때 멀티모달 사전학습에서 MoE보다 일관되게 더 좋은 성능을 보였고(Figure 16), MoE-4x와 상호보완적으로 결합될 수 있음도 보였다(Figure 17).

최근 연구들은 MoE를 feedforward layer를 넘어 attention mechanism까지 확장해 왔지만(Wang et al., 2023; Shen et al., 2024; Liu et al., 2024b), 본 논문의 접근법은 여러 중요한 측면에서 다르다. CogVLM (Wang et al., 2023)은 텍스트 출력 생성에 제한되어 있는 반면, MoT는 이미지와 텍스트 생성이 모두 가능하다. 본 연구와 동시에 진행된 Playground v3 (PGv3) (Liu et al., 2024b)는 DiT-style image transformer를 Llama3-8B text backbone과 global self-attention으로 통합하여, 텍스트 및 이미지 생성에서 state-of-the-art 성능을 달성했다. PGv3는 학습 중 text LLM을 freeze하고 image transformer component만 업데이트한다.

반면, CogVLM과 PGv3가 사전학습된 LLM 위에서 멀티모달 학습을 수행하는 것과 달리, 본 논문은 MoT를 scratch부터 학습할 수 있는 일반적인 sparse architecture로 제시한다. 또한 MoT는 layer normalization을 포함하여 transformer layer 전반의 모든 비임베딩 파라미터를 모달리티별로 분리하는 반면, 이전 접근법들은 layernorm parameter를 공유한다. 이러한 결과는 MoT가 멀티모달 사전학습을 위한 유연하고 확장 가능한 방법이며, MoE 기반 architecture를 보완하고 대규모 멀티모달 모델을 더 계산 효율적으로 학습하는 방향을 제공할 수 있음을 보여준다.

8 Conclusion

본 연구에서는 멀티모달 모델 사전학습의 계산적 어려움을 해결하기 위해 설계된 희소하고 확장 가능한 아키텍처인 Mixture-of-Transformers (MoT)를 제시했다. MoT는 모달리티별로 비임베딩 파라미터를 분리하면서도, 멀티모달 시퀀스 전체에 대한 global self-attention은 유지한다. 이를 통해 MoT는 모달리티별 처리를 최적화하면서도 cross-modal interaction을 보존한다.

실험 결과, MoT는 다양한 설정과 모델 규모에서 학습 비용을 크게 줄일 수 있음을 보였다. Chameleon 및 Chameleon+Speech 설정에서 MoT는 훨씬 적은 FLOPs를 사용하면서도 dense baseline과 비슷하거나 더 좋은 성능을 달성했다. 또한 텍스트와 이미지에 서로 다른 학습 objective를 적용하는 더 복잡한 Transfusion 설정에서도 이러한 개선을 유지했으며, 이미지 생성과 같은 task에서 일관된 효율성 향상과 성능 개선을 보였다.

FLOP 감소뿐 아니라, 시스템 프로파일링 결과는 MoT가 실제적인 이점도 제공함을 보여준다. 텍스트와 이미지 task 모두에서 wall-clock time이 줄어들었고, GPU 수를 늘려 scale을 키웠을 때 MoT의 개선 폭은 더 커졌다. 이는 MoT가 대규모 분산 학습 환경에 적합하다는 것을 시사한다.

또한 MoT와 Mixture-of-Experts (MoE-4x)를 결합한 초기 결과는, 계산 비용을 늘리지 않으면서 성능을 더 향상시킬 수 있는 hybrid model의 가능성을 보여준다. 이러한 결과들은 MoT가 향후 멀티모달 LLM을 위한 효과적인 프레임워크가 될 수 있으며, 다양한 모달리티에서 경쟁력 있는 성능을 유지하면서도 더 효율적인 대규모 학습을 가능하게 할 수 있음을 시사한다.

Reference

댓글

댓글 쓰기