Routing Entropy: A Hidden Self-Verifier for Free in Mixture-of-Experts LLMs [COLM 2026]

Paper: OpenReview

TL;DR: Routing Entropy is a free, training-free uncertainty signal inside MoE LLMs. It is more reliable than output confidence for self-verification, works best when combined with output confidence, and can regularize router post-training to improve uncertainty calibration without sacrificing generation quality.

Installation

Python 3.10+ and a recent CUDA-capable PyTorch environment are recommended for the paper models.

git clone https://github.com/tianyi-lab/routing-entropy.git
cd routing-entropy
python -m venv .venv
source .venv/bin/activate
pip install -e '.[all]'

The tested dependency range is declared in pyproject.toml. Large-model quantization and distributed execution are intentionally left to the user's Transformers/Accelerate setup.

Supported model interface

The scorer discovers modules named gate or router with a [num_experts, hidden_size] weight matrix. It captures the raw linear projection in a forward pre-hook, so it does not depend on generate() returning router logits. This covers the common Hugging Face implementations of:

• allenai/OLMoE-1B-7B-0924;
• Qwen/Qwen3-30B-A3B and instruct variants;
• DeepSeek-style remote-code models whose MoE gate exposes a two-dimensional weight.

Always inspect discovery before a large run:

routing-entropy inspect-routers \
  --model allenai/OLMoE-1B-7B-0924

For a custom implementation, pass a regular expression with --router-pattern. The matched modules must expose a linear router weight; gate_proj MLP projections should not be matched.

1. Extract OC and RE

Optionally download and format one of the six public evaluation benchmarks with the provided example template:

routing-entropy prepare-benchmark \
  --benchmark arc_challenge \
  --output data/arc_challenge.jsonl

Valid names are mmlu, hellaswag, piqa, arc_challenge, arc_easy, and winogrande. Dataset revisions should be pinned by the experiment runner when exact archival reproduction is required. The generated prompts ask for only an option label; this transparent template is not a claim about an unspecified private prompt used for the paper.

Create a JSONL input file:

{"id":"example-1","prompt":"Question: 2 + 2 = ?\nAnswer:","reference":"4"}
{"id":"example-2","prompt":"Write one sentence about routing."}

Run deterministic generation and full-sequence feature extraction:

routing-entropy score \
  --model allenai/OLMoE-1B-7B-0924 \
  --input examples/prompts.jsonl \
  --output runs/olmoe_features.jsonl \
  --token-window 5 \
  --deep-layer-fraction 0.5 \
  --max-new-tokens 64

Each output row contains the prompt, generated response, OC, output surprisal, router confidence, RE, and an optional correctness label. The built-in exact-match labeler is only a safe default. For MMLU, HellaSwag, PIQA, ARC, and WinoGrande, supply benchmark-specific answer parsing and write the resulting Boolean to is_correct before fitting the verifier.

Alignment convention

For a generated token at position t, OC is its probability under causal-LM logits at position t-1. RE is computed from router decisions while processing the generated token at position t. Generation is completed first; a single full-sequence forward pass then captures every router layer with consistent shapes.

2. Fit OC + RE fusion

Input rows must contain output_confidence, routing_entropy, and Boolean is_correct.

routing-entropy fit-fusion \
  --features runs/train_features.jsonl \
  --output runs/oc_re_xgboost.joblib \
  --folds 5

The command uses the paper's selected configuration:

Parameter	Value
max_depth	3
n_estimators	100
learning_rate	0.05
subsample	0.8
min_child_weight	5
colsample_bytree	1.0

Metrics are calculated from out-of-fold predictions and written beside the fitted model.

3. RE-regularized router post-training

Prepare JSONL rows with a text field containing the complete supervised sequence:

{"text":"Question: 2 + 2 = ?\nAnswer: 4"}

Then run:

routing-entropy train-router \
  --model allenai/OLMoE-1B-7B-0924 \
  --dataset examples/train.jsonl \
  --output-dir runs/olmoe_router \
  --re-lambda 0.05 \
  --learning-rate 2e-5 \
  --max-steps 1000

The command freezes every parameter except discovered routers and saves only router_state.pt, plus a manifest and training metrics. It does not duplicate multi-gigabyte base-model weights.

The equations in the paper regularize ground-truth SFT sequences toward low entropy. The code implements that stated objective. It does not add an unstated incorrect-response/high-entropy loss.

Repository layout

src/routing_entropy/
  capture.py       router discovery and differentiable logit capture
  benchmarks.py    public benchmark download and example prompt adapters
  metrics.py       NumPy reference metrics
  scoring.py       generation and OC/RE extraction
  fusion.py        XGBoost fusion and cross-validated evaluation
  training.py      router-only RE-regularized SFT
  cli.py           command-line entry points
configs/           paper defaults and model examples
examples/          small JSONL schemas (not paper data)
tests/             lightweight unit tests

Validation

pip install -e '.[dev]'
pytest -q
python -m compileall -q src

GPU integration tests are not run in GitHub Actions because the paper models are large. Before reporting results, record the exact model revision, Transformers version, dtype, device mapping, prompt template, answer parser, and dataset revision.