MoE Routing — Top-k gating, và bài toán load balancing

Mixtral 8x7B và DeepSeek-V3 đang làm nên kỳ tích: sở hữu capacity tương đương model hàng trăm tỷ parameters, nhưng mỗi token chỉ "chạm" đến vài tỷ parameters — như một thư viện khổng lồ chỉ mở 2-3 kệ sách mỗi lần tra cứu. Nhưng để làm điều này, họ phải giải một bài toán cực kỳ tinh vi: làm sao để "thủ thư" (router) biết chính xác kệ nào chứa sách cần tìm, mà không để 90% độc giả đổ xô vào cùng một kệ?

Vấn đề

Trong kiến trúc Dense truyền thống, mọi token đều phải "đọc" qua toàn bộ Feed-Forward Network (FFN) — dù token đó là dấu cách, một từ hiếm, hay một đoạn code Python. Điều này lãng phí compute khủng khiếp: bạn đang dùng 70B parameters để xử lý cả những tác vụ chỉ cần 7B.

MoE (Mixture of Experts) giải quyết điều này bằng cách thay thế một FFN duy nhất bằng hàng chục "chuyên gia" (experts) nhỏ hơn. Lý thuyết là: token về toán học đi đến expert toán, token về ngôn ngữ đi đến expert ngôn ngữ. Nhưng thực tế khắc nghiệt hơn nhiều.

Vấn đề "expert collapse": Nếu không can thiệp, router sẽ rơi vào vòng lặp "người giàu càng giàu". Giả sử expert #3 học nhanh hơn một chút về Python, router bắt đầu gửi thêm code đến đó. Expert #3 nhận được nhiều gradient hơn → giỏi hơn → router gửi thêm code. Chỉ sau vài nghìn steps, 90% token đổ về 1-2 experts, còn lại "chết đói" vì không có gradient. Bạn có 64 experts nhưng chỉ dùng 2, biến MoE thành mô hình dense đắt đỏ vô ích.

Ý tưởng cốt lõi

MoE routing là một "công tắc differentiable" — nó học cách chọn expert qua backpropagation, không phải quyết định cứng nhắc.

Thay vì bắt buộc mỗi token chọn 1 expert duy nhất (hard routing), ta dùng Top-k Gating: router là một linear layer nhỏ nhận hidden state, output ra logits qua softmax, rồi chọn k experts có score cao nhất (thường k=2). Output cuối cùng là weighted sum: output = 0.8 × Expert_A + 0.2 × Expert_B.

Đây là phần quan trọng nhất: gradient vẫn chảy về cả hai experts, dù expert B chỉ có trọng số 0.2. Nếu expert A hoàn toàn sai, lỗi sẽ khiến gate score của A giảm và B tăng ở lần sau. Router học cách "tin tưởng" dần dần, không phải all-or-nothing.

Tại sao k=2 là "sweet spot"?
Với k=1 (Switch Transformer), nếu router chọn sai expert, token đó bị xử lý bởi chuyên gia hoàn toàn không phù hợp, không có backup. K=2 tạo ra differentiable ensemble: expert thứ hai đóng vai trò "lưới an toàn". Ngay cả khi chỉ đóng góp 20%, nó vẫn học được từ gradient và giảm thiểu thiệt hại khi router "ngáo". Đây là implicit curriculum learning — router học từ sai lầm mà không làm hỏng hoàn toàn output.

Bài toán load balancing
Để tránh "rich get richer", ta thêm Auxiliary Loss — một "bàn tay quản lý thị trường" ép router phải khám phá các experts ít được dùng. Loss này tính tích của (tỷ lệ token gửi đến expert) và (trung bình score router cho expert đó). Nếu một expert nhận quá nhiều token, loss tăng cao, ép router "điều chỉnh phân bổ". Ngoài ra còn có Z-Loss để phạt logits quá lớn, tránh router trở nên quá tự tin và gây softmax overflow.

That's it. Bạn có một hệ thống mà 47B parameters được chia vào 8 experts, mỗi token chỉ kích hoạt ~12B, nhưng vẫn đảm bảo mọi expert đều được học đều đặn.

Tại sao nó hoạt động

Cơ chế học của router dựa trên gradient flow qua weighted output. Khi backpropagation chạy qua lớp gating, gradient được phân bổ tỷ lệ thuận với gate score. Nếu expert A được chọn với weight 0.8 và expert B với 0.2, lỗi sẽ backpropagate mạnh hơn vào A, nhưng B vẫn nhận được tín hiệu để cải thiện. Điều này tạo ra hiệu ứng "tự sửa lỗi": nếu A thực ra không giỏi về loại token này nhưng được chọn do may mắn, lỗi sẽ nhanh chóng điều chỉnh gate score xuống.

Về mặt toán học, auxiliary loss được tính như sau:

aux_loss = α × Σ(fraction_i × mean_score_i)

Trong đó fraction_i là tỷ lệ token batch được gửi đến expert i, mean_score_i là trung bình router score cho expert đó. Khi một expert chiếm quá nhiều token, tích này tăng, tạo áp lực buộc router phải "giảm giá trị" của expert đó và tăng giá trị các expert khác.

Capacity Factor là một chi tiết tinh tế: thay vì cứng nhắc giới hạn mỗi expert xử lý đúng (batch_size × k / num_experts) tokens, ta thêm buffer 1.0–1.25 để cho phép imbalance tạm thời. Điều này tránh crash khi một batch có quá nhiều token Python đổ về expert code cùng lúc.

Ý nghĩa thực tế

Benchmark thực tế: Mixtral 8x7B đạt chất lượng ngang LLaMA 2 70B nhưng tốc độ inference ngang model 12B. Switch Transformer T5-Base đạt tốc độ pretraining nhanh gấp 4 lần dense baseline với cùng ngân sách compute.

Ai đang dùng: Mistral AI (Mixtral), DeepSeek (V3), Qwen3, Grok, và GPT-4 (được đồn đoán sử dụng kiến trúc tương tự).

Limitations:

• Memory Wall: Bạn vẫn phải load TẤT CẢ experts vào VRAM dù chỉ dùng 2. Mixtral 8x7B cần 90GB VRAM để chỉ tính toán 12B parameters — lãng phí bộ nhớ để đổi lấy tốc độ.
• Fine-tuning Instability: Khi fine-tune trên dataset nhỏ, experts dễ "quên" cách phối hợp với nhau, dẫn đến overfit nhanh hơn dense models.
• Communication Overhead: Trong distributed training, việc gửi token đến đúng GPU chứa expert (all-to-all communication) có thể trở thành bottleneck nếu không dùng InfiniBand/NVLink.

Đào sâu hơn

Paper gốc:

• Switch Transformers — Bài báo gốc từ Google giới thiệu Top-k routing và auxiliary loss cho T5.
• DeepSeekMoE — Cải tiến với shared experts và fine-grained routing để tránh collapse.

Cùng cụm (architecture-moe):

• mixture-of-experts — Tổng quan về kiến trúc MoE, so sánh capacity vs compute.
• decoder-only — Tại sao decoder-only là nền tảng cho mọi MoE hiện đại.
• temperature-sampling — Điều chỉnh ngẫu nhiên trong sampling có liên quan đến exploration trong routing.

Đọc tiếp:

• sequence-modeling — Nền tảng Attention và Transformer mà MoE dựa trên.
• training-efficiency — Các kỹ thuật tối ưu khác như gradient accumulation, pipeline parallelism để train MoE quy mô lớn.
• new-architectures — Các kiến trúc thay thế như Mamba, linear attention cho long-context.