MoE, MoA

|

MoE(Mixtures of Experts)

MoE는 모든 네트워크가 활성화되는 dense model과는 달리 “experts”라고 불리는 여러 specialized subnetwork를 사용하여 입력에 따라 관련된 expert만 활성화시켜 학습과 추론을 떠 빠르고 효율적으로 만든다.

MoE는 Mixtral-8x7B이 공개되며 크게 주목 받았다. Mixtral-8x7B은 MoE를 사용하여 당시 가장 높은 성능을 보인 모델이다. GPT-4도 MoE가 적용되었을 가능성이 있다는 소문이 있는데, 이는 dense model에 비해 OpenAI가 GPT-4를 저렴하게 운영할 수 있게 하기 때문이다. (여러 특수 모델들을 Merge 시키면 되니까 큰 모델을 한번에 학습하지 않아도 됨.)

Core Component

MoE model의 핵심 요소는 아래 두 가지이다:

  • Sparse MoE Layers (Experts) : transformer architecture의 dense feed-forward network layer를 대체하는 layer이다. (dense feed-forward network layer 여러개라고 볼 수 있음.)각 MoE layer는 expert라고 불린다. expert는 특정 유형의 정보를 처리하는데 특화된 독립적인 신경망(모델)이다. 입력에 대해 일부 expert만 활성화된다.
  • Gate Network / Router : Router는 token의 특성과 Router가 학습한 parameter에 따라 어떤 expert에게 token을 보낼지 결정한다. 즉 어떤 token이 어떤 expert에 의해 처리될지, 어떤 expert를 활성화시킬지 결정한다. 입력의 각 부분이 가장 적합한 expert에 의해 처리되도록 한다. 이는 expert를 잘 선택할 수 있도록 학습된 모델이다.

아래 이미지에서 왼쪽은 일반적인 GPT 계열 모델(dense model)의 구조이고, 오른쪽은 GPT 계열 모델 구조에서 feedforward network가 Router + Exterts로 대체된 구조이다. 오른쪽 이미지를 보면 현재 time step에서의 token은 Router를 통해 Expert2와 Expert4이 활성화되어 이를 통과하고 있다.

GPT 계열의 모델들은 모두 Transformers의 Decoder 구조를 기반으로 한다. 아래 그림에서는 layer normalization 수행 시점이 Decoder 구조와는 다르게 각 block 앞에 위치해 있다. 이는 GPT3부터 변경된 구조로, 최근 나오는 많은 모델들이 이 구조를 따르고 있다.

GPT3 : Language Models are Few-Shot Learners

Illustrated by the author


하지만 위와 같이 복잡해진 구조로 인해 두 가지 문제점이 발생한다.

  • MoE를 통해 모델을 병합한 후의 fine-tuning이 어렵다. tuning하는 동안 expert usage의 균형을 맞추어 적절히 학습시켜야 하는데 이것이 쉬운 일이 아니다.
  • 추론 시 많은 Memory가 요구된다. 각 expert들이 포함된 모델을 memory에 올려야 하기 때문에 VRAM 용량이 많이 필요하다.

Essential Parameter

  • Number of experts:
    • huggingface/transformers의 mixtral config 기준 num_local_experts
    • 해당 parameter는 expert 수를 결정하는 parameter이다. 예를 들어 Mixtral은 8개의 expert를 가지고 있다.
    • expert가 많을수록 모델 크기가 커져 memory 요구량이 많아진다.
  • Number of experts/token
    • huggingface/transformers의 mixtral config 기준 num_experts_per_tok
    • 이는 하나의 token에 대해 활성화되는 expert 수이다. 예를 들어 Mixtral은 2개이다.
    • 많은 expert를 선택하면 정확도는 올라가겠지만 연산 속도가 느려질 수 있다. 따라서 빠른 추론과 좋은 성능을 충족시키는 적절한 expert 수 선택이 중요하다.

Routing techniques

  1. Top-k Routing
    • 입력 token에 대해 가장 높은 affinity(친화성) score를 가진 상위 $k$개의 expert를 선택하는 방법이다.
    • token-choice routing이다.
    • affinity(친화성) score는 일반적으로 Softmax를 통해 산출한다.
    • 이때 선택된 $k$개의 expert 이외의 나머지 expert들은 음의 무한대로 치환하여 무시한다.
  2. Expert Choice Routing
    • token이 expert를 선택하는 방법이 아닌 expert가 스스로 가장 잘 처리할 수 있는 token을 선택하는 방법이다.
  3. Sparse Top-k Routing
    • 각 token에 대해 일부 expert만 활성화되어 network에 sparsity를 만드는 방법이다. Sparse routing은 모든 expert가 각 token에 대해 활성화 되는 dense routing에 비해 계산량이 적다.
    • token-choice routing 방식으로 작동된다.
    • Mixtral-8x7B는 Sparse Routing 방식을 사용하는데, 이 방법론은 base model에서 layer normalization layer와 self-attention layer의 weight을 복사한 후 각 expert의 FFNN layer weight를 복사하여 사용한다. 이는 FFNN을 제외한 모든 parameter가 공유된다는 뜻이다. 그래서 Mixtral-8x7B가 $8*7B = 56B$ parameter가 아닌 46B parameter를 가지는 이유이다. 그리고 Mixtral-8x7B의 num_experts_per_tok는 2개인데 한 token당 추론에 소요되는 속도(FLOPs)가 dense model 기준 $2$x$7B = 14B$ 정도가 아닌 dense model 12B 정도의 속도가 나는 이유이기도 하다.

아래는 토큰별로 어떤 expert를 통과했는지 시각적으로 나타내는 이미지이다.

Mixtral of Experts


Router initialization(Mergkit)

original MoE 방법은 expert와 router를 함께 학습하지만 최근에는 병합되는 기존 모델들을 upcycling하여 나중에 router를 초기화하는 방식으로 병합이 이루어진다.

  1. Random weight를 랜덤하게 초기화. 이 방법은 동일한 expert가 계속 선택될 가능성이 있어 사용 시 이 점을 주의해야 한다.
  2. Cheap embed 입력 token의 raw embedding을 direct하게 사용하고 모든 layer에 동일한 transformation을 적용한다. 이 방법은 계산량이 적어 낮은 하드웨어 스펙에서도 잘 작동한다는 장점이 있다.
  3. Hidden positive와 negative prompt list의 hidden representation을 생성하여 LLM의 마지막 layer에서 추출한 후 이를 평균화-정규화를 하여 gate를 초기화한다. 이 방법은 token을 관련성 높은 expert에 올바르게 routing하기에 가장 적합하고 효율적이다.

Mixtral Moe

아래는 Mixtral-8x7B 구조에서 MoE layer와 관련된 class들이다.

class FeedForward(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.w1 = nn.Linear(args.dim, args.hidden_dim, bias=False)
        self.w2 = nn.Linear(args.hidden_dim, args.dim, bias=False)
        self.w3 = nn.Linear(args.dim, args.hidden_dim, bias=False)

    def forward(self, x) -> torch.Tensor:
        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x)) # 시그모이드

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):
        inputs_squashed = inputs.view(-1, inputs.shape[-1]) # 입력된 tensor 크기를 input tensor의 마지막 차원으로 통일
        gate_logits = self.gate(inputs_squashed) # gate linear layer 통과
        weights, selected_experts = torch.topk(
            gate_logits, self.args.num_experts_per_tok
        ) # gate_logits에서 top-K개의 expert 뽑기
        weights = nn.functional.softmax(
            weights,
            dim=1,
            dtype=torch.float,
        ).type_as(inputs) # top-K개의 expert에 대한 weight를 softmax에 넣어 score 구하기
        results = torch.zeros_like(inputs_squashed) # 나중에 결과 담을 inputs_squashed tensor와 같은 크기의 초기화된 tensor 생성
        for i, expert in enumerate(self.experts):
            batch_idx, nth_expert = torch.where(selected_experts == i) # selected_experts에서 값이 i인 요소들의 index를 반환.  
            results[batch_idx] += weights[batch_idx, nth_expert, None] * expert(
                inputs_squashed[batch_idx]
            ) # 결과 반영
        return results.view_as(inputs)


class TransformerBlock(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.attention = Attention(args)
        self.feed_forward = MoeLayer(
            experts=[FeedForward(args=args) for _ in range(args.moe.num_experts)],
            gate=nn.Linear(args.dim, args.moe.num_experts, bias=False),
            moe_args=args.moe,
        )
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.args = args

    def forward(
        self,
        x: torch.Tensor,
        freqs_cis: torch.Tensor,
        positions: torch.Tensor,
        mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        r = self.attention.forward(self.attention_norm(x), freqs_cis, positions, mask)
        h = x + r
        r = self.feed_forward.forward(self.ffn_norm(h))
        out = h + r
        return out

위 코드에서 FeedForward는 아래 Mixtral strcture의 block_sparse_moe를 보면 이해가 쉽다.

MixtralForCausalLM(
  (model): MixtralModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0-31): 32 x MixtralDecoderLayer(
        (self_attn): MixtralAttention(
          (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): MixtralRotaryEmbedding()
        )
        (block_sparse_moe): MixtralSparseMoeBlock(
          (gate): Linear4bit(in_features=4096, out_features=8, bias=False)
          (experts): ModuleList(
            (0-7): 8 x MixtralBLockSparseTop2MLP(
              (w1): Linear4bit(in_features=4096, out_features=14336, bias=False)
              (w2): Linear4bit(in_features=14336, out_features=4096, bias=False)
              (w3): Linear4bit(in_features=4096, out_features=14336, bias=False)
              (act_fn): SiLU()
            )
          )
        )
        (input_layernorm): MixtralRMSNorm()
        (post_attention_layernorm): MixtralRMSNorm()
      )
    )
    (norm): MixtralRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

Dense-MoE Hybrid Transformer

최근 open_llm_leaderboard를 보면 snowflake model들이 종종 보이는데 이는 Dense-MoE Hybrid 구조를 가진다.

https://www.snowflake.com/blog/arctic-open-efficient-foundation-language-models-snowflake/


MoA(Mixture-of-Agents)

“집단 지성”으로부터 착안된 방법론으로 볼 수 있다. 단일 LLM만 사용하는 것이 아니라 다수의 LLM이 팀을 이루어 각 LLM의 집단적 전문성을 활용하여 여러 LLM이 힘을 모아 응답하는 방법으로, 2024년 06월 Mixture-of-Agents Enhances Large Language Model Capabilities을 통해 제안되었다.

Mixture-of-Agents Enhances Large Language Model Capabilities


위 이미지는 MoA 작동 구조이다. MoA는 여러 layer로 구성되어 있다. 각 MoA layer에는 여러 LLM들이 있는데 위 그림에서 ‘A’들이 다 LLM이다. 각 LLM은 MoA layer에 입력된 prompt를 입력 받아 response를 생성한다. MoA layer 1의 LLM들의 응답은 layer 2에 전달된다. layer 2에도 LLM들이 있는데 layer 1의 LLM과 동일할 수도 아닐 수도 있다. 이런식으로 MoA layer를 다 거친 후 final layer에 들어가는데 final layer에는 하나의 LLM이 있다. final layer의 LLM은 input prompt + 이전 layer들에서 수집된 response를 입력 받아 최종 response를 만들어낸다. final layer에서 사용되는 prompt를 해당 논문에서는 “Aggregate-and-Synthesize prompt”라고 부르는데, 각 layer들의 응답들을 삽입하고 모델에게 이를 평가하여 최종 응답을 도출하도록 유도하는 prompt이다.

Mixture-of-Agents Enhances Large Language Model Capabilities


MoE에서는 각 모델들의 일부를 병합하지만 MoA에서는 각 모델 통으로 사용한다.

Reference

MoE

Mixtral of Experts

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

A REVIEW OF SPARSE EXPERT MODELS IN DEEP LEARNING

https://towardsdatascience.com/create-mixtures-of-experts-with-mergekit-11b318c99562

MoA

Mixture-of-Agents Enhances Large Language Model Capabilities

https://github.com/togethercomputer/MoA