기존 MoE 모델은 모달리티를 구분하지 않고 토큰을 처리했지만, MoMa는 모달리티별 전문화를 통해 효율성과 처리 속도를 동시에 향상시켰다고 하며, 1조 토큰을 아키텍처별로 학습시켜본 뒤 실험 내용을 보고함.
구현 참고자료
1. 서론
Auto-regressive mixed-modal foundation models은 이미지와 텍스트를 포함한 다양한 모달리티를 처리하고 생성하는데 중요한 역할을 합니다. 이런 모델들은 비주얼 질문 응답 및 멀티모달 콘텐츠 생성에서 향상된 성능을 보여주었습니다. 기존의 모델들은 모달리티 특화 인코더 또는 디코더를 결합하는 방식으로 설계되었으나, 이 방법은 모달리티 간 정보 통합을 제한할 수 있습니다. 따라서 이 연구에서는 모달리티별로 특화된 모듈을 사용하여 정보 처리 효율을 개선하려는 새로운 아키텍처인 MoMa를 제안합니다.
[문제 정의]
[해결 방법]
[선행연구 대비 차별점]
2. 모델
이 연구에서 제안하는 MoMa 아키텍처는 모달리티별 전문가 그룹을 통해 입력 토큰을 처리합니다. 각 그룹은 학습된 라우팅 메커니즘을 통해 토큰을 최적의 전문가에게 할당하며, 이는 모델의 효율성을 극대화합니다.
모달리티별 라우팅은 다음 수식으로 모델링됩니다.
\[\text{Routing}(x) = \text{softmax}(W_g x + b_g)\]$x$는 입력 토큰, $W_g$는 라우팅을 위한 가중치, $b_g$는 바이어스를 나타내며, 각 토큰을 가장 적합한 전문가 그룹에 할당하는 확률을 계산합니다.
[본 논문의 방법]
[실험 개요]
MoMa 아키텍처는 기존 모델들과 비교하여 향상된 효율성과 높은 성능을 보이며, 모달리티별 처리를 최적화함으로써, 대규모 멀티모달 데이터셋에서도 높은 효율성을 유지할 수 있다고 주장하며 Mixed-Modal, Early-Fusion에 대해 탐색합니다.
2.1 Early-Fusion
MoMa 모델은 Chameleon 아키텍처의 Early-Fusion 구조를 기반으로 합니다. 이 구조는 이미지와 텍스트를 통합된 트랜스포머 내에서 이산 토큰 시리즈로 처리하며, 셀프 어텐션 메커니즘을 통해 이미지와 텍스트 토큰의 결합된 시퀀스를 처리합니다. 이를 통해 모델은 모달리티 내부 및 모달리티 간의 복잡한 관계를 포착할 수 있습니다. Chameleon의 모델은 다음 토큰 예측 목표를 사용하여 훈련되며, 텍스트와 이미지 토큰을 자동 회귀 방식으로 생성하는 방법을 학습합니다.
Chameleon 참조 논문 - Chameleon: Mixed-Modal Early-Fusion Foundation Models
카멜레온은 이미지를 학습된 이미지 토크나이저를 사용하여 512x512 이미지에서 1024개의 이산 토큰으로 토크나이즈하며, 텍스트는 BPE 토크나이저를 사용하여 처리됩니다. 이 통합 토크나이징 방식은 모델이 이미지와 텍스트 토큰의 임의의 시퀀스를 처리할 수 있게 합니다.
2.2 너비 조절: 모달리티 인식 전문가의 혼합
너비 조절 접근 방식은 피드-포워드 모듈에서 모달리티 인식 블록 희소성을 도입하여 표준 전문가의 혼합 아키텍처를 확장하며, 이 방식의 핵심 인사이트는 다양한 모달리티의 토큰이 각기 다른 특성과 정보 밀도를 가지게 한다는 것입니다. 각 모달리티별로 별도의 전문가 그룹을 생성함으로써, 모델은 전문화된 처리 경로를 개발할 수 있으며, 모달리티 간 정보를 통합할 수 있는 능력을 유지합니다.
각 입력 토큰 $x$에 대해, 라우팅 메커니즘은 두 단계에서 작동합니다.
이 과정은 각 토큰을 가장 적합한 전문가에게 효율적으로 할당합니다.
[실험]
실험을 통해 MoMa 아키텍처는 표준 전문가의 혼합 아키텍처와 비교하여 효율성과 성능을 크게 향상시키는 것으로 나타났습니다. 이를 증명하기 위해, 대규모 토큰 데이터셋을 사용하여 모델을 훈련시켰고, 모델의 인퍼런스 성능과 훈련 중 FLOPs 절감을 측정했습니다.
3. 효율성 최적화
본 절에서는 모달리티 인식 전문가의 혼합(MoMa)을 분산 훈련하기 위해 Fully Sharded Data Parallel (FSDP)을 사용하는 과정에서 발생하는 효율성에 대한 도전과 이를 해결하기 위한 전략에 대해 논의합니다.
3.1 부하 균형
시스템 내에서 제약 없이 진행될 경우, GPU와 반복 사이에서 텍스트와 이미지 토큰의 비율이 크게 달라질 수 있어 부하 불균형이 발생할 수 있습니다. 이 불균형은 후속 레이어의 가중치 프리패칭과 이전 레이어의 그래디언트 해제를 지연시키는 캐스케이딩 스트래글러 효과를 유발할 수 있습니다. 이는 배치 내 모든 GPU에서 텍스트와 이미지 전문가를 처리하는데 필요한 최대 시간으로 훈련 지연을 결정합니다.
이런 문제를 완화하기 위해, 각 GPU에서 텍스트-이미지 데이터 비율이 전문가 비율과 일치하는 균형 잡힌 데이터 믹스를 개발했습니다. 이 접근 방식은 기대에 따라 부하 균형을 보장합니다. 각 FFN 레이어에서 토큰 재분배와 같은 대안적 재조정 알고리즘이 가능하지만, 추가 통신 오버헤드를 도입할 수 있습니다.
3.2 효율적인 전문가 실행
다양한 모달리티에 대해 전문가를 효율적으로 실행하기 위한 여러 전략을 탐구했습니다. 첫 번째 접근 방식은 모달리티 간에 동질의 전문가를 제한하고 텍스트 토큰을 이미지 전문가에게 라우팅하지 않도록 합니다. 이 방법은 모든 전문가가 동일한 토큰 수를 공유하는 경우, 모든 토큰과 모달리티를 동시에 처리할 수 있습니다. 또 다른 방법으로, 블록 희소성을 사용하여 완벽한 전문가 토큰 균형을 요구하지 않고도 첫 번째 접근 방식과 유사한 이점을 제공할 수 있습니다. 추가로, 모달리티 수가 제한적인 경우 서로 다른 모달리티의 전문가를 순차적으로 실행하는 방법을 고려했습니다. 이 접근 방식은 이전 모달리티 전문가의 계산과 후속 모달리티 전문가의 가중치 프리패칭 간에 더 나은 중첩을 허용하여 메모리 압력을 완화합니다. 또한 전문가 부하 균형에 대한 가정을 제거합니다.
실험에서 충분히 큰 토큰 수를 GPU당 처리하기 때문에, 하드웨어 활용은 주요한 관심사가 아닙니다. 따라서 현재 규모에서 실험 환경에 대해 순차적 접근 방식이 깔끔하고 유연한 선택으로 판단됩니다.
3.3 추가 최적화
처리량을 더욱 향상시키기 위해 여러 최적화를 구현했습니다. 이런 최적화에는 기울기 통신 양자화, 자동 GPU 커널 퓨전 및 torch.compile을 통한 그래프 최적화가 포함됩니다. 또한, CPU와 GPU 간 동기화를 최소화하기 위해 다양한 레이어에서 모달리티 토큰 인덱스를 재사용하는 MoMa 특화 최적화를 개발했습니다. 각 레이어의 통계 통신을 통합하고 이런 작업을 훈련의 중요 경로에서 이동시켰습니다.
이 섹션에서는 MoMa 아키텍처의 효율적인 분산 훈련을 위한 다양한 도전과 전략에 대해 자세히 설명하였습니다. 이런 최적화는 모델의 효율성과 성능을 극대화하는 데 중요한 역할을 합니다.
4. 실험
4.1 설정
Chameleon Team (2024)의 사전 훈련 데이터셋과 전처리 방법을 그대로 사용합니다.
모델의 스케일링 성능을 평가하기 위해, 1조 개 이상의 토큰을 사용하여 모든 모델을 훈련하였습니다. 특별히 명시하지 않는 한, 4096개의 토큰 길이와 모델 병렬 크기 1을 사용합니다. 훈련은 최고 학습률 \(1 \times 10^{-4}\), 4000단계의 워밍업 기간 및 학습률을 최고치의 1%까지 선형 감소시키는 방법을 포함했습니다.
모든 MoE 아키텍처에서, 모든 레이어에 MoE를 구현하고, 각 전문가의 훈련 용량 \(k_e\)를 \(\frac{b_{\mathcal{M}}}{\\|\mathcal{E}_{\mathcal{M}}\\|}\)로 설정하여 토큰당 FLOPs를 기본 밀집 모델과 비교하여 유사하게 유지하였습니다.
MoD 아키텍처에서는 0번 레이어부터 시작하여 번갈아 가며 MoD를 구현하고, 레이어 용량 계수 \(c_d\)를 25%로 설정합니다. 기본 밀집 모델과 FLOPs가 동등하도록 총 레이어 수를 늘리면서 히든 임베딩 차원을 일정하게 유지합니다.
4.2 성능의 계산적 스케일링
다양한 계산 수준에서 모델의 스케일링 성능을 제시합니다. 90M, 435M, 그리고 1.4B 파라미터의 세 가지 밀집 모델 크기와 그에 해당하는 희소 모델 변형을 FLOPs와 맞추어 비교합니다. 훈련 손실과 사전 훈련 속도 향상 계수 \(\eta\), 즉 희소 모델이 동일 FLOPs의 밀집 베이스 라인 모델과 동등한 사전 훈련 손실을 달성하는 데 필요한 전체 FLOPs의 \(\frac{1}{\eta}\)만 사용한다는 두 가지 주요 지표를 보고합니다.
4.3 전문가 수의 스케일링
전문가 수를 조정하는 영향을 조사하기 위한 추가적인 제거 실험을 수행했습니다. (1) 각 모달리티별로 동일한 수의 전문가를 할당하는 것과 (2) 서로 다른 수의 전문가를 할당하는 두 가지 시나리오를 탐구했습니다. 밀집 모델과 8개의 혼합 모달 전문가를 가진 전문가 선택 MoE (moe_8x)를 베이스 라인으로 포함했습니다.
4.4 업사이클링
2.3B MoD 모델에서 업사이클링의 영향을 추가적으로 조사했습니다. 특히, mod_moe_4t4i 모델을 스크래치에서 훈련시키는 것과 mod_moe_1t1i 체크포인트에서 초기화하여 훈련시키는 것을 비교했습니다. 공정한 비교를 위해 데이터 로더와 훈련 단계를 조정하여 mod_moe_1t1i에서 이미 완료된 훈련 단계 수를 고려하여 동등한 훈련 FLOPs를 유지했습니다.
4.5 처리량 분석
희소 모델은 추가된 동적성과 관련된 데이터 균형 문제로 인해 즉시 성능 향상을 이끌어내지 못할 수 있습니다. 제안된 방안이 훈련 효율성에 미치는 영향을 정량화하기 위해 여러 아키텍처의 훈련 처리량을 비교하는 통제 실험을 수행했습니다.
4.6 인퍼런스 시간 성능
OBELICS (Laurençon et al., 2023) 및 Shutterstock7의 서브셋에서 언어 모델링 데이터와 downstream 작업에 대해 모델을 평가했습니다. 또한, 사전 훈련된 언어 모델을 벤치마킹하는 데 일반적으로 사용되는 일련의 상식 인퍼런스 작업에 대한 0-샷 성능을 보고합니다.
5. 관련 연구
Early-Fusion 시각-언어 모델
Early-Fusion 기술은 처리의 시작부터 모달 간 상호 작용을 포착할 수 있다는 장점으로 멀티모달 학습에서 주목받고 있습니다. PerceiverIO(Jaegle et al., 2021)는 텍스트와 이미지를 포함한 다양한 입력에 대해 모달리티에 구애받지 않는 잠재 공간에서 작동하는 완전히 어텐션 기반의 읽기-처리하기-쓰기 아키텍처를 도입했습니다. NÜWA(Wu et al., 2021)는 텍스트, 이미지, 비디오를 다양한 조합으로 이해하고 생성할 수 있는 3D-어텐션 트랜스포머를 제시했습니다. CM3(Aghajanyan et al., 2022)는 인터넷의 혼합 모달 문서에서 학습할 수 있는 인과적 마스크 트랜스포머를 채택하여 이 패러다임의 대규모 사전 훈련 가능성을 보여주었습니다. 기본 트랜스포머 아키텍처로 Chameleon을 채택하고 모달리티 인식 희소성이 그 스케일링 성능을 효과적으로 향상시킬 수 있음을 입증했습니다.
멀티모달 표현 학습
BEIT-3(Wang et al., 2022b)는 시각 및 시각-언어 작업 모두에서 최고 성능을 달성하는 범용 멀티모달 기초 모델을 소개했습니다. 이는 모달리티별 전문가 모듈을 사용하는 Multiway 트랜스포머 아키텍처를 채용하고 이미지, 텍스트, 이미지-텍스트 쌍에 대한 사전 훈련을 위해 통합된 마스크 데이터 모델링 목표를 사용합니다.
희소 멀티모달 언어 모델
최근 희소 모델링 기술의 발전은 언어 모델을 효율적으로 확장하는 데 있어 유망한 결과를 보여주었습니다. VL-MoE(Shen et al., 2023)는 통합된 프레임워크 내에서 시각 및 언어 피드-포워드 네트워크에 희소 전문가의 혼합(MoE)을 적용하여 사전 훈련 목표를 사용하는 시각-언어 모델을 제안했습니다. VL-MoE는 밀집 모델보다 적은 파라미터를 사용하면서 시각-언어 작업에서 강력한 성능을 달성하며, 멀티모달 설정에서 MoE 훈련의 추가 증거를 제공합니다.
6. 한계점
현재 MoMa 구현은 데이터셋의 토큰 혼합 비율과 모델의 전문가 혼합 비율을 일치시켜 GPU 간의 부하 균형을 유지하는 데 의존하고 있습니다. 그러나 배치가 각 GPU 당 반복 수준에서 그 비율에서 벗어나는 것을 엄격히 제한하지 않기 때문에 소규모 불균형이 여전히 발생할 수 있습니다. 이런 개선을 위한 추가 연구가 필요합니다.
전문가 선택 라우팅
전문가 선택 라우팅은 훈련 중 전문가 부하 균형 문제를 완화하지만 인퍼런스 시 자동 회귀 언어 모델(LMs)에 추가적인 챌린지를 제시합니다(Zhou et al., 2022).
보조 라우터
보조 라우터는 네트워크 파라미터의 소수만을 차지하지만 그 역할은 중요하며, 전체 네트워크 훈련을 완료한 후 보조 라우터를 훈련시켰고, 이 과정을 몇 천 단계로 제한했습니다. 하지만 이전 연구에서는 전체 네트워크와 함께 이런 모듈을 공동으로 훈련할 가능성을 보여주었습니다(Raposo et al., 2024).
향후 연구는 보조 라우터의 아키텍처와 훈련 기법을 탐구하여 성능 병목 현상을 방지하고 다양한 데이터 분포에서의 일반화 가능성을 보장해야하며 특히, 보조 라우터와 원래 모델을 포함한 깊이의 혼합 아키텍처를 효과적으로 훈련하여 인과 인퍼런스 시나리오에서 효과적인 성능을 보장하는 것이 필요합니다.
MoD
MoD의 바닐라 수식과 MoE와의 단계적 통합만을 실험했습니다. 모달리티 인식 MoD를 포함한 다른 MoD 변형의 조사는 향후 작업으로 남겨두었습니다. 또한, 깊이의 혼합(MoD)을 사용한 일괄 시퀀스 생성은 표준 시퀀스 생성과 달리 동적 형태와 각 레이어에서 특정 시퀀스와 레이어가 다른 토큰을 건너뛸 때 KV 캐시의 동적 업데이트가 있기 때문에 자명하지 않습니다. MoD 모델의 인퍼런스 구현을 최적화하기 위한 추가적인 과제가 남아 있습니다.
7. 결론
이 작업에서는 Early-Fusion, 혼합 모달 기초 모델을 위한 일련의 모달리티 인식 희소 아키텍처를 도입했습니다. 접근 방식은 도메인 특수성을 활용하면서 모달 간 지식 공유와 기능 상호 작용을 보존합니다. 너비 차원(모달리티 인식 전문가의 혼합을 통해)과 깊이 차원(깊이의 혼합을 통해)에서 희소성을 통합함으로써 효율적이고 효과적인 모델 아키텍처를 개발했습니다.
최고 아키텍처인 Chameleon-MoMa는 최신 베이스 라인을 크게 개선하며, 복잡도 제어 실험을 통해 최대 3.7배의 전체 FLOPs 요구 사항을 줄이며, 모달리티 인식 희소 아키텍처가 경험적 스케일링 법칙을 유지한다는 것을 확인합니다.
이 특성은 접근 방식이 즉각적인 성능 개선을 보일 수 있으며, 혼합 모달 기초 모델 훈련의 미래 발전을 위한 확장 가능한 프레임워크를 제공한다는 것을 시사합니다.
Chameleon-MoMa을 통해 더 정교한 라우팅 메커니즘 탐색, 모달리티 간의 다양한 희소 패턴의 영향 조사, 그리고 접근 방식을 더 넓은 범위의 모달리티와 작업에 확장하는 것 등을 추가로 탐색해볼 필요가 있음을 언급합니다.