ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

• URL: http://arxiv.org/abs/2606.01509v1
• Date: Mon, 01 Jun 2026 00:16:36 GMT
• Title: ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
• Abstract summary: Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token.<n>We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets.
• Score: 32.864092506780004
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token. However, training such models remains challenging because top-$k$ routing is discrete and non-differentiable, requiring gradient estimators for expert selection whose design remains a central open problem. We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets and formulates routing as probabilistic inference in this discrete subset space. We first propose ProbMoE Exact-$k$ routing, which samples $k$-expert subsets in the forward pass, and the backward pass uses gradients through each expert's exact marginal probability as a tractable surrogate for the true gradient. ProbMoE naturally generalizes to a dynamic-$k$ routing setting, where both training and inference constrain the routing cardinality to the same predefined range, allowing adaptive expert allocation per token. Across benchmarks and model backbones, ProbMoE Exact-$k$ achieves strong performance compared to competitive baselines, with improved expert utilization and routing diversity; ProbMoE Dynamic-$k$ achieves comparable performance with fewer activated experts.

Related papers

• Coverage-Maximizing Multinomial Subset Routing under Operational Constraints [6.499454916436427]
  We introduce Multinomial Subset Routing (MSR), a new online routing framework over $K$ experts.<n>At each round, the learner samples $M$ experts i.i.d. from the multinomial policy, and the resulting set of distinct sampled experts forms the routed subset.<n>The reward depends only on the best-performing expert(s) in the routed subset.
  arXiv  Detail & Related papers  (2026-08-17T10:30:20Z)
• Share First, Route What Remains: A Unified Framework for Token-Adaptive MoE Computation [49.8459545291965]
  Mixture-of-experts (MoE) models have recently moved beyond routing a fixed number of complete experts.<n>We study this dependency by decomposing sparsely upcycled feed-forward experts into key-value channels.<n>We instantiate it in UniF-MoE, a unified framework for token-adaptive MoE computation.
  arXiv  Detail & Related papers  (2026-08-11T02:40:58Z)
• Elbow-Based MoE Routing: A Training-Free Inference Time Plugin for Expert Selection [0.0]
  Mixture-of-Experts (MoE) models enable model scaling while maintaining low inference-time compute.<n>We introduce elbow-based routing, a training-free inference-time modification that dynamically adjusts the number of experts on a per-token basis.<n> Experiments on a state-of-the-art MoE model demonstrate an average latency reduction of 5.3% while maintaining accuracy across six benchmarks.
  arXiv  Detail & Related papers  (2026-08-05T03:10:17Z)
• Beyond Routing Weights: Faithful Response-Level Interpretation of Mixture-of-Experts Reward Models via Contribution Contrast [22.088665451247753]
  Mixture-of-Experts (MoE) reward models seek to improve interpretability by routing prompts to specialized experts.<n>We propose $textbfCo$ntribution-$textbfCo$ntrast ($textbfCo$) response-level interpretation.<n>CoCo yields more coherent, faithful, and specialized interpretations than router-based, score-based, and sparse autoencoder-based alternatives.
  arXiv  Detail & Related papers  (2026-07-31T17:43:30Z)
• Selective Ensemble Based on Preference-Directed Multi-Objective Bandits [90.75513823660775]
  We formalize a sequential decision problem under partially specified linear preferences.<n>We then propose a preference-directed upper confidence bound (PrefUCB) algorithm.<n> Experiments on large pre-trained model selective ensemble tasks and online asset allocation under institutional mandates validate our method.
  arXiv  Detail & Related papers  (2026-06-20T07:52:46Z)
• Expert Routing for Communication-Efficient MoE via Finite Expert Banks [3.8673630752805437]
  Resource-efficient machine learning increasingly uses sparse Mixture-of-Experts (MoE) architectures.<n>We treat the gate as a channel and use $I(X;T)$ to quantify the routing information available to the selected expert.<n>The proposed framework provides a practical tool for analyzing resource-aware MoE inference systems.
  arXiv  Detail & Related papers  (2026-05-06T16:00:46Z)
• Path-Constrained Mixture-of-Experts [23.80197713200086]
  MoE routing selects each layer's experts independently, creating NL possible expert paths -- for N experts across L layers.<n>We propose pathmoe, which shares router parameters across consecutive layers.<n>Experiments on 0.9B and 16B parameter models demonstrate consistent improvements on perplexity and downstream tasks over independent routing.
  arXiv  Detail & Related papers  (2026-03-18T21:35:53Z)
• Token-Level LLM Collaboration via FusionRoute [60.72307345997823]
  FusionRoute is a token-level multi-LLM collaboration framework.<n>It selects the most suitable expert at each decoding step and contributes a complementary logit that refines or corrects the selected expert's next-token distribution.<n>It outperforms both sequence- and token-level collaboration, model merging, and direct fine-tuning.
  arXiv  Detail & Related papers  (2026-01-08T16:53:16Z)
• Sparsity-Controllable Dynamic Top-p MoE for Large Foundation Model Pre-training [30.589225478300023]
  DTop-p is a sparsity-controllable dynamic Top-p routing mechanism.<n>We show that DTop-p consistently outperforms both Top-k and fixed-threshold Top-p baselines.<n>DTop-p exhibits strong scaling properties with respect to expert granularity, expert capacity, model size, and dataset size.
  arXiv  Detail & Related papers  (2025-12-16T01:28:57Z)
• RouteMark: A Fingerprint for Intellectual Property Attribution in Routing-based Model Merging [69.2230254959204]
  We propose RouteMark, a framework for IP protection in merged MoE models.<n>Our key insight is that task-specific experts exhibit stable and distinctive routing behaviors under probing inputs.<n>For attribution and tampering detection, we introduce a similarity-based matching algorithm.
  arXiv  Detail & Related papers  (2025-08-03T14:51:58Z)
• Token-Level Prompt Mixture with Parameter-Free Routing for Federated Domain Generalization [51.562474873972086]
  Federated domain generalization (FedDG) aims to learn a globally generalizable model from decentralized clients with heterogeneous data.<n>Recent studies have introduced prompt learning to adapt vision-language models (VLMs) in FedDG by learning a single global prompt.<n>We propose TRIP, a Token-level prompt mixture with parameter-free routing framework for FedDG.
  arXiv  Detail & Related papers  (2025-04-29T11:06:03Z)
• Mixture of Efficient Diffusion Experts Through Automatic Interval and Sub-Network Selection [63.96018203905272]
  We propose to reduce the sampling cost by pruning a pretrained diffusion model into a mixture of efficient experts. We demonstrate the effectiveness of our method, DiffPruning, across several datasets.
  arXiv  Detail & Related papers  (2024-09-23T21:27:26Z)
• Patch-level Routing in Mixture-of-Experts is Provably Sample-efficient for Convolutional Neural Networks [74.68583356645276]
  In deep learning, mixture-of-experts (MoE) activates one or few experts (sub-networks) on a per-sample or per-token basis. We show for the first time that pMoE provably reduces the required number of training samples to achieve desirable generalization.
  arXiv  Detail & Related papers  (2023-06-07T00:16:10Z)
• Soft Merging of Experts with Adaptive Routing [38.962451264172856]
  We introduce Soft Merging of Experts with Adaptive Routing (SMEAR) SMEAR avoids discrete routing by using a single "merged" expert constructed via a weighted average of all of the experts' parameters. We empirically validate that models using SMEAR outperform models that route based on metadata or learn sparse routing through gradient estimation.
  arXiv  Detail & Related papers  (2023-06-06T15:04:31Z)
• On the Representation Collapse of Sparse Mixture of Experts [102.83396489230375]
  Sparse mixture of experts provides larger model capacity while requiring a constant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such a routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse.
  arXiv  Detail & Related papers  (2022-04-20T01:40:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.