SMART: Self-Learning Meta-strategy Agent for Reasoning Tasks

## Abstract

Abstract Tasks requiring deductive reasoning, especially those involving multiple steps, often demand adaptive strategies such as intermediate generation of rationales or programs, as no single approach is universally optimal.
While Language Models (LMs) can enhance their outputs through iterative self-refinement and strategy adjustments, they frequently fail to apply the most effective strategy in their first attempt. This inefficiency raises the question: Can LMs learn to select the optimal strategy in the first attempt, without a need for refinement? To address this challenge, we introduce SMART ( S elf-learning M eta-strategy A gent for R easoning T asks), a novel framework that enables LMs to autonomously learn and select the most effective strategies for various reasoning tasks. We model the strategy selection process as a Markov Decision Process and leverage reinforcement learning-driven continuous self-improvement to allow the model to find the suitable strategy to solve a given task. Unlike traditional self-refinement methods that rely on multiple inference passes or external feedback, SMART allows an LM to internalize the outcomes of its own reasoning processes and adjust its strategy accordingly, aiming for correct solutions on the first attempt.
Our experiments across various reasoning datasets and with different model architectures demonstrate that SMART significantly enhances the ability of models to choose optimal strategies without external guidance (+15 points on the GSM8K dataset). By achieving higher accuracy with a single inference pass, SMART not only improves performance but also reduces computational costs for refinement-based strategies, paving the way for more efficient and intelligent reasoning in LMs. https://github.com/kumar-shridhar/SMART/

## 1 Introduction

When people first encounter complex reasoning tasks, such as solving mathematical problems, they often make mistakes or approach them inefficiently []. However, with experience, humans tend to improve their performance by replacing ineffective or incorrect strategies with more effective ones, using a mix of strategies tailored to the specific task [, , , , inter alia].

Language Models (LMs) similarly struggle with reasoning tasks, sometimes producing incoherent results [, , ]. A common remedy is to resample the output, a process known as *refinement*. This refinement may involve reusing the same reasoning approach [] or adopting an entirely new one []. In addition, providing feedback on initial results has proven beneficial during resampling [, , , , inter alia].
This raises a critical question: *Can LMs be taught to optimize their choice of reasoning strategy for specific tasks overtime on the first trial, much like humans do?*

To address this question, we propose a novel framework called SMART (Self-learning Meta-strategy Agent for Reasoning Tasks), which allows LMs to learn optimal strategy selection through a continuous self-learning approach. We model the task of identifying the optimal strategy as a *Markov Decision Process* (MDP) [, ], where the agent (LM) starts with its pre-trained knowledge and iteratively improves its performance by learning from its own outputs and strategy choices. By integrating the LM’s reasoning abilities with reinforcement learning-driven self-improvement, the agent can simulate different reasoning strategies, evaluate their effectiveness based on past outcomes, and adjust its strategy choice accordingly.

Our approach differs from traditional methods by focusing on iterative reward-based learning, which encourages the agent to produce the correct inference on the first attempt without resampling. This not only improves cost efficiency - only one sampling step is required during inference - but also results in a more generalizable model capable of adapting its strategy selection based on the specific task. We validate SMART on a variety of reasoning datasets and LM architectures and show that our method significantly improves the ability of LMs to select optimal strategies on the first try, outperforming baseline models that rely on traditional self-refinement techniques in both accuracy and computational efficiency.
On three mathematical datasets (GSM8K [], SVAMP [], ASDiv []) over three LLM agents (Llama3 8B [], Gemma 7B [], and Mistral 7B []), we demonstrate the effectiveness of our approach. On iterative refinement with SMART, we achieve gains of up to +15 points (a relative gain of +35%) on the GSM8K dataset without the need for refinement. In addition, we improve refinement accuracy by +16 points over baselines.

## 2 Methodology

Figure: Figure 1: Our proposed methodology: In the first step (initial sampling), an agent (LM) chooses a strategy and solves the given task with it. If it is correct, the process ends successfully. If an incorrect strategy is chosen, the agent iteratively refines its strategy, taking previous strategies into account. The process stops when a correct strategy is chosen to solve a task, or when a stopping criterion such as the number of attempts is reached. All correct strategies are used to further refine the model, and the process is repeated. During testing, we sample once from LMt without refinement.
Refer to caption: x1.png

Let $q$ be a problem best solved with multi-step reasoning. An example is presented in the form of a mathematical word problem in Figure 1. An agent or language model (LM) can approach it using various strategies, such as solving it step by step (Chain of Thought, CoT []), decomposing it into subproblems, and solving each one (Least to Most, L2M []), or writing a program to solve it programmatically (Program of Thought, PoT []), among others. A common method is to prompt the LM with a specific strategy for solving the task. However, LMs are error-prone but can fix their answers when asked to do so, a process called *refinement*.
In the refinement process, LMs can either stick to the same strategy [] or switch to a more effective reasoning strategy []. Ideally, LMs could learn to choose the best strategy and reasoning path on the first try, minimizing the need for costly refinement.

Our objective:
The primary goal of our work is to enable language models (LMs) to autonomously learn and select the most effective strategies for various reasoning tasks on their first attempt, thereby improving both efficiency and accuracy. Unlike traditional self-refinement methods that require multiple inference passes or external feedback, our approach aims to internalize the learning process within the LM, allowing it to adjust its strategy selection based on past experience. This mirrors how humans learn to choose optimal strategies through experience when faced with complex tasks.

### 2.1 SMART : Self-learning Meta-strategy Agent for Reasoning Tasks

We model the strategy selection process as a *Markov Decision Process* (MDP), where the LM acts as an agent that interacts with the environment (the reasoning tasks) by selecting strategies and observing the outcomes. Using reinforcement learning techniques, the LM can learn a policy that maximizes the expected reward, effectively learning to choose the optimal strategy for each task. This framework allows the LM to simulate different reasoning strategies, evaluate their effectiveness based on past outcomes, and adjust its strategy choice accordingly.

In the following sections, we formalize the problem formulation, define the agent’s policy, and describe the learning objective. We then present our two-stage process with iterative refinement, which allows the agent to learn from its own reasoning processes and improve its strategy choice over time.

MDP Setup:
We model the strategy selection framework as a Markov Decision Process (MDP) given by the tuple $\langle\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\mu\rangle$. Here, $\mathcal{S}$ represents the state space encapsulating all possible states of the environment, where each state $s\in\mathcal{S}$ is the current problem statement or the subsequent LM response to it. The initial state distribution is given by $\mu$. The action space,
$\mathcal{A}$, is the set of all strategies available to the agent, where each action $a_{t}\in\mathcal{A}$ corresponds to the choice of a particular strategy at time $t$. The transition function $\mathcal{P}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S}$ defines the probability of transitioning to a new state $s_{t+1}$ after applying strategy $a_{t}$ in state $s_{t}$.
Particularly for our case, the transition function is non-deterministic as the next state is sampled by the agent (see Algorithm 1 later).
The reward function $\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ assigns a scalar reward based on the correctness of the result after applying the chosen strategy.

We start with the initial sampling step, where the LM chooses a strategy to solve the given task. In other words, given a problem statement $s_{1}\sim\mu$, the agent draws a strategy $a_{1}$ to solve the problem according to its policy (see below). Given the initial problem $s_{1}$ and strategy $a_{1}$, the environment transitions to the next state $s_{2}\sim\mathcal{P}(\cdot|s_{1},a_{1})$ with transition probability $\mathcal{P}$ and receives a reward $r_{1}$ indicating the correctness of the response. If a correct strategy is chosen and the LM solves the task using that strategy, the process terminates. Otherwise, the agent chooses another strategy $a_{2}$ to solve the problem, and the process repeats until the problem is solved. A realization of this stochastic process is a trajectory $\tau=\left(\left(s_{t},a_{t},r_{t}\right)_{t=1}^{T-1},s_{T}\right)$ up to a finite number of steps $T$. We denote a partial trajectory (history up to time $t$) as $\tau_{1:t}=\left(\left(s_{t},a_{t},r_{t}\right)_{t=1}^{t-1},s_{t}\right)$.

Agent’s policy: We start by defining a history $h_{t}\coloneqq\tau_{1:t}$ that includes the past actions up to time $t-1$ and the observed states up to time $t$. Then, we model the agent using a non-Markovian stochastic policy $\pi_{\theta}(a_{t}|h_{t})$ parameterized by $\theta$, which models the probability of choosing the action $a_{t}$ given the history $h_{t}$:

$$ $\pi_{\theta}(a_{t}|h_{t})=\Pr(a_{t}=a\,|\,h_{t};\theta)$ $$

Objective function: We want to optimize the following objective:

$$ $\centering\theta^{\star}=\arg\max_{\theta}J(\pi_{\theta})=\arg\max_{\theta} \mathbb{E}_{s_{1},a_{1}\sim\mu_{\text{test}},\pi_{\theta}(a_{1}|s_{1})}[r_{1}( s_{1},a_{1})]\@add@centering$ (2.1) $$

where, $\mu_{\text{test}}$ represents the state distribution over the test data.

### Two-Stage Process with Refinement

### Stage 1 (Initial Sampling)

- 1.
The process begins with a problem statement $s_{1}\sim\mu$ where $\mu$ represents the initial state distribution.
- 2.
The agent samples an action $a_{1}$ from the policy $\pi_{\theta}(a_{1}|h_{1})$, where $h_{1}=s_{1}$ for the step 1.
- 3.
The agent then generates an output based on strategy $a_{1}$
- 4.
The agent receives reward $r_{1}(h_{1},a_{1})$ based on it correctness as follows:
$r_{1}=\begin{cases}1,&\text{if the output is correct}\\
0,&\text{otherwise}\end{cases}$
- 5.
Termination Check: If $r_{1}=1$, the process terminates successfully.

### Stage 2 (Iterative Refinement, if r 1 = 0 subscript 𝑟 1 0 r_{1}=0 italic_r start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = 0 )

For each time step $t=2$ to $T$:

- 1.
The agent observes history $h_{t}$.
- 2.
Samples an action $a_{t}\sim\pi_{\theta}(a_{t}|h_{t})$.
- 3.
Generates an output based on strategy $a_{t}$.
- 4.
Receives reward $r_{t}(s_{t},a_{t})$ :
$r_{t}=\begin{cases}1,&\text{if the output is correct}\\
0,&\text{otherwise}\end{cases}$
- 5.
Termination Check: If $r_{t}=1$, terminate the process.
- 6.
Otherwise, transition to the next state $s_{t+1}$, which includes the reasoning choice from all the previous steps.

Trajectory and Reward Structure:
The resulting trajectory $\tau$ with reward has the following structure:

$$ $\tau=\left((s_{1},a_{1},r_{1}),(s_{2},a_{2},r_{2}),\ldots,(s_{T^{\prime}},a_{T ^{\prime}},r_{T^{\prime}}),s_{T^{\prime}+1}\right)$ $$

where $T^{\prime}<T$ is the time step at which the process terminates (either by solving the problem or by reaching the maximum steps). We keep the number of trajectories one less than the number of strategies in our work, since each time the model samples with a different strategy than the one present in its history. For simplicity, we omit $r_{t}=0$ rewards in the trajectory above.

The total reward for the trajectory is:

$$ $R(\tau)=\sum_{t=1}^{T^{\prime}}r_{t}$ $$

Given that $r_{t}=0$ for $t<T^{\prime}$ and $r_{T^{\prime}}=1$, if successful, total reward would be $R(\tau)=1$. However, in some cases, when a task is never solved, the total reward for that trajectory can be 0.

As typical in RL, the agent aims to learn a policy that maximizes the expected cumulative reward [, ]:

$$ $\theta^{\star}=\arg\max_{\theta}J(\pi_{\theta})=\arg\max_{\theta}\mathbb{E}_{ \tau\sim f(\tau;\pi_{\theta})}[R(\tau)]$ (2.2) $$

### Policy Gradient Update

The gradient of the expected reward with respect to the policy parameters $\theta$ is:

$$ $\nabla_{\theta}J(\pi_{\theta})=\mathbb{E}_{\tau}\left[\sum_{t=1}^{|\tau|-1} \nabla_{\theta}\log\pi_{\theta}(a_{t}|h_{t})\cdot r_{t}\right]$ (2.3) $$

The gradient is computed iteratively, and a step is only taken at time $t$ only if all the previous attempts were incorrect.

Implicit bias for Stage 1 to choose the right strategy in the first attempt: As described in (2.1), our goal is to maximize the reward received as early as possible. To implicitly bias the model towards selecting the correct action in earlier steps, we adjust the dataset $\mathcal{D}$ by replacing the sequences of unsuccessful actions with the final successful action taken at the initial state. Specifically, for trajectories where the problem is solved at time $T^{\prime}$ ($T^{\prime}>1$), we replace the samples $(h_{1},a_{1},r_{1}),\ldots,(h_{T^{\prime}},a_{T^{\prime}},r_{T^{\prime}})$ with $(h_{1},a_{T^{\prime}},r_{T^{\prime}})$. This encourages the model to learn to take the correct strategy $a_{T^{\prime}}$ for the problem statement $s_{1}$ in the first step.

Since we only update the model based on correct outputs, the policy update can be viewed as maximizing the likelihood of the correct actions given the history:

$$ $\theta^{\star}=\arg\max_{\theta}\sum_{(h_{i},a_{i})\in D}\log\pi_{\theta}(a_{i }|h_{i})$ (2.4) $$

See Algorithm 1 for a complete algorithm that illustrates our methodology in detail.

## Algorithm: SMART : Self-learning Meta-strategy Agent for Reasoning Tasks

Figure: Algorithm 1 Training Procedure for SMART

## 3 Experimental Details

SMART operates within a *self-learn* framework, where we prompt a pre-trained model with 8-shot examples to collect the initial training data. Prompts used for various strategies are provided in Subsection 8.1. The 8-shot pre-trained model also serves as a baseline comparison for our method. We use the MetaMath dataset [], a variant of the GSM8K [] training set, which contains 110K reasoning problems. We evaluate our methodology on the GSM8K test set, which contains 1,319 samples. To show the generalization ability of our method, we also test on two out-of-distribution datasets: the SVAMP dataset [] with 1,000 samples where the questions are made more challenging by altering them in a non-trivial way, and the ASDiv dataset [] with 2,300 samples, which consists of diverse mathematical problems from elementary to middle school.

**Table 1: Test Accuracy (maj@1) comparison between different baselines using three strategies CoT, L2M, PoT with our approach SMART on the GSM8K dataset. Baselines used 8-shot in-context examples to generate the output. Additionally, we report refinement accuracy obtained through the Oracle verifier for both the baselines and our approach. Refinements for the baselines can follow the same strategy as in [] and a different strategy than the initial one as in []. Results are presented for three models: Gemma 7B, Mistral 7B, and Qwen2 7B. The best results among the baseline are underlined while the best overall results are in bold.**
| Model | Method | Test Accuracy (%) | Refinement |  |
| --- | --- | --- | --- | --- |
| Strategy | Accuracy (%) |  |  |  |
|  | Baseline (*Using 8-shot examples*) |  |  |  |
|  | Chain of Thought (CoT) | 40.0 | same | 44.6 |
|  | different | 49.4 |  |  |
|  | Least to Most (L2M) | 34.9 | same | 43.4 |
|  | different | 48.7 |  |  |
|  | Program of Thought (PoT) | 40.4 | same | 46.1 |
|  | different | 51.3 |  |  |
| Gemma 7B |  |  |  |  |
|  | SMART (*Proposed Approach*) |  |  |  |
|  | *Iteration 1* | 46.5 | SMART | 64.6 |
|  | *Iteration 2* | 50.6 | SMART | 64.0 |
|  | *Final Iteration - Iteration 5* | 55.6<br>($\uparrow$ +15.2) | SMART | 67.5<br>($\uparrow$ +16.2) |
|  | Baseline (*Using 8-shot examples*) |  |  |  |
|  | Chain of Thought (CoT) | 50.6 | same | 59.0 |
|  | different | 67.5 |  |  |
|  | Least to Most (L2M) | 52.4 | same | 56.6 |
|  | different | 67.2 |  |  |
|  | Program of Thought (PoT) | 56.9 | same | 61.3 |
|  | different | 70.1 |  |  |
| Mistral 7B |  |  |  |  |
|  | SMART (*Proposed Approach*) |  |  |  |
|  | *Iteration 1* | 63.8 | SMART | 74.1 |
|  | *Iteration 2* | 66.3 | SMART | 76.4 |
|  | *Final Iteration - Iteration 3* | 67.9<br>($\uparrow$ +11.0) | SMART | 78.0<br>($\uparrow$ +7.9) |
|  | Baseline (*Using 8-shot examples*) |  |  |  |
|  | Chain of Thought (CoT) | 81.9 | same | 90.4 |
|  | different | 89.3 |  |  |
|  | Least to Most (L2M) | 80.0 | same | 84.9 |
|  | different | 88.7 |  |  |
|  | Program of Thought (PoT) | 76.9 | same | 86.0 |
|  | different | 88.6 |  |  |
| Qwen2 7B |  |  |  |  |
|  | SMART (*Proposed Approach*) |  |  |  |
|  | *Iteration 1* | 84.5 | SMART | 91.4 |
|  | *Iteration 2* | 85.4<br>($\uparrow$ +3.5) | SMART | 91.9<br>($\uparrow$ +1.5) |
|  | *Final Iteration - Iteration 3* | 85.2 | SMART | 91.2 |

We ran all our experiments on models with 7-8 billion parameters, specifically Gemma 7B [], Mistral 7B [], Qwen2 7B [], and Llama3 8B []. This is important for our *self-learn* setup, as the model needs a basic understanding of the task to begin the process. Smaller models often have difficulty starting the process, while very large models are too expensive to iterate over multiple times.

We employed three reasoning strategies: Chain of Thought (CoT) [], Least to Most (L2M) [], and Program of Thought (PoT) [] in our work due to their effectiveness on the multi-step reasoning tasks. We take the best out of these strategies as our baseline (underlined in Table 1). Since we used a different reasoning strategy during refinement following Shridhar et al. [], the majority of our experiments are limited to two trajectories as the third trajectory would make the remaining strategy the obvious choice. However, we also explore our approach beyond 3 strategies later (in Section 5).
We report the top-1 accuracy (maj@1) for all experiments, as we want to test the accuracy of the output on the first try.
We used a temperature of $0.7$ to generate samples at each iteration.
Since the models are already trained on the downstream datasets during pre-training, we chose to train only part of the model and used LoRA [] with rank 16, alpha 32, and a starting learning rate of 2e-4 with decay. We used the Unsloth library [] for training and we did the inference using the VLLM library []. Finally, we run SMART for multiple iterations until our accuracy gains no longer justify the cost of training and sampling. In our experiments, 3 iterations worked well for most models, although we trained Gemma 7B for 5 iterations before performance plateaued.

**Table 2: Test accuracy (maj@1) comparison between different baselines using three strategies CoT, L2M, PoT with our proposed approach SMART on the out-of-domain datasets of SVAMP and ASDiv. Baselines used 8-shot in-context examples to generate the output. Results are presented for three models: Gemma 7B, Mistral 7B, and Qwen2 7B. The best baseline results are underlined, while the best overall results are in bold.**
| Model | Method | SVAMP Acc (in %) | ASDiv Acc (in %) |
| --- | --- | --- | --- |
|  | Baseline (*Using 8-shot examples*) |  |  |
|  | Chain of Thought (CoT) | 65.6 | 68.0 |
|  | Least to Most (L2M) | 63.4 | 62.7 |
| Gemma 7B | Program of Thought (PoT) | 51.6 | 54.2 |
|  | SMART (*Proposed Approach*) |  |  |
|  | *Final Iteration* | 68.2<br>($\uparrow$ +2.6) | 70.9<br>($\uparrow$ +2.9) |
|  | Baseline (*Using 8-shot examples*) |  |  |
|  | Chain of Thought (CoT) | 65.8 | 71.3 |
|  | Least to Most (L2M) | 67.9 | 67.3 |
| Mistral 7B | Program of Thought (PoT) | 71.7 | 70.1 |
|  | SMART (*Proposed Approach*) |  |  |
|  | *Final Iteration* | 70.9<br>($\downarrow$ -0.8) | 72.3<br>($\uparrow$ +1.0) |
|  | Baseline (*Using 8-shot examples*) |  |  |
|  | Chain of Thought (CoT) | 89.3 | 89.5 |
|  | Least to Most (L2M) | 87.5 | 84.2 |
| Qwen2 7B | Program of Thought (PoT) | 82.8 | 80.8 |
|  | SMART (*Proposed Approach*) |  |  |
|  | *Final Iteration* | 90.3<br>($\uparrow$ +1.0) | 90.5<br>($\uparrow$ +1.0) |

## 4 Results

SMART significantly improves results on in-distribution dataset:
We compared SMART with baselines on the GSM8K dataset, which we also consider to be an in-distribution dataset since the train and test sets have the same distribution. Table 1 shows that SMART outperformed the baseline in its first iteration on the GSM8K dataset, achieving a gain of +6 points for both the Gemma 7B and Mistral 7B models (40.4 $\rightarrow$ 46.5 and 56.9 $\rightarrow$ 63.8, respectively). Although Qwen2’s performance is already very strong on the GSM8K dataset, we still observed a gain of +2.6 points (81.9 $\rightarrow$ 84.5) in the first iteration. After a few more iterations, we saw a total gain of +15 points for Gemma 7B (40.4 $\rightarrow$ 55.4), +11 points for Mistral 7B (56.9 $\rightarrow$ 67.9), and +4 points for Qwen2 7B (81.9 $\rightarrow$ 85.4).

SMART serves as a great refinement strategy:
Since SMART involves iterative refinement to find the optimal action given the trajectory, we also compare against two refinement baselines: refinement with the same strategy [] and refinement with a strategy change []. We used the Oracle Verifier, which identifies incorrect samples and refines them either with the same strategy or by choosing a different one. (^1^11Note that the Oracle Verifier is used to set the upper bound for all the methods compared and cannot be used during inference. It is only to compare the capabilities of different approaches.) Table 1 compares the refinement accuracy with SMART  and our proposed methodology shows significant improvements over the baselines. Gemma 7B gains over +16 points (48.9 $\rightarrow$ 67.5) compared to the best refinement baseline, Mistral 7B gains +8 points (66.5 $\rightarrow$ 78.0), and Qwen2 7B gains +1.5 points (86.9 $\rightarrow$ 91.9).

SMART generalizes well to out-of-distribution dataset:
We test our trained checkpoints using our proposed approach SMART on two out-of-distribution datasets: ASDiv and SVAMP. These are referred to as out-of-distribution because the model was not trained on these datasets but only evaluated on them. Table 2 compares the results of SMART with the baselines (created using 8-shot in context examples) across the three reasoning strategies: CoT, L2M, and PoT. Among the three models, Gemma 7B showed the most improvement, with a gain of +2.6 points on the SVAMP dataset (65.6 $\rightarrow$ 68.2) and +2.9 points on the ASDiv dataset (68.0 $\rightarrow$ 70.9). Mistral 7B, on the other hand, gained +1 point on the ASDiv dataset (71.3 $\rightarrow$ 72.3) but performed worse on the SVAMP dataset. We suspect that the test dataset for SVAMP is different from GSM8K, and since the model was not optimized for SVAMP-style questions, this could explain the drop in performance. Finally, for the Qwen2 7B model, we observed a modest gain of +1 point across both datasets (89.3 $\rightarrow$ 90.3 on SVAMP and 89.5 $\rightarrow$ 90.5 on ASDiv).

## 5 Discussion

How the strategy distribution changes over iterations: With SMART, the goal is to select the desired strategy at the first attempt, i.e., over iterations the LM should learn to select the appropriate strategy for a given task.
Figure 2 shows the changes in strategy distribution over iterations for the Gemma 7B model on the GSM8K dataset. As indicated by the baseline in Table 1, PoT emerges as the best strategy for the Gemma 7B model, followed by CoT, with L2M being the least effective. A similar pattern is observed in Figure 2, where the model increasingly favors PoT (indicated by the upward trend in the red line) while decreasing its preference for the other two strategies. Correspondingly, the accuracy for PoT improves the most, followed by CoT and L2M, demonstrating that over iterations the model learns to select the optimal strategy for a given task. A qualitative example where over iterations the model learned to choose the right strategy in its first attempt is provided in Figure 4. Initially, the model chose the wrong strategy but was fixed during refinement, and over iterations, the model picked the correct strategy in its initial sampling, without the need for refinement.

Figure: Figure 2: Strategy distribution change over iterations for Gemma 7B model on GSM8K dataset.
Refer to caption: x2.png

Extending the strategies beyond three strategies in SMART:
We want to test if it is possible to go beyond the three strategies explored in this paper by extending our approach to other strategies. Although we could not find a strategy with performance comparable to those explored in this paper, we introduced a <Unsolvable> tag for questions that the model could not answer during either sampling or refinement attempts and tried to use it as a fourth strategy. However, extending this strategy did not yield promising results due to the limited number of samples in the <Unsolvable> category, as the model was able to solve the task using one of the strategies during refinement. Similarly, we experimented with a <Answer Only> strategy, where the model directly predicts the answer to very simple questions without any intermediate reasoning. As with the <Unsolvable> strategy, the skewed data distribution led to poorer performance.

Is the strategy selection effective?:
We compare the strategy selection made by SMART during inference with a fixed strategy (CoT) on the GSM8K dataset using the Gemma 7B model, as shown in Table 3. Across all five iterations, the results highlight the importance of selecting the appropriate strategy, with the SMART approach outperforming the fixed strategy (CoT) by up to 6 points when the optimal strategy is selected. While the fixed strategy shows improvements over iterations, illustrating the *self-learning* effect over multiple generations, our work demonstrates that performance can be further improved by selecting the most appropriate strategy for each task.

**Table 3: Comparison of fixed strategy (CoT) vs SMART based strategy during inference for GSM8K dataset using Gemma 7B model.**
| Strategy | Iteration 1 | Iteration 2 | Iteration 3 | Iteration 4 | Iteration 5 |
| --- | --- | --- | --- | --- | --- |
| Fixed (CoT) | 45.5 | 47.5 | 48.5 | 50.6 | 49.7 |
| SMART | 46.5 | 50.6 | 52.9 | 55.0 | 55.6 |

SMART with better starting samples:

Figure: Figure 3: Figure showing a comparison of the effects of different starting data points for the Gemma 7B model. SMART is compared against two baselines: the pre-trained Gemma 7B model and the fine-tuned Gemma 7B model on Llama3 8B data.
Refer to caption: x3.png

Since we start with the model-generated samples, for a weaker model the starting samples can be improved if those samples come from the stronger model.
We investigated whether the SMART based *self-learning* approach could be extended to a setup where the initial data points are collected from a stronger model to initiate training for a weaker model. This is particularly beneficial for weaker models that cannot independently initiate the self-learning process using SMART due to their limited capabilities on a given task. Initially, we collected data using 8-shot in-context examples generated by the Llama3 8B model []. This approach yielded an average accuracy of about 80% across all three strategies, significantly higher than Gemma 7B’s initial accuracy of about 40%. Figure 3 illustrates the comparison between different iterations of SMART and the pre-trained Gemma 7B baseline at 40.4%, as well as the fine-tuned Gemma 7B baseline at 68.7%. Using SMART resulted in an additional improvement of +3.1 points (from 68.7% $\rightarrow$ 71.8%).

## 6 Related Work

Refinement in LLMs:
Refinement refers to the process of improving the initial output of large language models (LLMs) through iterative adjustments. This refinement can be achieved by following the same method used initially [], by incorporating feedback while using the same approach [, , , ], or by using an alternative method [, ]. However, recent studies have shown that naively applying self-correction can sometimes degrade performance [, , ], highlighting the need for more effective strategies. Supervised fine-tuning with feedback from larger models [, ], or using an ensemble of models [], has produced notable results. Nevertheless, relying on larger or multiple models for feedback presents challenges. In contrast, our method learns the optimal strategy during training and applies it correctly on the first attempt, eliminating the need for refinement.

Self-Training in LLMs:
Self-training is a semi-supervised learning method in which the model’s own predictions are used as additional data to improve its performance [, ]. This technique has been applied to NLP tasks, such as machine translation [, , ]. We use self-learning principles to generate new data and continually update the policy in an on-policy fashion. This approach can be viewed as an on-policy counterpart to self-imitation learning [], where the policy learns from prospective successful trajectories in its initial stages, rather than imitating past successful behavior.

## 7 Conclusion

Figure: Figure 4: Qualitative example demonstrating that Gemma 7B model learnt to refinement strategy in its initial sampling stage, removing the need for refinement.

We present SMART: Self-learning Meta-strategy Agent for Reasoning Tasks, a solution to the challenges LMs face in selecting strategies for complex reasoning tasks. By modeling the strategy selection process as a Markov decision process and leveraging reinforcement learning, SMART enables LMs to autonomously learn and apply the most effective reasoning strategies on the first trial, thereby reducing the reliance on iterative self-refinement. Our proposed approach not only improves the accuracy of LMs, as demonstrated by significant performance gains across multiple datasets but also enhances computational efficiency by minimizing the need for multiple inference passes.

## 8 Appendix

### 8.1 Prompts

Figure: Figure 5: 8-shot Chain of Thought demonstration.

Figure: Figure 6: 8-shot Least to Most demonstration.

Figure: Figure 7: 8-shot Program of Thought demonstration.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The method introduced appears intricate, particularly with its reliance on the MDP. Could you provide more concrete evidence or reasoning to support the necessity of this complex setup over simpler alternatives, such as supervised learning with strategy labels or existing reinforcement learning baselines?

2. Experimental details on configurations and hyperparameters are sparse. You should specify the settings for each dataset and model, especially for reward design and training.

3. The generalization to OOD datasets shows limited gains. What factors might limit this, and could SMART be adapted to improve performance on diverse datasets?

4. Only 7B parameter models were tested. Can you scale SMART to larger models, like 70B parameters?

5. The strategy selection process is crucial but lacks detailed evidence of consistency across tasks. Could you provide quantitative results showing its reliability on different reasoning tasks?
soundness: 3
strengths: 1. SMART provides a novel framework by modeling strategy selection in LMs as a Markov Decision Process, adding a unique layer of self-learning.
2. The paper demonstrates high-quality experimentation, validating SMART’s effectiveness across multiple datasets and model architectures, such as GSM8K, SVAMP, and ASDiv.
3. The paper  is clearly articulated, particularly in its two-stage process design, making the approach easy to understand and apply.
4. By enabling LMs to select optimal strategies without iterative refinements, the work significantly reduces computational costs.
summary: The paper introduces SMART, a novel reinforcement learning framework enabling language models (LMs) to autonomously select optimal reasoning strategies for complex tasks on the first attempt, reducing reliance on iterative refinements. By modeling strategy selection as a Markov Decision Process, SMART applies continuous self-improvement, enabling LMs to adapt strategies dynamically. Experiments across datasets show significant performance gains and computational efficiency, with SMART achieving up to +15 points in accuracy on the GSM8K dataset, outperforming traditional refinement methods.
weaknesses: 1. The method is overly complex without sufficient justification for why simpler methods cannot achieve similar results. There should be a more detailed baseline comparison.

2. The experimental section lacks clarity, particularly regarding the setup and specific configurations for each dataset, making the results difficult to interpret and replicate. Besides, the paper's discussion on the generalization capability of SMART across different datasets is weak, as it does not convincingly demonstrate robustness beyond minor accuracy gains.

3. The use of only 7B models raises concerns about the scalability of the proposed method, with no evidence provided for its effectiveness on larger models.

4. Theoretical claims about strategy optimization are vague and insufficiently supported by quantitative evidence, especially regarding the selection mechanism's consistency across different reasoning tasks.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 3
questions: 1. **Comparison with Existing Baselines**: How does SMART compare against other self-training baselines like ReST and ReST^{EM}? A direct comparison with these methods would provide stronger support for SMART’s effectiveness.
2. **Ensuring Strategy Selection**: How does SMART ensure that during data generation, one of the three strategies (L2M, PoT, CoT) is selected? Without proper constraints or instructions, LLMs might struggle to consistently use the intended strategies. Does SMART simply try each of these strategies until a correct solution is found, or is there a more sophisticated mechanism in place?
3. **Algorithm Implementation**: In the algorithm description, should the implicit biasing step occur before the policy update? This order seems more intuitive and would ensure the update incorporates the biased data correctly.
soundness: 2
strengths: 1. **Relevance**: The problem of efficient self-training is important for improving post-training capabilities of large language models (LLMs). A successful self-training method can enhance LLM performance on complex tasks, such as multi-step reasoning, making the research direction highly relevant.
2. **Clarity**: The overall presentation of SMART, including the RL framework for strategy selection, is clearly explained, and the experimental methodology is easy to follow.
3. **Diverse Empirical Evaluation**: Experiments are conducted using several LLM architectures (e.g., Gemma 7B, Mistral 7B) and multiple benchmarks, which adds credibility to the empirical results.
summary: This paper proposes SMART, a self-training framework that aims to improve problem-solving performance by searching for solutions through self-reflection and subsequently fine-tuning the model to predict the correct solution directly. Specifically, SMART addresses the problem of automatic selection of problem-solving strategies, such as Least-to-Most (L2M) and Chain-of-Thought (CoT), by leveraging reinforcement learning (RL). Empirical results demonstrate that SMART outperforms baseline methods that do not employ self-training.
weaknesses: 1. **Limited Novelty**: The core idea of SMART resembles existing self-training methods that utilize rejection sampling (https://arxiv.org/pdf/2308.01825). The concept of generating correct samples through self-reflection has already been widely explored, such as [ReST](https://arxiv.org/pdf/2308.08998) (Gulcehre et al., 2023) and [Re-ReST](https://arxiv.org/pdf/2406.01495), with the only new aspect being the emphasis on strategy selection rather than general solution generation. This difference does not appear significant enough to warrant a new framework.
2. **Lack of Related Work Discussion**: The paper does not adequately discuss related self-training approaches such as [ReST](https://arxiv.org/pdf/2308.08998), [ReST^{EM}](https://arxiv.org/pdf/2312.06585), and other comparable methods like STaR. Including a detailed discussion of these related works would help contextualize SMART’s contributions and clarify how it advances beyond these methods.
3. **Unfair Empirical Comparison**: SMART is compared only against baselines without self-training, which makes the comparisons somewhat unfair. Including direct comparisons against other self-training methods would offer a more meaningful assessment of SMART’s effectiveness.
4. **Contribution of Strategy Selection**: The paper lacks a detailed analysis of the impact of strategy selection. Previous self-training methods, even when limited to a single strategy like CoT, have shown significant improvements. Therefore, it is unclear how much of SMART’s gains can be attributed specifically to strategy selection as opposed to other aspects of the self-training process.
5. **Limitation to LoRA Results**: The experiments only report results from LoRA fine-tuning, which may not fully capture the performance gains that could be achieved by fine-tuning the entire model. While this is acceptable given possible resource constraints, including results from full-model fine-tuning could strengthen the empirical evaluation.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: The paper presents an approach to improving reasoning in language models through self-learning strategies. The experimental results seem significant. However, the key details of the method are missing and I cannot fully understand the contribution of this paper.

1. In the methodology section, the authors have emphasized the formulation of MDP, which is a crucial aspect of SMART. However, there is a noticeable absence of critical methodological details. For instance, the refinement process lacks clarity on the specific method (or prompts) used to facilitate the model's self-improvement. Whether additional training conducted to enhance the model's evaluation capabilities? How rewards are derived?  How do data from multiple refinements contribute to improving the model's success rate on the first attempt?

2. In the experimental section, the authors refer to three methods: Chain of Thought (CoT), Least to Most (L2M), and Program of Thought (PoT). How these methods are integrated and utilized within the SMART framework? Why does Table 3 only provide a comparison with CoT and not with the other methods?
soundness: 3
strengths: see questions
summary: This paper introduces SMART, a novel framework that enables language models to autonomously learn and select the most effective strategies for reasoning tasks. It models strategy selection as a Markov Decision Process and employs reinforcement learning for continuous self-improvement, aiming to achieve correct solutions on the first attempt. The approach significantly enhances the ability of models to choose optimal strategies without external guidance, improving both accuracy and computational efficiency.
weaknesses: see questions

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: 1.	Equation (1) has a mistake. It is s_1 that samples on \mu, not a_1. 
2.	In Algorithm, step 2 makes step 31 totally meaningless. I am also confused if Implicit Biasing should happen before Policy Update or not?
soundness: 2
strengths: 1.	The topic is interesting since it will reduce the number of attempts that LM needs before achieving successful solution. 
2.	Even though there is some details that need to be clarified, the approach is generally easy to understand.
summary: The paper tries to solve a problem for language model: Can LMs learn to select the optimal
strategy in the first attempt, without a need for refinement? The strategy I believe indicates the prompt strategy like CoT, PoT, step-by-step, etc. The authors design a strategy policy that selects the strategy based on current state and optimize it to select the best strategy on the first attempt.
weaknesses: 1.	Authors formulate the problem as MDP, but the policy is designed based on history of states and actions. It conflicts with conventional understanding of MDP in RL field. If the problem has Markov state, there is sufficient to make decision based on state. Otherwise, the state is mistakenly designed and it contains no full information for decision making. 
2.	The algorithm is named with “meta-strategy”. People might wonder the meaning of “meta”. Does the policy adjust to different tasks or it just choose appropriate strategy action based on current state of the problem?
3.	The objective is modified to use (4) for policy optimization. Only current outputs are used for update. It becomes a supervised learning with maximum likelihood to reproduce the pattern of correct samples, in stead of reinforcement learning through trial and error. 
4.	Following the above question, if the policy is trained on only correct (h_1, a_1/a_t), how it is able to make correct refinement in evaluation where h_t is not seen in the training set. One possible reason I guess might because the authors force “the model samples with a different strategy than the one present in its history” (Line 197). So perhaps randomly selecting the rest strategies in the refinement of SMART may still improve the performance. Authors can refute my guess by adding a comparison experiment. 
5.	Does SMART increase the LM parameters and computation because an additional policy LM is needed, in addition to the original pretrained LM to produce problem solutions? If so, I think a discussion is necessary in the paper.

Meta-review: This paper, while presenting an interesting framework (SMART) that models strategy selection in LMs as a Markov Decision Process, ultimately falls short of being acceptable for publication.  Although the authors highlight the novelty of their approach, the clarity of their experimentation, and the computational efficiency gains, several significant weaknesses undermine its contribution. The complexity of SMART seems unwarranted without a more rigorous comparison to simpler baselines, raising doubts about its necessity. Further, the experimental section is poorly detailed, hindering both replicability and a clear understanding of the results. The claims of generalization across datasets are weak, with limited evidence of robust performance beyond marginal improvements. Concerns also arise regarding the scalability of the method, given its testing on only 7B models, and the lack of theoretical grounding for some of the central claims. Therefore, despite the paper's strengths in its innovative concept and clear articulation of the two-stage process, the methodological concerns and insufficient empirical support lead to its rejection.  The authors are encouraged to address these issues in future work.

Decision: Reject