전문가 혼합(MoE) 설명

전문가 혼합(MoE)은 트랜스포머 모델이 밀집 모델보다 훨씬 빠르게 사전 학습될 수 있게 하고, 전체 파라미터 수에 비해 더 빠른 추론 속도를 달성합니다. 이는 밀집 피드포워드 네트워크(FFN) 레이어를 희소 MoE 레이어로 교체함으로써 이루어지며, 해당 레이어는 주어진 입력 토큰에 대해 소수의 “전문가”만 활성화합니다.

전문가 혼합의 핵심 아키텍처

MoE 모델은 표준 FFN 레이어를 희소 MoE 레이어와 라우팅 메커니즘의 조합으로 교체합니다. 이 아키텍처는 두 가지 주요 구성 요소로 이루어집니다:

  • Sparse MoE Layers: 단일 밀집 레이어 대신, 이 레이어들은 여러 개의 “전문가”를 포함하며, 각 전문가는 일반적으로 신경망(종종 FFN)입니다. 이러한 전문가들은 자체가 MoE인 계층 구조로 배열될 수 있습니다.
  • Gate Network (Router): 학습된 라우터가 토큰을 어떤 전문가에게 보낼지 결정합니다. 라우터는 네트워크의 나머지와 함께 사전 학습되며, 토큰을 하나 이상의 전문가에게 보낼 수 있도록 구성할 수 있습니다.

MoE는 계산 효율성을 제공하지만 특정 트레이드오프를 동반합니다. 모든 전문가가 메모리에 로드되어야 하기 때문에 높은 VRAM이 필요합니다—단일 순전파 동안 활성화되는 전문가 수는 적지만 말이죠. 예를 들어, Mixtral 8x7B는 47B 파라미터 모델을 담을 수 있는 충분한 VRAM이 필요하지만, 토큰당 두 명의 전문가를 사용할 경우 추론 계산(FLOPs)은 12B 모델에 가까워집니다.

희소성 및 라우팅 메커니즘

희소성은 조건부 계산을 사용하여 모델이 크기를 확장하면서도 계산량이 비례적으로 증가하지 않게 합니다—예시마다 시스템의 일부만 활성화합니다.

게이팅 함수

라우팅은 어떤 전문가 $E$를 활성화할지 결정하는 게이팅 네트워크 $G$에 의해 처리됩니다. 전통적인 설정에서는 softmax 함수가 사용됩니다. Noisy Top-k Gating과 같은 보다 진보된 접근법은 조정 가능한 노이즈를 도입하고 상위 $k$ 값을 선택하여 학습 중 더 나은 탐색과 부하 균형을 보장합니다.

부하 균형 및 전문가 용량

소수의 전문가가 과도하게 사용되고 다른 전문가가 무시되는 “승자 독식” 상황을 방지하기 위해 MoE는 다음을 사용합니다:

  • Auxiliary Loss: 라우터가 토큰을 모든 전문가에게 균등하게 분배하도록 장려하는 추가 손실 함수.
  • Expert Capacity: 단일 전문가가 처리할 수 있는 토큰 수를 제한하는 임계값. 전문가가 용량에 도달하면 토큰이 삭제되거나 잔차 연결을 통해 다음 레이어로 전달될 수 있습니다.

MoE 구현의 진화

GShard와 Switch Transformers

Google의 GShard는 top-2 게이팅을 사용하고 무작위 라우팅(두 번째 전문가가 가중치에 비례하여 선택됨)을 도입하여 트랜스포머를 6000억 파라미터 이상으로 확장했습니다.

Switch Transformers는 단일 전문가 전략을 도입하여 이를 더욱 단순화했습니다. 각 토큰을 하나의 전문가에게만 라우팅함으로써 Switch Transformers는 라우터 계산을 감소시키고, 전문가 배치 크기를 절반으로 줄이며, 통신 비용을 낮추면서도 품질을 유지했습니다. 이 아키텍처는 T5-XXL 대비 4배의 사전 학습 속도 향상을 달성했습니다.

안정성 및 전문화

학습 안정성은 Router z-loss(ST-MoE에서 도입)와 같은 기법을 통해 관리되며, 이는 게이팅 네트워크에 들어가는 큰 로짓을 패널티하여 반올림 오류를 감소시킵니다.

전문가 전문화에 대한 연구는 인코더 전문가가 종종 얕은 개념(예: 구두점이나 고유명사)에 특화되는 반면, 디코더 전문가는 덜 특화된다는 것을 보여줍니다. 다국어 설정에서는 부하 균형 요구사항 때문에 전문가가 단일 언어에 특화되는 경우는 드뭅니다.

파인튜닝 및 인스트럭션 튜닝

희소 모델은 밀집 모델보다 과적합에 더 취약합니다. 효과적인 파인튜닝 전략은 다음과 같습니다:

  • Higher Regularization: 희소 레이어 내에서 높은 드롭아웃 비율을 사용합니다.
  • Selective Freezing: MoE 레이어를 고정하고 비-MoE 레이어를 업데이트하면 품질을 유지하면서 학습 속도를 높일 수 있습니다.
  • Hyperparameter Adjustment: 희소 모델은 종종 작은 배치 크기와 높은 학습률에서 이점을 얻습니다.

최근 연구에 따르면 MoE는 밀집 모델보다 인스트럭션 튜닝으로부터 훨씬 큰 혜택을 받습니다. MoE와 동일한 T5 모델을 비교했을 때, MoE의 인스트럭션 튜닝에 의한 성능 향상(Flan-MoE vs MoE)이 밀집 모델의 향상(Flan T5 vs T5)보다 더 컸습니다.

계산 효율성 및 병렬성

병렬 전략

MoE의 방대한 파라미터 수를 처리하기 위해 다양한 병렬 기법이 사용됩니다:

  • Expert Parallelism: 전문가들을 서로 다른 워커에 분산합니다. MoE 레이어에서는 토큰이 해당 전문가가 위치한 워커로 라우팅되고, 비-MoE 레이어에서는 데이터 병렬과 동일하게 동작합니다.

하드웨어 최적화

  • Capacity Factor: 품질과 통신 비용 사이의 트레이드오프. 높은 용량 팩터는 품질을 향상시키지만 활성화 메모리와 디바이스 간 통신 메모리를 증가시킵니다.
  • Block-Sparse Operations: MegaBlocks와 같은 프로젝트는 전통적인 배치 행렬 곱셈을 블록-희소 연산으로 대체하여 토큰을 삭제하지 않고도 불균형 토큰 할당을 처리하며, 이는 상당한 속도 향상을 가져옵니다.

비교: 희소 MoE vs. 밀집 모델

특징 희소 MoE 밀집 모델
사전 학습 계산 보다 효율적 / 빠름 덜 효율적
추론 속도 빠름 (전체 파라미터 대비) 느림 (전체 파라미터 대비)
VRAM 요구량 높음 (모든 전문가를 로드해야 함) 낮음 (전체 파라미터 대비)
최적 사용 사례 고처리량, 다수의 머신 저처리량, 제한된 VRAM

오픈 소스 MoE 생태계

여러 프레임워크와 모델이 MoE를 오픈 커뮤니티에 도입했습니다:

  • 프레임워크: MegaBlocks, Fairseq, OpenMoE.
  • 모델: Google's Switch Transformers (최대 1.6T 파라미터), Meta's NLLB MoE (54B), 그리고 Mistral의 Mixtral 8x7B는 Llama 2 70B보다 빠른 추론으로 성능을 능가합니다.

Sources