κ-SwiGLU: Khi MoE Router Tự Biết Mình Đang Confident Đến Đâu

Nếu bạn đang làm việc với các model kiến trúc Mixture-of-Experts (MoE), chắc bạn đã quen với cái flow cơ bản: router nhận input, tính logits, chọn top-k experts, rồi mỗi expert xử lý token theo phần của mình. Nghe đơn giản, nhưng cái phần "gate" bên trong từng expert cụ thể là SwiGLU activation lại đang bị bỏ qua một thông tin cực kỳ có giá trị: router đang confident đến mức nào khi routing cái token đó.

Đây chính xác là vấn đề mà paper κ-SwiGLU (Confidence-Adaptive SwiGLU) giải quyết.

SwiGLU Trong MoE Đang Bỏ Phí Thông Tin Gì?

SwiGLU là activation function phổ biến trong các LLM hiện đại LLaMA, Mistral, Mixtral đều dùng. Công thức cơ bản:

1SwiGLU(x, W, V) = SiLU(xW) ⊙ (xV)

Trong MoE, mỗi expert có SwiGLU gate riêng của nó. Vấn đề là cái gate này hoàn toàn không biết gì về context routing nó không biết router đang "chắc chắn" hay "phân vân" khi gửi token này đến mình.

Hãy nghĩ thế này: khi router có distribution logits rất peaked (confident cao, một expert rõ ràng là winner), token đó có khả năng cao là "đúng chuyên môn" của expert được chọn. Ngược lại, khi logits flat (router phân vân giữa nhiều experts), token đó có thể là borderline case, expert nào cũng xử lý được một phần.

Theo kinh nghiệm của mình khi debug các MoE model, cái "confidence" của router này thực ra chứa rất nhiều signal về tính chất của token nhưng trước giờ nó chỉ được dùng để chọn expert, sau đó bị throw away hoàn toàn.

κ-SwiGLU Làm Gì Khác?

Ý tưởng của paper khá elegant. Họ introduce một scalar κ (kappa) được derive từ router logits để điều chỉnh độ "sharpness" của SwiGLU gate bên trong từng expert.

Cụ thể hơn, κ được tính từ confidence score của router (ví dụ: entropy của softmax distribution, hoặc gap giữa top-1 và top-2 logit). Khi router confident cao → κ lớn → gate sắc nét hơn, expert "chuyên tâm" hơn vào đặc trưng của nó. Khi router phân vân → κ nhỏ → gate mềm hơn, expert linh hoạt hơn.

Mình thấy cái này hay ở chỗ nó tạo ra một feedback loop implicit giữa routing decision và expert computation mà không cần thêm bất kỳ component nào phức tạp. Không có auxiliary loss mới, không có module riêng, chỉ là tận dụng thông tin đã có sẵn.

So Sánh Với Các Hướng Tiếp Cận Khác

| Approach | Thay đổi Architecture | Thêm Params | Overhead | Cải thiện
|---|---|---|---|---|
| Standard SwiGLU trong MoE | Không | 0 | 0 | Baseline |
| Expert-specific gating | Có | Đáng kể | Trung bình | Varies |
| Auxiliary routing losses | Không | 0 | Training only | Load balancing |
| κ-SwiGLU | Minimal | Negligible | Nhỏ | Pretraining perplexity ↓ |
| Binary Expert Masking (BEAM) | Có | Nhỏ | Trung bình | Efficiency |

Nhìn vào bảng này thì rõ ràng κ-SwiGLU đang chơi ở một góc khá khác biệt nó không phải efficiency optimization, không phải load balancing fix, mà thuần túy là cải thiện quality của expert computation với cost thêm vào gần như bằng 0.

Tại Sao "Negligible Parameters" Lại Quan Trọng?

Anh em làm ML infra chắc hiểu cái pain khi một paper propose cải tiến nhưng kèm theo +10-20% parameter count. Với các MoE model đang scale lên hàng trăm tỷ params, mỗi phần trăm tăng thêm là một vấn đề thực sự về memory, về checkpoint size, về serving cost.

κ-SwiGLU chỉ thêm một scalar per expert (hoặc thậm chí shared across experts trong một số config), nên về mặt thực tế nó gần như free. Đây là điểm mình đánh giá cao nhất của paper không phải cứ thêm complexity là tốt hơn, đôi khi cái insight đúng chỗ mới là thứ tạo ra impact.

Context Rộng Hơn: MoE Research Đang Đi Về Đâu?

Nhìn vào cluster các paper liên quan được recommend cùng thời điểm này, có thể thấy rõ một pattern:

• BEAM (Binary Expert Activation Masking) dynamic routing với binary mask
• φ-Balancing cân bằng load training cho MoE
• DOT-MoE dùng optimal transport cho MoEfication
• Post-Trained MoE Skip Half Experts pruning experts sau training
• Geometric Coupling in Sparse MoE routers học geometry của experts

Cộng đồng đang tấn công MoE từ nhiều góc: routing algorithm, training stability, inference efficiency, và giờ thêm κ-SwiGLU là góc expert computation quality. Điều thú vị là các hướng này không loại trừ nhau về lý thuyết bạn có thể stack κ-SwiGLU lên trên một routing scheme tốt hơn và thu được benefit từ cả hai.

Mình đặc biệt tò mò về interaction giữa κ-SwiGLU và các load balancing technique. Nếu router được train để balance hơn (logits flat hơn theo thiết kế), κ sẽ luôn nhỏ liệu đó có phải trade-off không? Paper chưa explore nhiều về chỗ này và đây có thể là một hướng follow-up thú vị.

Practical Takeaway Cho Ai Đang Build MoE

Nếu bạn đang pretraining một MoE model từ đầu đây là loại change mà bạn nên thử ngay vì risk gần như bằng 0. Không cần refactor architecture, không cần điều chỉnh training recipe phức tạp, chỉ cần modify cái SwiGLU implementation để nhận thêm router confidence signal.

Nếu bạn đang fine-tune một MoE pretrained sẵn (như Mixtral hay các open-source MoE khác) câu chuyện phức tạp hơn vì pretrained weights đã được train với standard SwiGLU, inject κ vào lúc này có thể cần careful initialization để không phá vỡ learned representations.

Anh em lưu ý: paper này focus vào pretraining performance, nên kết quả trên fine-tuning hay downstream tasks cụ thể vẫn cần được validate thêm. Đây là điều bình thường với research papers claim được support bởi perplexity improvement trong pretraining, còn generalization ra ngoài thì cần thêm work.

Nhìn chung, κ-SwiGLU là một ví dụ đẹp của "elegant engineering" trong ML research tìm ra thông tin đang bị bỏ phí trong system hiện tại và tận dụng nó theo cách minimal nhất có thể. Không phải mọi cải tiến đều cần một module mới hoành tráng.