MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

기술 요약: MMOE: 효율적인 전문가 설계를 통한 확산 트랜스포머의 현대화

문제 정의
대규모 언어 모델(LLM)은 희소 혼합 전문가(sparse Mixture-of-Experts, MoE) 및 효율성 메커즘(예: 제로-연산 전문가, 게이트-잔차 라우팅)을 사용하여 전체 파라미터 용량과 활성화된 연산량을 분리함으로써 성공적으로 규모를 확장해 왔습니다. 반면, AIGC 파운데이션 모델(AFM), 특히 확산 트랜스포머(diffusion-transformer) 백본은 이와 다른 궤적을 따라왔습니다. 확산 트랜스포머에 MoE를 적용하려는 현재의 노력(예: DiT-MoE, Race-DiT)은 주로 전체 전문가 풀과 희소 비율을 늘리는 데 집중되어 있으며, LLM의 스케일링을 실용적으로 만든 효율성 메커니즘을 도입하지는 못했습니다. 이러한 접근 방식은 모델을 자원이 제한된 장치에서 학습 및 배포하기 어렵게 만들며, 생성 품질과 학습 및 배포 비용 사이의 균형을 맞추는 데 실패합니다. 본 논문은 다음과 같은 질문을 던집니다: "효율적인 LLM 스케일링의 구조적 원칙을 AFM에 적응시켜, 유효 연산량이나 학습 부담을 비례적으로 늘리지 않고도 생성 품질을 개선할 수 있는가?"

방법론: ModernMOE (MMOE)
저자들은 SiT 스타일의 확산 트랜스포머를 현대화한 **ModernMOE (MMOE)**를 제안하며, 이는 LLM에서 입증된 네 가지 효율성 메커니즘을 확산 백본에 체계적으로 통합합니다. 결정적으로, MMOE는 외부 인터페이스(디노이징 목적 함수, 컨디셔닝, 잠재 표현, 샘플러)는 변경하지 않고 내부 블록 수준의 연산만을 수정합니다.

이 아키텍처는 다음 구성 요소들을 도입합니다:

1. MoE++ 전문가 레이어: 표준 라우티드 MLP 대신, 전문가 풀은 다음을 포함합니다:
   • 헤비 전문가(Heavy Experts): 표준 MLP.
   • 경량 제로-연산 전문가(Lightweight Zero-Computation Experts): "복사(Copy)"(항등), "제로(Zero)"(출력 0), "상수(Constant)"(학습된 벡터 혼합) 전문가. 이를 통해 모델은 무거운 업데이트가 불필요한 토큰에 대해 전체 MLP 실행을 건너뛸 수 있습니다.
2. 게이트-잔차 라우팅(Gate-Residual Routing): 라우터는 이전 전문가 레이어의 게이트 결정으로부터 잔차(residual)를 포함하여, 라우팅 경향이 레이어 전반에 걸쳐 상속되고 결정의 안정성을 높이도록 합니다.
3. 공유 전문가(Shared Experts): 토큰 의존적 특성화를 위한 라우티드 전문가와 함께, 공통 연산 경로를 제공하기 위해 공유 밀집 전문가가 항상 활성화됩니다.
4. 어텐션-잔차 집계(Attention-Residual Aggregation): 모델은 이전에 완료된 블록 상태의 정보를 재사용합니다. 어텐션 및 MLP 서브 레이어 이전에, 모델은 학습된 의사 쿼리(pseudo-query)를 사용하여 완료된 블록 상태 목록으로부터 표현(representation)을 집계함으로써, 새로운 생성 목적 함수 없이도 크로스-뎁스(cross-depth) 정보 재사용을 가능하게 합니다.

순전파(forward pass) 과정에서는 이 집계와 라우팅을 용이하게 하기 위해 순환 상태(완료된 블록 상태, 현재 부분 상태, 게이트 잔차)를 유지합니다.

주요 기여

• AFM 현대화의 정식화: 본 논문은 AFM 스케일링을 단순한 파라미터 수나 희소 비율의 팽창이 아니라, LLM의 효율적 설계를 도입하는 현대화 문제로 정의합니다.
• 통합 아키텍처: 라우티드 전문가, 공유 전문가, 경량 제로-연산 전문가, 게이트-잔차 라우팅, 어텐션-잔차 집계를 단일 SiT 스타일 백본에 통합한 MMOE를 제 제안합니다.
• 통제된 경험적 연구: 저자들은 동일한 접근 가능한 학습 예산(단일 8-GPU H100 노드, 배치 크기 256, 400k 스텝) 하에서 현대화 경로(Dense SiT $\to$ Standard MoE $\to$ Shared MoE $\to$ MoE++ $\to$ AttnRes-MoE $\to$ MMOE)에 대한 엄격하고 통제된 비교를 제공합니다.

결과
Class-conditional ImageNet-256 생성에 대해 평가한 결과:

• 수렴 및 품질: MMOE는 모든 기록된 체크포인트에서 밀집(dense) SiT 및 중간 단계의 희소 전문가 베이스라인보다 낮은 FID를 달s성했습니다. 또한 학습 스텝당 더 빠르게 수렴하여, 400k 스텝에서 밀집 베이스라인의 FID 5.20에 비해 3.75의 FID를 달성했습니다.
• 효율성: 표준 MoE 변형들은 라우팅 오버헤드로 인해 실제 벽시계(wall-clock) 학습 시간을 증가시킨 반면, MMOE는 희소 변형들 중 최적의 품질-비용 균형을 달성했습니다. 가장 비용이 많이 드는 변형인 AMOE의 학습 시간을 120시간에서 67시간으로 단축하면서도 우수한 FID를 유지했습니다.
• 확장성: 개선 사항은 다양한 모델 규모(S/2, B/2, L/2)와 해상도(ImageNet-512) 전반에서 유지되었으며, 이는 이점이 XL/2 설정에만 국한되지 않음을 나타냅니다.
• 라우팅 분석: 수렴된 모델의 분석 결과는 다음과 같습니다:
  • 안정적인 특성화: 균일한 사용보다는 강력한 블록 단위 전문가 특성화(높은 부하 Gini 계수)를 보였습니다.
  • 깊이 의존적 경량 사용: 경량 전문가(copy/zero/constant)는 초기 블록에서 집중적으로 사용되었으며(44.6%의 경로), 후기 블록에서는 감소했습니다(23.4%). 이는 초기 디노이징 단계가 저렴한 변환으로부터 이득을 얻는다는 가설을 뒷받침합니다.
  • 낮은 전환율: 인접한 디노이징 스텝 사이에서 상위 2개 전문가 세트가 변하는 토큰은 약 1.4%에서 2.7%에 불과하여, 안정적인 라우팅 결정을 나타냈습니다.
• 메모리 절감: 경량 전문가를 사용하는 것은 경량 전문가가 없는 attention-residual MoE와 비교했을 때 블록당 순전파 활성화 메모리를 약 20%, 역전파 메모리를 약 32% 감소시켰습니다.

의의 및 주장
본 논문은 AFMs가 단순히 총 파라미터와 희소 비율을 늘리는 것이 아니라, 입증된 효율적 설계를 도입함으로써 LLM의 "균형 잡힌 스케일링 경로"를 따를 수 있다고 주장합니다. MMOE의 의의는 다음을 입증했다는 점에 있습니다:

1. 품질-비용 균형: 단순히 더 큰 전문가 풀에 의존하기보다, 희소 라우팅과 저렴한 전문가 경로 및 표현 재사용을 결 조합함으로써 고품질 생성을 위한 더 나은 품질-비용 트레이드오프를 달성할 수 있습니다.
2. 접근 가능한 예산에서의 실현 가능성: 이러한 개선은 단일 머신 예산(8x H100) 내에서도 달성 가능하며, 이는 대규모 멀티 노드 클러스터와 수백만 번의 반복 학습이 필요한 이전 연구들에 비해 고급 확산 트랜스포머 연구를 더욱 접근 가능하게 만듭니다.
3. 메커니즘 전이: 특정 LLM 효율성 메커니즘(제로-연산 전문가, 게이트-잔차, 어텐션-잔차)은 확산 트랜스포머로 효과적으로 전이되어 수렴 속도와 특성화 안정성을 향상시킵니다.

저자들은 자신들의 주장에 대해 겸손한 태도를 유지하며, 훨씬 더 큰 규모(예: 7M iterations)로 학습된 방법들과 비교하여 SOTA(State-of-the-art) FID를 달성했다고 주장하는 것이 아니라, 통제된 단일 노드 예산 하에서의 우수한 성능과 통제된 현대화 연구 자체를 기여로 강조합니다. 또한 class-conditional 생성에 집중된 점과 현재 통신 병목 현상이 있는 분산 학습 효율성에 대한 추가 연구가 필요하다는 한계점도 인정하고 있습니다.