Mixture of Universal Experts: Scaling Virtual Width via Depth-Width Transformation

이 논문은 MoE 아키텍처의 확장성 한계를 극복하기 위해 고정된 토큰 활성화 예산 하에서 깊이를 가상 너비로 변환하는 '가상 너비'라는 새로운 차원을 도입한 '범용 전문가 혼합 (MOUE)'을 제안하며, 이를 통해 기존 MoE 모델 대비 성능을 향상시키고 확장 가능한 새로운 아키텍처를 제시합니다.

핵심

1. 문제: "왜 모든 층에 똑같은 전문가를 두나요?"

기존의 AI 모델 (Transformer) 은 마치 고층 빌딩과 같습니다.

• 층 (Layer): 빌딩의 1 층부터 100 층까지 있습니다.
• 전문가 (Expert): 각 층마다 특정 업무를 처리하는 전문가 팀이 따로 있습니다. 1 층의 팀은 '문법'을, 100 층의 팀은 '추론'을 담당합니다.

문제점:
이 방식은 비효율적입니다.

1. 공간 낭비: 1 층의 '문법 전문가'와 100 층의 '문법 전문가'가 사실은 같은 일을 할 수 있는데, 각각 따로 고용하고 급여 (메모리) 를 줍니다.
2. 확장의 한계: 더 똑똑해지려면 빌딩을 더 높이 쌓거나 (깊이), 층마다 전문가를 더 많이 고용해야 (너비) 합니다. 하지만 빌딩을 너무 높게 올리면 구조가 불안정해지고, 전문가를 너무 많이 고용하면 비용이 천문학적으로 늘어납니다.

2. 해결책: "공유된 전문가 풀 (Universal Expert Pool)"

이 논문은 **"층마다 전문가를 따로 두지 말고, 빌딩 전체가 공유할 수 있는 '유니버설 전문가 풀'을 만들자"**고 제안합니다.

• 비유: 각 층마다 별도의 식당을 짓는 대신, 빌딩 중앙에 거대한 '공유 식당'을 하나만 짓고, 모든 층의 사람들이 필요할 때 그곳으로 가서 음식을 시켜 먹는 방식입니다.
• 핵심 아이디어 (가상의 너비):
  • 물리적으로 전문가 (식당) 는 그대로지만, **깊이 (층수) 를 활용해서 가상의 넓은 공간 (Virtual Width)**을 만들어냅니다.
  • 같은 전문가가 1 층에서 쓰이다가 50 층에서도 다시 쓰일 수 있습니다.
  • 결과적으로 전문가 수는 그대로인데, 모델이 처리할 수 있는 아이디어의 조합 (경로) 은 기하급수적으로 늘어납니다.

3. 새로운 기술 3 가지 (어떻게 구현했나?)

이 아이디어를 실현하려면 세 가지 중요한 기술이 필요합니다.

① 계단식 회전 구조 (Staggered Rotational Topology)

• 문제: 모든 층이 모든 전문가를 동시에 볼 수 있게 하면, AI 가 "누구를 선택할지" 고민하다가 망설여집니다 (경로가 너무 많아서).
• 해결: 계단식 회전 방식을 씁니다.
  • 비유: 빌딩의 층을 3 층씩 묶어서 '구역'을 만듭니다. 13 층은 'A 구역 전문가들'만 볼 수 있고, 46 층은 'B 구역 전문가들'만 봅니다. 하지만 A 구역과 B 구역은 겹치는 전문가들이 있어서 자연스럽게 연결됩니다.
  • 이렇게 하면 AI 는 매번 모든 전문가를 다 볼 필요 없이, 현재 층에 맞는 전문가들만 골라 효율적으로 일할 수 있습니다.

② 공정한 분배 시스템 (Universal Expert Load Balance)

• 문제: 공유 식당이 생기면, 인기 있는 식당은 너무 붐비고 (과부하), 다른 식당은 아예 손님이 없습니다. 기존 AI 는 "누가 많이 쓰였나"만 보고 불공정하게 처리합니다.
• 해결: "접근 기회"를 고려한 공정한 분배를 합니다.
  • 비유: "이 식당은 100 층에서나 볼 수 있어서 자연스럽게 손님이 많을 수밖에 없는데, 그걸 '불공정'하다고 처벌하면 안 됩니다."
  • 대신, **"자신이 접근할 수 있는 층들 안에서 얼마나 골고루 쓰였는지"**를 기준으로 평가합니다. 이렇게 하면 모든 전문가가 고르게 일하게 되어 시스템이 안정됩니다.

③ 기억력이 있는 길잡이 (Universal Router)

• 문제: AI 가 층을 올라갈 때, 매번 "지금 내가 어디에 왔지?"를 잊어버리면 같은 일을 반복하게 됩니다.
• 해결: **작은 메모 (상태)**를 가지고 다니는 길잡이를 둡니다.
  • 비유: 식당을 갈 때, "아까 1 층에서 A 전문가를 만났으니, 50 층에서는 B 전문가를 만나야겠다"라고 이전 경로를 기억하고 다음 선택을 합니다.
  • 이렇게 하면 AI 는 단순한 반복이 아니라, 복잡한 추론 과정을 거칠 수 있게 됩니다.

4. 결과: 더 적은 비용, 더 큰 능력

이 방법을 실험해 보니 놀라운 결과가 나왔습니다.

• 성능 향상: 같은 크기의 기존 모델보다 최대 4.2% 까지 성능이 좋아졌습니다.
• 비용 절감: 전문가 수를 늘리지 않아도, 공유 시스템을 통해 가상의 넓은 능력을 얻었습니다.
• 유연성: 이미 만들어진 기존 AI 모델을 이 방식으로 쉽게 업그레이드할 수 있습니다. (기존 모델을 재활용해서 '공유 식당'을 추가하는 식)

요약

이 논문은 **"빌딩을 더 높이 쌓거나 (깊이), 층마다 사람을 더 많이 고용하는 (너비) 대신, 기존 전문가들을 clever하게 공유해서 (재사용) 가상의 넓은 공간을 만들어내자"**는 혁신적인 아이디어입니다.

마치 한 명의 요리사가 빌딩 전체를 돌며 각 층의 주문을 받아 요리하는 것처럼, 자원은 적게 쓰면서 훨씬 더 다양한 요리를 해내는 효율적이고 똑똑한 AI를 만든 것입니다.

기술 요약

1. 문제 정의 (Problem)

기존의 혼합 전문가 (Mixture-of-Experts, MoE) 아키텍처는 모델의 전체 파라미터 수와 토큰당 활성화되는 파라미터 수를 분리하여 효율적인 확장을 가능하게 했습니다. 그러나 MoE 는 여전히 물리적 차원 (깊이와 너비) 에 의해 확장성이 제한받는 두 가지 주요 한계를 가지고 있습니다.

1. 깊이 (Depth) 의 비효율성: 기존 MoE 는 각 레이어마다 별도의 전문가 집합을 고정적으로 쌓는 방식에 의존합니다. 이는 복잡한 재귀적 계산이나 깊은 계층 간의 기능적 중복을 활용하기 어렵게 만듭니다.
2. 너비 (Width) 의 확장 비용: 성능을 높이기 위해 전문가 (Expert) 의 수를 늘리면, 물리적 메모리 사용량과 시스템 오버헤드가 비례하여 급증합니다.
3. 핵심 질문: "모델의 깊이를 재사용하여 모델 용량을 확장하면서도, 추가적인 계산량이나 메모리 오버헤드를 최소화할 수 있는 아키텍처는 존재하는가?"

2. 방법론 (Methodology)

저자들은 가상 너비 (Virtual Width) 라는 새로운 확장 차원을 도입한 범용 전문가 혼합 (Mixture of Universal Experts, MOUE) 을 제안합니다. 이는 레이어에 구애받지 않는 공유된 '범용 전문가 (Universal Experts, UEs)' 풀을 여러 레이어에서 재사용함으로써, 깊이를 유효한 너비로 변환하는 구조입니다.

핵심 구성 요소

MOUE 는 두 가지 주요 도전 과제 (라우팅 경로 폭발, 노출 편향) 를 해결하기 위해 세 가지 핵심 컴포넌트를 도입했습니다.

1. 계단식 회전 토폴로지 (Staggered Rotational Topology):

   • 문제: 모든 레이어가 모든 전문가에 접근하면 라우팅 공간이 기하급수적으로 커져 최적화가 불안정해집니다.
   • 해결: 전문가 풀을 계층적 링 구조로 조직화합니다. 인접한 레이어 그룹 (Connectivity Group) 은 동일한 접근 가능한 전문가 집합을 공유하지만, 그룹이 깊이에 따라 이동함에 따라 접근 가능한 범용 전문가의 창 (Window) 이 링을 따라 계단식 (Staggered) 으로 회전합니다.
   • 효과: 지역적 전문화 (Layer-local experts) 와 제어된 재사용 (Cross-layer reuse) 사이의 균형을 유지하며, 라우팅 검색 공간을 국소화하여 최적화 안정성을 높입니다.

2. 범용 전문가 로드 밸런싱 (Universal Expert Load Balance, UELB):

   • 문제: 범용 전문가 (UE) 는 여러 레이어에서 접근 가능하므로, 기존 로드 밸런싱 로스 (Load Balancing Loss) 는 UE 가 과도하게 사용된 것으로 오인하여 이를 억제하는 편향을 만듭니다.
   • 해결: 접속성 정규화 (Connectivity Normalization) 원리를 도입합니다. 각 전문가의 노출 횟수 (Topological Degree, $c_j$) 로 나누어, 레이어별 평균 활용도를 기준으로 균형을 맞추도록 로스를 수정합니다.
   • 효과: 아키텍처적으로 설계된 노출 증가를 '과도한 사용'으로 잘못 판단하는 것을 방지하여, UE 풀이 효과적으로 활용되도록 합니다.

3. 범용 라우터 (Universal Router):

   • 문제: 재귀적 재사용은 각 레이어의 결정이 독립적이지 않고, 전체 계산 경로 (Trajectory) 에 의존해야 함을 의미합니다.
   • 해결: 의미적 경로 (Semantic Pathway) 와 맥락적 경로 (Contextual Pathway) 로 구성된 듀얼 경로 라우팅을 도입합니다. 맥락적 경로는 경량화된 상태 행렬 (Fast Weights) 을 사용하여 이전 단계의 라우팅 결정과 토큰 흐름을 추적하고, 이를 바탕으로 일관된 라우팅 결정을 내립니다.
   • 효과: 깊은 계층을 거치며 일관된 전문가 조합을 선택할 수 있게 하여, 재사용의 잠재력을 극대화합니다.

4. 점진적 웜스타트 (Progressive Warm-Start):

   • 기존 MoE 체크포인트를 MOUE 로 변환할 때, 범용 전문가 풀을 초기화하고 로지트 억제 (Logit Suppression) 커리큘럼을 통해 점진적으로 재사용을 활성화하여 학습 안정성을 확보합니다.

3. 주요 기여 (Key Contributions)

• 새로운 확장 차원 (Virtual Width) 의 제안: 물리적 파라미터 수를 늘리지 않고, 깊이를 재사용하여 유효한 모델 용량 (가상 너비) 을 기하급수적으로 확장하는 새로운 패러다임을 제시했습니다.
• 구조적 최적화 및 학습 전략: 재사용으로 인한 최적화 난이도를 해결하기 위한 계단식 토폴로지, 노출 편향을 교정하는 UELB, 그리고 경로 일관성을 위한 상태 기반 라우터를 통합한 프레임워크를 완성했습니다.
• 하위 호환성 및 점진적 변환: MOUE 는 기존 MoE 의 완전한 초집합 (Superset) 으로, 재사용을 비활성화하면 기존 MoE 로 돌아갑니다. 또한, 기존 MoE 체크포인트를 손실 없이 MOUE 로 변환하여 성능을 즉시 향상시킬 수 있는 방법을 제시했습니다.

4. 실험 결과 (Results)

다양한 규모 (MoE-160M, MoE-700M, JetMoE, OLMoE 등) 와 설정에서 MOUE 는 기존 MoE 베이스라인을 일관되게 능가했습니다.

• 너비 확장 (Width Expansion): 활성화 파라미터와 총 파라미터를 고정하고 가상 너비만 늘린 경우, MoE 대비 최대 1.3% 의 상대적 성능 향상을 기록했습니다.
• 깊이 확장 (Depth Expansion): FFN 파라미터를 공유하여 깊이를 늘린 경우, MoE 대비 2.5% ~ 4.2% 의 성능 향상을 보였습니다. 특히 MOUE L36 은 MoE 64A8 L16 보다 적은 파라미터로 더 높은 정확도를 달성했습니다.
• 점진적 변환 (Progressive Conversion): 기존 MoE 체크포인트를 MOUE 로 변환하여 추가 학습을 진행한 결과, 평균 4.2% 의 상대적 개선을 보였으며, 범용 전문가 풀이 클수록 성능 향상이 더 컸습니다.
• 학습 안정성: UELB 와 계단식 토폴로지는 라우팅 편향 (Routing Skew) 을 효과적으로 제어하여 학습 중 안정적인 수렴을 보장했습니다.

5. 의의 및 결론 (Significance)

이 논문은 MoE 아키텍처의 확장에 있어 물리적 자원 (메모리/계산량) 의 증가 없이도 깊이를 활용하여 모델 능력을 확장할 수 있음을 증명했습니다.

• 효율성: 기존 MoE 가 직면한 '파라미터 수 증가 = 메모리/계산량 증가'의 제약을 깨고, Depth-to-Width Transformation을 통해 더 넓은 모델 용량을 실현합니다.
• 실용성: 기존 모델의 체크포인트를 재사용하여 성능을 즉시 향상시킬 수 있어, 대규모 모델 학습 비용 절감에 기여합니다.
• 미래 방향: MOUE 는 MoE 의 확장 한계를 새로운 차원 (가상 너비) 으로 재정의하며, 재귀적 계산과 구조적 재사용을 결합한 차세대 언어 모델 아키텍처의 방향성을 제시합니다.

요약하자면, MOUE 는 공유된 범용 전문가 풀을 계단식 토폴로지와 지능형 라우팅을 통해 깊이 전반에 걸쳐 재사용함으로써, 제한된 활성화 예산 내에서 모델의 표현력을 극대화하는 혁신적인 프레임워크입니다.