안녕하세요. 이번에는 Multimodal MoE(mixture of experts)에 관련된 논문을 리뷰하려고 합니다. 본 리뷰는 이전에 작성한 리뷰보다 먼저 작성했으나, 업로드 과정에서 문제가 있어 순서가 바뀌었네요. 본 글에서는 MoE의 개념을 먼저 다룰 예정이니 이 점을 참고해 주셨으면 좋겠습니다. 그럼 바로 리뷰를 시작하겠습니다.
Background
Mixture of Experts(MoE)는 원래 앙상블 기법 중 하나입니다. 이 개념은 복잡한 문제를 해결하기 위해 여러 개의 expert 모델이 각기 다른 하위 태스크에 특화되어 학습되도록 설계되었습니다. 일반적으로 앙상블 학습은 모든 모델이 동일한 전체 데이터셋을 학습하고, 그 출력 값들을 평균, 가중 평균, 또는 다수결 방식으로 결합합니다.하지만 MoE는 각 expert가 전체 데이터 중 자기한테 가장 잘 맞는 일부 데이터만을 학습하도록 설계되어 있습니다. 쉽게 말해, 하나의 큰 문제를 여러 개의 작은 하위 문제로 나누고, 각 하위 문제를 해결할 수 있는 expert 모델을 따로 학습시키는 구조입니다. 이 구조 덕분에 각 expert 모델은 자신이 잘 다룰 수 있는 데이터만 집중적으로 학습하니 훨씬 더 정밀하고 효율적인 예측이 가능해지겠죠.
MoE의 구성요소는 크게 3가지로 구성되어있습니다.
첫 번째, Expert model은 특정 sub-task에 특화된 모델들입니다. Expert가 sub-task를 수행하는 방식에는 각 expert가 특정한 데이터의 패턴을 학습해서 처리할 수도 있고, 각 expert를 모듈화 시켜서 멀티태스킹을 수행할 수도 있습니다. 예를 들어, NLP 모델에 각 언어에 특화된 MoE를 적용시킨다고 하면 각 Expert는 하나의 언어를 담당하는 모듈로 구성할 수 있습니다. 그리고 멀티모달 모델에 MoE를 적용시킨다고 하면 이미지,텍스트,오디오를 각각 처리하는 expert를 설정하여 모델을 구성할 수 있습니다.
두 번째 구성 요소는 Gating Network로, 흔히 Router라고도 불립니다. Gating Network는 MoE에서 핵심적인 역할을 하며, 입력을 받아 어떤 expert가 처리할지를 결정합니다. 예를 들어, 10개의 expert가 있을 경우 softmax를 이용해 상위 k개의 expert만 활성화시켜 메모리 사용량과 연산 비용을 최적화할 수 있으며, 특정 expert가 특정 역할을 수행하도록 유도할 수 있습니다. 하지만 Gating Network는 단순히 어떤 expert를 선택하는 역할만 수행하지 않습니다. 추가적인 기능으로 Clustering이 있으며, 이는 학습 과정에서 입력 데이터가 특정 expert로 라우팅되는 과정을 통해 데이터를 특성별로 그룹화하는 역할을 합니다. 예를 들어, 비슷한 유형의 문장은 동일한 expert에 배정되고, 유사한 이미지도 같은 expert에 배정되어 특화된 표현을 학습할 수 있습니다. 이 덕분에 단일 모델보다 더 다양한 데이터를 효과적으로 학습할 수 있는 장점이 있습니다
마지막은 Pooling method인데, experts가 각각 값을 출력했으면 그것들을 하나로 통합하는 과정이 필요하겠죠. 이때 Gating Network의 가중치를 활용해 experts의 출력을 조합하는게 Pooling Method입니다. 여기서 중요한 점은 Gating Network와 experts이 함께 학습된다는 것이며, 두 요소는 서로 영향을 주고받으며 전체 loss를 최소화하는 방향으로 최적화됩니다. 즉 Gating Network는 더 적절한 expert를 선택하는 방법을 학습하고, expert는 자신에게 주어진 task를 더 잘 학습하도록 학습되는 것입니다.
MoE는 모델 성능을 향상시키고 계산 비용을 유지할 수 있어 매우 이상적이지만, 학습 과정에서는 몇 가지 문제점이 존재합니다.
첫 번째는 Load Balancing 문제입니다. 이는 MoE에서 가장 큰 문제로 특정 Expert만 과도하게 선택되어 다른 Expert는 사용되지 않는 현상을 말합니다. 이렇게 되면 모델의 효율성이 떨어지고 과적합 위험이 커져 학습이 제대로 이루어지지 않을 수 있습니다.
두 번째는 Token Dropping 현상입니다. Token Dropping은 MoE 모델에서 특정 입력 토근이 어떤 Expert에도 배정되지 않는 현상을 의미합니다. 즉 Gating Network가 특정 토큰을 무시하거나 선택하지 않음으로써, 그 토큰이 모델에서 처리되지 않는 문제가 발생하는 것입니다. 이러한 현상이 발생하는 이유는 Expert Capacity와 Sparse Activation 때문입니다. MoE에서 일반적으로 한 Expert가 처리할 수 있는 토큰 개수(Expert Capacity)는 정해져있습니다. 그렇기 때문에 해당 Expert가 처리할 수 있는 토큰의 개수가 한도를 넘긴다면 나머지 토큰들은 drop해야만 하기에 정보 손실로 이어질 수 있습니다.
한편, MoE에는 Dense MoE와 Sparse MoE가 있습니다. Dense MoE는 모든 입력 토큰이 모든 Expert를 돌시에 활성화하는 방식으로 데이터 분배가 균등해져 학습이 안정적이고 일반화 성능이 우수하다는 장점이 있지만 그만큼 연산량도 크고 특정 Expert가 필요없는 경우에도 연산을 수행하기 떄문에 비효율적이라는 단점이 있습니다. 그렇기 때문에 MoE는 주로 Sparse MoE에서 연구가 진행되고 있습니다. Sparse MoE는 말 그대로 모든 Expert를 활성화시키는게 아니라 특정 Expert만 활성화시켜 계산 비용은 유지한채로 모델을 확장하는 방식입니다. 하지만 앞에서 말했다시피 각 Expert가 처리할 수 있는 용량은 한정적이므로 모델이 처리할 수 있는 토큰의 개수는 더 줄어들어 정보 손실로 인한 성능 저하가 발생할 수 있습니다. 이 밖에도 더 다양한 문제가 존재하지만 여기서는 일반적으로 다루어지는 문제만 기술하고 넘어가도록 하겠습니다.
1. Introduction
이제 본격적으로 해당 논문의 내용을 살펴보겠습니다. 논문에서 제안하는 모델은 LIMoE (Language-Image MoE)입니다. 이름에서 알 수 있다시피 멀티 모달 모델에 MoE를 적용한 방법론인데, 이전까지는 주로 하나의 모달(ex: 텍스트, 이미지)에서만 연구가 되었고 처음으로 멀티 모달에 MoE를 적용시킨 논문입니다. 저자는 멀티모달 학습 과정에서 발생하는 불안정한 현상을 분석하고, 이를 해결하기 위해 Entropy Regularization을 도입하여 안정적인 학습이 가능하도록 하였습니다. 이를 통해 동일한 연산량으로 Dense 모델보다 7~13% 높은 성능을 기록하였고, 단순히 모델을 개발하는 것을 넘어 각 구성 요소의 역할과 성능을 분석한 결과를 함께 제시합니다.
저자의 Contribution은 다음과 같습니다.
- 처음으로 Large-scale multimodal mixture of experts models 제안
- 기존 MoE 모델의 정규화 방식이 멀티모달 학습에서는 한계를 가지며, 이를 해결하기 위해 새로운 entropy-based regularisation 기법을 도입하여 학습을 안정화하는 방법을 제안
- LIMoE가 다양한 아키텍처 규모에서 일반화되며, 동일한 연산량 조건에서 Dense 모델 대비 ImageNet zero-shot 정확도가 7%~13% 향상됨을 보여줌. 더 확장된 모델인 LIMoE-H/14는 84.1%의 zero-shot ImageNet 정확도를 달성/
2. Multimodal Mixture of Experts
멀티모달의 학습은 일반적으로 모달리티별로 독립적인 인코더를 사용하는 투 타워(Two-Tower) 구조가 사용됩니다. 이 방식에서는 모달리티 𝑚에 따라 각 입력을 처리하는 별도의 모델이 존재합니다. 예를 들어, 이미지 𝑖는 𝑓_image(𝑖), 텍스트 𝑡는 𝑓_text(𝑡)와 같은 각각의 전용 인코더를 통해 개별적으로 처리하는 방식입니다.
하지만 이러한 구조는 모달 간 정보 공유가 어렵고, 확장성이 제한되는 단점이 있습니다. 이를 해결하기 위해 본 연구에서는 One-Tower 구조를 사용하는데, One-Tower 구조는 이미지와 텍스트를 비롯한 모든 모달리티를 하나의 모델에서 처리하는 방식으로, 투 타워 구조와 달리 모든 입력이 동일한 네트워크를 거쳐 학습이 됩니다.
이러한 접근 방식은 다음과 같은 장점을 가집니다. 하나의 모델로 다양한 모달리티를 처리할 수 있어 보다 유연한 학습이 가능해져 범용성(Generality)이 향상됩니다. 새로운 모달리티를 추가할 때도 기존 모델을 크게 변경하지 않고 쉽게 확장할 수 있다는 장점을 가지고 있습니다. 마지막으로 이미지와 텍스트 등의 정보를 동일한 모델에서 함께 학습하면서, 각 모달리티 간의 관계를 더 효과적으로 이해할 수 있어 모달 간 지식의 공유 가능성이 증가합니다.
2.1 Multimodal contrastive learning
모델의 학습에 있어서는 저희가 일반적으로 사용하는 contrastive learning으로 학습을 진행합니다.
2.2 The LIMoE Architecture
LIMoE는 One-Tower Transformer 아키텍처를 기반으로 이미지와 텍스트 두 가지 모달리티를 처리합니다. 각 모달리티는 linear layer를 통해 차원을 맞추게 되는데, 텍스트는 SentencePiece 방식으로 벡터로 변환되고, 이미지는 ViT 방식을 통해 임베딩됩니다. 그런 다음, 입력된 임베딩은 Dense Encoder Layer를 거쳐 기본적인 특징을 추출하고, MoE Encoder Layer에 전달됩니다. 이 과정에서 Self-Attention과 Router가 작동하여 입력 토큰을 E개의 Expert 중 K개만 선택해 처리합니다. 최종적으로, 모델의 출력은 Average Pooling을 사용해 각 모달리티를 하나의 벡터로 결합합니다
2.2.1 Challenges for multimodal MoEs
하지만 이전 Introduction에서도 말했듯이 각 Expert가 처리할 수 있는 용량은 정해져 있습니다. 이를 초과하면 일부 토큰은 버려지게 되고 그 출력은 모두 0이 됩니다. (이때 토큰이 처리된 비율을 success rate라고 하며 이 비율을 통해 모델의 학습 안정성을 평가하는 지표로 활용이 됩니다.) 만약 라우터가 대부분의 토큰을 한 Expert로만 보낸다면 용량이 초과되어 토큰이 버려지게 되고, 나머지 Expert는 훈련되지 않기 때문에 전처 모델의 성능 저하로 이어질 수 있습니다.
또한 데이터셋의 구성을 살펴보면 대부분 이미지 토큰이 훨씬 많고 텍스트는 적은 경우가 많습니다. 이 경우 이미지 토큰은 여러 Expert로 잘 분산될 수 있지만 텍스트는 한 Expert로 몰려 용량이 초과되는 문제가 발생합니다. 이러한 경우 이미지 토큰은 잘 분산되고 텍스트 토큰은 대부분 버려져도 Loss는 0으로 수렴할 수 있기 때문에 학습이 잘 되는 것처럼 보이지만 실제로는 전혀 학습이 안될 수도 있습니다. 따라서 저자는 모달리티 기반의 Local entropy loss와 Global entropy loss를 제안합니다.
Local entropy loss는 각 토큰 단위로 확률 분포의 엔트로피를 계산한 뒤 평균을 내는 방식으로 토큰들이 다양한 Expert로 선택되도록 유도하는 함수입니다. Global entropy loss는 배치 내 모든 토큰의 학률 분포에 대해 엔트로피를 계산해 모든 Expert가 골고루 사용되도록 유도하는 함수입니다.
그림 3에서는 이러한 Loss 설계 방식에 대한 성능을 보여주고 있습니다. 그림 3.a 그래프는 각각 ImageNet 10-shot , ImageNet zero-shot , 그리고 COCO text-to-image retrieval 성능을 나타냅니다. 기본 설정인 “classic”과 비교했을 때, local entropy loss를 추가한 “classic + ent local”은 성능이 향상되었고, 여기에 global entropy loss까지 포함한 “classic + ent local + global”에서는 가장 높은 성능을 기록했습니다. 이를 통해 Auxiliary losses을 추가할수록 모델의 성능이 점진적으로 개선됨을 확인할 수 있습니다.
그림 3.b에서는 Image routing success rate, Text routing success rate, Layer 5 text specialisations의 결과를 보여주고 있습니다. 주목할 점은 모든 Loss를 추가해줬을때 학습의 불안정성을 없앨 수 있다는 것이고 C2,C3를 비교했을때 텍스트 토큰이 여러 expert에게 고르게 분산되는 효과가 있다는 것을 확인할 수 있네요.
3. Experiments
이제 실험 부분을 살펴보겠습니다. 저자는 동일한 규모의 dense Transformer 모델과 LIMoE를 비교하고, LIMoE를 확장하면 더 높은 성능을 얻을 수 있다는 점을 보여줍니다. 학습 데이터셋은 웹에서 크롤링한 36억 개의 이미지와 Alt-텍스트로 구성되어 있습니다. LIMoE-H/14는 더 큰 규모의 모델로, JFT-4B 데이터셋을 추가해 학습됩니다. 평가 방법으로는 zero-shot 방식의 이미지 분류와 cross-modal retrieval을 진행하며, LIMoE와 추출한 표현에 linear adaptation을 적용하여 10-shot accuracy 성능을 평가합니다.
3.1 Controlled study across scales
저자는 Dense Transformer를 Baseline으로 설정하여 비교를 수행합니다. Dense Transformer와 LIMoE는 Layer 수를 각기 다르게 설정하여( S,B,L 순) 학습시키고 LIMoE는 1개의 Expert만 활성화시켜 학습과 추론을 진행합니다.
그림 4에서는 FLOP(연산량) 대비 성능을 비교합니다. 여기서 S, B, L은 각각 Layer 수의 크기를 나타냅니다. 동그라미는 LIMoE를, 세모는 베이스라인을 나타냅니다. 결과를 보면, 비슷한 연산량에도 불구하고 모든 모델에서 성능이 크게 향상된 것을 확인할 수 있습니다. 첫 번째와 두 번째는 zero-shot과 10-shot에 대한 성능을 나타내며, 평균 accuracy가 각각 10%, 12% 증가했는데 모델의 효율성이 크게 향상되었음을 알 수 있네요.
3.2 Scaling up LIMoE
저자는 기존 LIMoE L/16 (24 layer)의 크기를 증가시킨 32-layer LIMoE-H/14 모델로도 실험을 진행합니다. 이는 12개의 Expert Layer를 가지며 각 Layer에는 32개의 Expert가 구성되며 이중 1개의 Expert만 활성화됩니다. 표 1에는 당시 SOTA contrastive models들과의 성능을 비교한 결과가 나와있습니다.
다른 모델들과 비교해 봤을때 상대적으로 작은 모델 크기를 가졌음에도 불구하고 크기가 큰 모델과 성능이 맞먹는 모습을 보이고 있습니다. 또한 LIMoE만 One-Tower 구조라는 것을 생각해보면 각 모달리티를 독립적으로 처리하는 것이 아닌 하나의 모델에서 두 모달리티를 처리하므로 서로 다른 모달리티 간의 상호 작용을 잘 학습할 수 있으며 메모리와 계산 비용에서도 큰 이점을 가져가는 것 같습니다.
3.3 Model Analysis
마지막으로 정성적 결과를 살펴보고 마무리하겠습니다. 이 결과는 CoCo 데이터셋에서 각 Expert가 처리한 이미지 패치를 모은 것입니다. 보시면 Expert 2는 식물 관련 패치를, Expert 7은 눈 관련 패치를, Expert 19는 바퀴 관련 패치를 처리하는 모습을 확인할 수 있습니다. 이를 통해 각 토큰에 가장 적합한 Expert를 선택함으로써, 특정 범주에 대한 정보를 선별적으로 학습하는 방식이 잘 드러난다고 할 수 있습니다.
감사합니다.
