DOT-MoE: Biến model dense thành MoE bằng Optimal Transport

Nếu bạn đang follow mảng LLM efficiency thì chắc không lạ gì với câu chuyện muôn thuở: model càng lớn thì performance càng tốt, nhưng inference cost cũng tăng theo kiểu không thể chịu nổi. Mixture of Experts (MoE) sinh ra để giải quyết đúng cái tension này thay vì activate toàn bộ parameters cho mỗi token, bạn chỉ route token đó qua một vài "chuyên gia" nhất định. Tổng số params có thể rất lớn, nhưng active params mỗi lần inference thì nhỏ hơn nhiều.

Vấn đề là train MoE từ đầu cực kỳ unstable và tốn compute. Nên hướng đi thực tế hơn là lấy một model dense đã được train tốt, rồi convert nó sang MoE hay còn gọi là MoEfication. Và đây chính là sân chơi của paper DOT-MoE vừa được publish.

Tại sao các cách cũ không ổn?

Các phương pháp MoEfication hiện tại chủ yếu dựa vào heuristic để chia FFN (Feed-Forward Network) thành các expert. Kiểu như: cluster các neurons có activation pattern giống nhau vào cùng một expert, hoặc đơn giản hơn là... random split. Nghe có vẻ thô, nhưng đó là thực tế.

Vấn đề với heuristic clustering là nó static bạn quyết định một lần rồi thôi, không có feedback loop nào để optimize cái assignment đó tốt hơn. Random split thì khỏi nói, nó nhanh nhưng performance drop khá đáng kể. Cả hai cách đều không tận dụng được gradient signal từ quá trình fine-tuning sau đó.

Mình thấy đây là một điểm rất quan trọng mà nhiều người hay bỏ qua: việc phân chia neurons vào experts và việc học routing policy là hai bài toán liên quan chặt chẽ nhau, nhưng các phương pháp cũ lại giải chúng hoàn toàn tách biệt.

DOT-MoE formulate lại bài toán từ đầu

Thay vì dùng heuristic, DOT-MoE nhìn nhận việc assign neurons vào experts như một Optimal Transport (OT) problem. Cụ thể hơn là balanced transport mỗi expert phải nhận đúng một lượng neurons nhất định, không được để một expert quá tải trong khi expert khác rảnh rỗi.

Ý tưởng core ở đây là: bạn có một "phân phối" neurons ở phía source, và một "phân phối" experts ở phía target. OT tìm cách vận chuyển neurons sang experts sao cho tổng "chi phí" là nhỏ nhất chi phí ở đây được định nghĩa dựa trên sự tương đồng về chức năng của các neurons.

Nhưng bài toán OT thuần túy không differentiable, nên bạn không thể backprop qua nó. DOT-MoE giải quyết điều này bằng Sinkhorn-Knopp iterations một kỹ thuật biến bài toán OT thành một dạng entropy-regularized version có thể differentiate được. Nếu bạn đã từng đọc về Wasserstein distance hay Earth Mover's Distance trong context của generative models thì đây cùng họ hàng.

Cái hay là sau khi có differentiable assignment rồi, vẫn còn một vấn đề: assignment cuối cùng phải là discrete (neuron X vào expert Y, không phải 30% vào Y và 70% vào Z), nhưng discrete operations không có gradient. DOT-MoE dùng Straight-Through Estimator (STE) để handle cái này forward pass dùng discrete assignment, backward pass pretend nó continuous để gradient chạy qua được.

End-to-end learning đây mới là điểm mấu chốt

Kết hợp lại, DOT-MoE cho phép joint learning của hai thứ cùng lúc:

1. Neuron-to-expert assignment: neurons nào vào expert nào
2. Token-to-expert routing: mỗi token được route đến experts nào

Hai thứ này được optimize cùng nhau thông qua gradient descent. Đây là sự khác biệt căn bản so với các phương pháp trước thay vì quyết định assignment trước rồi mới học routing, bạn học cả hai đồng thời và chúng có thể influence lẫn nhau.

Theo kinh nghiệm của mình khi làm với các bài toán optimization phức tạp, cái joint learning này thường cho kết quả tốt hơn hẳn so với giải tuần tự, dù training phức tạp hơn. Và kết quả của paper cũng confirm điều này.

Số liệu thực tế

Giữ được 90% performance của dense model trong khi chỉ dùng 50% active parameters là một con số khá ấn tượng. Để so sánh: nếu bạn prune thẳng model xuống còn 50% params, performance drop thường sẽ tệ hơn nhiều vì bạn mất hẳn những neurons đó. Với MoEfication, bạn vẫn giữ tất cả params, chỉ là không activate hết chúng cùng một lúc.

Practical implications cho anh em

Anh em lưu ý một điểm quan trọng: DOT-MoE không phải là cách để train MoE từ đầu. Nó là một post-training conversion framework bạn cần có một dense model đã được train tốt trước, rồi mới apply DOT-MoE để convert sang sparse MoE.

Flow thực tế sẽ là:

1. Train hoặc lấy một dense pretrained model (ví dụ LLaMA, Mistral, v.v.)
2. Apply DOT-MoE để convert FFN layers thành MoE layers
3. Fine-tune thêm một chút để recover performance
4. Deploy model MoE với inference cost thấp hơn đáng kể

Cái này có nghĩa là nếu bạn đang có một dense model đang chạy production và muốn giảm inference cost mà không muốn train lại từ đầu, đây là một hướng đáng explore.

Mình thấy cái này hay ở chỗ nó tận dụng được toàn bộ "kiến thức" đã được học trong quá trình pretraining của dense model, thay vì throw away và train MoE từ scratch. Pretraining tốn rất nhiều compute, nên việc reuse nó là điều rất có giá trị về mặt thực tế.

Optimal Transport không phải lần đầu xuất hiện trong ML

Nếu bạn chưa quen với Optimal Transport, đây là một framework toán học có lịch sử từ thế kỷ 18 (bài toán Monge gốc), nhưng gần đây được dùng khá nhiều trong ML đặc biệt trong domain adaptation, generative models (Wasserstein GAN), và point cloud matching. Việc DOT-MoE apply nó vào bài toán neuron assignment là một cách dùng khá sáng tạo.

Sinkhorn-Knopp algorithm mà paper dùng về cơ bản là một iterative procedure bạn alternately normalize rows và columns của một matrix cho đến khi nó converge về doubly stochastic matrix (mỗi row và column sum to 1). Cái này enforce được balanced assignment constraint một cách tự nhiên mà không cần hard constraints phức tạp.

Mình không đi sâu vào math ở đây vì sẽ dài, nhưng nếu bạn muốn hiểu rõ hơn thì search "Sinkhorn distance" hoặc "regularized optimal transport" có rất nhiều tài liệu tốt.

Còn những gì chưa biết?

Paper này mới publish đầu tháng 6/2026 nên community chưa có nhiều thời gian reproduce và verify. Một vài câu hỏi mình còn để ngỏ:

• Overhead của Sinkhorn iterations trong quá trình conversion training là bao nhiêu? Nó có đáng để đổi lấy performance gain không?
• Performance retention 90% được đo trên benchmarks nào? Có generalize tốt không khi apply vào tasks cụ thể hơn?
• Scale behavior ra sao khi model size tăng lên 7B, 13B, 70B?

Những câu hỏi này thường sẽ được answer trong các follow-up works hoặc khi community bắt đầu experiment với code (nếu authors release).

Dù sao thì đây vẫn là một hướng tiếp cận thú vị và có cơ sở lý thuyết vững chắc hơn so với heuristic clustering. Với trend hiện tại của cộng đồng là chuyển dần sang sparse architectures để optimize inference, những paper như DOT-MoE sẽ ngày càng relevant hơn đặc biệt khi anh em cần deploy LLMs với budget inference hạn chế.