SHAPE of Chain-of-Thought in Math Reasoning

AI summary

Overview

• Research area: Large language model reasoning, specifically Chain-of-Thought (CoT) analysis for mathematical problem solving, with links to post-training (reinforcement learning with verifiable rewards, RLVR).
• Technical level: Intermediate. The framework rests on entropy-based metrics and a taxonomy of heuristics, but the paper explains both in plain terms; the appendix-level math is not required to follow the argument.
• Scope: The paper introduces SHAPE (Semantic-space and Heuristic Analysis for Problem-solving Evolution), an automated pipeline that annotates LLM CoT trajectories with mathematical heuristics and semantic spaces, and then uses those annotations to diagnose reasoning patterns, audit post-training effects, and improve training.

What This Paper Is About

Final-answer accuracy says little about how a language model organizes a mathematical solution. Existing CoT analyses look at surface properties—length, self-revision words like "wait" or "aha", broad episode labels, graph/tree structure—but none track the specific mathematical interpretation under which the model is operating at each step. SHAPE borrows two concepts from mathematics education research, semantic spaces (the solver's current mathematical interpretation of a problem, e.g. algebraic or trial-based counting) and heuristics (purposeful mathematical actions taken within that interpretation, e.g. simplifying, working backward), and turns them into an automated annotation pipeline and a set of metrics for diagnosing LLM reasoning.

Key Contributions

1. The SHAPE framework and pipeline. A theoretically grounded representation of LLM CoT as a sequence of semantic spaces (with New, Return, Maintain transitions) and heuristic labels drawn from a consolidated 11-family taxonomy (H1–H11, plus non-heuristic labels N1–N4), operationalized as a three-stage automated pipeline: content-unit segmentation, heuristic tagging, and semantic-space tracking.
2. A validation that heuristic-level features beat existing CoT representations. Heuristic frequency features predict whether a trajectory reaches the correct answer with higher AUROC than CoT length, reasoning-token features, self-revision markers, or ThinkARM episode labels.
3. A diagnostic audit of post-training. SHAPE shows that RL post-training induces mode-seeking in heuristic usage (concentration into the dense core of the base model's heuristic distribution) and that hard-perturbed problems cause structural changes in heuristics and semantic-space dynamics that nonetheless fail to converge on correct solutions.
4. A training method motivated by the diagnosis. HA-Plan-GRPO (Heuristic-Augmented, Planning-guided GRPO) injects eleven mathematical heuristics into the rollout prompt of a planning-based RLVR setup and improves performance over a planning baseline that lacks heuristic information.

Main Findings

• Heuristics predict correctness better than surface CoT features. Under 5-fold stratified cross-validation on the 100-problem Omni-MATH subset used by ThinkARM (trajectories from 15 models), SHAPE (H+N) reaches AUROC 0.664 ± 0.02 with 12 features, and SHAPE (H) reaches 0.653 ± 0.02 with 11 features. Baselines: CoT length 0.504 ± 0.03 (1 feature), length + reasoning 0.503 ± 0.03 (3 features), self-revision markers 0.618 ± 0.03 (3 features), ThinkARM 0.618 ± 0.02 (8 features).
• Correct reasoning concentrates effort in fewer semantic spaces; incorrect reasoning scatters. Across most models in Table 2, incorrect trajectories show higher effective number of semantic spaces (N_space^eff), higher effective number of transitions (N_trans^eff), and higher transition ratio (ρ = N_trans^eff / N_space^eff) than correct ones. For example, QwQ-32B goes from N_space^eff 1.74 (correct) to 2.74 (incorrect), and ρ from 0.32 to 0.60. The paper notes GPT-4o is an exception, with a higher N_space^eff for correct (1.68) than incorrect (1.40) trajectories.
• Extended reasoning changes the kind of traversal, not just its length. Open-source reasoning models with full traces show N_space^eff of 1.81–2.53 and ρ of 0.40–0.51, whereas instruction-tuned models without extended reasoning sit at N_space^eff of 1.37–1.71 and ρ of 0.19–0.32.
• High ρ on incorrect trajectories suggests unproductive revisiting. The paper links the higher ρ of incorrect trajectories to the overthinking phenomenon, describing it as revisiting the same semantic spaces more intensively rather than making forward progress.
• Hard perturbations change strategy but not successfully. On the 115-problem MATH-Perturb test split, Pass@1 drops sharply under hard perturbation for all four post-trained models studied: Qwen3-8B .96 → .94 (simple) / .84 (hard); Qwen3-32B .92 → .90 / .76; Nemotron-Cascade-8B .92 → .95 / .71; Olmo-3-7B-Think-RLVR .97 → .94 / .76. The Jensen–Shannon divergence of heuristic frequency distributions, the change in N_space^eff, and the change in ρ are all larger for hard than for simple perturbations (one-sided paired Wilcoxon signed-rank tests, p < .05). Models open more semantic spaces and revisit them more often, yet fail to commit to a successful interpretation. This structural divergence appears early—within the first five and ten content units—not as a late correction.
• Post-training narrows the heuristic distribution (mode-seeking). Using Density and Coverage (k = 3 nearest neighbors, cosine distance) on heuristic frequency vectors of successful trajectories, every post-trained model shows Density above 1.0 and Coverage below 1.0: Qwen3-1.7B-Base → Qwen3-1.7B-GRPO gives Density 1.220, Coverage 0.871 (834 vs 886 trajectories); Olmo-3-7B → Olmo-3-7B-Think-RL-Zero gives 1.250 / 0.707 (1229 vs 1307); Olmo-3-7B → Olmo-3-7B-Think-RLVR gives 1.032 / 0.531 (1507 vs 1600). The cross-model baseline of two unrelated base models shows low Density (0.520) and low Coverage (0.437), confirming the pattern is specific to base–post-trained pairs. A PCA projection shows Olmo-3-7B-Think-RLVR concentrating around the peak of the base distribution while leaving the left tail uncovered.
• Injecting heuristics into RL training improves accuracy. On the MATH-Perturb test split with Qwen3-1.7B-Base, HA-Plan-GRPO outperforms both the base model and Plan-GRPO across all three conditions: Original Avg@64 36.80 / Pass@64 80.00; Simple 35.80 / 79.13; Hard 17.72 / 62.61. For comparison, Qwen3-1.7B-Base scores 23.54 / 77.40, 23.10 / 79.13, and 11.84 / 57.39; Plan-GRPO scores 30.00 / 80.00, 29.86 / 78.26, and 14.52 / 61.74.

Methodology in Plain English

The authors first build a gold-standard annotation set. They take 48 CoT trajectories from the MATH-Perturb dataset generated by four models spanning thinking and non-thinking modes and different sizes (Qwen3-30B-A3B-Instruct, Qwen3-30B-A3B-Thinking, Qwen3-8B, Nemotron-Cascade-8B), segment them into content units—the smallest spans that can be coherently annotated—yielding 1,598 units from 8,334 sentences, and have four authors, including a graduate researcher in mathematics education, annotate by consensus discussion until they agree. Because manual annotation does not scale, they select annotator models by comparing candidates against this gold set with weighted F1 for unit-level agreement and macro F1 for class-level agreement; Grok-4.1-Fast and Qwen3.5-27B are the best closed-source and open-source options respectively, with Qwen3.5-27B used by default.

The automated pipeline has three stages: (1) segment the CoT into heuristic-bearing content units, (2) assign one or more heuristic labels from the taxonomy to each unit, and (3) track semantic-space state, where the model assigns New, Return, or Maintain to units containing a representation-changing heuristic (H1, H2, H3, H5, H8, H11), and Return triggers a joint lookup of the target space ID from a memory buffer in a single call.

From the annotated sequences they derive two distributions of heuristic activity: a space distribution that merges repeated visits to the same semantic space, and a segment distribution that keeps each contiguous visit separate. Each is summarized by its "effective number"—the exponential of its entropy—giving N_space^eff and N_trans^eff (with the latter minus one). A transition ratio ρ = N_trans^eff / N_space^eff captures how often the model oscillates per space, and a heuristic frequency distribution u(h) captures which heuristics are used and how often. The paper illustrates the computation on the Figure 1 example (S = (1,2,1), N_space^eff = 2, N_trans^eff ≈ 1.83).

For the diagnostic studies, they train logistic regression classifiers (ℓ1 or ℓ2, selected by inner 5-fold cross-validation on AUROC) to predict answer correctness from each feature set, and they measure perturbation effects using Jensen–Shannon divergence between heuristic frequency distributions plus changes in N_space^eff and ρ. For post-training analysis, they compare base and post-trained models using Density and Coverage. For training, they compare Plan-GRPO and HA-Plan-GRPO, which share the same reward, verifier, optimizer, and evaluation prompt and differ only in the rollout prompt (HA-Plan-GRPO adds eleven mathematical heuristics), trained on Qwen3-1.7B-Base with the MATH training split.

Why This Matters

Impact on research. SHAPE offers a process-level vocabulary for reasoning analysis that sits between raw accuracy and surface-form statistics. It shows that the mathematical interpretation a model is operating under is measurable at scale, that these measurements carry stronger correctness signal than length or lexical markers, and that RL post-training's diversity collapse—previously documented through output diversity or surface-form similarity—also appears at the level of strategic heuristic choice. It also supplies a diagnostic that leads directly to a training intervention.

Real-world applications:

• Model evaluation and auditing: replacing or supplementing accuracy-only benchmarks with process-level diagnostics that reveal whether a model genuinely reorganizes its approach when a problem changes.
• Tutoring and education technology: the metrics are imported from mathematics education research on how students organize problem solving, so the same lens could support analysis of both student and model solution traces.
• Data curation and post-training design: knowing that RL concentrates successful trajectories into a narrow region of the base model's heuristic distribution can guide diversity-preserving training objectives.
• Robustness testing: the MATH-Perturb setup isolates whether a model adapts strategy when problems look similar but require different solution methods, a scenario that matters for deploying models on unseen variants.

Industry relevance. The findings that post-training narrows strategy usage and that silently adding heuristic guidance to the rollout prompt alone yields measurable gains (36.80 vs 30.00 Avg@64 on the original split, and 17.72 vs 14.52 under hard perturbation) are directly actionable for teams training reasoning models, since the intervention changes only a prompt rather than the reward, verifier, or optimizer. The pipeline is also cheap enough to be practical: the paper notes annotating 445 trajectories with Grok-4.1-Fast costs approximately $100, which it calls impractical for large-scale annotation, motivating the use of Qwen3.5-27B as the default annotator.

Future Directions

• Extending beyond mathematics. The authors state that SHAPE analyzes observable CoT traces and that current validation is limited to mathematical benchmarks, leaving other domains as future work.
• Understanding why hard perturbations fail. Models demonstrably change their heuristics and semantic-space dynamics yet do not converge on correct solutions; the paper does not identify the mechanism that prevents commitment to a successful interpretation.
• Preserving strategic diversity during RL. Since post-training concentrates successful trajectories into the dense core of the base distribution and leaves parts of it uncovered, a natural next step is training objectives that retain heuristic coverage without sacrificing accuracy.
• Testing heuristic augmentation more broadly. HA-Plan-GRPO is evaluated on Qwen3-1.7B-Base and the paper labels the results preliminary (Table 5); scaling the method to larger models, other post-training recipes, and larger problem sets remains open.

Target Audience

Researchers and engineers working on LLM reasoning, Chain-of-Thought interpretability, and reinforcement learning post-training for math, as well as mathematics education researchers interested in computational models of how solvers organize problems. Readers who only need benchmarking numbers will find the process-level metrics less directly applicable, but those designing or diagnosing reasoning training pipelines will find both the framework and the concrete training intervention useful.

Authors’ abstract

Large language models (LLMs) achieve strong performance on mathematical reasoning benchmarks, yet the mathematically meaningful skills underlying their reasoning remain underexplored. We introduce \texttt{SHAPE}, a framework that analyzes Chain-of-Thought (CoT) trajectories through two lenses developed in mathematics education: (1) semantic spaces: the model's evolving mathematical interpretations of a problem (e.g., algebraic, geometric), and (2) heuristics: the specific mathematical actions taken within those spaces (e.g., simplifying the problem, working backward). We first use \texttt{SHAPE} to analyze the reasoning patterns of various models. Our findings reveal that the mathematical heuristics employed by a model better explain final answer correctness than traditional CoT features. Furthermore, models are likely to reach correct solutions by concentrating their reasoning effort within a few semantic spaces rather than exploring many disparate ones -- a pattern consistent with human behavior. Next, we utilize the \texttt{SHAPE} lens to evaluate whether post-training truly enhances mathematical proficiency. We find that reinforcement learning induces mode-seeking in heuristic usage. Lastly, we post-train LLMs by promoting diverse heuristics and demonstrate its effectiveness in improving accuracy. Overall, \texttt{SHAPE} provides a theoretically-grounded diagnostic framework for decoding LLM reasoning and offers a new path toward post-training LLMs for math reasoning. The code for our model is available at https://github.com/holi-lab/SHAPE-of-CoT