Mixture-of-Experts with Expert Choice Routing

하임·2024년 11월 27일

MoE

목록 보기

4/9

간단 요약

이 논문에서는 Mixture-of-Experts (MoE) 모델의 효율성을 향상시키기 위한 Expert Choice Routing이라는 새로운 라우팅 방식을 제안합니다. MoE 모델은 대규모 언어 모델의 효율성을 높이는 데 유용하지만, 기존 방법들에는 여러 한계가 있었습니다. 아래에서는 이 논문의 핵심 내용을 쉽게 설명하겠습니다.

1. 문제점: 기존 MoE 모델의 라우팅 한계

기존의 MoE 모델은 각 입력 토큰이 가장 적합한 전문가(Expert)를 선택하여 할당하는 방식(토큰 선택 라우팅)을 사용합니다. 하지만 이 방식은 다음과 같은 문제점을 야기합니다.

로드 불균형: 특정 전문가에게 토큰이 집중되거나 거의 할당되지 않는 경우가 발생할 수 있습니다. 이로 인해 일부 전문가는 과부하되고, 다른 전문가는 활용되지 않아 효율성이 떨어집니다.
미세 조정의 어려움: 토큰이 고정된 수의 전문가에게 할당되기 때문에, 각 토큰의 복잡도나 중요도에 따라 자원을 유동적으로 할당하기 어렵습니다.

이를 해결하기 위해 논문에서는 Expert Choice Routing 방식을 도입하였습니다.

2. 해결책: Expert Choice Routing 방법론

Expert Choice Routing에서는 토큰이 전문가를 선택하는 대신, 전문가가 토큰을 선택하도록 변경합니다. 즉, 각 전문가는 자신에게 가장 적합한 토큰을 선택하는 방식으로 작동하며, 다음과 같은 장점을 제공합니다.

로드 균형 보장: 전문가가 처리할 토큰의 수를 일정하게 유지하여 로드 불균형 문제를 해결합니다.
유동적 자원 할당: 각 토큰은 필요에 따라 여러 전문가에 의해 처리될 수 있어, 복잡한 토큰에 대해 더 많은 계산 자원을 할당할 수 있습니다.

동작 방식

전문가-토큰 친화도 계산: 각 토큰과 전문가 간의 친화도(연결 강도)를 계산하여, 전문가가 자신에게 적합한 토큰을 선택합니다.
최적 토큰 선택: 각 전문가는 자신에게 가장 적합한 (k)개의 토큰을 선택합니다. 이렇게 선택된 토큰은 해당 전문가에 의해 처리됩니다.

이 과정은 로드 균형을 자동으로 유지할 수 있게 하며, 모델이 효율적으로 학습하도록 돕습니다.

3. 실험 결과: 기존 방법과의 비교

논문에서는 Expert Choice Routing이 기존의 MoE 라우팅 방식보다 훈련 효율성 및 성능에서 뛰어나다는 것을 입증하였습니다.

훈련 속도 향상: Switch Transformer와 GShard와 비교했을 때, 제안된 방식이 훈련 수렴 속도가 2배 이상 빨랐습니다.
다운스트림 작업 성능 향상: GLUE와 SuperGLUE 벤치마크에서 테스트한 결과, 제안된 방법이 기존 방법보다 더 높은 성능을 보였습니다.

4. 결론 및 한계

이 논문에서는 MoE 모델의 라우팅 문제를 해결하기 위해 Expert Choice Routing을 제안하였고, 이를 통해 모델의 로드 균형을 유지하며, 더 빠르고 효율적인 훈련이 가능함을 보여주었습니다. 다만, 이 방법은 현재 미래 토큰까지 활용하여 최적화를 진행하기 때문에, 오토리그레시브 텍스트 생성과 같은 작업에는 바로 적용하기 어렵다는 한계가 있습니다.

Expert Choice Routing의 주요 단계

토큰-전문가 친화도 계산:
- 각 토큰 ( $x_i$ )와 전문가 ( $E_j$ ) 사이의 친화도(affinity score)를 계산합니다. 이 친화도는 토큰과 전문가 간의 연결 강도, 즉 해당 토큰을 전문가가 얼마나 잘 처리할 수 있는지를 나타냅니다.
- 친화도 계산식은 일반적으로 내적(dot product) 또는 코사인 유사도(cosine similarity)를 사용하여 각 토큰과 전문가 간의 점수를 산출합니다.
전문가가 선택할 토큰 수 결정:
- 각 전문가는 자신이 처리할 수 있는 토큰의 수를 미리 설정합니다. 예를 들어, 각 전문가는 최대 ( k )개의 토큰을 선택할 수 있다고 가정합니다.
- 이렇게 하면 로드 균형(load balancing)이 유지됩니다. 모든 전문가가 일정한 수의 토큰을 처리하게 되어, 특정 전문가에 토큰이 몰리는 것을 방지할 수 있습니다.
최적 토큰 선택:
- 각 전문가는 모든 토큰과의 친화도를 바탕으로 자신에게 가장 적합한 ( $k$ )개의 토큰을 선택합니다. 이를 통해 각 전문가는 자신이 처리할 수 있는 최적의 토큰 집합을 확보하게 됩니다.
- 이 과정에서 친화도가 낮은 토큰들은 선택되지 않으며, 친화도가 높은 토큰들이 각 전문가에게 선택되어 배정됩니다.
선택된 토큰 처리:
- 선택된 각 토큰은 자신을 선택한 전문가에게 전달됩니다. 각 전문가는 이 토큰을 처리하며, 이 과정에서 MoE 모델은 병렬로 계산을 수행하여 효율을 극대화합니다.
- 최종적으로 각 전문가가 산출한 결과가 모여 최종 출력이 됩니다.

Expert Choice Routing의 예시

다음은 이 방법이 실제로 어떻게 작동하는지에 대한 예시입니다.

예시 상황: 모델에는 총 4개의 토큰 ( $T_1, T_2, T_3, T_4$ )와 3개의 전문가 ( $E_1, E_2, E_3$ )가 있습니다. 각 전문가는 최대 2개의 토큰을 선택할 수 있습니다.
1단계 - 친화도 계산: 각 토큰이 전문가와의 친화도를 갖습니다.
- 예를 들어, ( $T_1$ )의 ( $E_1$ )에 대한 친화도가 0.9, ( $E_2$ )에 대한 친화도가 0.7, ( $E_3$ )에 대한 친화도가 0.3이라고 할 수 있습니다.
2단계 - 최적 토큰 선택: 각 전문가는 자신과 가장 높은 친화도를 가진 토큰 2개를 선택합니다.
- 예를 들어, ( $E_1$ )는 친화도가 높은 ( $T_1$ )과 ( $T_3$ )을 선택하고, ( $E_2$ )는 ( $T_2$ )와 ( $T_4$ )를 선택합니다.
3단계 - 토큰 처리: 각 전문가가 선택한 토큰을 처리하고, 최종 결과를 산출합니다.

이 방식은 기존의 방식보다 유연하고 효율적으로 토큰을 분배하며, 전문가 간의 로드 균형을 유지할 수 있게 합니다.

장점 요약

로드 균형 개선: 각 전문가가 균등한 수의 토큰을 선택하여 처리함으로써, 특정 전문가에 과부하가 걸리는 문제를 방지합니다.
효율적 계산: 선택된 토큰만을 전문가가 처리하기 때문에, 필요하지 않은 계산을 줄이고 모델 효율성을 높입니다.
성능 향상: 전문가들이 자신의 강점에 맞는 토큰을 선택하기 때문에, 모델이 더 정확하고 신뢰성 있는 결과를 제공합니다.

요약

희소 활성화된 Mixture-of-Experts (MoE) 모델은 특정 토큰 또는 샘플에 대해 계산량을 일정하게 유지하면서도 파라미터 수를 크게 늘릴 수 있습니다. 하지만, 잘못된 전문가 라우팅 전략은 특정 전문가가 충분히 훈련되지 않거나 지나치게 특화되게 만드는 문제를 야기할 수 있습니다. 기존 연구에서는 서로 다른 토큰의 상대적 중요성을 고려하지 않고 각 토큰에 고정된 수의 전문가를 할당하는 top-k 함수를 사용했습니다. 이를 해결하기 위해 우리는 전문가 선택 방식을 도입한 이질적 Mixture-of-Experts를 제안합니다. 토큰이 top-k 전문가를 선택하는 대신, 전문가가 top-k 토큰을 선택하는 방식으로, 각 토큰이 가변적인 수의 전문가로 라우팅될 수 있으며, 각 전문가는 고정된 버킷 크기를 가질 수 있습니다. 우리는 Switch Transformer의 top-1과 GShard의 top-2 게이팅과 동일한 계산 자원을 사용하여 사전 훈련 속도를 체계적으로 연구했으며, 우리의 방법이 훈련 수렴 시간을 2배 이상 개선함을 확인했습니다. 동일한 계산 비용에서, 우리의 방법은 GLUE와 SuperGLUE 벤치마크의 11가지 선택된 과제에서 더 높은 성능을 보여주었습니다. 더 적은 활성화 비용으로도, 우리의 방법은 11개 과제 중 7개에서 T5 밀집 모델을 능가했습니다.

1. 서론

모델 용량, 데이터셋 크기, 훈련 시간을 확장하는 것은 컴퓨터 비전 아키텍처와 신경 언어 모델의 성능을 크게 향상시키는 데 성공적인 접근 방식임이 입증되었습니다. 최종 모델 품질은 데이터 양, 모델 크기, 연산 시간과의 거듭제곱 법칙 관계를 따르며, 최첨단 시스템을 뛰어넘는 모델 품질을 달성하는 데 사용되는 총 연산량을 뜻하는 훈련 효율성도 중요한 요소로 주목받고 있습니다.

희소 게이팅 MoE는 고정된 계산 비용에서 모델 용량을 확장하는 효과적인 방법을 제공하며, 최근 대규모 언어 모델의 훈련 효율성 향상에 중요한 역할을 해왔습니다. MoE는 여러 전문가를 서브 네트워크로 채택하고, 각 입력 토큰에 대해 한두 개의 전문가만 활성화하여 작동합니다. 라우팅 네트워크는 각 토큰을 가장 적합한 전문가에게 라우팅하기 위해 최적화되어야 합니다. 예를 들어, 최근 연구에서는 k-평균 클러스터링, 토큰-전문가 친화도를 극대화하는 선형 할당, 또는 해싱을 통해 희소 라우팅을 구현했습니다. 많은 이전 연구는 토큰 선택에 기반한 라우팅 전략을 사용하여 각 토큰이 가장 적합한 한두 전문가를 선택하는 방식을 따릅니다.

우리는 이전 연구에서 독립적으로 토큰이 전문가를 선택하는 방식이 전문가 간의 로드 불균형을 초래해 훈련 비효율성과 모델의 최적 훈련이 이루어지지 않는다고 주장합니다. 이러한 문제를 완화하기 위해 희소 게이팅 네트워크는 많은 토큰이 단일 전문가로 라우팅되는 것을 방지하기 위한 추가 보조 손실을 도입하지만, 그 효과는 제한적입니다. 최근 접근법은 사전 훈련에만 집중하며, 다운스트림 과제에서의 성능 향상을 보여주지 못했습니다. 또한, 이전의 방법들은 각 토큰의 중요도에 따라 가변적인 수의 전문가를 할당하는 것을 고려하지 않았으며, 이는 유익할 수 있습니다.

우리는 간단하면서도 효과적인 Expert Choice라는 라우팅 방법을 제안합니다. 기존 MoE에서는 토큰이 점수가 높은 한두 개의 전문가를 선택하는 것과 달리, 우리의 방법은 각 전문가가 top-k 토큰을 선택하도록 합니다. 이 방법은 완벽한 로드 균형을 보장하며, 각 토큰에 대해 가변적인 수의 전문가를 허용하고, 훈련 효율성과 다운스트림 성능에서 큰 향상을 보여줍니다. 우리의 주요 기여는 다음과 같습니다:

기존 MoE의 로드 불균형 문제를 식별하고(3.1절), 학습된 토큰-전문가 중요도에 기반한 유동적 모델 파라미터 할당을 제공하는 이질적 Expert Choice 방법을 제안합니다. 이 방법은 보조 손실을 추가하지 않고도 본질적으로 로드 균형을 보장합니다.
우리는 Switch Transformer와 GShard의 top-1 및 top-2 게이팅과 비교하여 8B/64E(80억 활성화 파라미터, 64명의 전문가) 모델에서 훈련 수렴 속도를 2배 이상 빠르게 한다는 것을 보여줍니다.
우리의 방법이 전문가 수를 16에서 128로 증가시킬 때 강력한 확장성을 보여주며, 훈련 퍼플렉시티를 평가합니다.
GLUE와 SuperGLUE에서 평가한 모든 규모에서 강력한 다운스트림 성능을 보여줍니다. 특히, 8B/64E 모델은 T5 11B 밀집 모델을 평가된 11개의 과제 중 7개에서 능가했습니다.

2. 관련 연구

확장성(Scaling)

신경망 성능을 높이기 위해 모델 용량을 확장하는 다양한 접근법이 제안되었습니다. 최근 연구에서는 모델 병렬화(model parallelism)를 통해 수십억 개의 파라미터로 모델을 확장했습니다. 모델 병렬화는 가중치와 텐서를 여러 코어에 나누어 처리하고, 파이프라인 병렬화는 각 레이어를 장치에 분할하여 마이크로 배치를 다른 레이어로 전송하는 방식입니다. 이러한 신경망 확장을 지속하기 위해서는 모델 훈련 및 서비스 효율성 개선이 중요한 연구 주제로 부각되고 있습니다.

조건부 계산(Conditional Computation)

입력에 따라 계산을 동적으로 결정하는 조건부 계산이 제안되었으며, 예를 들어 특정 파라미터나 계산을 상황에 맞게 활성화하여 신경망의 용량을 증가시키는 방식입니다. 조건부 컨볼루션 레이어는 작업별 게이팅을 통해 연속 학습에서 치명적 망각을 방지합니다. 게이팅 결정은 이진 또는 희소하고 연속적일 수 있으며, 확률적(stochastic) 또는 결정적(deterministic)일 수 있습니다.

Mixture of Experts (MoE)

희소 게이팅을 사용하는 MoE는 모델 용량, 훈련 시간 또는 모델 품질에서 큰 개선을 보여주었습니다. Switch Transformer는 각 토큰에 대해 상위 전문가 하나만 선택하는 간소화된 게이팅 방식을 사용하며, 이전 연구보다 더 나은 확장성을 보였습니다. 기존의 연구는 로드 균형을 유지하기 위해 보조 손실(auxiliary loss)을 추가해야 하며, 하드 용량 제한을 적용해도 많은 토큰이 MoE 레이어에서 처리되지 않을 수 있습니다. Hard MoE는 해시 태그 예측 작업에서 효과적으로 학습되었으며, Base Layers는 토큰-전문가 친화도를 극대화하여 각 전문가에게 동등한 수의 토큰을 배정합니다. THOR는 훈련 및 추론 시에 무작위로 전문가를 활성화하며, 번역 작업에서 강력한 성능을 보였습니다. 본 연구는 기존 연구와 달리 이질적 MoE를 가능하게 하는 학습 방법으로, 다운스트림 작업에서의 성능 향상을 효과적으로 지원합니다.

3. 방법론

우리는 기존 MoE 모델의 라우팅 방법의 몇 가지 문제점을 먼저 파악하고, 이를 해결하기 위해 Expert Choice 방식을 제안합니다.

3.1 토큰 선택 라우팅의 문제점

기존 MoE 모델은 각 토큰을 최적의 전문가에게 할당하기 위해 독립적으로 상위 ( k ) 전문가를 선택하는 토큰 선택 라우팅 방식을 사용합니다. 하지만, 이러한 전략은 다음과 같은 문제를 야기할 수 있습니다.

로드 불균형: 토큰 선택 라우팅은 전문가 간 로드 불균형을 초래할 수 있습니다. 일부 전문가는 대부분의 토큰으로 과부하 상태가 되고, 다른 전문가들은 사용되지 않아 모델의 용량이 낭비될 수 있습니다. 이전 방법들은 이 문제를 완화하기 위해 로드 균형을 맞추는 보조 손실을 추가하지만, 이는 초기 훈련 단계에서 불균형을 해결하지 못할 수 있습니다.
미숙한 전문화: 이상적인 MoE 레이어는 유사한 토큰을 동일한 전문가에게 라우팅해야 하지만, 부적절한 전략은 중복되거나 미숙하게 특화된 전문가를 초래할 수 있습니다. 특히, 로드 균형을 중시하여 보조 손실을 강하게 부과할 경우 라우팅의 효율성이 저하될 수 있습니다.
고정된 토큰 계산량: 토큰 선택 전략은 모든 토큰이 동일한 수의 전문가에게 라우팅되므로 동일한 계산량을 차지하게 됩니다. 모델은 입력의 복잡성에 따라 유연하게 자원을 할당할 수 있어야 합니다.

3.2 Expert Choice를 통한 이질적 MoE

Expert Choice 방법에서는 각 전문가가 상위 ( k )개의 토큰을 독립적으로 선택하여 처리하는 방식을 채택합니다. 이 방식은 로드 균형을 완벽하게 유지하며, 각 토큰이 가변적인 수의 전문가로 라우팅될 수 있어 유연한 계산 자원 할당이 가능합니다.

토큰-전문가 친화도 계산: 토큰과 전문가 간의 친화도를 계산하고, 각 전문가가 가장 적합한 토큰을 선택합니다.
가변 전문가 할당: 모델은 각 토큰에 대해 가변적인 수의 전문가를 할당할 수 있습니다.

입력 토큰 표현 ( $X \in R^{n \times d}$ )에 대해, 토큰-전문가 할당은 세 개의 출력 행렬 ( I ), ( G ), ( P )로 나타내어집니다. 여기서 ( I )는 선택된 토큰의 인덱스 행렬이며, 게이팅 행렬 ( G )는 각 선택된 토큰에 대한 전문가 가중치를 나타내고, ( P )는 각 전문가에 대한 토큰을 모으는 데 사용되는 원핫 행렬입니다.

이와 같은 Expert Choice는 Switch Transformer 및 GShard와 유사하게, Transformer 기반 네트워크에서 가장 계산 비용이 높은 밀집 피드포워드(FFN) 레이어에서 사용됩니다. 각 전문가의 출력을 효율적으로 계산하여 최종 FFN 레이어 출력을 얻습니다.

3.3 추가 제약이 있는 Expert Choice

각 토큰에 할당되는 전문가의 최대 수를 제한하여 라우팅을 정규화하는 추가 제약도 고려합니다. 이 제약은 학습과 미세 조정 결과를 개선할 수 있는지 분석하는 데 도움이 됩니다.

3.4 모델 아키텍처

전체적으로 이 연구에서는 희소 활성화된 Mixture-of-Experts (MoE) 모델의 개념을 채택하였고, Transformer 아키텍처를 사용하여 최근의 연구 방식에 따라 매 다른 Transformer 레이어의 피드포워드 컴포넌트를 MoE 레이어로 대체했습니다. 이렇게 일반 Transformer 레이어와 MoE 레이어를 교차(interleaving)하여 배치하면 모델 성능과 훈련 효율성이 개선되며, 이는 MoE 레이어 간에 일부 공유 컴포넌트를 강제하여 토큰 누락의 부정적 영향을 완화하는 데 도움이 되는 것으로 나타났습니다.

추가적으로, 최근 연구에서 채택된 몇 가지 수정 사항을 실험에 적용하였습니다. 예를 들어, 표준 위치 임베딩을 레이어별 상대 위치 바이어스로 대체하였고, 비 MoE 피드포워드 서브레이어에서는 첫 번째 선형 프로젝션과 활성 함수 대신 Gated Linear Unit (GLU)을 사용했습니다.

각 MoE 레이어는 독립적인 피드포워드 네트워크의 그룹으로 이루어져 있으며, 이를 "전문가"라고 합니다. 이 레이어에서 각 토큰이 가장 적합한 전문가들에게 할당될 확률을 계산하기 위해 softmax 활성화 함수를 사용하는 게이팅 함수가 사용됩니다. 훈련 과정에서는 이 게이팅 네트워크가 토큰 차원에서 상위 ( k )개의 전문가를 선택하도록 학습됩니다.

주요 단계 및 추가 요소

Shuffle 단계와 Unshuffle 단계가 MoE 레이어에 추가되어, 첫 번째 단계에서는 토큰이 지정된 전문가로 이동하고 두 번째 단계에서는 원래 순서로 되돌아옵니다.
MoE 레이어는 더 많은 파라미터를 포함하지만, 각 토큰에 대해 활성화되는 전문가가 제한적이기 때문에 밀집 레이어와 유사한 계산 비용을 가질 수 있습니다.

본 연구에서는 전문가의 활성화 수를 토큰마다 다르게 하면서도, 식 (1)에서의 용량 계수 ( c )를 고정하여 계산 비용을 기본 아키텍처와 동일하게 유지했습니다. 기본적으로 ( $c = 2$ )로 설정하여, 우리의 방법을 GShard의 top-2 토큰 선택 게이팅과 직접 비교할 수 있게 하였습니다.

4. 실험

4.1 설정

Table 1에 다양한 MoE 모델의 하이퍼파라미터 설정이 요약되어 있으며, 각 토큰당 활성화되는 파라미터 수와 유사한 밀집 모델 구성도 포함됩니다. 전문가 수를 확장하면서 전문가별 크기는 100M 파라미터로 고정하여, 모델 용량 증가가 훈련에 미치는 영향을 연구했습니다.

데이터셋: 다양한 자연어 사용 사례를 대표하는 1.6조 개의 토큰이 포함된 GLaM 데이터셋을 사용했습니다. 품질 분류기를 사용하여 고품질의 웹페이지 콘텐츠만을 선별해 책, Wikipedia 페이지, 대화, 포럼, 뉴스와 결합하여 최종 데이터셋을 생성했습니다.
모델 훈련: GLaM 설정을 따르며 최대 시퀀스 길이를 1024 토큰으로 하여 Adafactor 옵티마이저를 사용했습니다. 훈련의 첫 10K 단계 동안 학습률을 일정하게 유지하고 이후에는 역제곱근 스케줄을 사용해 감쇠시켰습니다.
모델 평가: GLUE와 SuperGLUE 벤치마크의 11개 과제에서 미세 조정 성능을 평가했습니다.

4.2 훈련 효율성

훈련 효율성 및 수렴 속도를 비교한 결과, GShard top-2와 비교하여 우리의 EC-CF2 방법이 훈련 속도를 2배 이상 빠르게 했습니다. EC-CF2는 GShard top-2와 동일한 퍼플렉시티에 도달하는 데 절반 이하의 단계만을 필요로 했고, GShard top-2의 각 단계가 우리의 방법보다 20% 더 느린 것으로 나타났습니다.

4.3 전문가 수 확장

전문가 수를 늘리는 것이 모델 용량을 증가시킬 수 있음을 확인했으며, 더 많은 전문가가 훈련 퍼플렉시티 개선에 지속적으로 기여함을 실험에서 확인했습니다.

4.4 GLUE와 SuperGLUE에서의 미세 조정

GLUE와 SuperGLUE의 11개 과제에서 미세 조정을 수행한 결과, EC-CF2가 모든 관련 방법보다 우수한 성능을 보였고, 특히 8B/64E 설정에서 평균 정확도가 2% 이상 증가했습니다.

4.5 이질성의 중요성 (Heterogeneity Matters)

제한된 전문가 선택 (Capped Expert Choice)

토큰당 최대 전문가 수를 제한하는 방식으로 Expert Choice를 정규화하여 Section 3.3에서 설명한 방식을 적용했습니다. Table 4에 따르면, EC-CAP2는 각 토큰에 대해 최대 두 명의 전문가만 할당하는 방식으로, 평균적으로 0.8 포인트의 미세 조정 정확도 감소가 발생했습니다. 반면, EC-CAP3는 각 토큰에 최대 세 명의 전문가를 허용했으며, 기본 Expert Choice 방식과 유사한 결과를 얻었습니다. 이 실험은 각 토큰에 가변적인 전문가 수를 허용하는 것이 도움이 된다는 점을 확인해줍니다.

토큰별 가변 전문가 할당 (Variable Experts per Token)

토큰에서 전문가로 라우팅되는 통계를 계산한 결과, 대부분의 토큰은 한두 명의 전문가에게 라우팅되었고, 약 23%의 토큰이 세 명 또는 네 명의 전문가로 라우팅되었으며, 3%의 토큰만이 네 명 이상의 전문가에게 라우팅되었습니다. 이러한 분포는 중요한 토큰에 대해 가변적인 전문가 수를 할당하는 방법이 효과적이라는 우리의 가설을 뒷받침합니다.

4.6 해시 레이어와의 비교 (Comparison with Hash Layer)

해시 레이어 방식과 우리의 Expert Choice 방식을 비교했습니다. 해시 레이어는 토큰 ID를 전문가 ID로 매핑하기 위해 모드 연산을 사용하여 로드 균형을 유지하고 전문화된 전문가를 생성합니다. 하지만, 해시 라우팅 방식은 평균 점수와 분산 면에서 Expert Choice보다 성능이 떨어졌습니다. 이는 로드 균형만으로는 모든 장점을 얻을 수 없음을 의미합니다.

4.7 소거 연구 (Ablation)

용량 계수 (Capacity Factor)

용량 계수는 각 토큰이 평균적으로 얼마나 많은 전문가에게 라우팅될 수 있는지를 결정하는데, 이는 전문가당 버킷 크기 ( k )를 의미합니다. 모든 이전 실험에서 용량 계수 2를 사용했으며, 이는 GShard 방식의 top-2 게이팅과 계산 비용이 일치합니다. 용량 계수를 1로 줄여 Switch Transformer의 top-1 게이팅과 공정하게 비교한 결과, 더 작은 용량 계수는 더 높은 퍼플렉시티를 초래했지만 여전히 top-1 게이팅보다 높은 성능을 보였습니다. 용량 계수를 0.5로 낮추어도 top-1 게이팅보다 성능이 뛰어난 것으로 나타났습니다.

밀집 모델과의 사전 훈련 비교

사전 훈련 단계에서 밀집 모델과 비교한 결과, 우리의 방식이 퍼플렉시티와 수렴 시간 면에서 일관되게 뛰어났습니다. 100M 파라미터의 작은 전문가 크기에서 희소 게이팅의 이점이 더욱 두드러졌습니다. 또한, 전문가 수를 늘리면 모델 성능이 개선됨을 확인했습니다.

5. 결론

이 연구에서는 희소 활성화된 MoE 모델을 위한 새로운 라우팅 방법을 제안했습니다. 이 방법은 기존 MoE 방법의 로드 불균형과 전문가 미활용 문제를 해결하며, 토큰당 다른 수의 전문가를 선택할 수 있게 합니다. 우리의 모델은 GShard와 Switch Transformer 모델 대비 2배 이상의 훈련 효율성 향상을 보여주었으며, GLUE와 SuperGLUE 벤치마크의 11개 데이터셋에서 강력한 성능 향상을 달성했습니다.

6. 한계 (Limitations)

자동회귀 텍스트 생성에의 적용: 현재 구현은 과거와 미래 토큰을 함께 사용해 top-k 선택을 수행하기 때문에 자동회귀 텍스트 생성에 즉각적으로 적용되기 어렵습니다. 이를 해결하기 위해 대량의 입력 시퀀스를 수집하고, 동일한 시퀀스의 토큰을 별도의 그룹으로 나누어 그룹별로 Expert Choice 라우팅을 수행할 수 있습니다.
작은 배치 크기에서의 제한: 추론 시 배치 크기가 매우 작아질 경우 Expert Choice 방식은 제한적으로 적용될 수 있습니다. 이 경우, 글로벌 top-k 방식을 선택하고 각 전문가 또는 토큰이 선택되는 횟수를 제한하는 방식이 가능할 수 있습니다.
큰 메모리 요구량: MoE의 오랜 문제 중 하나는 높은 메모리 요구량입니다. 희소 게이팅을 통해 계산 비용을 줄일 수 있지만, 전문가 수가 증가할수록 하드웨어 자원이 많이 필요합니다.

하임

NLP 공부합니당

이전 포스트

MoRAL: MoE Augmented LoRA for LLMs’ Lifelong Learning

다음 포스트