Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Bayesian Mixture of Experts For Large Language Models

Maryam Dialameh, Hossein Rajabzadeh, Weiwei Zhang, Walid Ahmed, Hyock Ju Kwon

TL;DR

The paper addresses unreliable calibration in fine-tuned large language models by introducing Bayesian-MoE, a post-hoc Bayesian framework that applies a structured Laplace approximation to the second linear layer of each MoE expert. By enforcing a block-diagonal, Kronecker-factored curvature and using memory-efficient low-rank approximations, it estimates posterior uncertainty without altering the training process or adding adapters. Empirical results on Qwen1.5-MoE-A2.7B and DeepSeek-MoE-16B show improved calibration (lower ECE) and competitive negative log-likelihood, across in- and out-of-distribution benchmarks, often outperforming deep ensembles and other baselines with only a single checkpoint. The method highlights the value of focusing Bayesian treatment on decisive MoE components, offering scalable, robust uncertainty quantification for large-scale MoE LLMs and guiding future work on modeling expert correlations and extending Bayesian inference to other network parts.

Abstract

We present Bayesian Mixture of Experts (Bayesian-MoE), a post-hoc uncertainty estimation framework for fine-tuned large language models (LLMs) based on Mixture-of-Experts architectures. Our method applies a structured Laplace approximation to the second linear layer of each expert, enabling calibrated uncertainty estimation without modifying the original training procedure or introducing new parameters. Unlike prior approaches, which apply Bayesian inference to added adapter modules, Bayesian-MoE directly targets the expert pathways already present in MoE models, leveraging their modular design for tractable block-wise posterior estimation. We use Kronecker-factored low-rank approximations to model curvature and derive scalable estimates of predictive uncertainty and marginal likelihood. Experiments on common-sense reasoning benchmarks with Qwen1.5-MoE and DeepSeek-MoE demonstrate that Bayesian-MoE improves both expected calibration error (ECE) and negative log-likelihood (NLL) over baselines, confirming its effectiveness for reliable downstream decision-making.

Bayesian Mixture of Experts For Large Language Models

TL;DR

The paper addresses unreliable calibration in fine-tuned large language models by introducing Bayesian-MoE, a post-hoc Bayesian framework that applies a structured Laplace approximation to the second linear layer of each MoE expert. By enforcing a block-diagonal, Kronecker-factored curvature and using memory-efficient low-rank approximations, it estimates posterior uncertainty without altering the training process or adding adapters. Empirical results on Qwen1.5-MoE-A2.7B and DeepSeek-MoE-16B show improved calibration (lower ECE) and competitive negative log-likelihood, across in- and out-of-distribution benchmarks, often outperforming deep ensembles and other baselines with only a single checkpoint. The method highlights the value of focusing Bayesian treatment on decisive MoE components, offering scalable, robust uncertainty quantification for large-scale MoE LLMs and guiding future work on modeling expert correlations and extending Bayesian inference to other network parts.

Abstract

We present Bayesian Mixture of Experts (Bayesian-MoE), a post-hoc uncertainty estimation framework for fine-tuned large language models (LLMs) based on Mixture-of-Experts architectures. Our method applies a structured Laplace approximation to the second linear layer of each expert, enabling calibrated uncertainty estimation without modifying the original training procedure or introducing new parameters. Unlike prior approaches, which apply Bayesian inference to added adapter modules, Bayesian-MoE directly targets the expert pathways already present in MoE models, leveraging their modular design for tractable block-wise posterior estimation. We use Kronecker-factored low-rank approximations to model curvature and derive scalable estimates of predictive uncertainty and marginal likelihood. Experiments on common-sense reasoning benchmarks with Qwen1.5-MoE and DeepSeek-MoE demonstrate that Bayesian-MoE improves both expected calibration error (ECE) and negative log-likelihood (NLL) over baselines, confirming its effectiveness for reliable downstream decision-making.

Paper Structure

This paper contains 25 sections, 32 equations, 1 figure, 8 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Mixture of Experts (MoE)
  4. Laplace Approximation
  5. Methods
  6. Experiments
  7. Experimental Setups
  8. Results
  9. Ablation: Layer-Wise Sensitivity of Bayesian-MoE
  10. Conclusion
  11. Posterior Estimation under Laplace Approximation
  12. Monte Carlo Sampling from the Laplace Predictive Posterior
  13. Algorithms
  14. Efficient Low-Rank Factorization via Randomized SVD
  15. Marginal Likelihood Optimization for Bayesian-MoE
  16. ...and 10 more sections

Figures (1)

  • Figure 1: Layer-wise Ablation of Bayesian-MoE on Qwen1.5-MoE-A2.7B and DeepSeekMoE-16B. Each line shows the performance when excluding one of the four quarters of transformer layers from Bayesian-MoE treatment. The results indicate that earlier layers (Q1 and Q2) contribute more significantly to model calibration and likelihood. Notably, excluding the first quarter leads to the sharpest degradation in both Expected Calibration Error (ECE) and Negative Log-Likelihood (NLL), suggesting that early experts encode more uncertainty-relevant information. These trends are consistent across both model architectures.