LLM & NLP

Mixture of Experts (MoE): model khổng lồ mà mỗi token chỉ dùng một phần

Tóm tắt

MoE chia model thành nhiều expert, mỗi token chỉ đi qua vài expert, nên dung lượng lớn mà tính toán thấp. Tốc độ theo tham số hoạt động nhưng bộ nhớ theo tổng tham số, điểm hay bị hiểu sai nhất khi tự host.

Một model có hàng trăm tỷ tham số nhưng mỗi token chỉ "dùng" vài chục tỷ, nghe như mâu thuẫn. Đó chính là ý tưởng của Mixture of Experts (MoE): thay vì mọi tham số cùng tham gia xử lý mỗi token, model có nhiều "chuyên gia" và một bộ định tuyến chọn vài chuyên gia phù hợp cho từng token. Nhờ vậy model có dung lượng lớn mà chi phí tính toán mỗi token vẫn thấp. Nhiều LLM hiện đại dùng kiến trúc này.

Cơ chế

  • Trong mỗi lớp Transformer, khối feed-forward được thay bằng nhiều expert (mỗi expert là một mạng feed-forward).
  • Một router (gating network) chấm điểm các expert cho từng token và chọn top-k.
  • Output là tổng có trọng số của các expert được chọn; các expert còn lại không tính toán gì cho token đó.
token → router → [E1: 0.6] [E5: 0.4]  (các expert khác bỏ qua)
output = 0.6·E1(token) + 0.4·E5(token)

Ba cột mốc

Nghiên cứuĐóng góp
Sparsely-Gated MoE (Shazeer và cộng sự, 2017)Lớp MoE với hàng nghìn expert, kích hoạt thưa; tăng dung lượng model hơn 1000 lần với chi phí tính toán tăng ít
Switch Transformer (Fedus và cộng sự, 2021)Đơn giản hoá: mỗi token chỉ tới 1 expert; mở rộng tới hàng nghìn tỷ tham số, tăng tốc tiền huấn luyện đáng kể so với model dày cùng compute
Mixtral 8x7B (Mistral AI, 2024)Model mở: 8 expert mỗi lớp, chọn 2 cho mỗi token; tổng 47 tỷ tham số nhưng chỉ 13 tỷ hoạt động mỗi token, ngang hoặc vượt Llama 2 70B trên nhiều benchmark

Dense và MoE

DenseMoE
Tham số dùng mỗi tokenTất cảMột phần nhỏ
Tốc độ suy luậnTỷ lệ với tổng tham sốTỷ lệ với tham số hoạt động
Bộ nhớ GPUChứa toàn bộVẫn phải chứa toàn bộ expert
Độ khó huấn luyệnỔn định hơnPhải cân bằng tải giữa các expert

Rủi ro và lưu ý

  • "13 tỷ hoạt động" không có nghĩa chạy được trên máy yếu: toàn bộ 47 tỷ tham số vẫn phải nằm trong bộ nhớ. Đây là hiểu lầm phổ biến nhất khi tự host.
  • Cân bằng tải: nếu router dồn token vào vài expert, các expert khác lãng phí và hệ thống chậm; huấn luyện cần thêm loss phụ để cân bằng.
  • Expert không chuyên theo chủ đề như tên gọi: nghiên cứu về Mixtral cho thấy việc định tuyến thiên về cú pháp nhiều hơn là theo lĩnh vực.
  • Fine-tune phức tạp hơn: MoE có thể kém ổn định khi tinh chỉnh so với model dày.

Áp dụng thế nào?

  • Dev: Khi chọn model mở để tự host, tính VRAM theo tổng tham số, còn ước lượng tốc độ theo tham số hoạt động.
  • Tech Lead: MoE hợp với phục vụ lưu lượng cao trên cụm GPU lớn; với một GPU đơn lẻ, model dày nhỏ thường thực tế hơn.
  • PM: Khi nhà cung cấp quảng bá "X tỷ tham số", hỏi rõ tổng hay hoạt động; con số này ảnh hưởng trực tiếp tới chi phí hạ tầng.

Nguồn & tham khảo

  1. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer — Shazeer et al. (arXiv)
  2. Switch Transformers — Fedus et al. (arXiv)
  3. Mixtral of Experts — Mistral AI (arXiv)