MoE là gì? Kiến trúc mô hình thưa cho LLM hiệu quả

Mixture of Experts (MoE) là kiến trúc mô hình học sâu cho phép mạng neural chỉ kích hoạt một phần tham số mỗi lần infer, giúp tăng quy mô mô hình mà không làm tăng chi phí tính toán tuyến tính. Trong khi mô hình dense truyền thống dùng toàn bộ tham số cho mỗi token, MoE dùng router network để chọn top-k experts phù hợp, từ đó cải thiện hiệu năng và tốc độ suy luận.

MoE ra đời từ nghiên cứu năm 1991 của Robert Jacobs, nhưng đến năm 2017 Google mới áp dụng thành công vào transformer với GShard và GLaM. Ngày nay GPT-4, Mixtral 8x7B, DeepSeek-V2 đều dùng MoE với hàng chục tỷ tham số, trong đó chỉ vài tỷ được kích hoạt mỗi lần forward pass.

Nguyên lý hoạt động

Mô hình MoE gồm nhiều expert networks (thường là feed-forward layers độc lập) và một router (gating network). Router nhận đầu vào, tính trọng số cho từng expert, rồi chỉ chọn top-k expert có điểm cao nhất. Các expert chưa được chọn không tính toán gì, tiết kiệm FLOPS đáng kể.

# Ví dụ routing trong MoE layer
def moe_forward(x, experts, gate, top_k=2):
    gate_logits = gate(x)           # [batch, seq, num_experts]
    topk_vals, topk_idx = torch.topk(gate_logits, top_k, dim=-1)
    out = torch.zeros_like(x)
    for i in range(top_k):
        expert_mask = (topk_idx == i)
        out += expert_mask.unsqueeze(-1) * experts[i](x)
    return out

Ưu điểm và thách thức

Ưu điểm chính là khả năng mở rộng quy mô mô hình với chi phí tính toán thấp hơn dense model. Ví dụ Mixtral 8x7B có tổng 46.7B tham số nhưng chỉ dùng ~13B tham số mỗi lần infer, cho chất lượng gần ngang GPT-3.5 với tốc độ nhanh hơn.

Thách thức lớn nhất là load balancing: nếu router cứ chọn cùng vài expert, những expert còn lại bị “chết đói”. Các giải pháp như auxiliary loss, expert capacity factor, và token-dispatch routing giúp cân bằng tải. Ngoài ra, fine-tuning MoE khó hơn dense model do đường dẫn gradient bị phân mảnh qua các expert.

Ứng dụng thực tế

Ngôn ngữ lớn: Mixtral 8x7B, DeepSeek-V2, GPT-4 (rumored)
Đa phương thức: MoE cho vision-language models như Google PaLI
Speech: Conformer-Transducer với MoE layers giảm latency cho ASR
Search: Retrieval-augmented MoE kết hợp knowledge retrieval và generation

MoE đang định hình lại cách chúng ta scale AI: thay vì train một mô hình dense khổng lồ, chúng ta train một mô hình thưa với hàng trăm expert, chỉ dùng vài chục mỗi lần suy luận. Đây là hướng đi tất yếu cho AGI với chi phí hợp lý.

Để tìm hiểu sâu hơn về router design và load balancing, xem bài viết Switch Transformers trên arXiv hoặc giới thiệu MoE từ Hugging Face.