기술 요약: DOT-MoE: MoEfication을 위한 미분 가능한 최적 운송(Differentiable Optimal Transport)

1. 문제 정의

대규모 언어 모델(LLM)의 급격한 규모 확장은 추론 효율성 측면에서 심각한 병목 현상을 야기했습니다. Mixture-of-Experts(MoE) 아키텍처는 토큰당 일부 파라미터만을 활성화함으로써 모델 크기와 추론 비용을 분리하지만, MoE를 처음부터 학습시키는 것은 매우 불안정하고, 많은 데이터가 필요하며, 계산 비용이 많이 듭니다.

결과적으로, 사전 학습된 밀집(dense) 모델을 희소(sparse) MoE로 변환하는 MoEfication이 유망한 대안으로 떠올랐습니다. 그러나 기존의 밀집 Feed-Forward Network(FFN)를 희소 전문가(expert)로 변환하는 방법들은 헤우리스틱 전략(예: 무작위 분할, 가중치 기반 클러스터링 또는 활성화 기반 클러스터링)에 크게 의존합니다. 이러한 접근 방식은 다음과 같은 세 가지 주요 한계를 가집니다:

1. 관심사의 분리(Separation of Concerns): 이들은 일반적으로 뉴런 할당(FFN 분할)과 라우터 학습을 별개의 순차적 단계로 취급하여, 두 과정이 상호 적응(co-adapt)하는 것을 방지합니다.
2. 대리 최적화(Proxy Optimization): 뉴런의 실제 출력 기여도가 아닌 중간 통계량(가중치 또는 활성화)을 기준으로 최적화하므로, 밀집 모델의 동작을 최적으로 재구성하는 데 실패하는 경우가 많습니다.
3. 미분 불가능성(Lack of Differentiability): 뉴로의 이산적인(discrete) 할당 특성은 엔드 투 엔드(end-to-end) 그래디언트 흐름을 방해하여, 구조적 분해와 라우팅 정책을 공동으로 최적화하는 것을 어렵게 만듭니다.

2. 방법론: DOT-MoE

저자들은 밀집 레이어를 전문가로 분해하는 과정을 미분 가능한 최적 운송(Differentiable Optimal Transport, DOT) 문제로 정식화하는 DOT-MoE 프레임워크를 제안합니다. 이 접근 방식은 전문가 분해와 라우팅 메커니즘 모두를 엔드 투 엔드로 학습할 수 있게 합니다.

2.1. 최적 운송 정식화

핵심 통찰은 뉴런 할당을 "질량(mass)"(뉴런)이 소스 분포(밀집 FFN 뉴런)에서 타겟 분포(전문가)로 엄격한 용량 제약 조건 하에 이동해야 하는 균형 잡힌 운송 문제로 보는 것입니다.

• 소스(Source): 각각 단위 질량을 가진 $dffn$개의 뉴런.
• 타겟(Target): 정확히 $s$개의 뉴런을 필요로 하는 $E$개의 전문가 ($s \cdot E = dffn$).
• 비용 함수(Cost Function): 고정된 비용 행렬 대신, 이 방법은 학습 가능한 친화도 행렬(affinity matrix) $A$를 사용합니다. 여기서 $A_{i,e}$는 뉴런 $i$를 전문가 $e$에 할당하는 친화도를 나타냅니다. 목표는 주변부 제약 조건을 만족하면서 총 친화도를 최대화하는 것입니다.

2.2. Sinkhorn-Knopp를 통한 미분 가능한 완화

이산적인 할당 문제를 미분 가능하게 만들기 위해, DOT-MoE는 **엔트로피 정규화(Entropic Regularization)**를 채택합니다.

• 최적 운송 문제는 행과 열을 반복적으로 정규화하는 Sinkhorn-Knopp 알고리즘을 사용하여 완화됩니다.
• 이를 통해 기대값 관점에서 용량 제약을 만족하는 소프트 할당 행렬 $M_{soft} \in [0, 1]^{dffn \times E}$을 얻습니다.
• 온도 매개변수 $\tau \to 0$이 됨에 따라, 솔루션은 이산 최적해에 도달합니다.

2.3. Straight-Through Estimators (STE)를 이용한 공동 최적화

배포 시에는 이산적인 할당이 필요하고, 라우팅에는 이산적인 top-k 선택이 포함되는데, 둘 다 미분 불가능하다는 점이 중요한 과제입니다. DOT-MoE는 이를 **Straight-Through Estimators (STE)**를 사용하여 해결합니다.

• 순전파(Forward Pass): 하드 결정(할당을 위한 $M_{soft}$의 그리디 반올림, 라우팅을 위한 top-k 선택)을 사용합니다.
• 역전파(Backward Pass): 소프트 대응물($M_{soft}$ 및 라우터 로짓 $P$)을 통해 그래디언트를 전달합니다.
• 이를 통해 라우터 가중치($W_r$)와 할당 친화도 행렬($A$)이 학습 중에 서로 적응할 수 있습니다.

2.4. 학습 목적 함수

프레임워크는 밀집 모델의 가중치를 고정한 상태에서 **정렬 단계(Alignment Phase)**를 통해 학습됩니다. 모델은 다음의 조합된 손실을 최소화합니다:

1. 출력 재구성(Output Reconstruction): 밀집 교사(teacher) 모델과 희소 학생(student) MoE 모델 출력 간의 KL 발산(divergence)으로 측정됩니다.
2. 언어 모델링(Language Modeling): 언어 모델링 목적 함수에 대한 교차 엔트로피 손실입니다.
3. 안정성 보조 항(Stability Auxiliaries): 로짓 폭발을 방지하기 위한 라우터 $z$-loss 및 전문가 붕괴를 방지하기 위한 로드 밸런싱(load-balancing) 손실입니다.

3. 주요 기여

1. 최적 운송 프레임워크: 본 논문은 밀집 모델을 희소 Mo로 변환하는 것을 균형 잡힌 운송 문제로 정식화하고, 엄격한 전문가 용량 제약을 강제하기 위해 미분 가능한 Sinkhorn 반복법을 사용하는 새로운 프레임워크를 도입했습니다.
2. 이중 수준의 공동 적응(Dual-Level Co-Adaptation): 라우터를 학습시키기 전에 할당을 고정하는 기존 방식과 달리, DOT-MoE는 상보적인 STE를 사용하여 뉴런-전문가 분해와 토큰-전문가 라우팅을 공동으로 최적화합니다.
3. 출력 인지형 구축(Output-Aware Construction): 가중치나 중간 활성화와 같은 대리 지표를 최적화하는 대신, 최종 출력에 대한 뉴런의 기여도를 직접 모델링함으로써 더 높은 재구성 충실도를 달acia합니다.

4. 실험 결과

저자들은 세 가지 모델 제품군(LLaMA-2, LLaMA-3, Qwen2.5)과 여섯 가지 벤치마크(ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ)에 대해 DOT-MoE를 평가했습니다.

• 프루닝 대비 성능: 50% 파라미터 예산을 가진 LLaMA-2 7B에 대해, DOT-MoE는 WikiText-2에서 7.99의 퍼플렉시티(perplexity)를 달성하여, 최첨단 구조적 프루닝 방법인 DISP-LLM(9.84)을 크게 앞질렀습니다.
• 기존 MoEfication 대비 성능: 제로샷(zero-shot) 설정에서, DOT-MoE는 CMoE 및 LLaMA-MoE-v2와 같은 방법들을 실질적으로 능가했습니다. 예를 들어, LLaMA-2 7B에서 DOT-MoE는 CMoE의 44.5% 대비 **61.5%**의 평균 정확도를 달성했습니다.
• 효율성: DOT-MoE는 활성 파라미터를 50% 줄이면서도 원래 밀집 모델 성능의 약 **90%**를 유지합니다.
• 미세 조정(Fine-Tuning): 1.2B 토큰의 연속 미세 조정을 통해, DOT-MoE 모델은 밀집 모델에 근접한 성능을 회복했습니다(예: LLaMA-2 7B에서 66.6% vs 76.6%). 이는 유사하거나 더 큰 데이터 예산을 요구하는 베이스라인 모델들보다 우수한 성능입니다.
• 절제 연구(Ablation Studies):
  • 그래뉼러리티(Granularity): 성능은 포화 지점까지 전문가 그래뉼러티가 증가함에 따라 향상되지만, 그 이후에는 이득이 감소합니다.
  • 추론 속도: 전문가 수가 많아지더라도, 전체 퓨즈드 중간 차원이 일정하게 유지되므로 퓨즈드 MoE 커널(예: vLLM)을 사용할 때 추론 처리량은 안정적으로 유지됩니다.
  • 희소성(Sparsity): 더 높은 희소성(75%) 수준에서 학습된 모델은 추론 시 훨씬 더 극단적인 희소성 영역에서도 견고하게 일반화됩니다.

5. 의의 및 주장

본 논문은 DOT-MoE가 밀집 모델을 희소 Mo로 변환하는 원칙적이고 미분 가능한 방법을 제공함으로써 LLM 추론 효율성에 있어 중요한 진전을 이루었다고 주장합니다.

• 패러다임 전환: 헤우리스틱하고 정적인 분할에서 벗어나, 전문가의 기능적 일관성을 존중하는 동적이고 학습 가능한 분해 방식으로 전환합니다.
• 실용성: 이 방법은 전체 재학습 없이도 사전 학습된 모델을 효율적으로 사후 희소화(post-hoc sparsification)할 수 있게 하며, 활성 파라미터를 획기적으로 줄이면서도 높은 용량을 유지합니다.
• 강건성: 본 프레임워크는 FFN의 충실도를 보존하기 위해서는 뉴런의 출력을 명시적으로 모델링해야 함을 보여주며, 이는 기존의 대리 지표 기반 방법들이 포착하지 못하는 원칙입니다.

저자들은 결론적으로 DOT-MoE가 희소 모델을 위한 우수한 초기화를 제공하며, 최소한의 추가 학습 데이터로 밀집 아키텍처와 희소 아키텍처 사이의 성능 격차를 줄여준다고 밝히고 있습니다.