Publications
2026
Breaking the Factorization Barrier in Diffusion Language Models
Ian Li*, Zilei Shao*, Benjie Wang, Rose Yu, Guy Van den Broeck, Anji Liu (* equal contribution)
International Conference on Machine Learning (ICML) 2026
Diffusion language models generate many tokens at once but assume they are independent, which costs coherence. CoDD replaces that factorized output with a lightweight tractable inference layer, keeping quality even at very few decoding steps.
Breaking the Factorization Barrier in Diffusion Language Models
Ian Li*, Zilei Shao*, Benjie Wang, Rose Yu, Guy Van den Broeck, Anji Liu (* equal contribution)
International Conference on Machine Learning (ICML) 2026
Diffusion language models generate many tokens at once but assume they are independent, which costs coherence. CoDD replaces that factorized output with a lightweight tractable inference layer, keeping quality even at very few decoding steps.
ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
Heng Zhao, Zilei Shao, Guy Van den Broeck, Zhe Zeng
International Conference on Machine Learning (ICML) 2026
Top-k expert routing in Mixture-of-Experts models is discrete and non-differentiable. ProbMoE treats expert selection as a distribution over size-constrained expert subsets and performs routing as probabilistic inference, giving tractable gradients and adaptive per-token budgets.
ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
Heng Zhao, Zilei Shao, Guy Van den Broeck, Zhe Zeng
International Conference on Machine Learning (ICML) 2026
Top-k expert routing in Mixture-of-Experts models is discrete and non-differentiable. ProbMoE treats expert selection as a distribution over size-constrained expert subsets and performs routing as probabilistic inference, giving tractable gradients and adaptive per-token budgets.
Rethinking Probabilistic Circuit Parameter Learning
Anji Liu, Zilei Shao, Guy Van den Broeck
International Conference on Artificial Intelligence and Statistics (AISTATS) 2026 Oral
Full-batch EM for probabilistic circuits is too slow on large datasets, and mini-batch variants overfit each batch. anemone is a mini-batch EM algorithm with an implicit per-parameter adaptive learning rate that converges faster and reaches better likelihoods.
Rethinking Probabilistic Circuit Parameter Learning
Anji Liu, Zilei Shao, Guy Van den Broeck
International Conference on Artificial Intelligence and Statistics (AISTATS) 2026 Oral
Full-batch EM for probabilistic circuits is too slow on large datasets, and mini-batch variants overfit each batch. anemone is a mini-batch EM algorithm with an implicit per-parameter adaptive learning rate that converges faster and reaches better likelihoods.
2025
Zero-Variance Gradients for Variational Autoencoders
Zilei Shao, Anji Liu, Guy Van den Broeck
NeurIPS Workshop on Structured Probabilistic Inference & Generative Modeling (SPIGM) 2025 Best Poster Award
Training a VAE normally requires noisy Monte Carlo gradient estimates. This work shows that with a suitably restricted decoder the expected ELBO can be computed analytically, giving gradients with zero estimation variance and better early optimization.
Zero-Variance Gradients for Variational Autoencoders
Zilei Shao, Anji Liu, Guy Van den Broeck
NeurIPS Workshop on Structured Probabilistic Inference & Generative Modeling (SPIGM) 2025 Best Poster Award
Training a VAE normally requires noisy Monte Carlo gradient estimates. This work shows that with a suitably restricted decoder the expected ELBO can be computed analytically, giving gradients with zero estimation variance and better early optimization.
Adversarial Tokenization
Renato L. Geh*, Zilei Shao*, Guy Van den Broeck (* equal contribution)
Annual Meeting of the Association for Computational Linguistics (ACL), Main Conference, 2025
The same string can be tokenized in exponentially many valid ways, and LLMs still understand the non-canonical ones. Simply re-tokenizing a prompt, without changing a single character, is enough to bypass safety alignment in state-of-the-art models.
Adversarial Tokenization
Renato L. Geh*, Zilei Shao*, Guy Van den Broeck (* equal contribution)
Annual Meeting of the Association for Computational Linguistics (ACL), Main Conference, 2025
The same string can be tokenized in exponentially many valid ways, and LLMs still understand the non-canonical ones. Simply re-tokenizing a prompt, without changing a single character, is enough to bypass safety alignment in state-of-the-art models.
2024
Bounding-Box Inference for Error-Aware Model-Based Reinforcement Learning
Erin J. Talvitie, Zilei Shao, Huiying Li, Jinghan Hu, Jacob Boerma, Rory Zhao, Xintong Wang
Reinforcement Learning Journal (RLC), vol. 5, pp. 2440-2460, 2024
In model-based RL, planning with an inaccurate learned model can destroy the policy. This paper propagates uncertainty by reasoning over bounding boxes around sets of possible states, letting the agent plan only where its model can be trusted.
Bounding-Box Inference for Error-Aware Model-Based Reinforcement Learning
Erin J. Talvitie, Zilei Shao, Huiying Li, Jinghan Hu, Jacob Boerma, Rory Zhao, Xintong Wang
Reinforcement Learning Journal (RLC), vol. 5, pp. 2440-2460, 2024
In model-based RL, planning with an inaccurate learned model can destroy the policy. This paper propagates uncertainty by reasoning over bounding boxes around sets of possible states, letting the agent plan only where its model can be trusted.