Bayesian Mixture-of-Experts: Towards Making LLMs Know What They Don't Know
This work addresses reliability issues in LLMs for AI safety and robustness, though it is an incremental improvement on existing MoE architectures.
The paper tackles the problem of overconfidence and miscalibration in Mixture-of-Experts (MoE) Large Language Models by proposing a Bayesian MoE routing framework, which improves routing stability, in-distribution calibration, and out-of-distribution detection in a 3-billion parameter model.
The Mixture-of-Experts (MoE) architecture has enabled the creation of massive yet efficient Large Language Models (LLMs). However, the standard deterministic routing mechanism presents a significant limitation: its inherent brittleness is a key contributor to model miscalibration and overconfidence, resulting in systems that often do not know what they don't know. This thesis confronts this challenge by proposing a structured \textbf{Bayesian MoE routing framework}. Instead of forcing a single, deterministic expert selection, our approach models a probability distribution over the routing decision itself. We systematically investigate three families of methods that introduce this principled uncertainty at different stages of the routing pipeline: in the \textbf{weight-space}, the \textbf{logit-space}, and the final \textbf{selection-space}. Through a series of controlled experiments on a 3-billion parameter MoE model, we demonstrate that this framework significantly improves routing stability, in-distribution calibration, and out-of-distribution (OoD) detection. The results show that by targeting this core architectural component, we can create a more reliable internal uncertainty signal. This work provides a practical and computationally tractable pathway towards building more robust and self-aware LLMs, taking a crucial step towards making them know what they don't know.