안녕하세요 이번에 리뷰할 논문은 2025년 Apple에서 발표한 논문으로 CLIP 모델에 Mixture-of-Experts(MoE)를 적용하여 다양한 실험을 통해 CLIP 모델 확장 과정에서 발생할 수 있는 문제점을 분석하고 새로운 학습 전략을 제안한 논문입니다. 저 또한 비디오 도메인에서 CLIP을 backbone으로한 CLIP4clip 모델에 MoE를 적용시키고자 하고 있기에 이번 논문에서 얻어갈 것이 많아 리뷰를 하게 되었습니다. 그럼 바로 시작하겠습니다.
1. Introduction
CLIP 모델은 대규모의 웹 데이터로 사전학습되어, 다양한 테스크에서 높은 성능을 보여주고 있습니다. 이에 따라 여러 분야에서 CLIP을 활용하는 연구가 늘고 있고 더 나은 성능을 위해 모델을 확장하는 것이 점점 더 중요해지고 있습니다. 하지만 지금까지 CLIP의 확장은 대부분 dense 모델의 크기를 키우는데 초점을 맞춰왔습니다. 여기서 Dense 모델이란, 모든 레이어의 파라미터가 활성화되어 있는 일반적인 신경망 구조를 의미합니다. 그러나 이러한 Dense 모델은 모든 파라미터를 매번 활성화하기 때문에 계산량이 많고 대규모 모델로 확장할 때 비효율적이라는 단점이 있습니다.
그래서 모델을 확장하는 다른 방식으로는 sparse 모델링 기법이 있습니다. sparse 모델링이란 모델의 일부 파라미터만 선택적으로 활성화하여 연산 효율을 높이고, 필요한 부분에만 집중해 학습과 추론을 수행하는 방법을 말합니다. 이 중 대표적인 방법이 Mixture-of-Experts(MoE)입니다. MoE는 여러 개의 작은 전문가 네트워크(Experts) 중에서 입력에 가장 적합한 일부만 선택해 활성화함으로써, 전체 모델을 크게 키우면서도 계산 효율을 유지할 수 있게 해주는 특징을 가지고 있습니다. 하지만 CLIP과 같은 대규모 모델을 처음부터 학습하는 데는 여전히 많은 시간과 계산 비용이 들며, MoE 모델의 경우 학습 과정에서 불안정성이 발생하기 쉬워 이를 안정화하기 위한 auxiliary loss 설계와 최적화가 중요한 문제로 남아 있습니다.
제가 이전에 리뷰한 논문 중 LIMOE는 텍스트와 이미지 인코더가 공유된 백본을 사용해 Multi Modal으로의 첫 확장 연구를 시도한 연구입니다. LIMOE는 기존 Dense model 보다 성능에서는 큰 이점을 가졌지만 계산 자원 비용 측면에서는 1.35배 더 많은 학습 FLOPs를 필요로 한다는 단점도 가지고 있었습니다. 또한 추가적인 auxiliary loss를 도입해야 했기에 학습 과정이 더 복잡해지고 하이퍼파라미터 조정이 어려워지는 문제점도 가지고 있습니다.
따라서 저자는 대안적인 학습 전략으로 사전 학습된 CLIP 모델을 Spare MoE 아키텍처로 upcycling 하는 기법을 제안합니다. 여기서 Sparse upcycling은 사전학습된 모델을 활용하여 MoE layer가 포함된 모델로 활용하는 방식을 의미합니다.
위의 그림 1을 보시면 여러 가지 학습 전략을 통해 모델을 평가한 모습을 확인할 수 있습니다. 저자는 총 4가지 학습 전략을 제시하는데(Recipe 1,2,3,4,), 자세한 셋팅은 이후 섹션에서 말씀드리고 결론부터 말씀드리면 텍스트와 이미지 인코더를 분리시키고 Sparse upcycling 기법을 적용 시켰을때 성능이 가장 많이 오르는 것을 확인할 수 있습니다. 또한 모델이 얼마나 많은 계산을 수행하는지 나타내는 ZFLOPS에서도 처음부터 학습시키는 Scratch training에 대해 값을 4.2에서 3.7로 줄여 학습의 효율성을 개선할 수 있음을 보여주었습니다. 따라서 저자가 제안하는 기법의 Contribution은 다음과 같이 요약할 수 있습니다.
1. 저자가 제안하는 CLIP-UP은 사전학습된 가중치로부터 sparse upcycling을 활용하여 Mixture-of-Experts (MoE) CLIP 모델을 구축하였으며, 복잡한 auxiliary loss 없이도 성능을 낼 수 있음을 증명함
2. CLIP-UP은 다양한 Text-Image retrieval benchmarks에서 기존 dense모델보다 훨씬 나은 성능을 보여줌
3. CLIP-UP은 B/32부터 L/14까지 다양한 모델 크기에 확장 가능하며, 설계 요소에 대한 포괄적인 분석을 제시함
2. Sparse Upcycling with Mixture-of-Experts
그럼 이제 본격적으로 저자가 제안하는 CLIP-UP 모델의 아키텍쳐 구조를 살펴보도록 하겠습니다. CLIP의 구조와 학습 방식은 생략을 하고, CLIP 모델을 MoE layer로 대체하는 방식과 Sparse Upcycling Training 접근 방식으로 나누어 설명드리겠습니다.
2.1 CLIP with Mixture-of-Experts upcycling
CLIP-UP에서는 전체 Transformer layer의 절반에서 dense MLP들을 sparsely-gated MoE layer로 교체하여, dense-sparse-dense의 구조로 설계했습니다. 이러한 구조는 모델의 forwarding 과정에서 안정성을 높일 수 있고, 초기 학습 시 gradient oscillation을 방지하는 데 도움을 줄 수 있습니다. 각 MoE layer는 E개의 expert와 하나의 router로 구성됩니다. expert는 입력에 따라 선택적으로 활성화되는 MLP이며 router는 각 입력 token에 대한 gating logit을 계산하여 총 E개의 expert 중에서 top-K expert를 선택하는 역할을 수행합니다. 구체적으로는 X_j를 MoE layer에 대한 j번째 token의 입력이라고 할때, G_j 는 E개의 expert에 대한 gating logit, W_e는 expert e에 대한 router weight matrix를 나타내며 최종 출력MoE(X_j) 는 아래 수식과 같이 계산됩니다.
기본적으로 MoE의 Expert는 한 번에 처리할 수 있는 용량이 정해져있습니다. 이를 buffer capacity라고 부르는데, 이 같은 제약 조건을 두는 이유는 expert가 과부화되는 것을 방지하고 hardware나 memory resource를 관리하기 위한 이유가 있고 동시에 계산의 효율성도 향상시키기 위해 이 같은 조건을 설정해둡니다. 하지만 이 같은 효과를 얻기 위해서는 모델의 학습이 더 어려워진다는 단점이 존재합니다. 왜냐하면 MoE를 도입함으로써 이루고자하는 목표는 입력된 토큰들에 대해 모든 expert들이 균형있게 사용되는 것을 목표로 합니다. 하지만 이러한 균형이 깨져버리면 특정 expert만 과도하게 사용이 되고 그 expert에 배정된 토큰들의 개수가 buffer capacity를 초과한다면 이는 정보 손실로 이루어져 최종적으로는 성능에 악영향을 끼칠 수 있습니다. 그렇기 때문에 이전 MoE의 연구에서는 Auxiliary loss를 잘 설계하여 이러한 문제를 방지하는데 초점을 맞춰왔습니다.
한편 저자는 추가적인 Auxiliary loss를 설계하는 대신 이전 연구에서의 loss를 활용하여 모델을 설계합니다. 먼저 저자는 2022년에 발표된 ST-MoE의 설계에 따라 각 modality에 대해 load balance loss와 router z-loss로 구성된 auxiliary loss를 도입하여 토큰이 균등하게 처리되도록 설계했습니다. Load balance loss는 gating mechanism에서 총 E개의 expert에 대해 token을 균일하게 할당하는 역할을 합니다. 이는 길이가 S인 sequence에 대해 다음과 같은 수식으로 정의할 수 있습니다.
여기서 R_e는 expert e에 대한 token 할당 비율을 나타내고, P_e는 expert e의 평균 router probability를 나타냅니다.
Router z-loss는 router logit의 출력을 정규화하여 gating 연산을 안정화하는 역할을 수행합니다. 수식으로는 아래와 같이 표현됩니다.
따라서 최종 Loss는 아래와 같이 contrastive loss와 auxiliary loss에 대해 scaling factor를 곱해 표현할 수 있습니다.
ST-MoE와 마찬가지로 2022년에 발표된 LIMOE 또한 새로운 Auxiliary loss를 제안을 합니다. 이는 하나의 expert로 집중을 유도하는 local entropy loss와 expert의 분산을 유도하는 global entropy loss 로 구성이 되어있습니다.
그리고 LIMoE의 Loss의 경우 각 모달리티별로 Loss가 설계된다는 특징을 가지고 있습니다. 하지만 학습 중 균형이 무너지면 한 모달리티가 너무 많은 expert를 사용하게 되고, 이로 인해 특정 expert에 과도하게 의존하게 되어 과적합이 발생할 수 있습니다. 이를 방지하기 위해 global entropy loss에 threshold를 적용하여, 최소한의 expert 수(m개)를 사용하도록 유도함으로써 모달리티 간 균형을 유지하고 학습의 안정성을 강화하는 부분을 추가합니다. 따라서 LIMOE의 Loss까지 추가한 최종 Loss는 다음과 같이 표현할 수 있습니다.
하지만 저자의 실험 결과에 따르면 LIMOE의 방식은 백본을 공유하는 구조에서는 성능을 확실하게 올려주기는 했지만 이미지와 텍스트를 따로 분리한 구조에서는 성능이 떨어졌다고 합니다. 따라서 저자는 LIMoE Loss를 제거한 수식 7번을 최종 Loss로 선택했습니다.
2.3 Sparse Upcycling Training
Sparse Upcycling Training은 이미 잘 훈련된 Dense CLIP 모델의 가중치를 재활용(upcycling)해서 MoE 구조로 바꾸고, 거기서부터 빠르고 안정적으로 훈련을 이어가는 방식을 의미합니다. 구체적으로는 각 expert는 dense CLIP 모델의 dense MLP 레이어에서 동일한 MLP 가중치를 사용하여 초기화가 되고 router 가중치는 무작위로 초기화됩니다. 그리고 Attention, embedding layer와 같은 다른 모든 layer는 원래 모델에서 복사를 하여 사용합니다. 이러한 학습 방식은 모델의 일부분만을 수정하도록 하며 대부분의 파라미터는 변경되지 않아 모델의 본래 기능은 유지됩니다. 초기화를 마친 후에는 batch 크기와 같은 원래 hyperparameter는 유지하고 학습률과 weight decay를 약간 줄여 추가 학습을 진행합니다.
3. Experiment
이제 실험 부분을 살펴보도록 하겠습니다. 먼저 모든 모델(initial dense CLIP, CLIP-UP, baseline dense CLIP)은 WIT-3000M, DFN-5B의 동일한 이미지-텍스트 쌍 데이터셋을 사용하여 학습하였고 ImageNet (classification), COCO (retrieval), Flickr30K (retrieval)의 벤치마크에서 평가하였습니다.
3.1 Recipe Study
먼저 shared backbone과 separated backbone의 성능 차이, 그리고 scratch training과 Sparse Upcycling 방식을 사용하는 것의 영향을 비교한 실험이고 모든 실험은 CLIP-B/16 모델을 기반으로 수행되었으며, 결과는 아래 Table 1에 정리되어 있습니다
결과는 separated backbone과 sparse upcycling의 조합이 전반적으로 가장 높은 성능을 보여주었습니다. separated backbone은 일반적으로 shared backbone보다 성능이 뛰어난데, 이는 각 modality에 할당된 추가 parameter 때문일 가능성이 크다고 저자는 분석하고 있습니다. 그리고 scratch vs sparse upcycling을 비교했을때 sparse upcycling이 모든 backbone의 성능을 향상시키는 것을 확인할 수 있는데 이는 sparse upcycling이 기존 모델의 표현력을 잘 살리면서 일부만 효율적으로 학습한 결과로 해석할 수 있을 것 같습니다.
그리고 저자는 LIMoE의 auxiliary loss가 성능에 미치는 영향에 대한 실험도 진행합니다. 결과는 위 그림 3에 나와있는데 처음부터 학습된 shared 백본에 LIMOE 손실을 적용하면 성능이 크게 향상되는 것을 확인할 수 있습니다. 그러나 다른 조합들과 비교 했을때는 낮은 성능을 보여주었고 Separated + Upcycle의 조합에 적용했을때는 retrieval 성능은 약간 향상되었지만 classification 성능은 오히려 떨어지는 모습을 보여주고 있습니다. 이러한 새로운 auxiliary loss를 추가하면 조정해야할 하이퍼파라미터가 늘어나기 때문에 모델의 학습이 불안정해질 수 있고, 저자 역시 리소스의 제약으로 인해 최적의 하이퍼파라미터를 찾는 대신 수식 7의 Loss를 최종 Loss로 사용했다고 언급하고 있습니다.
3.2 Final Model Evaluation and Baselines
이전 섹션의 결과에 따라 저자는 Separated Backbone + Sparse Upcycling 조합을 기본 전략으로 선택하고, 이 방법의 효과를 더 검증하기 위해 다양한 모델 크기(B/32부터 L/14까지)에서 CLIP-UP 모델을 학습시켰습니다. Table 2에서는 baseline dense CLIP 모델과 CLIP-UP 모델의 zero-shot classification , retrieval 성능을 비교하고 있습니다. Dense CLIP의 경우 440k step에서 790k step으로 늘려도 성능 향상은 적었지만, CLIP-UP은 모든 지표에서 성능이 향상되는 것을 보여주었습니다.
추가로 CLIP-UP B/32 모델이 추론 시 GFLOPS(연산량)의 47%만 사용하면서도 baseline dense CLIP B/16보다 COCO T2I recall@1에서 2.4% 더 높은 성능을 보였고, I2T recall@1 성능은 동등했습니다. 마찬가지로, CLIP-UP B/16은 inference GFLOPS의 31%만 사용하면서도 dense CLIP L/14보다 COCO T2I recall@1에서 1.9% 더 높고 I2T recall@1 에서 같은 수준의 성능을 유지했습니다
3.3 Training Efficiency
마지막으로 그림 4에서는 dense pretraining + upcycling과 scratch 학습의 과정을 비교하여 나타냈습니다. 이 결과를 통해 초기에 dense 모델을 먼저 잘 학습시켜두면, 그 위에 sparse upcycling 방식을 적용할 때 훨씬 더 안정적이고 효과적으로 성능을 끌어올릴 수 있는 것을 보여주고, 반대로 사전 학습된 dense 체크포인트를 활용하지 않고 처음부터 학습하면 CLIP-UP의 성능 수준에 도달하는 데 훨씬 더 많은 계산 비용이 드는 것을 확인할 수 있습니다.
