PaperlayerView exact revision on arXiv

Source-linked AI summary

A Strong MoE Is All You Need for Process Supervision

Guanzhi Deng, Zengyan Liu, Haibo Wang, Kuan Wu, Shing Yin Wong, Sichun Luo, Zhuoran Wang, Linqi Song

arXiv:2610.29498v1cs.CLcs.AI

TL;DR

Reliable process labels for multi-step reasoning are difficult to obtain, and many existing approaches rely on annotation, rollouts, or auxiliary evaluation. This paper turns native MoE routing into confidence scores for selective reasoning revision and a distilled standalone PRM. Both approaches consistently outperform strong baselines, and the PRM transfers across model architectures and tasks.

  • Problem

    Reliable process labels remain difficult to obtain, while existing approaches often rely on human annotation, rollouts, or auxiliary evaluation.

    Source: §1 · p. 1

  • Method

    The paper combines overall routing mass, head routing mass, and within-step variation to guide selective revision or label trajectories for a standalone PRM.

    Source: §1 · p. 2

  • Results

    RoCo-Rethink and RoCo-PRM consistently outperform strong baselines, with RoCo-Rethink best in all 18 settings and RoCo-PRM best or tied-best in all 30.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Expert routing provides an architecture-native source of process supervision, and routing-derived supervision transfers across MoE and dense target models.

    Source: §5.2 · p. 6, §7 · p. 10

Abstract

from arXiv · show

Process reward models (PRMs) provide fine-grained supervision for multi-step reasoning, yet obtaining reliable process labels remains a central bottleneck. Existing approaches typically rely on human annotation, rollout-based estimation, or auxiliary verification. In this work, we explore a different source of process supervision natively exposed by Mixture-of-Experts (MoE) models: expert routing. We show that simple routing statistics provide an intrinsic process-confidence signal that is competitive with conventional output-logit confidence. Building on this observation, we construct routing confidence by combining overall routing mass, head routing mass, and within-step routing variation, and introduce RoCo-Rethink, which identifies the least confident reasoning step and selectively regenerates the subsequent trajectory. We further introduce RoCo-PRM, which uses a strong MoE to automatically annotate reasoning trajectories with routing-derived step scores and distills them into a standalone PRM. Across multiple MoE architectures and diverse reasoning tasks, both RoCo-Rethink and RoCo-PRM consistently outperform strong baselines, while RoCo-PRM transfers effectively across both model architectures and tasks. Our findings suggest that expert routing is more than a computational dispatch mechanism: it provides an architecture-native source of process supervision for reasoning.

1 INTRODUCTION

The introduction asks whether MoE routing can provide process supervision without relying on the usual annotation, rollout, or auxiliary-verification sources. It proposes routing confidence and two ways to use it, reporting strong results for both.

1 INTRODUCTION

  • The paper positions process supervision as useful for evaluating intermediate reasoning states because errors can propagate through later steps.PRMs are used in reasoning verification, test-time search, self-correction, and reinforcement learning.Source: §1 · p. 1
  • The paper asks whether native MoE routing can serve as intrinsic process supervision, given the costs and dependencies of existing annotation, rollout, verification, and uncertainty-based approaches.Many automated approaches still require additional evaluation computation, outcome signals, or auxiliary evaluators.Source: §1 · p. 1
  • Routing confidence combines overall routing mass, head routing mass, and within-step routing variation, which relate systematically to step reliability across architectures and reasoning domains.Routing-based predictors often outperform conventional output-side confidence signals.Source: §1 · p. 2
  • RoCo-Rethink selects the least confident step and regenerates the subsequent trajectory, while RoCo-PRM distills routing-derived step scores into a standalone PRM.The PRM scores steps without access to routing information.Source: §1 · p. 2
  • RoCo-Rethink outperforms the strongest test-time baseline in all 18 evaluated settings, while RoCo-PRM achieves the best or tied-best performance in all 30 settings.Source: §1 · p. 2

2 RELATED WORK

Related work describes process supervision from human annotation and automated sources, emerging evidence that MoE routing reflects reasoning states, and test-time methods that use localized reliability signals.

2.1 PROCESS SUPERVISION AND PROCESS REWARD MODELS

  • Process-supervision methods use human labels, rollouts, outcome-derived signals, verification, or model uncertainty, but can require extra evaluation computation, outcome information, or auxiliary evaluators.This paper instead investigates routing dynamics produced during ordinary MoE inference.Source: §2.1 · p. 2

2.2 MOE ROUTING AS A REASONING-STATE SIGNAL

  • Prior studies report that MoE routing retains information about task context, trajectory state, and latent reasoning states, motivating its use as a step-level confidence signal.Related work also connects routing behavior with reasoning reliability and routing interventions.Source: §2.2 · pp. 2–3

2.3 TEST-TIME REASONING GUIDANCE

  • Test-time reasoning methods allocate computation to uncertain trajectory segments, commonly using output-side confidence or separately trained evaluators.RoCo-Rethink instead derives its guidance signal directly from native MoE routing.Source: §2.3 · p. 3

3 UNDERSTANDING MOE ROUTING AS A PROCESS SIGNAL

This section tests whether native MoE routing tracks step reliability and defines routing confidence for reasoning guidance and process supervision.

3 UNDERSTANDING MOE ROUTING AS A PROCESS SIGNAL

  • The section frames routing as a possible source of process guidance and supervision after examining whether it tracks step-level reliability.The proposed routing-confidence signal is intended for both reasoning guidance and process supervision.Source: §3 · p. 3

3.1 ROUTING SIGNALS TRACK STEP-LEVEL RELIABILITY

  • Higher overall and head routing mass are associated with higher step quality, while greater within-step routing variation is associated with lower quality across the evaluated settings.The figure reports within-trajectory correlations across three MoE models and four tasks, with GSM8K and MATH averaged over two math PRMs.Source: §3.1 · p. 3, Figure 1 · p. 4
  • Routing-based bad-step predictors often outperform their logit-based counterparts, supporting native routing as a competitive step-reliability signal.Figure 1(b) compares grouped-CV bad-step AUROC for routing-based and logit-based predictors.Source: §3.1 · p. 4, Figure 1 · p. 4

3.2 FROM ROUTING STATISTICS TO ROUTING CONFIDENCE

  • Routing confidence combines normalized overall and head routing mass with within-step variation, whose direction is reversed because variation is negatively associated with reliability.It uses quantities produced during ordinary MoE inference and requires neither additional rollouts nor an auxiliary verifier.Source: §3.2 · p. 4

4 METHOD

RoCo turns native MoE routing into test-time reasoning guidance and process supervision. RoCo-Rethink revises a trajectory from its least confident step; RoCo-PRM distills routing-derived labels into a standalone evaluator.

4 METHOD

  • RoCo is a framework with two uses for routing confidence: direct test-time guidance through RoCo-Rethink and a standalone process evaluator through RoCo-PRM.RoCo-PRM distills routing-derived supervision.Source: §4 · p. 4

4.1 ROCO-RETHINK

  • RoCo-Rethink selects the least confident reasoning step, preserves the preceding prefix, and regenerates the trajectory from that point onward.It recomputes confidence after each regeneration and stops when predictions stabilize or a maximum number of rounds is reached.Source: §4.1 · pp. 4–5

4.2 ROCO-PRM

  • RoCo-PRM trains on MoE-generated trajectories paired with routing-confidence scores for intermediate steps, then regresses predicted scores from each problem and reasoning prefix toward those targets.The resulting PRM need not be an MoE and does not require routing access from the target model.Source: §4.2 · p. 5
  • At test time, RoCo-PRM scores reasoning prefixes and uses the lowest-scoring step as the revision point, after which the target model regenerates the suffix.The same PRM can guide MoE and dense target models.Source: §4.2 · p. 5

5 EXPERIMENTS

Across the evaluated models and tasks, RoCo-Rethink improves test-time reasoning guidance, while RoCo-PRM provides strong process supervision that also transfers to dense models.

5.1 EXPERIMENTAL SETUP

  • The evaluation covers six benchmarks across mathematical, knowledge, code, instruction-following, and logical reasoning, with PRM training data disjoint from evaluation sets.RoCo-Rethink is tested on three MoE models; RoCo-PRM is evaluated on those models and two dense models.Source: §5.1 · p. 6
  • Comparisons include Vanilla decoding, Self-Consistency, and LECO for test-time guidance, and MC-PRM and CPMI-PRM for process supervision.Experiments are repeated with five random seeds and report mean and standard deviation.Source: §5.1 · p. 6

5.2 MAIN RESULTS

  • RoCo-Rethink achieves the best performance in all 18 model–task settings across three MoE models, improving over the strongest test-time baseline by 1.2–1.9 points on average.The reported average gains are 1.2 points on OLMoE, 1.9 on DeepSeek, and 1.5 on Qwen1.5.Source: §5.2 · p. 6, Table 1 · p. 7
  • RoCo-PRM is best or tied-best in all 30 model–task settings, with 29 strict wins and one tie, and improves over the strongest baseline on dense models.Average gains over the strongest baseline are 1.4 points on Qwen2.5 and 1.6 points on Llama.Source: §5.2 · p. 6, Table 1 · p. 7

6 FURTHER ANALYSIS

Further analyses find that composite routing confidence selects higher-value revision points, makes better use of additional rethink rounds, and supports RoCo-PRM performance through stronger teachers and cross-task transfer.

6.1 COMPOSITE ROUTING CONFIDENCE ENABLES BETTER REVISION-POINT SELECTION

  • M+H outperforms every individual routing statistic in 17 of 18 settings, and adding −D makes the full formulation best or tied-best in 17 of 18.Each individual statistic generally improves over random rethink.Source: §6.1 · p. 7, Figure 3 · p. 8
  • The full M+H−D selector lowers average oracle regret to 0.140 from 0.183 for random selection and raises hit rate to 0.583 from 0.511 across nine settings.It achieves the lowest regret and highest hit rate in all nine model–task settings.Source: §6.1 · p. 8, Figure 4 · p. 8

6.2 ROCO-RETHINK USES TEST-TIME RETHINK BUDGET MORE EFFECTIVELY

  • At the same rethink-round budget, RoCo-Rethink delivers larger gains than LECO and random rethink across all three MoE models.From one to three rounds, its average gains rise from 2.9 to 4.7 on OLMoE, 3.2 to 5.6 on DeepSeek, and 2.5 to 4.4 on Qwen1.5.Source: Figure 5 · p. 8, §6.2 · p. 9

6.3 ROCO-PRM IS MORE SENSITIVE TO TEACHER QUALITY THAN BACKBONE SCALE

  • With the PRM backbone fixed, Qwen3 routing supervision improves average score by 2.0 points over DeepSeek across 30 settings, while scaling the backbone from 1.5B to 7B adds 0.9 points.The teacher gain ranges from 1.4 to 2.7 points across target models; the backbone gain ranges from 0.6 to 1.5.Source: §6.3 · p. 9, Table 2 · p. 9

6.4 ROUTING-DERIVED PROCESS SUPERVISION TRANSFERS ACROSS TASKS

  • Leave-one-task-out RoCo-PRMs improve all 12 target-model–task settings by an average of 1.98 points despite excluding task-aligned routing supervision.Gains are positive across all three target models, including a smaller average gain of 0.87 points on MMLU.Source: §6.4 · p. 9, Figure 6 · p. 9

7 CONCLUSION

The conclusion presents RoCo as a framework that turns native MoE routing into reasoning guidance and process supervision. The accompanying appendix details its routing measures, step statistics, and trajectory-relative confidence.

7 CONCLUSION

  • RoCo-Rethink uses routing confidence for targeted test-time revision, while RoCo-PRM distills routing-derived scores into a standalone evaluator that transfers across MoE and dense models.The conclusion links reported gains to revision-point selection, stronger routing teachers, and cross-task transfer.Source: §7 · p. 10

AI USE STATEMENT

  • Generative AI assisted with manuscript preparation, research ideation, experimental work, analysis, and coding; the authors report reviewing and verifying AI-assisted outputs and take responsibility for the final work.The statement also says literature and citations suggested with AI assistance were checked against original sources.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The reproducibility materials include routing-feature definitions, data sources, training and decoding settings, evaluation protocols, and multi-seed and oracle-analysis details.The statement says code and scripts for data construction, routing extraction, training, and evaluation will be released.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 TOKEN-LEVEL ROUTING MASS AND LAYER NORMALIZATION

  • Token-level routing mass sums router probabilities for selected experts, is min–max normalized within each MoE layer over the trajectory, and is averaged across layers.The layer-normalized variant is used in the main RoCo experiments.Source: §A.1 · pp. 13–14

A.2 STEP SEGMENTATION AND ROUTING STATISTICS

  • Reasoning trajectories are segmented into newline-delimited non-empty steps, with token offsets assigning tokens to each step and its first h tokens.The main experiments use h = 3; shorter steps use all available tokens.Source: §A.2 · p. 14

A.3 TRAJECTORY-LEVEL NORMALIZATION AND ROUTING CONFIDENCE

  • Overall mass, first-token-region mass, and within-step variation are normalized across steps to form routing confidence, which is trajectory-relative rather than globally calibrated.RoCo-Rethink uses the score to order steps within a trajectory and select the lowest-confidence step.Source: §A.2 · p. 14, §A.3 · p. 14

A.4 ROUTING–RELIABILITY ANALYSIS

The routing predictor is evaluated against a dimension-matched output-logit predictor for identifying bad reasoning steps, and the appendix notes that each MoE model retains its native routing rule.

A.4 ROUTING–RELIABILITY ANALYSIS

  • The routing-based bad-step predictor outperforms the logit-based predictor in 13 of 18 model–task–verifier settings.The predictors share a pipeline and are evaluated with trajectory-level five-fold GroupKFold cross-validation; bad steps are those at or below the verifier-score 20th percentile.Source: §A.4 · pp. 14–15, Table 4 · p. 15

B.1 MODEL CONFIGURATIONS

  • Routing extraction uses each MoE model’s native top-k rule rather than imposing a common k across architectures.The appendix identifies this as an architectural configuration relevant to routing extraction.Source: §B.1 · p. 15

B.2 BENCHMARKS AND EVALUATION PROTOCOLS

The appendix specifies the benchmark and supervision setup, then details selective-revision baselines and analyses of routing components, revision-point quality, and rethink budgets. It also describes how routing-derived training examples are constructed.

B.2 BENCHMARKS AND EVALUATION PROTOCOLS

  • Evaluation spans six benchmarks across five reasoning domains, and all evaluation sets are disjoint from RoCo-PRM supervision data.Tables 5 and 6 summarize model configurations and benchmark/evaluation setup.Source: §B.2 · p. 15, Table 5 · p. 16, Table 6 · p. 16

B.3 PROCESS-SUPERVISION DATA SOURCES

  • Each teacher-generated trajectory yields one regression example per reasoning step, pairing a problem and prefix with its trajectory-normalized routing statistics and target C_i = f_Mi + e_Hi − e_Di.The examples are task-aligned and use supervision corpora disjoint from the corresponding evaluation sets.Source: §B.3 · p. 16

C.1 POST-HOC REVISION PROCEDURE

  • RoCo-Rethink replays a complete trajectory, selects its lowest-confidence step, preserves the preceding prefix, and regenerates the suffix; it can recompute the revision point after each round.For extractable answers, it may stop when successive answers match, subject to a minimum number of rounds; the main setup allows at most three rounds.Source: §C.1 · pp. 16–17

C.2 TEST-TIME BASELINES

  • LECO changes the confidence source to output logits, while random rethink selects a step uniformly; both use RoCo-Rethink’s suffix-regeneration machinery.Self-consistency samples eight completions by default; HumanEval’s main comparison uses pass@1 rather than pass@8.Source: §C.2 · p. 17

D.1 ROUTING-COMPONENT ABLATION

  • M+H−D is best or tied-best in 17 of 18 settings, while M+H beats each individual routing statistic in 17 of 18.Table 9 gives the complete five-seed routing-component ablation across three MoE models and six tasks.Source: §D.1 · p. 18, Table 9 · p. 18

D.2 MONTE CARLO REVISION ORACLE

  • The full M+H−D selector reaches the lowest regret and highest hit rate in all nine settings, reducing regret from 0.183 to 0.140 and raising hit rate from 0.511 to 0.583 versus random selection.Table 10 reports the complete Monte Carlo revision-oracle results; lower regret and higher hit rate are better.Source: §D.2 · p. 18, Table 10 · p. 19

D.3 RETHINK-BUDGET SCALING

  • Table 11 reports mean score gain over Vanilla as the maximum rethink-round budget B increases, averaged across GSM8K, MATH500, MMLU, and LogicBench.The passage describes the aggregation procedure but gives no gain values.Source: §D.3 · p. 19, Table 11 · p. 19

E.1 ROUTING-DERIVED TRAINING EXAMPLES

  • Training examples pair each problem and reasoning prefix with a routing-derived target; trajectories are min–max normalized, and prefixes are not truncated until tokenization.Empty steps and samples with invalid or missing routing features are removed.Source: §E.1 · p. 19

E.2 PRM ARCHITECTURE AND OPTIMIZATION

RoCo-PRM is trained as a scalar-regression evaluator of routing-derived confidence scores and then used to guide selective revision. The appendix also specifies its optimization and the supervision procedures used by comparison PRMs.

E.2 PRM ARCHITECTURE AND OPTIMIZATION

  • RoCo-PRM predicts routing-derived confidence from the problem and reasoning prefix using a Qwen2.5-7B-Instruct scalar-regression model trained with MSE.Training uses LoRA adapters, a trainable regression head, and one epoch by default; inputs longer than 1,024 tokens are truncated during training tokenization.Source: §E.2 · pp. 19–20

E.3 PRM-GUIDED RETHINK

  • At evaluation, the PRM scores each prefix, selects the lowest-scoring step, and guides suffix regeneration without requiring routing access from the target model.Scores are min–max normalized within each trajectory before selecting the revision point.Source: §E.3 · p. 20

F PRM BASELINE DETAILS

  • MC-PRM, CPMI-PRM, and RoCo-PRM use the same Qwen2.5-7B-Instruct regression backbone and training recipe unless stated otherwise, isolating how their supervision targets are constructed.The CPMI-derived targets also use the same LoRA and optimization configuration.Source: F PRM BASELINE DETAILS · p. 20, §F.2 · p. 22

F.1 MC-PRM

  • MC-PRM trains on the fraction of sampled continuations that pass task-specific correctness checks, using the same scalar-regression PRM architecture.Its default construction samples eight continuations per prefix with Qwen3-30B-A3B.Source: §F.1 · p. 21

F.2 CPMI-PRM

  • CPMI-PRM scores a step by its incremental evidence for a correct target relative to plausible incorrect alternatives, using mean token log-probability.Task-specific positive and negative targets are constructed for discrete-answer, code, and instruction-following tasks.Source: §F.2 · pp. 21–22

G.1 TEACHER AND BACKBONE ABLATIONS

The appendix compares routing-teacher and PRM-backbone changes, tests leave-one-task-out transfer, and documents task-specific evaluation procedures. Routing-teacher improvements are larger than those from scaling the PRM backbone, and leave-one-task-out PRMs improve results across evaluated settings.

G.1 TEACHER AND BACKBONE ABLATIONS

  • A Qwen3-30B-A3B routing teacher raises average scores by 2.0 points across 30 settings, versus 0.90 points from increasing the PRM backbone size across the same settings.Teacher gains range from 1.4 to 2.7 points across five target models; backbone gains range from 0.55 to 1.53 points.Source: §G.1 · p. 22, Table 14 · p. 23

G.2 LEAVE-ONE-TASK-OUT CROSS-TASK GENERALIZATION

  • Leave-one-task-out RoCo-PRMs improve all 12 target-model–task settings by 1.98 points on average despite excluding supervision aligned to each evaluation task.Average gains are +2.53 on MATH500 and GSM8K, +1.97 on HumanEval, and +0.87 on MMLU; MMLU gains remain positive across all three target models.Source: §G.2 · p. 23, Table 15 · p. 23

H.1 TASK-SPECIFIC PROMPTING

  • Within each benchmark, compared methods use the same base prompt; the appendix specifies task-specific prompt formats, including native chat templates where available.HumanEval uses the canonical continuation prompt without a chat template, while IFEval uses its original prompt wrapped in the model’s native chat template.Source: §H.1 · p. 24

H.2 REASONING-STEP SEGMENTATION

  • Natural-language reasoning steps are newline-delimited non-empty blocks, while HumanEval steps are non-empty source-code lines.Tokenizer offsets align generated tokens with their text spans for step-level aggregation.Source: §H.2 · p. 24

H.3 ANSWER EXTRACTION

  • Answer extraction is applied consistently across Vanilla, rethink-based, and PRM-guided methods, with task-specific rules for numerical, multiple-choice, and Boolean answers.For GSM8K and MATH500, the extraction order prioritizes the last well-formed boxed answer, then other specified markers and numerical expressions.Source: §H.3 · pp. 24–25

H.4 HUMANEVAL EVALUATION

  • HumanEval completions are evaluated with the official functional-correctness harness, and the reported metric is pass@1.For rethink methods, evaluation combines the preserved code prefix with the regenerated suffix before execution.Source: §H.4 · p. 25

H.5 IFEVAL EVALUATION

  • IFEval uses the official checker under strict and loose criteria, but the main paper reports loose prompt accuracy and uses strict compliance for supervision construction.The evaluator records prompt-level and instruction-level compliance outcomes.Source: §H.5 · p. 25

I QUALITATIVE EXAMPLES

The examples show both successful correction and a failure to escape an erroneous continuation. A successful trajectory takes more than one regeneration to correct, while a failed regeneration can reproduce the original suffix.

Successful correction (W2R)

  • In the GSM8K success example, RoCo-Rethink selects Step 2 at confidence −0.707; after an initial answer change from $360 to $500, a later revision reaches the gold answer $800.The preserved prefix precedes Step 2, and the regenerated reasoning identifies five $40 blue ties and ten $60 red ties.Source: I QUALITATIVE EXAMPLES · pp. 25–26
  • The successful example requires multiple revisions: the first continuation remains incorrect, and the stopping rule terminates after the answer stabilizes.The final rethink round preserves the correct answer.Source: I QUALITATIVE EXAMPLES · pp. 25–26

Failure case (W2W)

  • In the failure example, selecting Step 3 at confidence −0.569 reproduces the erroneous suffix and leaves the answer at $195,000 rather than the $70,000 gold profit.The authors note that identifying a low-confidence point does not guarantee regeneration will leave the original reasoning basin.Source: I QUALITATIVE EXAMPLES · p. 26, Table 17 · p. 26
Loading 2610.29498v1…