MoE란? - Mixture of Experts

https://huggingface.co/blog/moe

Mixture of Experts Explained

 

MoE란? 

모델의 규모가 크면 클수록 더 좋은 성과를 보인다.

그러나 MoE를 사용하면 훨씬 더 적은 컴퓨팅으로 모델을 학습할 수 있으므로 dense 모델과 동일한 컴퓨터 리소스로 모델이나 데이터 크기를 극적으로 확장할 수 있다. 특히 사전학습 중에 dense보다 훨씬 빠른 학습 속도를 가진다.

Tranformer 모델에서 dense feed-forward(FFN) 대신 sparse MoE layer를 사용한다. MoE는 일정한 수의 전문가가 있으며, FFN이거나 더 복잡한 계층일 수 있어 계층적 MoE가 된다.

이 문장이 헷갈리는 이유는 MoE (Mixture of Experts)가 구조적으로 FFN보다 더 복잡하게 느껴질 수 있기 때문이에요. 하지만 본질적으로 MoE 레이어에서의 각 "전문가"는 실제로 FFN입니다.

다음과 같이 설명할 수 있어요:

1. MoE의 Sparse 활성화: 일반적인 FFN은 모든 유닛이 활성화되지만, MoE는 입력에 따라 선택된 "전문가" 일부만 활성화됩니다. 이 선택적 활성화가 MoE를 sparse하게 만드는 핵심입니다.
2. MoE에서 FFN을 전문가로 사용하는 이유: 각 전문가는 보통 FFN으로 구성되는데, 이는 FFN이 특정 데이터 패턴을 효율적으로 처리할 수 있기 때문이에요. 다만, 이 전문가들은 이론상 FFN 외의 다른 네트워크 유형이나 아키텍처로 구성될 수도 있습니다. 중요한 점은 전문가가 특정 데이터 패턴을 전문적으로 처리할 수 있다는 것이죠.
3. 계층적 MoE: 계층적 MoE는 각 전문가가 또 하나의 MoE로 구성될 때 형성됩니다. 이 경우, 각 "전문가"가 자체 FFN 또는 또 다른 전문가들 중 일부를 선택하게 되죠. 이렇게 계층화된 구조는 모델이 더욱 세분화된 데이터 측면에 집중할 수 있게 해주어, 효율성 및 특화 정도가 더욱 향상됩니다.

이 디자인은 입력에 따라 동적으로 계산 자원을 할당해 주고, 복잡한 라우팅 시스템에서 더 단순한 FFN을 활용하여 매우 효율적인 처리가 가능하도록 합니다.

게이트나 라우터(router)가 각각의 토큰이 어떤 전문가에게 보내질지 정한다. 

사전 학습에서는 훨씬 효율적이었지만 fine-tuning 과정에서 일반화하는데 어려움이 있고, overfitting도 발생하고있다.

추론 과정에서는 모든 전문가들을 메모리에 넣어야 하므로 VRAM이 충분히 필요하지만 동일한 크기의 dense 모델보다 빠른 추론 속도를 가지고 있다.(MoE 모델은 모든 파라미터를 사용하는 것이 아니라 일부만을 사용하기 때문에)

 

MoE는 "Adaptive Mixtures of Local Experts"에서 처음 나왔습니다. 

그 이후 계속적으로 업데이트해 나갔고, 희소성을 도입해 대규모에서도 매우 빠른 추론을 유지하였지만 높은 소통 비용과 훈련 불안전성을 가지게 되었습니다.

 

Sparsity란? (희소성)

dense가 모델의 모든 파라미터가 모든 입력에 사용되지만 sparsity는 일부만 사용한다.
Sparsity를 활용하여 계산을 늘리지 않고, 모델의 크기를 확장하여 MoE 계층에서 수천 명의 전문가가 만들어진다.

일반적인 Gate는 전문가에게 보내기 위해 Softmax를 활용하여 가중치를 구하지만 Noisy Top-k Gating 방식을 활용해 상위 k만 값을 유지합니다.

k = 2로 한다면 빠른 추론 및 훈련이 가능하고, 1을 하지 않는 이유는 게이트가 다른 전문가에게 라우팅하는 방법을 배우려면 두 명 이상의 전문가에게 라우팅해야 했기 때문이다.

노이징이 들어간 이유는 load 밸런싱을 맞추기 위해서이다.

 

그런데 일부 전문가만 학습하게 되면 훈련이 엄청 오래 걸리게 된다!
Auxiliary Loss(보조 손실)을 추가하여 모든 전문가가 학습을 진행하도록 도와준다.

 

 

전문가 용량(Expert capacity)가 있어서 일정한 토큰수 이상이 되면 전문가가 더이상 처리하지 않고 넘겨버린다는데 이게 가능한지..... 해도 되는 건지...

MoE 모델에서 "Expert Capacity"는 각 전문가가 처리할 수 있는 토큰의 최대 수를 설정하는 개념입니다. 이를 통해 효율적으로 자원을 관리하면서 예측 가능한 성능을 유지할 수 있습니다.

1. Expert Capacity 설정 이유: MoE의 각 전문가가 얼마나 많은 토큰을 받을지 예측할 수 없기 때문에, 모델의 텐서 모양이 컴파일 시점에서 고정될 수 있도록 용량을 미리 설정합니다. 이로 인해 각 전문가가 처리할 수 있는 토큰 수에 제한을 두어 모델의 안정성을 높이는 것입니다.
2. 오버플로 처리: 설정된 용량을 초과한 토큰들은 오버플로로 간주됩니다. 이 경우 두 가지 방식으로 처리됩니다:
   • Residual Connection으로 넘기기: 초과된 토큰을 다음 레이어로 전달해, 원래 계산을 완전히 무시하지 않고 모델의 다른 레이어에 영향을 줄 수 있도록 합니다.
   • 버리기: 일부 프로젝트에서는 오버플로 토큰을 아예 버려, 모델이 중요한 입력에만 집중하도록 합니다.
3. 왜 필요한가? 이렇게 하는 이유는, 모델의 일관성을 유지하면서 동적으로 많은 양의 데이터를 처리할 수 있게 하기 위함입니다. 또한, 일정 용량을 넘어가는 데이터가 모델의 다른 레이어로 넘어가도록 설계하면, 학습이 집중될 수 있도록 유도하여 효율성을 높이는 데 도움이 됩니다.

이처럼 Expert Capacity 설정을 통해 효율성을 높이면서도 무한히 많은 데이터가 쏟아질 때 발생할 수 있는 리소스 낭비를 방지하게 됩니다.

 

만약 버리는 방식을 선택하면 뒤에 있는 토큰에서 중요한 내용이 있다 그래도 버려져서 악영향을 끼치는 거 아니야?

네, 맞습니다. 토큰을 버리는 방식은 뒤에 중요한 내용이 포함된 토큰이라도 버려질 위험이 있습니다. 이는 모델 성능에 부정적인 영향을 미칠 수 있죠. 특히, 중요 정보가 중간이나 후반에 있는 경우, 버림으로 인해 모델이 중요한 패턴을 놓칠 가능성이 큽니다.

이 문제를 최소화하려면 다음과 같은 접근 방식이 사용될 수 있습니다:

1. Residual Connection 사용: 오버플로된 토큰을 다음 레이어로 넘기는 방법은 중요한 정보를 완전히 무시하지 않게 해줍니다. 이렇게 하면 정보 손실을 줄이면서도 모델이 유연하게 처리할 수 있게 되죠.
2. Adaptive Capacity 설정: 일부 MoE 모델에서는 데이터의 특성을 분석해 전문가가 다룰 수 있는 용량을 동적으로 조절하기도 합니다. 이렇게 하면 특정 전문가에 너무 많은 토큰이 몰리는 상황을 줄일 수 있습니다.
3. 토큰 우선순위 지정: 토큰의 중요도를 평가해 우선순위를 두고 처리할 수 있는 모델도 있습니다. 이를 통해 더 중요한 토큰을 먼저 할당받게 하여 정보 손실을 줄일 수 있습니다.

따라서, 토큰을 버리는 방식은 정보 손실을 줄일 대안을 함께 설계해야 합니다.

 

Switch Transformer를 통해 1.6조개의 파라미터를 가진 MoE가 빠른 pre-train 속도를 가지게 되었다.

 

전문가의 수가 늘어날수록 샘플 요율성이 향상되고 속도도 빨라지지만, 많은 VRAM과 이점이 점점 줄어듭니다.