[NLP]Mixture of Experts (MoE)란?

Mixture-of-Experts

최근 LLM 연구에 의하면, 모델의 성능을 향상시키기 위한 중요한 요소 중 하나는 모델의 크기이다. 주어진 컴퓨팅 자원이 한정된 경우, 작은 모델을 여러 번 학습시키는 것보다 큰 모델을 적은 횟수로 학습하는 것이 더 효율적이다. Mixture of Experts(MoE)는 훨씬 적은 컴퓨팅 자원으로도 사전 훈련이 가능하여, 동일한 자원 내에서 더 큰 모델이나 더 방대한 데이터셋을 학습할 수 있게 해준다. 특히, MoE 모델은 사전 훈련 시 밀집(dense) 모델에 비해 훨씬 빠르게 동일한 성능을 달성할 수 있다.

Concepts of MoE

1) Sparse MoE layers

• 인공 신경망에서 가중치 그래프가 완전 그래프인 Dense Fully-connected Network (FFN)와는 다른 구조로, 트랜스포머의 핵심 구성 요소 중 하나이다.
• Sparse MoE 레이어는 하나의 FFN을 \(N\)개의 Expert로 나누어 사용하는 방식이다. FFN이 나누어져 있기 때문에 각 Expert는 특정 토큰을 처리하게 된다.

(Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity)[https://arxiv.org/abs/2101.03961]

2) Gate network(=Router)

• Gate network 또는 라우터는 각 Expert에게 어떤 토큰을 보낼지 결정하는 모듈이다. 예를 들어, 위 그림에서 “More”라는 단어는 라우터에 의해 2번 Expert에게 할당되고, “Parameter”라는 단어는 1번 Expert에게 할당된다.
• 이렇게 입력된 토큰을 각 Expert에게 분배하는 역할을 하는 것이 바로 Gate network(=라우터)이다. 경우에 따라, 하나의 토큰을 여러 Expert에게 동시에 보낼 수도 있다.
• MoE에서 중요한 결정 중 하나는 토큰을 어떤 Expert에게 할당할지 정하는 것이며, 라우터는 학습된 파라미터로 구성되어 있고, 네트워크의 나머지 부분과 함께 사전 훈련된다.

Sparsity란?

Sparsity는 딥러닝 모델에서 가중치 그래프를 통해 설명할 수 있다. 일반적으로, 트랜스포머는 역전파를 통해 모델의 파라미터를 업데이트하는 과정에서 밀집(dense) 모델로 작동한다. 즉, 다음 레이어의 가중치를 업데이트하기 위해서는 이전 레이어의 모든 가중치 값이 영향을 미친다는 뜻이다.

이를 수학적으로 표현하면, ‘Sparsity’는 조건부 연산의 개념을 활용한다고 볼 수 있다. Dense한 모델에서는 모든 입력에 대해 모든 파라미터 \(W_i\)가 사용되어 \(f(x) = W_ix\)와 같이 계산된다. 하지만 Sparse한 모델에서는 입력에 따라 일부 파라미터만 활성화되므로, \(f(x) = W_{\text{selected}}x\) 형태로 작동하게 된다. 이러한 방식은 계산량을 줄이면서 모델의 크기를 효과적으로 확장할 수 있게 해준다.

Shazeer가 번역 작업에서 Mixture of Experts(MoE)를 적용한 사례가 이러한 조건부 연산의 좋은 예시이다. MoE는 특정 입력에 대해 네트워크의 일부 Expert만 활성화되며, 이를 통해 계산 비용을 최소화하면서 모델의 크기를 확장할 수 있었다. 이러한 방식 덕분에 적게는 수십 개에서 많게는 수천 개의 Expert가 MoE 레이어에서 사용될 수 있었다.

먼저, 총 8개의 Expert가 있다고 가정하고, 하나의 Expert로 토큰을 보내는 문제는 코드로 간단하게 구현할 수 있다. Linear layer + Softmax의 조합을 사용하면 쉽게 어떤 Expert로 보낼지 결정할 수 있으며, 해당 토큰은 선택된 Expert로 간단하게 포워딩하면 된다.

import torch.nn as nn

# inputs는 dim 차원을 가진 tensor
dim, num_experts = 512, 8
gate = nn.Linear(dim, num_experts)
nth_expert = F.softmax(gate(inputs), dim=-1).argmax()

하지만 이러한 구조는 몇 가지 문제를 일으킬 수 있다. 특히, K개의 Expert에게 게이팅(gating)할 때 이러한 문제가 더욱 두드러지게 나타난다. 일반적인 트랜스포머 모델에서는 모든 데이터가 모든 파라미터 업데이트에 관여하기 때문에 배치 크기가 일정하게 유지된다. 그러나 MoE에서는 데이터가 각 Expert에 다르게 할당되므로, 각 Expert의 배치 크기가 달라지고, 고정된 배치 크기보다 작아질 수 있다.

예를 들어, 10개의 토큰이 주어졌을 때, 그중 5개는 하나의 Expert에 할당되고, 나머지 5개는 여러 Expert에 분배될 수 있다. 이 경우, 배치 크기는 (5, 2, 1, 1, 1)과 같이 Expert마다 달라지게 되어, 불균형한 배치 크기와 자원의 비효율적인 사용이 발생할 수 있다. 또한, 각 토큰이 여러 Expert에게 분배되면서 일부 Expert는 과부하 상태가 되고, 다른 Expert는 거의 사용되지 않는 문제가 생겨, 전체적인 자원 활용이 비효율적으로 이루어지고 학습 성능이 저하될 가능성이 있다.

이러한 불균형 문제를 해결하기 위해 Gate network는 Top-K Gating방식을 채택한다. 먼저 학습된 Gating network (\(G\))는 입력의 어느 부분을 각각 어떤 Expert(\(E\))에게 할당할지를 결정한다.

$$y = \displaystyle\sum_{i=1}^n G(x)_iE_i(x)$$

이 방식에서는 모든 Expert가 모든 입력에 대해 실행된다. 그러나 \(G\)가 0이 되는 경우, 모든 Expert가 모든 연산을 처리할 필요가 없어지므로 컴퓨팅 자원이 절약된다. Gating 함수는 앞서 언급한 것처럼 Softmax 함수이다.

$$G_{\sigma}(x) = \text{Softmax}(x \cdot W_g)$$

Top-K Gating 방식은 말 그대로 K개의 값을 계속해서 유지하는 것이다. 이 때, Shazeer는 noise를 더해주는 Noisy Top-K Gating 방식을 사용하였다.

1. 먼저 Gating 함수에 noise를 더해준다.

$$H(x)_i = (w \cdot W_g)_i \;+ \; \text{StandardNormal()} \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i)$$

1. 후 Top-K개를 뽑는다.

$$\text{KeepTopK}(v, k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v, \\ -\infty & \text{otherwise.} \end{cases}$$

1. 마지막으로 Softmax 함수를 적용하면 새로운 Gating 함수가 정의된다.

$$G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k))$$

이 sparsity는 몇 가지 특성을 지닌다. 충분히 낮은 \(k\) (예: 1 또는 2)를 사용하면 더 많은 Expert를 활성화하는 경우보다 훨씬 빠르게 학습하고 추론할 수 있다. 하나 이상의 Expert에 라우팅을 해야 게이트가 다양한 Expert로의 라우팅을 학습할 수 있기 때문에, 최소 두 개의 Expert가 선택되어야 한다.

예를 들어, 하나의 Expert만 선택하는 경우, 매번 같은 Expert로만 라우팅하게 되어 모든 입력 데이터가 특정 Expert 하나로만 전달된다. 이로 인해 다른 Expert는 학습에 참여하지 못하게 되고, 모델은 특정 패턴에만 치우쳐 학습되며, 다양성을 확보하지 못하게 된다. 결국, 모든 Expert의 학습 능력을 충분히 활용하지 못하는 결과를 초래한다.

반면에, 최소 두 개의 Expert를 선택하게 되면 게이트는 다양한 Expert로 라우팅하는 방법을 학습할 수 있다. 예를 들어, Expert A와 B가 있을 때, 일부 입력은 A로, 다른 일부는 B로 라우팅되어 두 Expert가 모두 학습에 참여할 수 있게 된다. 이렇게 라우팅이 다양해지면 모델은 더 다양한 데이터 패턴을 학습할 수 있고, 더 일반화된 성능을 기대할 수 있다.

따라서, 최소 두 개의 Expert가 선택되어야 게이트가 여러 Expert를 효과적으로 활용하고, 각 Expert가 고유한 데이터 패턴에 맞춰 학습할 수 있게 된다.

MoE Implementation (by Mistral-7B)

import dataclasses
from typing import List

import torch
import torch.nn.functional as F
from simple_parsing.helpers import Serializable
from torch import nn

@dataclasses.dataclass
class MoeArgs(Serializable):
    num_experts: int
    num_experts_per_tok: int

class MoeLayer(nn.Module):
    def __init__(self, experts: List[nn.Module], gate: nn.Module, moe_args: MoeArgs):
        super().__init__()
        assert len(experts) > 0
        self.experts = nn.ModuleList(experts)
        self.gate = gate
        self.args = moe_args

    def forward(self, inputs: torch.Tensor):
        # Step 1 : Expert로 보내기 위한 gate linear layer 통과
        gate_logits = self.gate(inputs)
        
        # Step 2 : gate logits에 대해 Top-K개 Expert 뽑기
        weights, selected_experts = torch.topk(gate_logits, self.args.num_experts_per_tok)

        # Step 3 : Top-K개의 experts에 대한 weights 구하기 (by softmax)
        weights = F.softmax(weights, dim=1, dtype=torch.float).to(inputs.dtype)
        results = torch.zeros_like(inputs)

        # N개의 experts 돌면서 순회
        for i, expert in enumerate(self.experts):
            # Step 4 : i_th expert에 해당하는 tokens 뽑기
            batch_idx, nth_expert = torch.where(selected_experts == i)
            
            # Step 5 : i_th expert에 해당하는 token들 i_th expert에 통과
            # Step 6 : 통과된 결과값에 expert weight 반영     
            results[batch_idx] += weights[batch_idx, nth_expert, None] * expert(
                inputs[batch_idx]
            )
        return results

Reference

[1] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
[2] FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models
[3] Blog: Mixture of Experts Explained
[4] Blog: What is MoE?
[5] Github: Mistral

NLP 카테고리 내 다른 글 보러가기