[LLM]Mixture of Experts(MoE)

LLM은 파라미터의 개수가 많으면 많을수록 뛰어난 성능을 가진다는 것을 여러 연구를 통해 증명함에 따라, LLM을 개발하는 Meta, Google, OpenAI 등의 기업들은 경쟁적으로 더 큰 모델을 출시하고 있다. 그러나, 모델이 커짐에 따라 이를 학습시키고 서빙하기 위한 컴퓨팅 자원 비용 역시 같이 증가하고 있다. 이런 상황에서 더 거대한 LLM을 더 적은 비용으로 학습하고, 서빙하기 위해서 나온 기술이 바로 Mixture of Experts 즉 MoE이다.

MoE를 채택한 모델은 대표적으로 Mistral AI의 Mixtral-8x7B가 있다. OpenAI의 GPT도 MoE 구조를 사용한다고 추측되지만, 정확한 구조는 공개되지 않았다. 이러한 MoE 모델들은 전체 파라미터 중 일부만을 사용하여 효율적으로 추론한다.

 

MoE를 구성하는 가장 기본적인 개념은 다음과 같다:

• 전문가 네트워크(Experts): MoE 모델은 여러 개의 독립적인 신경망으로 이루어져 있습니다. 이 각각의 신경망을 "전문가(Expert)"라고 부릅니다. 일반적으로 수백에서 수천 개의 전문가 네트워크가 포함됩니다.
• 게이트 네트워크(Router): 주어진 입력에 대해 어떤 전문가가 활성화될지를 결정하는 네트워크입니다. 게이트 네트워크는 입력을 받아 각 전문가의 출력을 가중치로 조합하여 최종 출력을 만듭니다. 보통 입력 데이터마다 몇 개의 전문가만 활성화되며, 나머지는 비활성화됩니다.

기존의 LLM의 Transformer는 Self-Attention -> Add & Norm -> FFN -> Add & Norm으로 이루어져 있다. 그러나 MoE는 Add & Norm -> FFN -> Add & Norm 부분에서 차이가 있는데, 기존 LLM과는 다르게 여러 개의 FFN을 가진다. 첫 번째 토큰이 들어와서 Self-Attention -> Add & Norm을 거친 이후 router가 지정해준 FFN을 거치고, 다음으로 들어온 토큰은 첫 번째 토큰과는 다른 FFN을 통과할 수 있다. 이런 식으로 router에 의해 여러 FFN 중 특정 FFN만 지나게 되는데, 1개 혹은 N개 이상을 지날 수 있다. 설정 값에 따라 Router가 지정한 top-k FFN을 설정하게 되면 해당 FFN을 독립적으로 통과한 이후 나온 출력을 가중치 합 혹은 여타 방법론을 활용하여 결합하여 최종 출력을 만든다. 여기서 각 FFN의 출력에 가중치를 부여하는 방식은 Router에서 계산된 확률을 바탕으로 할 수 있다.

MoE의 종류는 다음과 같다:

• Standard MoE: 모든 전문가의 출력을 결합하여 사용하며, 계산 비용이 높지만, 모든 전문가의 지식을 활용할 수 있습니다.
• Switch Transformer: 각 입력에 대해 하나의 전문가만 선택하며 매우 효율적이나, 정보 손실의 위험이 있습니다.
• GShard: Google에서 개발한 방식으로, 큰 규모의 모델에 MoE를 적용하며 다양한 작업에 대해 확장성이 뛰어납니다.
• BASE Layers: 전문가 간의 균형을 유지하면서 학습하는 방식이며, 전문가들이 고르게 활용되도록 합니다.
• HashMoE: 해시 함수를 사용하여 전문가를 선택하며, 계산 효율성이 높으나 유연성이 떨어질 수 있습니다.
• Soft MoE: 모든 전문가의 출력을 부드럽게 결합하며, 표준 MoE와 Switch Transformer의 중간 형태라고 볼 수 있습니다.
• ST-MoE (Sparse Top-k MoE): 입력에 따라 다수의 Experts 중 상위 k개를 선택적으로 활용하여 효율성과 성능을 높이는 대규모 언어 모델 기술입니다.

이 외에도 Expert Choice Routing, Mixture of Experts with Expert Choice (MoEC) 등의 방식이 있으며, 현재는 ST-MoE 방식이 널리 사용되고 있습니다.

대표적인 MoE 모델인 Mixtral-8x7B의 사용방법은 아래의 예시와 같다.

from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from mistral_common.protocol.instruct.messages import UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest

# Tokenization

mistral_models_path = "MISTRAL_MODELS_PATH"
 
tokenizer = MistralTokenizer.v1()
 
completion_request = ChatCompletionRequest(messages=[UserMessage(content="Explain Machine Learning to me in a nutshell.")])
 
tokens = tokenizer.encode_chat_completion(completion_request).tokens

# Inference by mistral_inference

from mistral_inference.transformer import Transformer
from mistral_inference.generate import generate

model = Transformer.from_folder(mistral_models_path)
out_tokens, _ = generate([tokens], model, max_tokens=64, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)

result = tokenizer.decode(out_tokens[0])

print(result)

# Inference with huggingface transformers

from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
model.to("cuda")
 
generated_ids = model.generate(tokens, max_new_tokens=1000, do_sample=True)

# decode with mistral tokenizer
result = tokenizer.decode(generated_ids[0].tolist())
print(result)