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의 적응성과 효율성 향상을 평가하기 위해 다음과 같은 멀티모달 시나리오에서 실험을 수행했다.

텍스트와 이미지 모두에 대한 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).
세 번째 모달리티로 음성 도입 (Chameleon: Text+Image+Speech).
Chameleon 설정에 세 번째 모달리티로 discrete speech token을 추가했을 때, MoT는 모든 모달리티에서 유사한 성능을 달성했으며, 음성 모달리티의 경우 더 적은 학습 FLOPs, 즉 **37.2%**만 필요했다(Figure 8). 이러한 결과 역시 다른 여러 모델 규모에서도 일관되게 나타났다(Figure 8, Appendix Figure 25).
텍스트에는 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와 비교한 효과를 평가했다.

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).
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).
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을 포착한다.

$$
\mathrm{GlobalAttn}(x, {\theta^m_{\mathrm{attn}}}) =
\left(
\mathrm{softmax}
\left(
\frac{QK^T}{\sqrt{d_k}}
\right)
V
\right)
W_O^{m_i}
$$

$$
Q_i = x_i W_Q^{m_i}, \quad
K_i = x_i W_K^{m_i}, \quad
V_i = x_i W_V^{m_i}
\tag{3}
$$

여기서 (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가 모달리티별로 쪼개진 것