Characterizing, Evaluating, and Optimizing Complex Reasoning
SAI replication review · Referee report
Summary
This paper proposes a unified framework for characterizing, evaluating, and optimizing the reasoning traces produced by large reasoning models. The conceptual move is to lift reasoning quality out of the linear, step-wise view that dominates existing PRM work and treat a trace as a two-dimensional object described by the ME² principle — macro vs. micro granularity crossed with efficiency vs. effectiveness. Traces are abstracted as DAGs whose macro-level linearization exposes branching/merging structure and whose micro-level dominant path exposes local step quality. Under this abstraction the authors train a lightweight Thinking Reward Model (TRM) on 103K DeepSeek-V3.2-labeled pairwise preferences drawn only from verified-correct traces, and plug it into Best-of-N test-time selection and GRPO-based RL with gated reward shaping. Empirically, TRM beats existing PRMs on validation-set preference accuracy (88.6% vs 62.5%/46.3%), lifts Best-of-N accuracy on AIME by up to 19 points at N=16, and adds 2–4 average points to RLVR fine-tuning across ten benchmarks and three policies. The two-axis quality principle plus a scalable DAG-based evaluation pathway is a genuinely valuable contribution, and casting the TRM addition as potential-based reward shaping is the right formalism for arguing the optimal-policy set is preserved. The main methodological limitation is that TRM's advantage is confounded along several axes at once — DAG structure, ME² prompting, pairwise Bradley–Terry training, and DeepSeek-V3.2 supervision — with no ablation isolating them, and because DeepSeek-V3.2 both labels the data and judges the downstream win rates, the reasoning-quality story has an evaluator-loop concern. Executed replication sharpens this assessment asymmetrically: the pairwise-evaluation half reproduces strongly (TRM 88.47%, Qwen-PRM near-random 45.9%, tie counts near-exact) and the Best-of-N trend and selector ordering reproduce, but the RL half — the paper's second headline — could not be run (~588 GPU-hours) and remains unverified, and the released code required substantial repair to execute.
Strengths
- Conceptual contribution. The ME² principle is a clean two-axis characterization (macro/micro × efficiency/effectiveness) that carves out four intuitive quality dimensions and is a substantially more principled starting point than the piecewise definitions in prior PRM work.
- DAG abstraction is scalable. The incremental DAG-construction procedure with a simplified attachment pool (main branch + representative branch endpoints) is a pragmatic, well-motivated compromise between full history and purely local context, and the collapse into super-nodes keeps downstream evaluation tractable.
- Decoupling reasoning quality from final-answer correctness. Training TRM only on verified-correct traces removes correctness as a training-label confounder and is what lets the reward compose cleanly with verifiable rewards in GRPO.
- Pairwise-evaluation result holds up under execution. Independent replication of Table 1 on the held-out 1500-pair split reproduced TRM at 88.47% (paper 88.6%) with near-identical counts, preserved the full method ranking, and confirmed Qwen2.5-Math-PRM-7B at chance (45.9%), giving high confidence in the paper's core evaluation claim.
- Empirical breadth. Table 2 spans three policies × three reward configurations × ten benchmarks, and the test-time experiments cover two mid-size and two large generators — a wider surface than typical PRM papers.
- Cost discipline (as reported). The TRM addition adds only ~12% GPU-hours over the Verifier baseline and stays cheaper than ReasonFlux, an important property for an auxiliary reward.
Weaknesses
- Confounded factors behind TRM's win. TRM differs from the PRM baselines along four axes simultaneously (DAG abstraction, ME² dimensions, pairwise Bradley–Terry supervision, DeepSeek-V3.2 labels). No ablation isolates them; in particular, a pairwise-preference reward model trained on flat traces with the same aggregated labels is missing, so the reader cannot tell whether the DAG machinery is what matters.
- Circular evaluator loop for the reasoning-quality claim. DeepSeek-V3.2 produces the TRM-Preference labels and is also the judge for the Section 6.1 pairwise win rates, so TRM-trained policies winning under the same judge is partly evidence of imitation success rather than an independent quality signal.
- Overstated framing of the headline numbers. The abstract's "up to 19.3%" and "up to 3.9%" are best-case maxima from different setups and reference points; the "orthogonal to answer-based verifiable signals" claim overstates decoupling from correctness (the test-time story exploits exactly that correlation, and Micro-Effectiveness penalizes local errors); and the italicized "leads to better final outcomes" infers causation from a correlational Best-of-N selection.
- Train/test distribution mismatch at selection time. TRM is trained only on correct-vs-correct pairs (and every prompt tells the judge both traces are correct), yet at Best-of-N it scores candidates that include incorrect traces — a shift the paper does not discuss.
- Micro abstraction restricted to the dominant path. Micro-Efficiency targets looping/redundant restarts that live in off-path branches, which the dominant-path abstraction discards; the two micro dimensions are conflated under one abstraction and the efficiency form is never separately tested.
- Theoretical gaps. The Theorem 5.2 bound is not sign-controlled ( can be negative); the "variance" interpretation of conflicts with the appendix expectation rendering; the sequence-level "final bonus only" identity requires but is used; a factor is dropped in the step; and is undefined. The optimal-policy-invariance theorem also implies the empirical gains are a finite-horizon optimization effect, which should be stated.
- Statistical support is thin. Table 2 gives single-point, apparently single-seed accuracies with many ~1-point gaps; AIME24/25 have 30 problems each; test-time scaling uses only five repeats, giving standard errors comparable to the N=16 TRM-over-PRM gaps. No confidence intervals or significance tests are reported.
- Uncontrolled distribution overlap. TRM training and RL fine-tuning both draw from WebInstruct-verified, so TRM's RL utility may partly reflect distribution match rather than a general reasoning-quality signal; a disjoint-prompts control is missing.
- Under-specified reliability and cost reporting. The reliability validation collapses three human annotators into one 90.0% number with no per-annotator/κ breakdown and no per-family stratification; the Table 3 GPU-hour units and hardware are ambiguous.
- Local presentation issues. Eq. 1 does not render in the main text; the super-node symbol overloads ; the Macro-Efficiency prompt penalizes absent verification (an effectiveness concern) under the efficiency axis; and the Case 2 narrative overstates Trace B's failure relative to the exhibit.
Reproducibility & code
- Pairwise-evaluation validation (Table 1) reproduces strongly. On the held-out 1500-pair split TRM came out 88.47% (paper 88.6%), counts 1327/1/172 (paper 1326/0/171), ranking preserved; Qwen2.5-Math-PRM-7B reproduced at 45.9% (paper 46.3%), confirming near-random behavior with the PRM scored truncated to its 4096-token window. The only divergence is ReasonFlux reproducing higher (67.9% vs 62.5%), which does not threaten the claim.
- Best-of-N reproduces as a trend but not in magnitude. The AIME24/Qwen3-8B curve rose monotonically (0.63 → 0.767) with the correct selector ordering (TRM > ReasonFlux > Qwen-PRM at N=16), but the peak gain was +13.7pp vs the paper's 19.3pp because the reproduced N=1 pass@1 was 0.63 vs 44.7% (a base-model/decoding difference, compounded by a forced halving of
max_new_tokens32768→16384). Only 1 of the 4 panels (AIME24/Qwen3-8B) was run. - The RL half was not reproduced. Full TRM-guided GRPO training (Table 2, Table 3, Figs 5–6; ~588 GPU-hours across three reward configs) was infeasible in the environment; only a 3-step, 1000-row smoke test ran and stalled during model materialization. The reward-fusion pipeline (general-verifier correctness fused with the TRM score per Eq. 1) was validated end-to-end and GRPO launched, but no trained checkpoints or benchmark accuracies were produced, so the second headline claim (up to 3.9% average gain) and the win-rate/training-dynamics figures remain unverified. The compute requirement is documented (Table 3), so this is an environmental limitation, not a paper defect.
- Released code did not run as-is (nine fixes). Two are genuine defects: the answer verifier uses a main-thread-only SIGALRM timeout that, called from worker threads, silently scores every sample incorrect (pass@1 = 0) despite correct boxed answers — a naive multi-threaded reproducer would obtain zeroed metrics with no error; and the reward/verifier backend hardcodes a private VPN IP (100.97.200.2). Remaining friction (flash-attn/CUDA mismatch, missing sglang
[srt]extra and dep pins, undeclaredverlruntime deps, placeholder model paths) indicates an incompletely packaged release. The reward-shaping coefficientalphais also a hard-coded literal rather than a config parameter. - Preference-construction pipeline is not reproducible from the release. Only the final HuggingFace dataset is usable; the trace-generation, four-dimension judgment, aggregation, and pair-subsampling stages were not exercised and the ~1.8M→103K bridge (verified-correct pass rate, "at most four pairs per prompt", validation split) is underspecified. TRM training itself was not re-run; a released checkpoint was used.
Recommended Changes
Essential
- Add an ablation isolating the DAG contribution. Train a variant on the same TRM-Preference labels but with flat, unstructured traces as input (no DAG, no macro/micro split) and report its Section 5.2 accuracy and Best-of-N gains, directly addressing the confounded-factors weakness in Table 1.
- Break the DeepSeek-V3.2 evaluator loop for Figure 5. Recompute the Section 6.1 pairwise win rates under at least one independent judge (Gemini 3.1 Pro / GPT-5.4 / Claude Opus 4.6, already used in App. G.4) and report side by side.
- Report per-seed variance and significance. Run ≥3 seeds per RL configuration for Table 2 (mean ± std) and add a paired-bootstrap test (or ≥10 repeats) for the AIME24/25 Best-of-N gaps, given the 30-problem benchmarks and five-repeat protocol.
- Repair the theoretical section. State an explicit alignment condition ensuring in Theorem 5.2; reconcile the variance-vs-expectation rendering of ; set (or carry ) so the "thinking reward bonus of the final response" identity is justified; restore the factor; define ; and note that Theorem 5.1 makes the gains a finite-horizon effect.
- Render Eq. 1 and fix the released code. Show the gated-reward formula and specify when contributes; fix the SIGALRM scoring bug so the shipped harness works multi-threaded, sanitize the hardcoded VPN IP, expose
alphaas a config parameter, and document the sglang/verlinstall so the pairwise-eval and RL paths run without private patches. - Release the preference-construction pipeline. Ship the trace-generation, four-dimension judgment, aggregation, and pair-subsampling scripts, and report the verified-correct pass rate and validation-split logic that bridge the raw candidate pool to the 103K pairs.
Suggested
- Reframe the headline numbers. Report the Best-of-N absolute lift and the TRM-over-PRM margin separately; soften "leads to better final outcomes" to "is predictive of"; and qualify "orthogonal" to "decoupled from final-answer supervision".
- Probe the selection-time distribution shift. Report TRM's score distribution on correct vs. incorrect candidates at Best-of-N, since TRM never saw incorrect traces in training.
- Add a distribution-overlap control between TRM training data and RL prompts (disjoint WebInstruct partitions) to separate distribution match from a general reasoning-quality effect.
- Test micro-efficiency on the full trace, not only the dominant path, or show the two forms agree.
- Discuss the scale-vs-gain trend explicitly, engaging the reward-model-capacity ceiling hypothesis for the large generators.
- Clarify the Macro-Efficiency prompt so the presence-of-verification criterion moves under Macro-Effectiveness (or justify blending).
- Break out the reliability numbers into per-annotator accuracies with an inter-annotator agreement statistic, stratified by generating-model family.
- Reconcile local presentation issues: rename the super-node symbol (e.g. ), disambiguate the Table 3 GPU-hour units and hardware, reconcile the "three evaluations per order" vs "four orderings" descriptions, and revise the Case 2 narrative to match the exhibit.