기술 요약: EMO - 등장적 모듈성을 위한 사전 학습 전문가 혼합 모델

문제 제기

대형 언어 모델 (LLM) 은 일반적으로 전체 모델을 로드해야만 하는 단일 구조 (monolithic) 시스템으로 배포됩니다. 이는 애플리케이션이 코드 생성이나 수학적 추론과 같은 좁은 범위의 기능만 필요로 할 때도 발생합니다. 이러한 패러다임은 불필요한 계산 비용과 메모리 사용량을 초래합니다. 전문가 혼합 (MoE) 아키텍처는 입력당 전문가의 하위 집합만 활성화함으로써 잠재적인 해결책을 제시하지만, 표준 MoE 는 모듈형 배포를 지원하지 못합니다. 기존 MoE 에서는 단일 입력 내의 토큰들이 종종 전치사나 구두점과 같은 저수준 어휘 패턴에 의해 주도되어 다양한 전문가를 활성화합니다. 따라서 특정 도메인에 대해 전문가의 특정 하위 집합으로 추론을 제한하면 성능이 심각하게 저하되어, 모델의 일부만 로드할 수 있는 메모리 제약 환경에서 MoE 를 실용적으로 만드는 것을 불가능하게 만듭니다.

방법론: EMO

저자들은 사전 학습 과정에서 모듈성을 1 순위 목표로 유도하도록 설계된 MoE 아키텍처인 EMO를 소개합니다. 핵심 목표는 인간이 정의한 도메인 레이블에 의존하지 않고도 특정 도메인에 대해 전문가 하위 집합을 독립적으로 사용할 수 있으며, 이를 강력한 범용 모델로 구성할 수 있는 모델을 훈련하는 것입니다.

핵심 메커니즘: 문서 수준 공유 풀 (Document-Level Shared Pools)

핵심 직관은 동일한 문서 내의 토큰들이 일반적으로 동일한 의미 도메인을 공유한다는 것입니다. EMO 는 이러한 구조를 강제하기 위해 약한 감독 신호로 문서 경계를 활용합니다:

1. 공유 전문가 풀: 각 문서에 대해 라우터는 라우팅된 전문가 전체 집합에서 공유 전문가 풀 ($D$) 을 선택합니다.
2. 제한된 라우팅: 해당 문서 내의 모든 토큰은 이 공유 풀에서만 활성 전문가를 선택하도록 제한됩니다.
3. 동적 세분성: 다양한 메모리 예산에 걸쳐 유연성을 보장하기 위해 전문가 풀의 크기 ($d$) 는 사전 학습 중 $k$ (토큰당 활성 전문가 수) 와 $n_r$ (전체 라우팅 전문가 수) 에서 균일 분포 $U\{k, n_r\}$ 로 샘플링된 무작위 변수로 처리됩니다.

훈련 목표 및 기술적 고려 사항

• 목표: 모델은 표준 자기회귀 언어 모델링 손실 ($L_{CE}$) 과 보조 손실을 결합하여 훈련됩니다.
• 로드 밸런싱: 중요한 기술적 기여는 전역 로드 밸런싱의 채택입니다. 표준 마이크로 배치 로드 밸런싱은 문서 수준 라우팅 제약 (마이크로 배치 내 토큰을 전문가 전체에 분산하도록 강제) 과 충돌합니다. EMO 는 데이터 병렬 그룹 전체에 걸쳐 라우팅 통계를 집계하여 로드 밸런싱 손실을 계산함으로써, 문서 내 일관성을 유지하면서 전체 데이터셋에 걸쳐 균일한 전문가 활용을 보장합니다.
• 아키텍처: 저자들은 1 조 개의 토큰으로 128 개의 전체 전문가 (127 개의 라우팅 전문가, 1 개의 공유 전문가) 를 가진 10 억 활성 파라미터, 140 억 총 파라미터 모델을 훈련했습니다. 이 아키텍처는 사전 정규화 (pre-norm), QK-정규화 제거, 공유 전문가 포함 등 이전 베이스라인 대비 개선을 포함합니다.

주요 기여

1. 등장적 모듈성: EMO 는 문서 경계를 유일한 신호로 사용하여 데이터에서 직접 모듈 구조가 등장할 수 있음을 입증하며, 사전 정의된 도메인 레이블이나 인간 선입견의 필요성을 제거합니다.
2. 선택적 전문가 사용: 이 모델은 도메인별 작업에서 거의 전체 모델 성능을 유지하는 작은 전문가 하위 집합 (예: 전문가의 12.5%~25%) 의 배포를 가능하게 하여, 표준 MoE 가 lacking 한 기능을 제공합니다.
3. 의미적 전문화: 저수준 구문 패턴에 전문화되는 표준 MoE 와 달리, EMO 전문가들은 수학, 코드, 생물의학, 정치와 같은 고수준 의미 도메인에 전문화되는 일관된 그룹으로 조직됩니다.
4. 파레토 프론티어 개선: EMO 전문가 하위 집합은 처음부터 훈련된 메모리 일치 모델 (밀집형 및 표준 MoE 모두) 보다 성능이 우수하여, 메모리 - 정확도 트레이드오프에 대한 새로운 파레토 프론티어를 확립합니다.

결과

• 전체 모델 성능: 완전한 모델로서 EMO 는 동일한 데이터 (1 조 토큰) 로 훈련된 표준 MoE 와 성능을 일치시켜, MMLU, MMLU-Pro, GSM8K 와 같은 벤치마크에서 경쟁력 있는 결과를 달성합니다.
• 하위 집합 성능:
  • 전문가의 25% 를 유지할 때, EMO 는 MMLU 정확도에서 1% 의 절대적 감소만 발생시키는 반면, 표준 MoE 는 10% 이상 저하됩니다.
  • 전문가의 12.5% 를 유지할 때, EMO 는 3% 감소하는 반면, 표준 MoE 는 15% 감소합니다.
  • EMO 하위 집합은 동일한 파라미터 수로 처음부터 훈련된 밀집형 모델 및 표준 MoE 보다 일관되게 우수한 성능을 보입니다.
• 강건성: EMO 의 성능은 다양한 전문가 선택 방법 (라우터 기반 대 Easy-EP 가지치기) 에 걸쳐 강건하며, 전문가 선택을 위한 극도로 작은 검증 세트에서도 효과적입니다.
• 분석: 클러스터링 분석은 동일한 문서에서 온 EMO 토큰이 "건강, 의료 및 웰빙"과 같은 동일한 의미 클러스터에 그룹화되는 반면, 표준 MoE 토큰은 표면적 특징에 기반하여 클러스터 전체에 분산됨을 보여줍니다.

중요성 및 주장

이 논문은 EMO 가 대규모 희소 모델의 모듈형, 메모리 효율적 배포를 위한 길을 제공한다고 주장합니다. 전문가 하위 집합의 독립적 사용과 구성을 가능하게 함으로써 EMO 는 MoE 확장과 관련된 메모리 병목 현상을 해결합니다. 저자들은 다음과 같은 함의를 강조합니다:

• 접근 가능한 배포: 사용자는 특정 작업이나 도메인과 관련된 전문가만 제공하여 메모리 요구 사항을 크게 줄일 수 있습니다.
• 세분화된 제어: 모듈성은 추론 시 스팸 또는 성인 콘텐츠와 관련된 전문가를 제외하는 등 기능의 선택적 활성화 또는 비활성화를 가능하게 합니다.
• 모듈형 개발: 이 구조는 작업별 전문가 하위 집합을 독립적으로 훈련한 후 나중에 재통합할 수 있는 모듈형 사전 학습을 위한 잠재적 패러다임을 시사합니다.
• 모니터링 가능성: 명확한 라우팅 패턴은 모델 행동을 감사하고 디버깅하기 위한 구조화된 신호를 제공합니다.

저자들은 모듈성이 사전 학습 자체에 내장될 수 있으며, 배포, 적응 및 검사가 더 쉬운 구성 가능한 아키텍처를 향해 단일 구조 시스템에서 벗어나고 있다고 결론지었습니다.