Mixture of Experts(MoE)에서의 라우팅은 현대 기계 학습 아키텍처의 핵심 메커니즘으로, 각 입력 토큰을 전문가 네트워크의 하위 집합에 동적으로 할당한다. 모든 입력에 대해 모든 매개변수를 활성화하는 대신, 라우터(또는 게이팅 네트워크)는 토큰을 가장 관련성 높은 전문가에게 보내도록 학습하여 모델 용량과 계산 효율성의 균형을 맞춘다. 이 접근 방식은 대규모 언어 모델이 추론 비용을 관리 가능한 수준으로 유지하면서 수조 개의 매개변수로 확장할 수 있게 하며, 토큰당 네트워크의 일부만 활성화된다.
이 개념은 1990년대 초 적응형 로컬 전문가 혼합에 대한 연구에서 비롯되었지만, 2010년대 딥러닝에서 희소 게이팅 MoE 레이어의 도입으로 두각을 나타냈다. 구글 딥마인드와 오픈AI의 모델과 같은 현대적 구현은 전문가에 대한 확률 분포를 출력하는 학습된 라우팅 함수에 의존하며, 종종 균형 잡힌 사용을 보장하기 위해 보조 손실을 통합한다. 라우팅은 최첨단 생성형 AI 시스템의 중요한 설계 선택이 되었으며, 훈련 효율성과 최종 모델 품질 모두에 영향을 미친다.
역사적 발전
입력을 전문 구성 요소로 라우팅한다는 아이디어는 딥러닝 이전부터 존재했으며, 앙상블 방법과 모듈식 신경망에 뿌리를 두고 있다. 1991년 마이클 조던과 로버트 제이콥스는 여러 피드포워드 네트워크의 출력에 가중치를 부여하는 게이팅 네트워크를 사용하는 계층적 전문가 혼합에 대한 기초 연구를 발표했다. 이 초기 프레임워크는 입력 공간을 분할하여 서로 다른 전문가가 서로 다른 영역을 처리하도록 학습한다는 핵심 원리를 확립했다.
2017년, 구글 딥마인드(당시 구글 브레인)의 연구자들은 Noam Shazeer가 주도한 논문에서 희소 게이팅 MoE 레이어를 도입했다. 이 설계는 각 토큰에 대해 top-k 전문가를 선택하기 위해 소프트맥스 게이팅 함수를 사용했으며, k는 일반적으로 1 또는 2로 설정되었다. 저자들은 MoE 레이어가 수백 개의 전문가로 확장될 수 있음을 입증했으며, 언어 모델링 및 기계 번역 벤치마크에서 최첨단 결과를 달성하면서 예시당 계산 비용을 줄였다. 이 연구는 토큰당 하나의 전문가만 선택하도록 라우팅을 단순화한 Switch Transformer(2021)와 대규모 병렬 훈련을 위한 GShard 프레임워크와 같은 이후 아키텍처에 직접적인 영향을 미쳤다.
라우팅 메커니즘
라우팅 함수는 여러 유형으로 분류할 수 있다. 가장 일반적인 것은 토큰 선택 라우팅으로, 각 토큰이 학습된 점수를 기반으로 독립적으로 top-k 전문가를 선택한다. 이는 토큰의 은닉 상태를 선형 투영한 다음 전문가 인덱스에 대한 소프트맥스를 적용하여 구현된다. 선택된 전문가는 토큰을 처리하고, 그 출력은 라우팅 확률에 따라 가중치가 부여되어 합산된다.
대안은 전문가 선택 라우팅으로, 각 전문가가 배치에서 top-k 토큰을 선택하여 모든 전문가가 최소 부하를 받도록 보장한다. Switch Transformer에서 대중화되고 Mixtral과 같은 모델에서 개선된 이 접근 방식은 부하 불균형 문제를 해결하지만 토큰-전문가 할당을 신중하게 처리해야 한다.
다른 변형으로는 학습된 매개변수 없이 결정적 해시 함수를 사용하여 토큰을 할당하는 해싱 기반 라우팅과, 2단계 게이팅 시스템이 먼저 전문가 그룹을 선택한 다음 그룹 내에서 특정 전문가를 선택하는 계층적 라우팅이 있다. 각 방법은 유연성, 계산 오버헤드, 훈련 안정성 간의 균형을 맞춘다.
부하 균형 및 보조 손실
MoE 라우팅의 주요 과제는 부하 불균형이다. 라우터가 항상 소수의 인기 있는 전문가를 선택하여 다른 전문가를 충분히 훈련시키지 못할 수 있다. 이를 완화하기 위해 대부분의 구현은 전문가 간 토큰 분포의 불균일성을 패널티로 주는 보조 부하 균형 손실을 추가한다. 2017년 Shazeer 논문에서 도입된 표준 공식은 각 전문가로 라우팅된 토큰의 비율을 계산하고 교차 엔트로피 항을 통해 균일성을 장려한다.
Switch Transformer의 부하 균형 손실과 같은 최근 접근 방식은 라우터의 평균 확률과 전문가당 실제 토큰 수 간의 내적을 곱하는 더 간단한 계수를 사용한다. GShard와 같은 일부 시스템은 각 전문가가 처리할 수 있는 토큰 수를 제한하는 용량 계수를 사용하여 라우터가 작업을 분산하도록 강제한다. 이러한 기술은 안정적인 훈련과 일부 전문가가 죽은 가중치가 되는 전문가 붕괴를 방지하는 데 필수적이다.
아키텍처 통합
라우팅은 일반적으로 트랜스포머 블록 내에 적용되며, 피드포워드 네트워크(FFN)를 MoE 레이어로 대체한다. 표준 트랜스포머에서 각 토큰은 멀티 헤드 어텐션 하위 레이어와 위치별 FFN을 통과한다. MoE 트랜스포머에서 FFN은 각각 고유한 매개변수를 가진 일련의 전문가 FFN과 활성화할 전문가를 선택하는 라우터로 대체된다.
이 통합을 통해 모델은 토큰당 활성 매개변수 수를 일정하게 유지하면서 총 매개변수 수를 크게 늘릴 수 있다. 예를 들어, 각각 1억 개의 매개변수를 가진 64개의 전문가가 있는 모델은 총 640억 개의 매개변수를 가지지만 k=2이면 토큰당 2억 개만 활성화한다. 이 속성은 앤트로픽 및 기타 연구소의 Mixture of Experts 모델과 같은 시스템에서 볼 수 있듯이 수조 개의 매개변수 모델로 확장하는 데 중요하다.
라우터 자체는 종종 소프트맥스가 뒤따르는 단일 선형 레이어인 작은 신경망이며, 전문가와 함께 공동으로 훈련되는 고유한 매개변수를 가진다. 일부 아키텍처는 레이어나 어텐션 헤드에 대해 별도의 라우터를 사용하며, 최근 연구는 토큰 유형이나 위치에 따라 적응하는 학습된 라우팅 정책을 탐구한다.
훈련 역학
라우팅으로 MoE 모델을 훈련하면 고유한 과제가 발생한다. 라우터의 이산적 결정(top-k 전문가 선택)은 미분 불가능하므로 기울기는 선택된 전문가의 출력에 라우팅 확률을 곱한 값으로만 흐른다. 이는 이동 표적 문제를 만든다. 전문가가 개선됨에 따라 라우터의 선호도가 바뀌어 잠재적으로 불안정성을 유발할 수 있다.
이를 해결하기 위해 연구자들은 훈련 중 탐색을 장려하기 위해 라우팅 로짓에 가우시안 노이즈를 추가하는 노이즈가 있는 top-k 게이팅과 같은 기술을 사용한다. 또 다른 접근 방식은 라우팅 결정에 대해 직선 추정기를 사용하여 선택을 하드 할당으로 처리하지만 확률 가중치를 통해 기울기를 전달하는 것이다. 또한 일부 방법은 훈련 중 소프트맥스의 온도를 점진적으로 낮추어 라우팅 결정을 선명하게 한다.
부하 균형 손실은 일반적으로 기본 작업 손실을 지배하지 않도록 작은 계수(예: 0.01)로 가중치를 부여한다. 실제로 MoE 모델을 훈련하려면 신중한 하이퍼파라미터 튜닝이 필요하며, 많은 시스템이 부하 균형과 라우터 신뢰도 모두에 대해 보조 손실을 사용한다.
대규모 언어 모델에서의 응용
MoE의 라우팅은 대규모 언어 모델에서 표준 기술이 되었다. 주목할 만한 예는 다음과 같다.
- Google의 Switch Transformer(2021)는 희소 MoE 아키텍처로 1.6조 개의 매개변수로 확장되었다.
- GShard(2020)는 수천 개의 TPU 코어에서 MoE 모델의 효율적인 훈련을 입증했다.
- Mistral AI의 Mixtral 8x7B(2023)는 top-2 라우팅과 함께 8개의 전문가를 사용하여 더 큰 밀집 모델과 유사한 성능을 달성했다.
- DeepSeek-V3(2024)는 256개의 전문가와 top-8 라우팅을 사용하는 세분화된 MoE를 사용한다.
- 오픈AI와 앤트로픽의 여러 모델이 MoE 레이어를 사용하는 것으로 알려져 있지만, 정확한 세부 사항은 종종 독점적이다.
이러한 모델은 라우팅이 총 매개변수가 증가해도 토큰당 계산이 일정하게 유지되므로 비용 효율적인 확장을 가능하게 한다는 것을 보여준다. 이로 인해 MoE는 아마존 웹 서비스, 애저, 구글 클라우드와 같은 클라우드 플랫폼을 포함한 프로덕션 환경에서 대규모 모델을 제공하는 데 선호되는 선택이 되었다.
효율성 및 하드웨어 고려 사항
라우팅은 분산 훈련 및 추론에서 통신 오버헤드를 도입한다. 토큰은 다른 장치에 있을 수 있는 적절한 전문가에게 전송되어야 하기 때문이다. 이 all-to-all 통신 패턴은 특히 전문가가 많을 때 병목 현상이 될 수 있다. 엔비디아와 AMD와 같은 하드웨어 공급업체는 MoE 연산에 최적화된 커널을 개발했으며, 그록 및 삼바노바의 특수 AI 가속기는 희소 활성화 패턴을 효율적으로 처리하도록 설계되었다.
통신을 줄이기 위해 일부 시스템은 전문가를 장치에 복제하는 전문가 병렬 처리 또는 데이터 및 전문가 병렬 처리를 결합한 하이브리드 접근 방식을 사용한다. 라우팅 세분성(토큰 수준 대 블록 수준) 선택도 효율성에 영향을 미친다. 토큰 수준 라우팅은 더 세밀한 제어를 제공하지만 통신이 더 많고, 블록 수준 라우팅은 토큰을 그룹화하여 오버헤드를 줄인다.
과제 및 향후 방향
성공에도 불구하고 MoE의 라우팅은 몇 가지 공개된 문제에 직면해 있다. 하나는 전문가 전문화와 일반화 간의 균형이다. 과도하게 전문화된 전문가는 새로운 작업에 잘 전이되지 않을 수 있다. 또 다른 문제는 라우터가 새로운 데이터 분포에 적응해야 할 수 있으므로 MoE 모델을 미세 조정하기 어렵다는 것이다. 토큰당 활성 전문가 수가 변하는 적응형 라우팅과 현재 토큰 너머의 맥락을 고려하는 학습된 라우팅 정책에 대한 연구가 진행 중이다.
향후 방향으로는 모델 가지치기 및 양자화와 같은 다른 효율성 기술과 라우팅을 결합하고, 다중 모달 입력을 처리할 수 있는 라우터를 개발하는 것이 포함된다. 모델이 계속 확장됨에 따라 라우팅은 용량과 계산의 균형을 맞추는 핵심 메커니즘으로 남을 것이며, 언어를 넘어 비전 및 강화 학습에도 잠재적으로 적용될 수 있다.
같이 보기
참고 문헌
- Shazeer, N., et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.
- Fedus, W., et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.
- Lepikhin, D., et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.
- Jiang, A. Q., et al. (2023). Mixtral of Experts.
- DeepSeek-AI (2024). DeepSeek-V3 Technical Report.