Nếu bạn theo dõi các mô hình ngôn ngữ lớn gần đây, bạn sẽ thấy cụm từ Mixture of Experts (MoE) xuất hiện khắp nơi. DeepSeek, Mixtral hay Grok đều dùng kiến trúc này. Trước khi MoE trở nên phổ biến, một bài báo năm 2021 của Google Brain đã thay đổi cách giới nghiên cứu nhìn về mô hình thưa (sparse model): Switch Transformer.

Vấn đề của MoE truyền thống

Ý tưởng MoE không mới. Từ những năm 1990, giới nghiên cứu đã biết rằng thay vì dùng một mạng nơ-ron duy nhất xử lý mọi tác vụ, có thể chia việc cho nhiều "expert" chuyên biệt. Mỗi expert là một mạng truyền thẳng nhỏ, còn mạng điều phối sẽ quyết định expert nào xử lý đầu vào nào. Vấn đề là các hệ thống MoE trước đó thường đưa mỗi token tới nhóm expert top-k, thường từ 2 expert trở lên, rồi kết hợp đầu ra theo trọng số. Cách làm này vừa tốn tài nguyên tính toán, vừa làm tăng lượng dữ liệu phải truyền giữa các thiết bị khi các expert nằm trên những GPU hoặc TPU khác nhau.

Top-1 routing: đơn giản đến bất ngờ

William Fedus, Barret Zoph và Noam Shazeer đặt câu hỏi: tại sao phải đưa token tới nhiều expert? Họ đề xuất Switch Routing, tức chỉ chọn một expert cho mỗi token. Nhiều người lo ngại k = 1 sẽ làm tín hiệu lan truyền ngược bị hạn chế, khiến mô hình khó huấn luyện. Nhưng thực tế chứng minh điều ngược lại. Mỗi token chỉ đi qua một expert, giúp giảm đáng kể chi phí tính toán ở mỗi bước. Quan trọng hơn, lượng dữ liệu truyền giữa các thiết bị giảm mạnh vì mỗi token chỉ cần gửi đến một nơi, thay vì phải tỏa ra nhiều nhánh rồi gom lại.

Expert parallelism và capacity factor

Switch Transformer phân bổ các expert trên nhiều TPU, mỗi thiết bị phụ trách một hoặc vài expert. Khi token được đưa tới expert nằm trên thiết bị khác, hệ thống chỉ cần truyền dữ liệu một lần thay vì nhiều lần như cách top-k. Tuy nhiên, nếu quá nhiều token dồn về cùng một expert, thiết bị đó sẽ bị quá tải. Lúc này, capacity factor phát huy tác dụng: nó quy định số token tối đa mỗi expert được xử lý trong một batch. Token vượt giới hạn sẽ đi thẳng qua residual connection mà không qua expert nào. Nhóm tác giả nhận thấy capacity factor từ 1,0 đến 1,25 cho kết quả tốt nhất, đủ để cân bằng giữa việc tận dụng expert và tránh lãng phí bộ nhớ.

Auxiliary loss giữ cân bằng tải

Một vấn đề khác là nếu mạng điều phối "ưu ái" một vài expert, chúng sẽ nhận phần lớn token, trong khi các expert khác gần như không có việc. Switch Transformer xử lý điều này bằng auxiliary load balancing loss, tức một hàm phụ khuyến khích phân bổ token đồng đều giữa các expert. Giá trị nhỏ nhất đạt được khi mọi expert nhận lượng token như nhau. Nếu thiếu cơ chế này, mô hình sẽ nhanh chóng suy biến và chỉ dùng lặp đi lặp lại một vài expert.

Kết quả và di sản

Với kiến trúc đơn giản này, nhóm Google Brain mở rộng mô hình lên 1,6 nghìn tỷ tham số và đạt tốc độ tiền huấn luyện nhanh gấp 7 lần T5 với cùng lượng tài nguyên tính toán. Trên các bộ đánh giá như GLUE, SuperGLUE và SQuAD, Switch Transformer cho kết quả tốt hơn T5 có cùng quy mô. Bài học lớn nhất: đôi khi giải pháp tốt nhất không phải phức tạp hơn mà là đơn giản hóa đúng chỗ. Top-1 routing nghe có vẻ quá đơn giản để hiệu quả, nhưng chính sự đơn giản đó đã mở đường cho thế hệ mô hình MoE phổ biến ngày nay.