---
title: "ProbMoE Fixes the Broken Gradient Problem Holding Back MoE Models"
subtitle: "ProbMoE replaces top-k's broken gradients with probabilistic subset sampling, delivering better expert diversity and adaptive compute per token"
url: "https://alphasignal.ai/news/probmoe-fixes-the-broken-gradient-problem-holding-back-moe-models"
author: "AlphaSignal Newsroom"
author_url: "https://alphasignal.ai/editorial-team"
company: "AlphaSignal"
company_url: "https://alphasignal.ai/company/alphasignal"
published: "2026-06-01T00:00:00.000Z"
updated: "2026-06-09T05:41:54.203Z"
type: "Paper"
topics:
  - "Llms"
  - "Training Infra"
subtopics:
  - "Mixture Of Experts"
  - "Distributed Training"
  - "Scaling Laws"
upvotes: 3074
source_url: "https://arxiv.org/abs/2606.01509"
comments: 0
truncated: true
---

# ProbMoE Fixes the Broken Gradient Problem Holding Back MoE Models

> ProbMoE replaces top-k's broken gradients with probabilistic subset sampling, delivering better expert diversity and adaptive compute per token

Mixture-of-Experts (MoE) models -- the architecture powering DeepSeek, Mixtral, and Qwen -- are built on a clever premise: instead of running every parameter for every token, route each token to a small subset of specialized sub-networks called experts. The efficiency gains are real, but training these models has always carried a hidden tax: the routing step is fundamentally broken from a gradient perspective.

A new paper from researchers introduces [ProbMoE](https://arxiv.org/abs/2606.01509), a probabilistic routing framework that attacks this problem at its root by reframing expert selection as probabilistic inference over discrete subsets -- and the results suggest it is a meaningful step forward for both training quality and compute efficiency.

## The Gradient Problem Nobody Solved

Standard MoE routing works by scoring each expert for a given token, then picking the top-k highest scorers. The argmax operation at the heart of this is not differentiable -- backpropagation requires gradients, and gradients require continuity, but argmax produces a discrete integer index that jumps discontinuously as inputs change. This means the router cannot be trained end-to-end in the usual sense.

The conventional workaround is the **straight-through estimator (STE)** \-- a trick where you use the hard discrete selection in the forward pass but pretend the operation was a softmax during the backward pass. The basic idea is to use two different computations for the forward and backward passes, accepting this inconsistency as a pragmatic compromise that enables training to proceed. It works well enough to train models like Mixtral and DeepSeek, but it introduces bias and can hurt expert specialization.

The downstream effects are real. Experts that lose routing competitions receive no gradient signal and no opportunity to refine their representation of that kind of input, creating strong selective pressure where experts must learn to be distinctively better than all others at some identifiable class of tokens. In practice, this often leads to uneven expert utilization -- some experts handle most tokens, others are nearly idle.

_Preview only. Full text requires an AlphaSignal Pro account._

## Sources

- [ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts](https://arxiv.org/abs/2606.01509)
- [SIMPLE: A Gradient Estimator for k-Subset Sampling](https://arxiv.org/html/2210.01941v2)
- [ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing](https://arxiv.org/abs/2412.14711)
- [Top-K Routing: Expert Selection in Mixture of Experts Models](https://mbrenndoerfer.com/writing/top-k-routing-mixture-of-experts-expert-selection)
- [DirMoE: Dirichlet-routed Mixture of Experts](https://arxiv.org/pdf/2602.09001)