PaperlayerView exact revision on arXiv

Source-linked AI summary

THOR: Tool-Integrated Hierarchical Optimization via RL for Mathematical Reasoning

Qikai Chang, Zhenrong Zhang, Pengfei Hu, Jun Du, Jiefeng Ma, Yicheng Pan, Jianshu Zhang, Quan Liu, Jianqing Gao

arXiv:2509.13761v3cs.AIcs.CL

TL;DR

LLMs struggle with high-precision mathematical tasks, while tool-integrated reasoning methods face gaps in data construction, fine-grained optimization, and inference. THOR combines policy-aligned TIRGen data, hierarchical RL, and tool-feedback self-correction; it reports strong math and code benchmark performance across models.

  • Problem

    LLMs struggle with high-precision tasks such as numerical computation and symbolic manipulation, and existing tool-integrated reasoning methods face gaps in data construction, fine-grained optimization, and inference.

    Source: §1 · p. 2

  • Method

    THOR uses multi-agent TIRGen data construction, episode- and step-level RL, and inference-time self-correction based on immediate tool feedback.

    Source: §1 · pp. 2–3, §4 · p. 10

  • Results

    THOR reports state-of-the-art performance among comparable-size models on mathematical benchmarks, reduced inference overhead, and gains on code benchmarks, with generalization across reasoning and non-reasoning models.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The authors describe THOR as promising reliable automated assistance for complex problems in mathematics, engineering, and the formal sciences.

    Source: ETHICS STATEMENT · p. 10

    The evaluation focuses on code-integrated mathematical problem solving; other tool types and larger-scale models were not explored.

    Source: E LIMITATION AND FUTURE WORK · p. 18

Abstract

from arXiv · show

Large Language Models (LLMs) have made remarkable progress in mathematical reasoning, but still continue to struggle with high-precision tasks like numerical computation and formal symbolic manipulation. Integrating external tools has emerged as a promising approach to bridge this gap. Despite recent advances, existing methods struggle with three key challenges: constructing tool-integrated reasoning data, performing fine-grained optimization, and enhancing inference. To overcome these limitations, we propose THOR (Tool-Integrated Hierarchical Optimization via RL). First, we introduce TIRGen, a multi-agent based pipeline for constructing high-quality datasets of tool-integrated reasoning paths, aligning with the policy and generalizing well across diverse models. Second, to perform fine-grained hierarchical optimization, we introduce an RL strategy that jointly optimizes for both episode-level problem solving and step-level code generation. This is motivated by our key insight that the success of an intermediate tool call is a strong predictor of the final answer's correctness. Finally, THOR incorporates a self-correction mechanism that leverages immediate tool feedback to dynamically revise erroneous reasoning paths during inference. Our approach demonstrates strong generalization across diverse models, performing effectively in both reasoning and non-reasoning models. It further achieves state-of-the-art performance for models of a similar scale on multiple mathematical benchmarks, while also delivering consistent improvements on code benchmarks. Our code will be publicly available at https://github.com/JingMog/THOR.

1 INTRODUCTION

LLMs struggle with high-precision mathematical tasks, and tool-integrated reasoning still faces gaps in data construction, optimization, and inference. THOR addresses these gaps with policy-aligned tool-use data, hierarchical RL, and inference-time self-correction.

  • LLMs remain weak on high-precision tasks such as numerical computation, equation solving, symbolic manipulation, and formal proofs.The introduction motivates combining LLM semantic reasoning with precise, verifiable execution by external code-based tools.Source: §1 · p. 2
  • Existing tool-integrated reasoning methods face challenges in constructing suitable data, optimizing individual reasoning steps, and using immediate tool feedback during inference.Prompt-only data generation can fail to elicit effective tool use from reasoning models, while rule-based insertion can struggle to find suitable positions.Source: §1 · p. 2
  • THOR combines policy-aligned TIRGen data, episode- and step-level reinforcement learning, and inference-time revision using immediate tool feedback.Step-level optimization targets execution-failure steps to improve code generation.Source: §1 · pp. 2–3
  • TIRGen uses a generator to produce natural-language reasoning and a refiner to test code executability and refine reasoning with an external executor.Source: §1 · p. 2
  • THOR reports state-of-the-art results among comparable-size models across mathematical benchmarks, reduced inference overhead, and consistent gains on code benchmarks.The approach is reported to generalize across reasoning and non-reasoning models.Source: §1 · pp. 2–3

2 METHODOLOGY

THOR models tool-integrated reasoning as an interaction trajectory, builds tool-augmented data with TIRGen, and uses cold-start training, hierarchical RL, and inference-time self-correction. Its optimization combines final-answer feedback with execution feedback on individual tool calls.

2.1 PROBLEM FORMULATION

  • A reasoning trajectory alternates thoughts, tool actions, and execution observations, which inform subsequent thoughts and actions until the final answer.Source: §2.1 · p. 3

2.2 TIRGEN: TIR DATA GENERATION PIPELINE

  • TIRGen iteratively has a Generator produce reasoning steps and a Refiner convert code-solvable operations into executable calls whose results update the reasoning.The Refiner observes isolated steps, aligning synthesized data with the Generator’s reasoning preferences; filtering removes invalid or trivial code and problems solvable by pure chain-of-thought.Source: Figure 2 · p. 3, §2.2 · p. 4

2.3.1 COLD START

  • Cold-start supervised fine-tuning on TIRGen data teaches the model fundamental tool-invocation patterns, which reasoning models may otherwise struggle to use effectively.Source: §2.3.1 · pp. 4–5

2.3.2 HIERARCHICAL REINFORCEMENT LEARNING

  • Episode-level RL uses final-answer correctness, while step-level optimization targets code generation for failed tool calls using execution correctness.The framework is motivated by the finding that intermediate tool-call success strongly predicts final-answer correctness.Source: §2.3.2 · pp. 5–6
  • The episode-level objective uses GRPO with final-answer correctness as its reward, assigning one for a correct answer and zero otherwise.Source: §2.3.2 · p. 5

2.4 SELF-CORRECTION DURING INFERENCE

  • Inference-time self-correction backtracks after failed execution, regenerating the reasoning suffix and action for up to Ncorr attempts rather than recomputing the whole trajectory.The authors describe the additional computational cost as minimal compared with the total cost.Source: §2.4 · pp. 6–7

3 EXPERIMENTS

THOR improves mathematical reasoning across reasoning and non-reasoning models, while its inference strategy uses tool-execution feedback for candidate selection and self-correction. Ablations and code benchmarks further examine the contributions of its training stages and generalization.

3.1 DATASET

  • Evaluation covers school- and competition-level math benchmarks plus basic and competition-level code-generation tasks.Math evaluation uses an LLM judge; code evaluation uses HumanEval+, MBPP+, and LiveCodeBenchv6.Source: §3.1 · p. 7

3.2 COMPARISON WITH STATE-OF-THE-ART METHODS

  • THOR substantially improves mathematical reasoning on both reasoning and non-reasoning models, reaching state-of-the-art results among comparable-size models while remaining competitive with larger systems at low inference cost.The evaluation compares THOR with TIR and chain-of-thought methods across the two model settings.Source: §3.2 · p. 8

3.3 SELF-REWARDED INFERENCE ENHANCEMENT

  • Self-rewarded best-of-N selection uses code-execution pass rate without an external reward model and outperforms single-path reasoning, especially on challenging AIME tasks.Candidates in this comparison do not use self-correction.Source: §3.3 · p. 8, Table 2 · p. 8
  • Self-rewarded selection matches or surpasses PRM-based selection, while combining execution feedback with PRMs yields further improvements.At large N, PRMs can help distinguish candidates with identical code-pass rates; the two reward signals are reported as complementary.Source: §3.3 · pp. 8–9

3.4 ABLATION STUDY

  • TIRGen cold-start data improves pass@16 and code invocation ratio, mitigating degradation from out-of-distribution samples relative to other TIR datasets.The comparison includes Nemotron Long CoT and ReTool Short CoT data.Source: §3.4 · p. 9, Figure 4 · p. 9
  • Episode-level TIR reinforcement learning improves over the baseline and outperforms standard chain-of-thought reinforcement learning; adding step-level optimization yields further gains across most datasets.The comparison identifies step-level optimization as important for code generation in a TIR setting.Source: §3.4 · p. 10
  • Self-correction using intermediate-code feedback with up to four attempts produces substantial performance gains.The ablation links successful code generation and execution to the final outcome.Source: §3.4 · p. 10

3.5 GENERALIZATION ON CODE BENCHMARKS

  • THOR models consistently improve pass@1 on HumanEval+, MBPP+, and LiveCodeBenchv6 in zero-shot evaluation without code-generation fine-tuning.The paper reports gains across all three benchmarks.Source: §3.5 · p. 10, Figure 5 · p. 10

4 CONCLUSION

THOR addresses data construction, hierarchical optimization, and inference enhancement with TIRGen and tool-feedback-driven reinforcement learning and self-correction. The conclusion reports improvements across diverse models and benchmarks at low inference cost.

4 CONCLUSION

  • THOR combines TIRGen data construction, episode- and step-level optimization, and tool-feedback-driven self-correction, with substantial improvements across diverse models and benchmarks at low inference cost.The authors state that they will release code and datasets to facilitate reproduction.Source: §4 · p. 10, REPRODUCIBILITY STATEMENT · p. 10

A LLM USAGE

  • The paper reports that LLMs were used only to polish manuscript grammar and clarity; the authors conducted the research, experiments, and analyses.Source: A LLM USAGE · p. 15

B RELATED WORKS

  • The related-work section groups prior reasoning methods into tool-free and tool-integrated approaches.Source: B RELATED WORKS · p. 15

B.1 TOOL-FREE REASONING

  • Tool-free search methods incur substantial inference overhead and do not directly optimize the model’s internal reasoning policy; RL methods often rely on sparse episode-level rewards.The paper notes that sparse rewards impede efficient policy learning on long reasoning chains.Source: §B.1 · p. 15

B.2 TOOL-INTEGRATED REASONING

  • Prior tool-integrated methods face generalization constraints from synthetic-data scope and quality, prompt-trigger reliance, and unexplored step-level tool feedback.Source: §B.2 · pp. 15–16

C.1 STATISTICAL VALIDATION

The analysis finds a statistically significant association between code-execution success and answer correctness, supporting the paper’s motivation for using tool feedback.

C.1 STATISTICAL VALIDATION

  • χ2 = 336.3 and p = 4.09 × 10^-75 reject independence between code-execution success and answer correctness on the test set.Source: §C.1 · p. 16, Table 4 · p. 16
  • The authors conclude that successful intermediate tool calls are a strong predictor of final-answer correctness.Source: §C.1 · p. 16

C.2.1 ABLATION ON EXECUTION FAILURE REPAIR STRATEGIES

The ablation compares three repair strategies for self-correction after execution failures, distinguishing code errors from possible reasoning-step inconsistencies.

C.2.1 ABLATION ON EXECUTION FAILURE REPAIR STRATEGIES

  • The controlled ablation compares action-only repair, step-suffix repair, and full re-repair after execution failures.The setup distinguishes syntax or runtime errors from deeper logical inconsistencies in the preceding reasoning step.Source: §C.2.1 · p. 16, Table 5 · p. 16

C.2.2 ANALYSIS OF SELF-CORRECTION MECHANISMS

The explicit self-correction mechanism improves accuracy and code pass rate while reducing failed tool executions relative to the compared settings.

C.2.2 ANALYSIS OF SELF-CORRECTION MECHANISMS

  • Explicit self-correction reduces failed executions, increases code pass rate, and consistently improves accuracy.The comparison covers accuracy, code ratio, failed and successful tool calls, and overall code pass rate.Source: §C.2.2 · p. 16, Table 6 · p. 17
  • The evaluation compares an uncorrected baseline, THOR-7B with emergent self-correction, and THOR-7B with the explicit mechanism.Source: §C.2.2 · p. 16

C.2.3 INFERENCE COST ANALYSIS

THOR generates fewer tokens and reduces inference latency versus baseline models, with the reported efficiency results including self-correction overhead.

C.2.3 INFERENCE COST ANALYSIS

  • TIRGen data identifies redundant computation in reasoning chains, and THOR learns to produce concise solutions that use tools to simplify inference computations.Source: §C.2.3 · p. 17
  • THOR consistently generates fewer tokens, reduces total inference latency, and improves FPS over baseline models, including self-correction overhead.Table 7 reports token consumption across benchmarks; runtime and FPS are measured in a unified vLLM environment.Source: §C.2.3 · p. 17, Table 7 · p. 17

D.1 TIR DATA CONSTRUCTION & COLD START

The appendix details cold-start TIR data construction, hierarchical RL training, inference self-correction settings, and the prompts used across model types.

D.1 TIR DATA CONSTRUCTION & COLD START

  • TIRGen produces 29,217 short-CoT samples from Qwen2.5-Math-7B and 57,598 long-CoT samples from Qwen3-8B for cold start.The source questions come from DAPO17k, DeepMath103k, and Deepscaler40k; Figure 6 shows the distribution of code invocation counts.Source: §D.1 · p. 17, Figure 6 · p. 18
  • Cold-start models are full-parameter fine-tuned for one epoch, with maximum context lengths of 20,000 tokens for reasoning models and 4,096 for non-reasoning models.Training uses a global batch size of 256 and a fixed learning rate of 2 × 10−6.Source: §D.1 · pp. 17–18

D.2 HIERARCHICAL REINFORCEMENT LEARNING

  • RL training uses off-policy GRPO and dynamic data filtering; early failed code steps are downsampled so step-level queries do not exceed episode-level queries.This prioritizes overall problem-solving ability while code execution success is relatively low.Source: §D.2 · p. 18

D.3 INFERENCE WITH SELF-CORRECTION

  • Inference allows up to four self-correction attempts, with backtracking lengths of 100 tokens for non-reasoning models and 500 for reasoning models.Source: §D.3 · p. 18

D.4 PROMPT SETTING

  • The prompt appendix provides separate TIRGen Generator and Refiner prompts and distinct tool-integrated reasoning prompts for reasoning and non-reasoning models.These are shown in Figures 7–10.Source: §D.4 · p. 18

E LIMITATION AND FUTURE WORK

The authors limit their investigation to code-integrated mathematical reasoning and report case studies of tool use in both non-reasoning and reasoning models.

E LIMITATION AND FUTURE WORK

  • The study focuses on code integration for mathematical problem solving; other tools remain unexplored, and 32B- and 72B-scale models were not tested due to computational constraints.The authors report validation across model sizes from 1.5B to 8B and identify larger models and multi-tool optimization as future work.Source: E LIMITATION AND FUTURE WORK · pp. 18–19

F CASE STUDY

Loading 2509.13761v3…