Heng Zhao

First-year Ph.D. Student @ TAKO Lab, UVA

[Heng Zhao @Hayouger]

MoE routing shouldn’t be a hard top-k guess.

ProbMoE turns TopK routing into probabilistic inference over expert subsets, giving routers more informative gradients, more exploration, and better expert utilization. Because it’s probabilistic, it naturally extends to Dynamic-k, learning both which experts to use and how many each token needs.

Moreover, ProbMoE is broadly applicable to TopK-routed MoEs, while keeping training efficient through sparse expert execution.

Sparse compute. Smarter routing. 🚀

[Heng Zhao @Hayouger]

Huge thanks to my amazing collaborators @zileishao @guyvdb @zhezeng0908 ! It’s been a great pleasure working together on ProbMoE.🥰

Full paper: https://t.co/Zv2AsIQ5Lt
Code: https://t.co/OaKbbh7XSA

[Heng Zhao @Hayouger]

MoE routing shouldn’t be a hard top-k guess.

ProbMoE turns TopK routing into probabilistic inference over expert subsets, giving routers more informative gradients, more exploration, and better expert utilization. Because it’s probabilistic, it naturally extends to Dynamic-k, learning both which experts to use and how many each token needs.

Moreover, ProbMoE is broadly applicable to TopK-routed MoEs, while keeping training efficient through sparse expert execution.

Sparse compute. Smarter routing. 🚀

[Heng Zhao @Hayouger]

A probabilistic router might sound expensive — but ProbMoE keeps the sparse execution benefits of TopK MoEs.

Our wall-clock results show that the optimized variant stays close to conventional routing in step time, total training time, and GPU memory, while delivering the benefits of better routing.

better routing without giving up efficiency

[Zilei Shao @zileishao]

Check out our most recent work on dLLM!