PaperlayerView exact revision on arXiv

Source-linked AI summary

Reasoning with Reinforced Functional Token Tuning

Kongcheng Zhang, Qi Yao, Baisheng Lai, Jiaxing Huang, Wenkai Fang, Dacheng Tao, Mingli Song, Shunyu Liu

arXiv:2502.13389v1cs.AI

TL;DR

Smaller LLMs face costly, difficult-to-scale access to high-quality reasoning data, and direct tree search can produce homogeneous paths across a large search space. RFTT trains learnable functional tokens using self-generated data, then uses them for autonomous tree exploration. Experiments report stronger mathematical reasoning results, with iterative updates and inference-time rollouts also improving performance.

  • Problem

    Training reasoning in smaller LLMs requires high-quality data, but obtaining it from stronger models or human annotators is costly and difficult to scale, while direct tree search faces homogeneous paths and a large search space.

    Source: §1 · p. 1

  • Method

    RFTT learns functional tokens through supervised fine-tuning on self-generated token-annotated reasoning paths, then samples those tokens for autonomous tree exploration during online reinforcement learning.

    Source: §1 · p. 2

  • Results

    LLaMA-3.1-8B-Instruct reaches 60.2% Pass@1 on MATH with RFTT, up from 50.6% with CoT prompting, while Qwen-2.5-7B-Instruct reaches 79.6% versus 80% for Qwen-2.5-14B-Instruct.

    Source: §4.3 · p. 6

  • Takeaways & Limitations

    RFTT's performance continues to improve with more inference-time search rollouts, and iterative updates show continuous enhancement of the initial policy.

    Source: §1 · p. 2, §4.3 · p. 8

    The work focuses primarily on mathematical reasoning, so its effectiveness on broader reasoning domains needs further exploration.

    Source: §5 · p. 10

Abstract

from arXiv · show

In this work, we propose Reinforced Functional Token Tuning (RFTT), a novel reinforced fine-tuning framework that empowers Large Language Models (LLMs) with self-play learn-to-reason capabilities. Unlike prior prompt-driven reasoning efforts, RFTT embeds a rich set of learnable functional tokens (e.g., <analyze>, <verify>, <refine>) directly into the model vocabulary, enabling chain-of-thought construction with diverse human-like reasoning behaviors. Specifically, RFTT comprises two phases: (1) supervised fine-tuning performs prompt-driven tree search to obtain self-generated training data annotated with functional tokens, which warms up the model to learn these tokens for reasoning; and (2) online reinforcement learning further allows the model to explore different reasoning pathways through functional token sampling without relying on prompts, thereby facilitating effective self-improvement for functional reasoning. Extensive experiments demonstrate the superiority of the proposed RFTT on mathematical benchmarks, significantly boosting Qwen-2.5-7B-Instruct (70.6% to 79.8%) and LLaMA-3.1-8B-Instruct (32.2% to 60.2%) on the MATH dataset. Moreover, the performance of RFTT consistently improves with more search rollouts at inference time. Our code is available at https://github.com/sastpg/RFTT.

1. Introduction

Learning to reason in smaller LLMs requires high-quality training data, but relying on stronger models or human annotators is costly and difficult to scale. RFTT trains models to reason with functional tokens through self-generated data and token-guided exploration.

1. Introduction

  • High-quality reasoning data is especially challenging for smaller LLMs, while data generated by stronger models or human annotators incurs substantial costs and limited scalability.Self-play methods aim to generate rationales without human-curated training data.Source: §1 · p. 1
  • Direct tree-search methods face homogeneous reasoning paths and an exponentially growing candidate space from the combinatorial vocabulary.rStar uses functional prompts to diversify and constrain search, but leaves the reasoning capabilities uninternalized and requires iterative traversal with redundant inference calls.Source: §1 · pp. 1–2
  • RFTT embeds learnable functional tokens in the model vocabulary and uses SFT on self-generated token-annotated paths followed by RL for autonomous token-guided tree exploration.The SFT phase builds reasoning trees with functional prompt-guided MCTS; RL samples functional tokens to expand trees and reinforce high-value pathways.Source: §1 · p. 2

2. Related Works

Related work covers prompting and training for mathematical reasoning, tree-structured search, and self-improvement from generated paths. RFTT is framed as learning functional tokens to guide reasoning-tree exploration.

Mathematical Reasoning

  • Mathematical-reasoning work spans inference-time prompting and approaches that train intrinsic reasoning ability through pretraining and post-training methods.Prompting includes chain-of-thought, multi-path reasoning, and divide-and-conquer; post-training approaches include supervised fine-tuning, preference optimization, and reinforcement learning.Source: §2 · p. 2

Tree-Search Reasoning

  • Tree-search methods explore reasoning trajectories, but direct approaches struggle with homogeneous paths and large search spaces; rStar uses external functional prompts to diversify and constrain search.RFTT instead introduces learnable functional tokens during training to support token-guided exploration and improve intrinsic reasoning ability.Source: §2 · p. 2

LLM Self-Improvement

  • Self-improvement methods use correct-answer rejection sampling, MCTS, or online random sampling, yet may not fully explore reasoning space, especially for smaller models.RFTT uses function tokens for automatic exploration through token-guided tree traversal.Source: §2 · p. 2

3. Methodology

RFTT models complex reasoning as tree search over intermediate steps and trains functional-token reasoning in two phases: supervised warmup followed by online reinforcement learning.

3.1. Problem Formulation

  • Functional tokens make actions more structured than vocabulary-level sampling, which the paper says would create a very large search space and reduce efficiency.The reasoning task is formulated as a tree whose root is the problem and whose paths are potential solution trajectories.Source: §3.1 · p. 3, Phase 2: Online Reinforcement Learning · p. 5

3.3. Functional Monte Carlo Tree Search

  • Functional prompts guide SFT tree search, while learned functional tokens guide RL exploration without prompts, addressing the homogeneity of direct-search trajectories.The paper describes functional behaviors such as clarifying, analyzing, verifying, refining, and outputting an answer.Source: §3.3 · pp. 3–4, §3.2 · p. 3

Step 1: Functional Tree Search

  • MCTS selects nodes with UCT, expands them using functional actions, simulates to an answer or depth limit, then backpropagates trajectory outcomes.Repeated rollouts yield multiple reasoning trajectories for each question.Source: Step 1: Functional Tree Search · p. 4

Step 2: Cross Verification

  • Cross-verification selects a correct trajectory with the highest average process reward and a nearby wrong trajectory, then prompts the model to identify the error.The wrong path is considered when it lies within a predetermined number of distinct steps from the correct path.Source: Step 2: Cross Verification · pp. 4–5

Step 3: Branch Merging

  • Branch merging joins reasoning paths with functional tokens to create structured SFT examples featuring self-verification and self-correction.Each intermediate step is concatenated with its corresponding functional token.Source: Step 3: Branch Merging · p. 5

3.4. Reinforced Functional Token Tuning

Phase 1: Supervised Fine-Tuning

  • During SFT, the model learns functional tokens embedded in its vocabulary from question–reasoning-path pairs, without relying on prompts.These tokens are treated as new special tokens.Source: Phase 1: Supervised Fine-Tuning · p. 5

Phase 2: Online Reinforcement Learning

  • During online RL, the policy samples functional tokens to expand reasoning trees, evaluates rollouts, and updates its parameters using process rewards or outcome rewards.The paper states that functional tokens can reduce the tree-search space.Source: Phase 2: Online Reinforcement Learning · p. 5

4. Experiments

RFTT is evaluated on mathematical benchmarks using self-generated training data and comparisons with fine-tuning, prompting, and tree-search baselines. The experiments report benchmark gains, benefits from inference-time search and iterative updates, and ablation and training-dynamics findings.

4.1. Experimental Setups

  • Training uses the MATH training set; evaluation measures Pass@1 on five mathematical benchmarks, using MATH-500 for MATH and full test sets elsewhere.Experiments use three open-source models under 10B parameters and compare RFTT with stronger models, fine-tuning, in-context learning, and tree-search methods.Source: §4.1 · p. 6

4.2. Implementation Details

  • Implementation uses MCTS to collect reasoning paths, then trains with functional-token SFT and functional-token-guided MCTS during RL.Search uses depth 15 and 16 rollouts per question; approximately 1,000 SFT examples are collected, followed by 10 SFT epochs and RL sampling of 16 paths per question.Source: §4.2 · p. 6

4.3. Main Results

  • 60.2% Pass@1 on MATH for LLaMA-3.1-8B-Instruct after RFTT exceeds its 50.6% CoT result; Qwen-2.5-7B-Instruct reaches 79.6% versus Qwen-2.5-14B-Instruct at 80%.RFTT also exceeds ReFT by an average margin of 5% across benchmarks; results remain effective on unseen problem sets despite training only on MATH data.Source: §4.3 · pp. 6–7, Table 1 · p. 7
  • More inference-time rollouts improve mathematical reasoning, with slower gains on MATH after 16 rollouts and steady gains on AMC; trained policies exceed o1-preview at 8 MATH and 20 AMC rollouts.The rollout comparison includes the policy model's Pass@1 accuracy at one rollout.Source: §4.3 · p. 7, Figure 3 · p. 8
  • RFTT dynamically selects high-value functional tokens and reports higher search-answer accuracy with substantially reduced time complexity versus other tree-search methods.The comparison measures average time per rollout and the proportion of questions answered correctly; RFTT outperforms rStar and LLaMA-Berry.Source: §4.3 · pp. 7–8
  • A second SFT round using correct paths from the first RL round yields continued policy improvement, with a 7.9% MATH gain versus 1.7% for ReST-MCTS*.The comparison reports iterative-update results on GSM8K and MATH; ReST-MCTS* uses an additional reward model to guide search.Source: Table 3 · p. 7, §4.3 · p. 8

4.4. Ablation Studies

  • Ablations report better performance with SFT warmup than pure RL, MCTS than random sampling, and PRM than ORM.Warmup gains are 5.7% for Qwen-2.5-7B-Instruct and 7.7% for LLaMA-3.1-8B-Instruct; MCTS yields mean 73.92 versus 72.34 across five benchmarks, while PRM averages about 2% higher than ORM on both model families.Source: §4.4 · pp. 8–9, Table 4 · p. 8

4.5. Evolution Analysis

  • During RL, training-set batch accuracy improves; with PRM, response length rises sharply in the first 50 steps and then more gradually.The analysis also reports improving search time, self-correction and diverse-solution behaviors with increasing search depth, and steadily improving average rewards and KL divergence.Source: §4.5 · p. 9, Figure 4 · p. 9

5. Conclusions

RFTT embeds learnable functional tokens to internalize human-like reasoning, while its prompt appendix specifies the behaviors those tokens guide. The work reports gains and inference-time search scalability, but focuses on mathematics and leaves broader reasoning domains for future exploration.

5. Conclusions

  • RFTT embeds learnable functional tokens in the model vocabulary to internalize human-like reasoning without external prompts.The framework uses self-play reasoning and reports performance gains with inference-time search.Source: §5 · pp. 9–10
  • The work focuses primarily on mathematical reasoning; effectiveness on broader reasoning domains remains to be explored.Source: §5 · p. 10

A. Functional Prompts

  • The <clarify>, <direct answer>, and <verify> prompts respectively restate the problem, complete a solution, and check subsequent steps for mistakes.Source: A. Functional Prompts · pp. 13–14
  • A clarification example restates the cylinder problem while preserving its given conditions and goal.Source: A. Functional Prompts · p. 13
  • The <analysis>, <subquestion>, and <next step> prompts guide insight, problem decomposition, and generation of an appropriate next reasoning step.The next-step prompt calls for the first step when no prior steps exist and continuation otherwise.Source: A. Functional Prompts · p. 13

B. Reasoning Cases

The reasoning cases illustrate functional-token sequences that support verification, reflection, and correction. The examples include checking a solution and revising an inverse-function derivation after identifying an error.

B. Reasoning Cases

  • The cases illustrate diverse reasoning behaviors, including double-checking, self-reflection, and self-correction.Source: B. Reasoning Cases · p. 14
  • A self-verification case reviews solution steps against a correct solution and checks for mistakes.Source: B. Reasoning Cases · pp. 14–15
  • Another case uses clarification, analysis, and subsequent reasoning steps to solve a complex-number problem about an equilateral triangle.Source: B. Reasoning Cases · p. 14
  • In the inverse-function example, the model notices an error in isolating y, refines the expression, and returns b = 2.Source: B. Reasoning Cases · pp. 15–16

C. Detailed MCTS

The detailed MCTS description treats reasoning steps as actions and generated trajectories as tree states. Search selects and expands actions, simulates continuations, and backpropagates rewards and visit counts.

C. Detailed MCTS

  • MCTS balances exploration and exploitation through selection, expansion, simulation, and backpropagation.Source: C. Detailed MCTS · p. 17
  • Selection uses the policy over functional tokens and UCT; expansion generates an unselected action subject to token dependencies.Simulations continue to a terminal state, after which backpropagation updates Q and N along the path to the root.Source: C. Detailed MCTS · p. 17
  • In LLM sampling, each action generates a reasoning step, which is appended to prior steps to form the next state.The state transition is deterministic, so the value can be stored at the resulting state node.Source: C. Detailed MCTS · p. 17

D. Experimental Detials

The experiments use selected MATH training questions for SFT and RL, with more Level 4 and Level 5 problems in training. Evaluation combines reported metrics from model technical reports with zero-shot API results when metrics are unavailable.

D. Experimental Detials

  • The SFT and RL phases use 1,000 and 3,994 MATH training questions, respectively, with an increased proportion of Level 4 and Level 5 problems.Source: D. Experimental Detials · p. 18
  • For stronger LLMs, evaluation uses relevant technical-report metrics and obtains unreported dataset accuracies through official APIs in zero-shot settings.The CoT evaluation prompt template is provided separately.Source: D. Experimental Detials · p. 18

E.1. How do Functional Tokens Benefit?

MCTS training faces exponential growth in candidate reasoning paths and difficulty training a fine-grained value model. Functional tokens constrain and diversify search while their changing probabilities can guide it.

E.1. How do Functional Tokens Benefit?

  • MCTS faces exponentially growing candidate reasoning paths in the combinatorial vocabulary space and difficulty training a fine-grained value model.Source: §E.1 · p. 18
  • Functional tokens diversify node exploration while constraining the search space, addressing MCTS’s expanding candidate paths.Their rich semantics provide distinct functional choices during search.Source: §E.1 · p. 18
  • Functional-token probabilities adjust dynamically during training and can guide search, potentially reducing the need for an extra value model.Functional tokens also divide reasoning into fine-grained steps that can integrate with PRM.Source: §E.1 · p. 18

E.2. Is SFT Warmup with Functional Tokens Needed?

Pure RL produced self-reflection at a low rate on Qwen-2.5-7B-Instruct and none observed on LLaMA-3.1-8B-Instruct. The authors argue that functional-token SFT warmup supports faster convergence and stable training.

E.2. Is SFT Warmup with Functional Tokens Needed?

  • About 1.5% of Qwen-2.5-7B-Instruct pure-RL cases showed self-reflection, while none were observed for LLaMA-3.1-8B-Instruct.The passage attributes smaller models’ difficulty reinforcing sophisticated reasoning patterns to their inherent preference for pretrained syntactic patterns.Source: §E.2 · p. 18
  • SFT warmup with functional tokens can play an important role in faster convergence and stable training.The authors contrast this with pure RL, where self-reflection was rare or unobserved in the reported models.Source: §E.2 · p. 18
  • Figure 5 presents statistical information about the training dataset used in the two-phase training.Source: Figure 5 · p. 18
  • Figure 6 presents training curves for RFTT without SFT warmup and without MCTS during RL.Source: Figure 6 · p. 19

E.3. Discussion of Response Lengths.

The response-length discussion links changes under RFTT without SFT warmup to KL divergence, while contrasting them with RFTT using PRM, whose lengths do not show the same downward trend.

E.3. Discussion of Response Lengths.

  • As response lengths dropped and then rebounded in RFTT without SFT warmup, KL divergence rose during t1 and fell during t2; the rebound was attributed primarily to the KL penalty, not sophisticated behavior exploration.The lengths remained below those of the original pretrained model.Source: §E.3 · p. 19
  • RFTT without PRM did not show a downward response-length trend, unlike the observed pattern discussed for RFTT without SFT warmup.Source: §E.3 · p. 19
  • The authors suggest functional-token-guided tree search provides diverse exploration, while PRM assigns fine-grained rewards to intermediate functional tokens and mitigates KL-penalty exploration limits.They identify functional tokens as a bridge between MCTS and PRM for learning to reason.Source: §E.3 · p. 19

Related papers

Loading 2502.13389v1…