PaperlayerView exact revision on arXiv

Source-linked AI summary

Satori: Reinforcement Learning with Chain-of-Action-Thought Enhances LLM Reasoning via Autoregressive Search

Maohao Shen, Guangtao Zeng, Zhenting Qi, Zhang-Wei Hong, Zhenfang Chen, Wei Lu, Gregory Wornell, Subhro Das, David Cox, Chuang Gan

arXiv:2502.02508v3cs.CLcs.AI

TL;DR

The paper asks whether search capabilities can be internalized in a single LLM rather than relying on externally guided, two-player test-time search. It introduces COAT and trains Satori through format tuning followed by reinforcement-learning self-improvement, reporting superior mathematical reasoning and strong out-of-domain transfer. The approach is presented as enabling autoregressive search without external guidance.

  • Problem

    Existing test-time search commonly uses extensive sampling guided by an external verifier, leaving search capabilities outside a single LLM.

  • Method

    The method combines COAT reasoning with small-scale format tuning and large-scale reinforcement-learning self-improvement using Restart and Explore techniques.

  • Results

    Satori shows superior in-domain mathematical reasoning, outperforms the instruct model built on the same base model, and transfers strongly to out-of-domain tasks.

  • Takeaways & Limitations

    Satori demonstrates that a single 7B LLM can perform autoregressive search without external guidance while using minimal supervision and large-scale self-improvement.

  • Takeaways & Limitations

    Long-horizon reasoning with sparse terminal rewards makes learning difficult because many correct steps are required before reward is received and failures restart from the initial state.

Abstract

from arXiv · show

Large language models (LLMs) have demonstrated remarkable reasoning capabilities across diverse domains. Recent studies have shown that increasing test-time computation enhances LLMs' reasoning capabilities. This typically involves extensive sampling at inference time guided by an external LLM verifier, resulting in a two-player system. Despite external guidance, the effectiveness of this system demonstrates the potential of a single LLM to tackle complex tasks. Thus, we pose a new research problem: Can we internalize the searching capabilities to fundamentally enhance the reasoning abilities of a single LLM? This work explores an orthogonal direction focusing on post-training LLMs for autoregressive searching (i.e., an extended reasoning process with self-reflection and self-exploration of new strategies). To achieve this, we propose the Chain-of-Action-Thought (COAT) reasoning and a two-stage training paradigm: 1) a small-scale format tuning stage to internalize the COAT reasoning format and 2) a large-scale self-improvement stage leveraging reinforcement learning. Our approach results in Satori, a 7B LLM trained on open-source models and data. Extensive empirical evaluations demonstrate that Satori achieves state-of-the-art performance on mathematical reasoning benchmarks while exhibits strong generalization to out-of-domain tasks. Code, data, and models are fully open-sourced.

1. Introduction

The paper targets the limitation of two-player test-time search by internalizing autoregressive search in a single LLM through COAT reasoning and two-stage training. Satori combines format tuning with reinforcement-learning self-improvement and reports strong mathematical and out-of-domain performance.

  • Motivation: Test-time search often relies on external feedback, creating two-player systems that do not internalize search capabilities in a single LLM.These methods use extensive sampling of complete solutions or intermediate steps guided by auxiliary reward models.
  • Approach: COAT enables autoregressive search through an extended reasoning process with self-reflection and self-exploration of new strategies.The mechanism allows LLMs to take various meta-actions during problem solving.
  • Approach: The training paradigm first uses small-scale format tuning to internalize COAT, then large-scale reinforcement learning with Restart and Explore techniques for self-improvement.The two stages are designed to combine reasoning-format acquisition with reinforcement-learning improvement.
  • Results: Satori is a 7B LLM trained on open-source base models and mathematical data, achieving superior performance on in-domain and out-of-domain tasks.The paper presents this outcome as the central result of the approach.
  • Results: Satori performs mathematical reasoning without external guidance as a single LLM and achieves this with minimal supervision and large-scale self-improvement.The paper characterizes this capability as autoregressive search.
  • Results: Satori outperforms the instruct model built on the same base model on in-domain mathematical reasoning tasks and transfers strongly to out-of-domain tasks.The paper describes both effectiveness on mathematics and generalizability beyond the training domain.

2. Related Work

The related work spans post-training for reasoning and methods that equip LLMs with search abilities. Satori is positioned as an RL-based, minimally supervised alternative to teacher distillation and prompting-based search.

  • Post-training LLMs for reasoning: Recent reasoning work commonly distills knowledge from stronger teacher models to enhance open-source LLMs.The paper contrasts this dominant concurrent approach with Satori’s reinforcement-learning perspective.
  • Post-training LLMs for reasoning: Satori differs from these approaches by using reinforcement learning and requiring minimal supervision, with only 10K format-tuning samples.This comparison is made in the context of concurrent work on open-source reasoning models.
  • Post-training LLMs for reasoning: Instruction-tuning approaches construct high-quality datasets but incur expensive annotation costs.The related-work discussion also identifies self-improvement as a response to these costs.
  • Enabling LLMs with searching abilities: Prompting-based methods guide search through error correction and alternative paths, but the paper states that they cannot fundamentally enhance LLM reasoning abilities.The discussion also notes reported difficulties with self-correction.

3. Preliminaries

The paper formulates mathematical problem solving as sequential decision making for a language-model policy trained with reinforcement learning. Its overview places format tuning before self-improvement in the Satori framework.

  • Problem formulation: The model πθ generates a solution ỹ matching ground truth y* for a mathematical problem prompt x.The formulation treats x, y, and y* as token sequences from a predefined dictionary.
  • Satori training framework: Satori’s training framework first learns COAT through imitation learning on small demonstration trajectories, then self-improves through large-scale reinforcement learning.The figure presents format tuning and self-improvement as consecutive stages.
  • Reinforcement learning: In reinforcement learning, the language model acts as a policy that selects actions across states, receives rewards, and continues until a terminal state.A trajectory is the sequence of states and actions, and training maximizes expected cumulative reward.
  • Reinforcement learning: Reinforcement learning optimizes expected cumulative rewards over the trajectory, whose length is denoted by L.The supplied formulation expresses the cumulative reward across steps l = 1 through L.

4. Method

The method equips an LLM with COAT reasoning, allowing it to continue, reflect, or explore during multi-step problem solving. It combines format tuning with reinforcement-learning-based restart-and-explore training for self-improvement.

  • 4. Method: COAT adds continue, reflect, and explore meta-actions to guide whether the model extends, checks, or revises its reasoning.Conventional CoT is a special case restricted to continuation.
  • 4. Method: The model treats reasoning as sequential decision-making, appending each generated step to the context until producing a final answer.The policy generates actions from an initial problem state and transitions to updated states.
  • 4.2. Self-improvement via Reinforcement Learning: RAE addresses long-horizon sparse-reward learning by restarting from intermediate correct or incorrect trajectories and triggering reflection.It broadens initial states beyond original problem prompts and encourages exploration without redundant recomputation.
  • 4. Method: Format tuning uses synthetic expert COAT trajectories to teach the base LLM how to follow the reasoning format and use meta-actions.A multi-agent framework uses a generator, critic, and reward model to construct demonstrations, with behavior cloning used for training.
  • 4.2. Self-improvement via Reinforcement Learning: PPO trains the format-tuned policy on original problems and partial trajectories, while reflection bonuses reward correcting errors and penalize unnecessary revisions.The procedure constructs restart buffers from correct and incorrect trajectories and adds a reflect token to intermediate states.

5. Experiment

Satori-Qwen-7B is evaluated on mathematical and diverse out-of-domain reasoning benchmarks against same-scale and larger baselines. It outperforms the same-base instruct model and shows continued gains after iterative self-improvement.

  • Main Results on Math Domain: Satori-Qwen-7B achieves state-of-the-art performance across five mathematical benchmarks, while Round 2 further improves performance on hard tasks.The comparison is against Qwen-2.5-Math-7B-Instruct, which uses the same base model.
  • Benchmark and Evaluation: The evaluation covers five mathematical benchmarks and diverse logical, code, commonsense, tabular, and domain-specific reasoning tasks.Math evaluation uses greedy decoding without tools and reports zero-shot pass@1 accuracy.
  • Main Results on Math Domain: Satori-Qwen-7B outperforms Qwen-2.5-Math-7B-Instruct despite using the same base model and significantly less supervision.Its training relies more on self-improvement and more reinforcement-learning data.
  • Out-of-Domain Transferability: Despite training only on math datasets, Satori-Qwen-7B transfers reasoning capabilities to diverse out-of-domain benchmarks and outperforms the same-scale instruct baseline.On BoardgameQA, it surpasses all baseline models of the same scale.
  • Results on Iterative Self-improvement: Round-2 training produces continued performance gains across most in-domain and out-of-domain benchmarks.The result supports iterative self-improvement as a way to further increase reasoning performance.

6. Analysis

Satori’s analysis examines how COAT reasoning and reinforcement learning support self-correction, test-time scaling, and generalization. Ablations indicate that COAT and large-scale RL outperform classical CoT and large-scale format tuning across supported evaluations.

  • Satori Exhibits Self-correction Capability: Satori-Qwen demonstrates stronger self-correction than Satori-Qwen-FT, and this capability extends from MATH500 and Olympiad to MMLUProSTEM.The analysis measures positive and negative answer changes after self-reflection.
  • RL Enables Satori with Test-time Scaling Behavior: More RL training-time compute increases policy accuracy and average generated-token length, indicating that Satori learns to allocate more reasoning time.Response length temporarily decreases early in training before increasing as reflection capabilities emerge.
  • RL Enables Satori with Test-time Scaling Behavior: Across MATH difficulty levels, Satori-Qwen uses more test-time compute than Satori-Qwen-FT and achieves consistent performance improvements.The comparison jointly evaluates response length and accuracy.
  • Large-scale FT v.s. Large-scale RL: Satori-Qwen trained with 10K FT data and 300K RL data outperforms the same base model trained with 300K FT data across math and out-of-domain benchmarks.Large-scale FT does not match the performance of small-scale FT combined with large-scale RL and performs significantly worse out of domain.
  • Distillation Enables Weak-to-Strong Generalization: Distilled models outperform format-tuned models when weaker Llama-3.1-8B and Granite-3.1-8B bases are trained from Satori-generated data.The study uses 240K synthetic examples for each comparison.

7. Concluding Remarks

The Satori training framework combines small-scale format tuning with large-scale reinforcement learning. The paper presents format tuning as a warm-up for internalizing the reasoning format and RL as important for intrinsic reasoning abilities.

  • Concluding Remarks: Small-scale format tuning internalizes a specific reasoning format, while large-scale reinforcement learning plays a crucial role in incentivizing intrinsic reasoning abilities.The conclusion identifies this two-stage framework as having significant potential for enhancing LLM reasoning capabilities.

A. Satori’s Demo Examples

The demonstrations show Satori verifying intermediate reasoning, detecting errors, and exploring alternative solutions across mathematical, logical, commonsense, scientific, economic, and tabular tasks. These examples illustrate self-reflection and solution revision in both math and out-of-domain settings.

  • Mathematical Reasoning: In mathematical examples, Satori verifies intermediate steps, identifies errors, and proposes alternative or simpler solutions.The examples include continuing after verification, switching solutions, and exploring a simpler approach.
  • Mathematical Reasoning: The vector example determines the largest possible projection magnitude by using the projection formula and maximizing |cos θ| at 1.The reasoning identifies θ = 0° as the maximizing case.
  • Logical Reasoning: In a boardgame reasoning example, Satori verifies an earlier solution, recognizes that it remains erroneous, and proposes an alternative correct solution.The example includes preferred rules and a subsequent corrected solution.
  • Out-of-domain Reasoning: Out-of-domain examples show Satori identifying potential mistakes and initiating another solution in commonsense, chemistry, and economics tasks.The supplied examples cover StrategyQA, MMLUPro Chemistry, and MMLUPro Economics.
  • Out-of-domain Reasoning: Satori verifies early intermediate steps, recognizes erroneous previous solutions, and proposes alternatives in additional out-of-domain demonstrations.These examples include domain-specific and broader reasoning cases.
  • Tabular Reasoning: A tabular reasoning example shows Satori selecting peaks, summing their prominences, and reporting an average prominence of 2606.25.The example states the selected peaks and calculates the average from the summed prominences.

B.1. Post-training LLMs for Reasoning

Prior work improves LLM reasoning through CoT data, self-training, prompting, and trajectory-level search, but these approaches rely on external supervision or face limitations in self-correction. The paper’s demonstration synthesis uses critic feedback and generator self-refinement to support search-oriented training.

  • Post-training LLMs for Reasoning: Post-training datasets use human-annotated or distilled CoT-like reasoning chains, while self-improvement trains models on data generated by themselves.Self-training reduces reliance on external resources but raises concerns about model collapse.
  • Post-training LLMs for Reasoning: CoT prompting improves complex-task reasoning through intermediate steps but lacks mechanisms to correct errors once committed.Later approaches add forward exploration, backtracking, and alternate paths through prompting or training.
  • Post-training LLMs for Reasoning: Trajectory-level search trains models to identify errors in complete previous responses and produce improved responses using human-annotated revisions or model-generated data.This line of work differs from methods that search within intermediate reasoning processes.
  • Post-training LLMs for Reasoning: The demonstration synthesis process samples multiple reasoning trajectories, obtains critic feedback, and uses that feedback for generator self-refinement.The critic verifies correct final answers and identifies mistakes or proposes alternatives for incorrect trajectories.

C. Details about Data Synthesis Framework

The data synthesis framework generates multiple reasoning trajectories, uses critic feedback and generator refinement, and filters high-reward paths into COAT demonstrations. It combines verification, correction, exploration, and meta-action formatting across several trajectory types.

  • Trajectory generation: The generator samples K initial reasoning trajectories and separates them by whether their final answers match the ground truth.Correctness determines whether trajectories receive verification or error-correction feedback.
  • Critic and refinement: For incorrect trajectories, the critic identifies the first error, produces a reflection and exploration hint, and the generator continues from that feedback.The continuation conditions on the original prefix plus the critic’s reflection and hint.
  • Critic and refinement: For correct trajectories, the critic verifies a randomly selected intermediate step before the generator produces subsequent reasoning steps.This creates verification-oriented trajectories rather than mistake-correction trajectories.
  • Trajectory generation: The procedure repeats for up to m iterations, checks the final answer after refinement, and inserts meta-action tokens indicating each reasoning step’s action type.The resulting feedback and refinements are used to synthesize demonstration trajectories.
  • Trajectory filtering: Top-k trajectories are selected by reward score to construct a diverse synthetic dataset containing short-cut CoT paths and multi-round self-reflection paths.The selected dataset includes both simplified reasoning and more complex reflective search.

D.1.2. SYNTHETIC DATA EXAMPLES

The synthetic examples illustrate COAT behaviors through verification, mistake correction, and exploration across mathematical and reasoning problems. They show how reflection changes an initially incorrect or incomplete solution before a final answer is produced.

  • Type I: No reflection: Type I contains a solution trajectory without critic feedback or reflection actions.It represents the no-reflection baseline among the synthetic trajectory types.
  • Type II-I: Verify correctness: The boat example verifies the arithmetic for five couples, computes 10 people, and rounds 10/4 up to 3 boats.The final distribution uses boats carrying 4, 4, and 2 people.
  • Trajectory taxonomy: The examples distinguish trajectories with verification, correction, and exploration actions, culminating in a Type III-II category for two rounds of self-reflection and self-exploration.The supplied Type III-II material also includes a reflection-bonus statement used in the training discussion.
  • Type II-II: Verify then correct: The Type II-II example reflects on a candidate perfect square, rejects 9801 because of its digit interpretation, and explores 961 instead.The revised candidate is presented as 31^2 and satisfies the stated conditions in the example.
  • Type III-I: Correct mistake: The Type III-I example corrects an interval solution by removing 0 because the variable belongs to the positive natural numbers.The corrected set A is {1, 2, 3}, which is then intersected with B.

D.4. Evaluation Details.

The evaluation covers mathematical reasoning and several out-of-domain reasoning settings using zero-shot pass@1 under greedy decoding. It includes datasets spanning arithmetic, competition mathematics, logic, tabular reasoning, code reasoning, and multi-hop question answering.

  • Mathematical benchmarks: MATH500 contains 500 uniformly sampled MATH test problems, while GSM8K contains 1319 test problems requiring multi-step elementary arithmetic reasoning.These datasets provide different scales and difficulty profiles for mathematical evaluation.
  • Out-of-domain benchmarks: The out-of-domain evaluation includes TableBench, MMLUProSTEM, FOLIO, BoardgameQA, CRUXEval, and StrategyQA.These benchmarks test tabular, STEM, logical, code, and multi-hop reasoning capabilities.

E.1. Ablation on Reflection Bonus

The ablations test the reflection bonus, restart-and-explore mechanism, and preference bonus within reinforcement learning. Removing these components produces degradation or suboptimal training behavior in the reported experiments.

  • Reflection bonus: Setting the reflection bonus to r_reflect = 0 slightly degrades performance on AMC2023 and AIME2024 compared with r_reflect = 0.5.The degradation is reported specifically on the challenging competition benchmarks.
  • Restart and explore: Removing RAE causes performance degradation, indicating that restart and explore is important to the reported reinforcement-learning training.RAE is described as encouraging self-reflection and diversifying initial states.
  • Preference bonus: Removing the Outcome Reward Model causes performance degradation, supporting a hybrid reward signal that combines rule-based and dense ORM-based rewards.The preference bonus is introduced to mitigate sparse rewards.
  • Restart buffer: Offline restart-buffer sampling is preferred because online collection produces an imbalance dominated by correct trajectories.Offline sampling balances intermediate states from correct and incorrect trajectories.
Loading 2502.02508v3…