Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

이 논문은 언어와 시각 토큰의 본질적 차이를 고려하여 조건부 및 프로토타입 라우팅을 통해 전문가 특화를 유도하는 'ProMoE' 프레임워크를 제안함으로써, 기존 확산 트랜스포머에 적용된 MoE 의 한계를 극복하고 ImageNet 에서 최첨단 성능을 달성했다고 요약할 수 있습니다.

핵심

🎨 1. 문제점: "모든 일을 한 사람이 다 하려고 하면 지친다"

지금까지 이미지를 만드는 AI 는 '모든 전문가 (Expert)'가 매번 모든 일을 다 하는 방식을 썼습니다.

• 비유: 마치 한 식당에서 **모든 요리사 (전문가)**가 손님이 "김치찌개"를 시키든 "스테이크"를 시키든, 모든 요리사가 동시에 모든 요리를 준비하는 것과 같습니다.
• 결과: 식당이 커질수록 (모델이 커질수록) 요리사들이 너무 많아서 주방이 붐비고, 전기세 (컴퓨팅 비용) 가 너무 많이 나옵니다.

그래서 사람들은 '모자이크 (MoE, Mixture of Experts)' 방식을 도입했습니다.

• 비유: "김치찌개"가 오면 김치찌개 전문 요리사만, "스테이크"가 오면 스테이크 전문 요리사만 나오게 하는 거죠. 나머지 요리사들은 쉬게 됩니다.
• 문제: 언어 (텍스트) AI 에서는 이 방식이 잘 먹혔는데, 이미지 AI 에서는 잘 안 먹혔습니다. 왜일까요?

🧐 2. 왜 안 됐을까? "이미지는 텍스트와 다르다"

저자들은 두 가지 이유를 발견했습니다.

1. 이미지는 너무 비슷하다 (공간적 중복성):
   • 비유: 텍스트는 "사과", "자동차"처럼 단어가 명확하게 다릅니다. 하지만 이미지는 픽셀들이 서로 너무 비슷합니다. (예: 하늘의 구름 조각들은 다 비슷해 보임). 그래서 AI 가 "어떤 요리사가 이 조각을 담당해야 할지" 헷갈려서, 모든 요리사가 똑같은 일을 하려고 합니다.
2. 역할이 섞여 있다 (기능적 이질성):
   • 비유: 이미지 생성 AI 는 두 가지 모드로 작동합니다.
     • 조건부 (Conditional): "고양이 사진을 그려줘"라고 할 때.
     • 무조건부 (Unconditional): "아무거나 그려줘"라고 할 때.
   • 기존 방식은 이 두 가지를 구별하지 않고 똑같이 처리해서, 요리사들이 혼란을 겪었습니다.

💡 3. 해결책: ProMoE (두 단계로 나누는 똑똑한 지시자)

이 문제를 해결하기 위해 ProMoE라는 새로운 시스템을 만들었습니다. 핵심은 **"명확한 지시 (Explicit Routing Guidance)"**를 주는 것입니다.

1 단계: 역할 나누기 (Conditional Routing)

• 비유: 식당 입구에 수위가 서 있습니다.
  • "고양이 그려줘"라고 하면? -> **고양이 전문 요리사 (Conditional Expert)**에게 보냅니다.
  • "아무거나 그려줘"라고 하면? -> **일반 요리사 (Unconditional Expert)**에게 보냅니다.
• 효과: 요리사들이 자신의 역할 (고양이 vs 아무거나) 을 명확히 구분해서 집중할 수 있게 됩니다.

2 단계: 세부 배정 (Prototypical Routing)

• 비유: 고양이 요리사들 사이에서도 "진한 갈색 고양이"를 그릴 사람과 "하얀 고양이"를 그릴 사람을 구분해야 합니다.
  • ProMoE 는 **학습된 '표본 (Prototype)'**이라는 가이드북을 사용합니다.
  • "이 그림 조각은 갈색 고양이와 비슷하네?" -> 갈색 고양이 전문 요리사에게 보냅니다.
  • "이건 하얀 고양이와 비슷하네?" -> 하얀 고양이 전문 요리사에게 보냅니다.
• 효과: 비슷한 것끼리 같은 전문가에게 몰려가게 되어, 각 전문가가 매우 특화된 기술 (전문성) 을 갖게 됩니다.

3 단계: 팀워크 강화 (Routing Contrastive Loss)

• 비유: 요리사들이 서로 너무 비슷해지지 않도록, 서로 다른 스타일을 유지하라고 독려하는 규칙을 만듭니다.
  • "너희는 서로 다른 맛을 내야 해!"라고 가르쳐서, 한 요리사가 모든 것을 다 하려는 것을 막고, 각자 고유한 영역을 지키게 합니다.

🚀 4. 결과: 왜 이것이 대단한가?

• 더 적은 비용으로 더 좋은 결과: ProMoE 는 기존 방식보다 **적은 요리사 (활성화된 파라미터)**를 쓰면서도, **더 많은 요리사 (전체 파라미터)**를 쓰는 기존 최강 모델보다 더 멋진 그림을 그립니다.
• 빠른 학습: 요리사들이 역할을 명확히 알기 때문에, 더 빨리 배우고 더 빨리 완성된 그림을 만들어냅니다.
• 다양한 상황 대응: 텍스트로 "파란 코끼리"를 요청하든, "달 위의 에펠탑"을 요청하든, 모든 전문가들이 제 역할을 잘 해냅니다.

📝 한 줄 요약

"이미지 AI 에게 '누가 무엇을 할지' 명확하게 지시하고, 비슷한 것끼리 묶어주니, 적은 인원으로 더 전문적이고 멋진 그림을 그릴 수 있게 되었습니다."

이 기술은 앞으로 더 크고 멋진 AI 이미지를 만들 때 필수적인 도구가 될 것입니다.

기술 요약

1. 문제 정의 (Problem)

• 배경: 혼합 전문가 (Mixture-of-Experts, MoE) 아키텍처는 대규모 언어 모델 (LLM) 에서 모델 용량을 확장하면서도 계산 효율성을 유지하는 데 성공적으로 적용되었습니다. 그러나 Diffusion Transformer (DiT) 에 MoE 를 적용하려는 시도는 LLM 에서만큼의 성능 향상을 거두지 못했습니다.
• 근본 원인: 언어 토큰과 시각 (이미지) 토큰의 본질적인 차이 때문입니다.
  1. 높은 공간적 중복성 (High Spatial Redundancy): 텍스트 토큰은 의미적으로 밀집되어 있고 토큰 간 차이가 뚜렷한 반면, 이미지 패치 (시각 토큰) 는 연속적이고 공간적으로 결합되어 있어 중복성이 매우 높습니다. 이로 인해 전문가 (Expert) 들이 동질적인 특징을 학습하게 되어 전문화가 어렵습니다.
  2. 기능적 이질성 (Functional Heterogeneity): 확산 모델은 Classifier-Free Guidance (CFG) 를 사용하며, 이는 '조건부 (conditional)'와 '비조건부 (unconditional)'라는 두 가지 기능적으로 다른 입력 토큰을 생성합니다. 기존 MoE 는 이를 구별하지 않고 균일하게 라우팅하여 전문가의 역할 분담을 방해합니다.
• 결과: 기존 DiT-MoE 방법들은 전문가 다양성 (Inter-expert Diversity) 과 일관성 (Intra-expert Coherence) 을 확보하지 못해 밀집 (Dense) 모델보다 성능이 낮거나 미미한 개선만 보였습니다.

2. 제안 방법: ProMoE (Methodology)

저자들은 ProMoE를 제안하며, 이는 **명시적 라우팅 가이드 (Explicit Routing Guidance)**를 갖춘 2 단계 라우터 (Two-step Router) 를 핵심으로 합니다.

A. 2 단계 라우터 (Two-Step Router)

1. 조건부 라우팅 (Conditional Routing):

   • 토큰의 기능적 역할에 따라 토큰을 먼저 분할합니다.
   • 비조건부 이미지 토큰 (null conditioning 하의 패치) 은 전용 '비조건부 전문가 (Unconditional Expert)'로 고정 라우팅됩니다.
   • 조건부 이미지 토큰 (특정 조건 하의 패치) 은 2 단계로 넘어가 세부 라우팅을 받습니다.
   • 이를 통해 기능적 이질성을 해결하고 전문가 간 기능 분리를 강제합니다.

2. 원형 라우팅 (Prototypical Routing):

   • 조건부 토큰에 대해 학습 가능한 원형 (Learnable Prototypes) 집합을 사용하여 세부 할당을 수행합니다.
   • 각 원형은 특정 전문가와 연결되며, 토큰 임베딩과 원형 간의 코사인 유사도를 계산하여 라우팅 가중치를 결정합니다.
   • 이는 잠재 공간 (Latent Space) 에서 의미적 유사성에 기반한 할당을 가능하게 합니다.
   • 활성화 함수: Softmax 대신 Identity 함수를 사용하여 계산 효율성과 안정성을 높였습니다.

B. 라우팅 대비 손실 (Routing Contrastive Loss, RCL)

• 문제: 기존 클러스터링 기반 라우팅은 토큰 수준의 유연성이 부족하거나 수동 레이블이 필요했습니다.
• 해결: RCL 을 도입하여 원형 라우팅을 명시적으로 강화합니다.
  • 인내 전문가 일관성 (Intra-expert Coherence): 의미적으로 유사한 토큰이 같은 전문가로 라우팅되도록 원형을 해당 토큰 군집의 중심 (Centroid) 쪽으로 당깁니다.
  • 전문가 간 다양성 (Inter-expert Diversity): 서로 다른 전문가로 할당된 토큰 군집은 원형이 서로 멀어지도록 밀어냅니다.
  • 이 과정은 수동 레이블 없이도 작동하며, 전통적인 로드 밸런싱 손실보다 더 효과적으로 전문가 전문화를 유도합니다.

C. 공유 전문가 (Shared Experts)

• 모든 토큰을 처리하는 공유 전문가를 포함하여 일반적인 지식을 학습하게 하고, 조건부/비조건부 토큰에 따라 선택적 전문가를 추가하여 활성화 파라미터 수를 일정하게 유지합니다.

3. 주요 기여 (Key Contributions)

1. 언어 - 시각 토큰 차이 분석: 시각 토큰의 공간적 중복성과 기능적 이질성이 MoE 의 비효율성을 유발함을 규명하고, 이를 해결하기 위한 명시적 가이드의 필요성을 제시했습니다.
2. ProMoE 아키텍처: 기능적 역할 (조건부/비조건부) 과 의미적 내용 (Prototypes) 에 기반한 2 단계 라우팅 메커니즘을 설계했습니다.
3. 라우팅 대비 손실 (RCL): 수동 레이블 없이 의미적 가이드를 강화하여 전문가 간 다양성과 일관성을 동시에 확보하는 새로운 손실 함수를 제안했습니다.
4. 범용성 및 확장성: Rectified Flow 와 DDPM 두 가지 학습 목표 모두에서 Dense 모델 및 기존 MoE SOTA 를 능가하는 성능을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋: ImageNet (256x256) 기준.
• 성능 비교:
  • Dense 모델 대비: 동일한 활성화 파라미터 수 (Activated Params) 를 가질 때, ProMoE 는 Dense DiT 보다 FID 를 크게 낮추고 IS 를 높였습니다. (예: ProMoE-L-Flow 는 Dense-DiT-XL-Flow 보다 적은 활성화 파라미터로 더 좋은 FID 달성).
  • 기존 MoE SOTA 대비: DiT-MoE, EC-DiT, DiffMoE 등 기존 방법들보다 우월한 성능을 보였습니다. 특히 ProMoE-L-Flow(총 파라미터 1.06B) 는 DiffMoE-L-Flow(총 파라미터 1.84B) 보다 더 적은 파라미터로 더 높은 성능을 기록했습니다.
  • 학습 목표: Rectified Flow 와 DDPM 모두에서 일관된 개선을 보였으며, Rectified Flow 에서 더 큰 이득을 얻었습니다.
• 확장성 (Scaling): 모델 크기 (S, B, L, XL) 와 전문가 수를 늘릴수록 성능이 선형적으로 향상되는 것을 확인했습니다.
• 계산 효율성: 추론 시간과 GFLOPs 면에서 기존 MoE 방법들과 유사하거나 더 낮으면서도 훨씬 높은 성능을 달성했습니다.
• 텍스트 - 이미지 생성: GenEval 벤치마크에서도 Dense 모델 및 Token-Choice MoE 대비 우수한 일반화 능력을 입증했습니다.

5. 의의 및 중요성 (Significance)

• DiT 확장성의 새로운 패러다임: MoE 가 DiT 에서 실패했던 이유를 명확히 분석하고, 시각 데이터의 고유한 특성 (중복성, 기능적 이질성) 에 맞춘 라우팅 전략을 제시함으로써 확산 모델의 대규모 확장을 가능하게 했습니다.
• 효율성과 성능의 균형: 더 적은 활성화 파라미터로 더 큰 Dense 모델이나 더 많은 파라미터를 가진 기존 MoE 모델보다 우수한 성능을 내어, 계산 자원이 제한된 환경에서도 고품질 생성 모델을 구축할 수 있는 길을 열었습니다.
• 시각 MoE 설계 원칙 정립: "명시적 라우팅 가이드"가 시각 MoE 의 성공에 필수적임을 증명하여, 향후 시각 모델의 MoE 설계에 중요한 지침을 제공합니다.

이 논문은 확산 모델의 효율적인 확장을 위해 MoE 아키텍처를 어떻게 재설계해야 하는지에 대한 이론적 통찰과 실증적 성과를 모두 제공한다는 점에서 의의가 큽니다.