Machine & Deep Learning

[yongggg's] Mixture of Experts (MoE) LLMS

Yonggg 2025. 4. 16. 15:43

MoE는 효율성과 성능을 동시에 향상시키기 위해 사용되는 신경망이며, 모델 architecture에 희소성을 도입하여 모델 크기를 크게 증가시키면서도 컴퓨팅 비용을 증가시키지 않는다는 장점을 갖고 있습니다.

 

이 내용을 공부하면서 간단한 아이디어이지만, 그 아이디어에서 나오는 문제까지 집요하게 해결할 수 있는 눈을 길러야 겠다는 생각을 했습니다. 지금부터 설명을 시작하겠습니다.

1. MoE를 공부하기 위한 기본 지식

1-1. LLM

MoE 기반 LLM은 Decoder 전용 transformer architecture를 기반으로 한다. 이 Decoder 전용 LLM의 한 Block에서는 Normalization(LayerNorm, RMSLayerNorm 등), Masked multi-headed self-attention, Feed-Forward(FF) transformation, residual connetion을 포함하며, 이 블록이 반복되는 구조로 LLM을 구성한다.

1-2. Process

 

text는 LLM의 입력이 되지만, 실제로 다음과 같은 과정을 거치니다. 먼저, text는 Vocab에 정의되어 있는 단어를 바탕으로 tokenizing 된다(이 Token 들은 words와 sub-words로 구성된다). text가 token으로 변환된 후, 입력의 각 token은 LLM의 token embedding layer를 가지며, 이 layer에서 각 token에 대한 vector를 검색하여 입력 matrix를 형성할 수 있다. 각 token embedding이 d 차원이고 C개의 token이 있다면, 입력의 최종 크기는 $(C \times d)$ 이다.

모든 transformer의 block 내의 layer은 위의 입력 크기를 유지한다. 결과적으로 transformers의 모든 Feed-Forward 또는 Attention module에 대한 입력 및 출력은 이와 동일한 크기의 Matrix가 된다는 것이다.

2. Mixture of Experts(MoE)

what is the Experts?

MoE는 decoder로만 구성되어 있는 transformers architecture의 Feed-Forward에서 단순한 수정으로 구성된다.

Feed-Forward Layer의 표준 architecture는 단일 Feed-Forward(일반적으로 두 개의 Feed-Forward layer와 그 사이에 비선형 activation layer로 구성)을 사용하며, 모든 token은 이 신경망을 통해 각각 위의 Layer로 전달된다.

위의 그림은 표준 Feed-Forward Transformation을 나타내는데, MoE는 이 구조를 약간 변경한다. Block의 Feed-Forward 구성 요소 내에 단일 FFN을 두는 대신에, 각각 독립적인 가중치를 갖는 여러 개의 FFN을 생성한다. 이 여러개의 Net을 "Experts"라고 한다. MoE 기반 LLM은 각 FF 하위 Layer에 $N$명의 독립적인 Experts를 포함한다고 부른다.

$$Experts : {E_{i}(\cdot)}_{i=1}^{N}$$

transformers layer 내의 Experts는 위와 같이 정의할 수 있으며, $N$ layer 안에 Experts가 있으며 i-th Experts는 $E_{i}$로 표기할 수 있다.

MoE 기반 transformers 생성.

MoE 기반의 decoder로만 구성되어 있는 transformers architecture를 생성하려면, transformers의 FF Layer를 MoE Layer로 변환하기만 하면 된다. MoE Layer 내의 각 Experts는 해당 Layer의 원래 FFN과 동일한 architecture를 갖는데, 단순히 원래 FFN의 여러 독립적인 복사본만 추가할 뿐이다.

그러나, 모든 FF Layer에 Experts를 사용하는 것이 아니다. 대부분의 MoE 기반 LLM은 P-Stride를 사용한다. 즉, 모든 P 번째 Layer는 Experts Layer로 변환되고 다른 Layer은 그대로 유지된다. 이를 "interleaved(끼워넣은)" MoE Layer라고 한다. 이 Interleaved MoE Layer 방식은 결과적으로 모델의 성능과 효율성 간의 균형을 이룰 수 있도록 한다.

Routing Algorithms

MoE 기반의 architecture의 주요 이점 중 하나는 효율성이지만, Experts만 사용한다고 해서 효율성이 향상되지 않는다. 실제로 모델의 각 Layer에 Experts를 추가하기 때문에, 모델의 총 parameter 수와 필요한 computing 용량이 크게 증가한다. architecture의 효율성을 높이려면 각 Layer에서 사용할 Experts를 분산적으로 선택해야 한다.

Selection Experts: 단일 token vector를 생각해보자. 목표는 이 token을 처리할 Experts의 부분집합을 선택하는 것이다. MoE 관련 문서에서는 일반적으로 token이 Experts에게 routing된다라고 말한다.

 

생각 해보면, 이는 간단하게 해결할 수 있다. token vector에 선형 변환을 적용하여 $N$(Experts 개수)의 벡터를 생성한 후, Softmax를 적용하여 token Experts 집단에 대한 확률 분포를 생성할 수 있다. 이 분포를 사용하여 분포 내 상위 Experts를 선택함으로써 token을 routing할 Experts를 선택할 수 있다.

이 routing 전략은 "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer"에서 사용됐다.

 

하지만, 위의 routing 메커니즘은 Experts의 균형 잡힌 선택이 불가능하게 된다. 즉, 모델은 Experts Layer를 완전히 균일하게 활용하는 대신에, 모든 token에 대해 동일한 소수의 Experts를 반복적으로 선택하는 상태로 수렴할 가능성이 있는데, 이를 "Routing collapse"라고 한다.

 

(Active Parameter: MoE Layer 내에서 각 toekn을 처리할 Experts의 Subset만 선택하기 때문에, "activate paramter"라는 개념이 있다. 즉, 각 MoE Layer에서 선택된 전문가가 제공하는 MoE 모델의 전체 parameter중, 일부의 parameter가 주어진 token을 처리할 때, 활성 parameter로 작용한다. 결과적으로 MoE가 수행하는 총 computation은 전체 parameter의 수가 아닌 active parameter의 수와 비례한다.)

해결 방법

Auxiliary Losses and Expert Load Balancing

학습 과정에서 균형 있는 Experts를 선택하기 위해, 각 Experts를 균등하게 활용하는 모델에 대한 보상을 제공하는 방법을 학습 loss에 추가할 수 있다. 이 추가적인 제약 조건은 각 Experts에 대한 중요도 점수를 정의하여 구현했다. 중요도 점수는 routing 메커니즘에서 각 Experts에 대해 예측한 확률을 기반으로 한다.

data batch가 주어지면, batch 내의 모든 token에 대해 각 Experts에게 할당된 확률의 합을 구하여 중요도를 계산한다. 그 다음, 이 확률이 균형을 이루는지 확인하기 위해 Experts 중요도 점수의 squared coefficient of variation(CV; 제곱 변동 계수)를 구할 수 있다. 즉, 모든 Experts의 중요도 점수가 비슷하게 되면, CV는 작은 값이 되고, 그 반대의 경우도 CV는 작은 값이 된다.

 

위 방법의 중요도 Loss를 표준 Language Modeling Loss에 더해 새로운 regularized training Loss를 구성할 수 있다. 이 추가적인 중요도 Loss 항은 MoE가 학습 과정 전반에 걸쳐 Experts에게 동일한 확률을 할당하도록 하는 데 도움이 된다.

Load balancing
위에서 설명한 중요도 손실은 유용하지만, Experts에게 동일한 중요도가 할당되었다고 해서 token이 동일하게 routing되는 것은 아니다. 예를 들어 다음과 같은 경우 Experts는 동일한 중요도를 갖는다.

  • 매우 높은 확률로 나오는 몇 개의 tokens
  • 낮은 확률로 배정되는 많은 수의 tokens

즉, 중요도 Loss를 사용하더라도 각 Experts에게 전송되는 tokens의 수는 매우 불균형할 수 있으며, 이는 과도한 Memory 사용, MoE의 전반적인 효율성 저하로 이어질 수 있다.

이 문제를 해결하기 위해, Experts의 중요도와 Load balancing을 모두 포괄하는 위의 단일 보조 Loss 항을 만들 수 있다.

이 항은 각 Experrts 간의 equal routing of tokens으로 정의된다. 이 접근법은 "Switch Transformers"에서 제안되었으며, 저자들은 두 가지 값을 고려하여 Loss를 만들었다고 한다.

  • 각 Experts에 할당된 router 확률의 비율
  • 각 Experts에 전송된 tokens의 비율

이 두 비율을 N차원 벡터에 저장하여 dot product를 취하면, 단일 Loss 항을 만들 수 있다. Experts가 균일한 확률과 부하 분산을 받을 때, 결과적으로 Loss가 최소화 되어 단일 auxiliary Loss 항 내에서 두 가지 goals를 이룰 수 있다.

Router z-loss

위에서 설명한 auxiliary load balancing Loss는 MoE에서 널리 사용되지만, "ST-MoE: Designing Stable and Transferable Sparse Expert Models"의 저자는 router z-loss라는 추가 auxiliary loss 항을 제안하여 학습 안정성을 더욱 향상시켰다.

router z-loss는 routing 메커니즘에서 예측한 logit(확률이 아닌 softmax를 적용하기 전의 값)의 size를 제한한다. 이상적으로는 이 logits은 매우 커질 수 있으며, router의 softmax 함수(exponential)를 통과할 때, 전체 precision(float32)을 사용하더라도 학습 process를 불안정하게 만들 수 있는 반올림 오류가 발생할 수 있다.

ST-MoE: router는 float32 정밀도로 Experts에 대한 확률 분포를 계산한다. 그러나 가장 큰 size에서는 신뢰할 수 있는 학습 결과를 얻기에 충분하지 않다.

router가 더 작은 logits을 예측하도록 유도하기 위해, 위에 제시된 Loss 항을 사용할 수 있다. 이 Loss는 router의 logits 정규화에만 초점을 맞추고 load balancing을 수행하지 않으므로, 일반적으로 auxiliary load balancing Loss와 router z-loss를 사용한다. 이 두 Loss는 LLM의 standard Language Modeling Loss에 추가 되었다.

Experts Capacity

MoE Layer에서 수행되는 계산은 학습 및 추론 과정에서 결정되는 routing decisio으로 인해 동적이 된다. 그러나 개부분의 sparse model의 구현을 살펴보면, 일반적으로 batch가 정적임을 알 수 있다. 이는 하드웨어 활용도를 높이는 데 유용하다. (Experts에 들어가는 batch를 추가적으로 구성하고자 함.)

각 Experts에 대해 설정하는 고정의 batch size를 공식화하기 위해 Experts의 capacity를 정의할 수 있다. Experts의 capacity는 다음과 같이 정의된다.

$$\text{Experts Capacity}= {\text{Total Tokens in the Batch} \over N} \cdot \text{Capacity Factor}$$

Experts Capacity는 각 Experts에 보낼 수 있는 batch 내의 최대 token 수로 정의한다. Experts에게 전송된 token 수가 Experts Capacity를 초과하면, 해당 token을 삭제한다. 더 구체적으로, 이러한 token에 대해서는 계산을 수행하지 않고, transformers의 residual connection을 통해 해당 token representation이 다음 Layer로 직접 전달되도록 한다.

ST-MoE: "하드웨어 활용도를 높이기 위해 대부분의 sparse model 구현은 각 Experts에 대해 고정된 배치 크기를 사용한다.
Experts Capacity는 각 Experts에게 routing될 수 있는 token 수를 나타낸다. 이 capacity를 초과하면 오버플로된 토큰은 residual connection을 통해 다음 Layer로 전달된다." 

Experts Capacity는 capacity factor를 통해 제어된다. capacity factor가 1이면 token이 Experts 간에 완벽한 균형 routing이 되고 있음을 의미한다. 반대로 capacity factor를 1보다 높게 설정하면, Experts 간 token 불균형을 수용할 수 있는 추가 buffer가 제공된다. 하지만, 이 경우 Memory 사용량 증가 및 효율성의 저하 등의 단점이 있다.

capacity factor 설정 방법: MoE 모델은 상대적으로 낮은 capacity factor로 좋은 성능을 보이는 경향이 있다. 그러나 학습 실행에 영향을 미치지 않도록 삭제되는 token 수가 너무 많지 않도록 해야한다(경험적으로 수행 가능). 학습 및 추론에 다른 capacity factor를 사용할 수 있다. 예를 들어 ST-MoE는 학습 중에는 1.25의 capacity factor를 사용하고, 평가 중에는 2.0의 capacity factor를 사용한다.

MoE Layer의 Output 계산

router의 output을 얻으면, 위와 같이 최종으로 output을 계산한다. 과정은 다음을 따른다.

  1. 선택된 Experts에 token을 보낸다.
  2. 이 token에 대한 Experts의 output을 계산한다.
  3. router가 각 Experts에게 할당한 확률인 weight를 적용하여 Experts output의 weighted mean을 구한다.

Shared Experts

shared Experts는 다음의 아이디어로 소개된 MoE 방식이다.

  • Experts 그룹은 shared Experts 그룹과 routed Experts 그룹 ,두 가지로 나뉜다.
  • 모든 token은 항상 shared Experts를 통해 전달된다.
  • token은 일반적인 MoE routing 메커니즘에 따라 routed Experts를 통해 전달된다.

위의 그림에서 shared Experts에는 Router가 적용되지 않고, 나머지 Exerts에 적용되는 것을 확인할 수 있다. 일반적으로 shared Experts의 수는 routed Experts보다 적어야 한다. shared Experts의 수를 늘리게 되면, MoE의 희소성 이점이 감소한다.

 

shared Experts를 사용하는 이유는 Experts 간의 중복 정보량을 최소화하기 위해서이다. shared Experts 집합을 사용하면 Network가 여러 Experts에게 동일한 정보를 복제하는 대신, shared Experts 내에 공유 정보를 저장할 수 있다. shared Experts가 있는 MoE Layer의 output을 계산하려면 shared Experts의 output을 일반적인 routing output에 더하기만 하면 된다.

Putting it all Together: Decoder-only LLMs with MoE Layers

위는 MoE Layer의 전체 structure이다. MoE에선 FFN를 Experts Layer로 대체하여 standard decoder-only transformers의 block 구조를 변경한다. 즉, 이 Experts Layer는 원래 FFN의 여러 독립적인 복사본을 포함한다. 주목할 점은 MoE Layer 내의 모든 구성요소(Normal Layer, Experts LAyers, Routing 메커니즘)이 경사하강법으로 공동 학습 된다는 것이다.

 

각 token에 대해 routing 메커니즘은 어떤 Experts를 사용할 지 선택할 수 있으며, 이는 일반적으로 token vector의 간단한 선형 변환을 통해 구현된다. 이를 종합하면, MoE의 수정된 Block 구조는 다음과 같다.

  • Self-Attention Layer
  • Residual Connection & Normalization 작업
  • Experts에게 token을 routing하는 Routing 메커니즘
  • 여러 개의 독립적인 FFN을 가지는 Experts Layer
  • 각 token의 Experts Layer의 최종 output에 적용되는 'final add' 및 Normalization 작업

수정된 block structure를 제외하면, transformers architecture는 동일하게 유지된다. 또한 transformers의 모든 $P$번째 block만 MoE 계층을 사용하도록 변환하고, 다른 block은 변경하지 않는다. 일부 MoE는 모든 Layer에 Experts를 사용하지만, 실제로는 2개, 4개 혹은 6개 등으로 설정하는 것이 일반적이다. 이 방법은 MoE LLM에서 사용되는 총 Parameter의 수를 제어하는 데 유용할 수 있다.

3. MoE 의 장단점

MoE의 장점

LLM의 확장성: 일반적으로 더 큰 model과 더 큰 dataset은 더 나은 성능을 가져온다. 하지만 LLM의 확장에는 비용이 더 많이 든다. MoE의 주요 이점 중 하나는 확장 시 발생하는 문제를 피할 수 있다는 것이다. 즉, token당 고정된 계산 비용으로 모델의 크기를 늘릴 수 있다. 이렇게 하면, 밀도가 높은 모델만 사용하는 경우보다 더 큰 모델을 학습할 수 있다. 언어 모델링 영역에서 이러한 sparse model의 추가 parameter와 표현 능력은 더 나은 model을 만든다.

 

MoE의 연산 이점은 추론 과정에서 가장 큰 영향을 미친다. MoE 모델은 총 parameter 수가 많기 때문에 이러한 parameter를 저장할 수 있는 충분한 개수의 GPU가 필요하다. 그러나 각 token을 처리할 때 이러한 parameter의 고정된 부분만을 사용하기 때문에, 계산 효율성이 크게 향상된다. 추론은 batch size가 작을수록 더 빠르고, batch size가 클수록 처리량이 더 많다. 또, MoE는 학습 효율성도 더 높다. Switch Transformer는 MoE architecture를 사용하면서 pre-train 속도가 7개 향상되었다고 한다.

MoE의 단점

단점은 다음과 같다.

  • 훈련 중 불안정해지기 쉽다.
  • fine-tuning이 어렵다(과도한 과적합 발생).
  • low / mixed precision 훈련 기술에 민감하다.
  • hyper parameter 설정(가중치 초기화 등)에 민감하다.

이와 같이, MoE를 최대한 활용하려면 더 많은 부가 기능이 필요하다. 때문에, MoE 기술은 모든 상황의 최선의 선택이 아닐 수 있다. 예를 들어, 특정 과제에 대한 LLM을 fine-tuning 하려는 경우 dense model이 더 쉬운 선택일 수 있다. 하지만, MoE를 제대로 활용한다면, 다양한 이점을 얻을 수 있다.