Rational Metareasoning for Large Language Models

Authors: C. Nicolò De Sabbata, Theodore Sumers, Thomas L. Griffiths
Primary area: foundation or frontier models, including LLMs
Keywords: LLM, Metareasoning, Problem solving, Chain of Thought, Inference Optimization, Value of Computation

Abstract:
Being prompted to engage in reasoning has emerged as a core technique for using large language models (LLMs), deploying additional inference-time compute to improve task performance. However, as LLMs increase in both size and adoption, inference costs are correspondingly becoming increasingly burdensome. How, then, might we optimize reasoning's cost-performance tradeoff?  This work introduces a novel approach based on computational models of metareasoning used in cognitive science, training LLMs to selectively use intermediate reasoning steps only when necessary. We first develop a reward function that incorporates the Value of Computation by penalizing unnecessary reasoning, then use this reward function with Expert Iteration to train the LLM. Compared to few-shot chain-of-thought prompting and STaR, our method significantly reduces inference costs (20-37\% fewer tokens generated across three models) while maintaining task performance across diverse datasets.

Decision: reject
Reviewer scores: 5, 5, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1729075128615

thread_type:
official_review

<thread_object>
creator: reviewer_mbzh

<content>
summary:
This paper presents an approach for optimizing small LLMs towards generating shorter chains of reasoning while maintaining capabilities. To this end, the authors designed a reward function that balances the log-likelihood of generating a target answer with the log normalized cost of the number of generating chains of thought tokens. Experiments on a number of question-answering datasets demonstrate that the proposed approach can reduce the number of output tokens while maintaining the LLM’s performance.

strengths:
A) The paper is clear and easy to follow. Furthermore, the method is simple enough and the paper detailed enough so that it would, I believe, be easy to reproduce. 

B) The problem of reducing the number of output tokens of CoT-like approaches has practical relevance.

C) While the LLMs that were investigated are on the small end (2.7B-8B), the authors demonstrate that their approach leads to reduced number of output tokens on multiple LLMs and datasets.

D) I appreciate the inclusion of insightful qualitative results in the Appendix.

weaknesses:
## Major
E) I appreciate the attempt to demonstrate the validity of the method for multiple chains of reasoning approaches (CoT and STaR). However, in my view, the paper is incremental, since the cost of generating chains of reasoning has already been addressed by follow up work like Zelikman et al. (2024). Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking. arXiv. https://doi.org/10.48550/arXiv.2403.09629 The authors should directly compare to Quiet-STaR, and investigate whether their approach can provide orthogonal gains.

F) In my view, the following important ablation is missing to ensure that the models are indeed retaining their ability to reason (and decompose a problem into subproblem) as opposed to relying simply on directly producing an answer. Looking at Figure 3 and Figure 4 in the Appendix, it seems to me the model learned to simply rely on its internalized knowledge (one conclusion from this could be that ARC and CommonsenseQA simply aren’t great benchmarks to assess complex reasoning capabilities). As a baseline, directly fine-tune the LLM to produce the correct answer. I also need to point out that the reduction in reasoning lengths is much less impressive on the harder reasoning benchmarks GSM8K and Proofwriter (see Table 5). Thus, my current hypothesis is that the proposed approach works well for problems that are not hard reasoning problems, but is much less effective on harder reasoning problems where we actually need more CoT-like methods.

G) Related to the above, I believe the authors should test on additional hard compositional problems such as AQUA-RAT (https://github.com/google-deepmind/AQuA).  

H) Another ablation I am missing is investigating to what extent RL is really needed here. What happens if you generate chains of reasoning on the training data, ask the LLM to summarize these chains, and then directly fine-tune on this synthetically generated data?

## Minor
I) LLMs are on the smaller scale, so it remains to be seen whether gains in terms of reduced number of output tokens persist for larger models.

questions:
I don’t have any questions, the paper was very clear.

What the authors would have to demonstrate to see an improved rating from me:
1. Compare directly to Quiet-STaR, and demonstrate whether their approach leads to orthogonal gains (E)
2. Add an ablation where the LLM is directly fine-tuned on the author’s training sets (F)
3. Demonstrate results on additional compositional reasoning benchmarks such as AQUA-RAT since ARC and CommonsenseQA aren’t informative (G)
4. Add an ablation where the LLM is used to summarize chains of reasoning and then fine-tuned on the resulting synthetic data (H)
</content>
</thread_object>

</thread>
<thread>
creation date:
1729673394299

thread_type:
official_review

<thread_object>
creator: reviewer_hjgt

<content>
summary:
This paper proposes a new fine-tuning approach for training LLMs to generate token-efficient and adaptive responses. The method encourages the LLM to produce longer reasoning chains only when necessary. Building on the STaR framework, it introduces a filtering step to fine-tune the model on correct answers with the highest reward, where the reward function penalises unnecessarily long responses. Experiments on a dataset constructed from a mixture of four benchmark datasets show that the method effectively reduces the token count of generated responses without sacrificing accuracy. The results also demonstrate that the method promotes adaptive behaviour, producing longer responses for more difficult questions and shorter ones for simpler questions.

strengths:
- The paper tackles an important and underexplored problem of lowering sampling costs at inference time without sacrificing model performance in an adaptive manner.
- The proposed method shows improved performance over a set of 3 baselines, including the STaR method which is a direct ablation of Rational Metareasoning.
- The method is evaluated on a good selection of benchmark datasets.

weaknesses:
- The link between VOC and the latter presented approach of modified expert iteration is unclear. Sections 2 and 3 are not well connected together.
- Incremental contribution. The method is an extension of existing work (STaR). More justification or analysis of this approach's novelty would strengthen the paper (see question 2).
- There a few places where writing could be improved (see the writing suggestions below).
- There are a few inconsistencies or unclear statements (see the clarity comments below).
- A potential limitation, not mentioned in the paper, is that while shorter responses may be preferred from the computational costs point of view, they may not necessarily be more human friendly. The paper would benefit from a small human study assessing the qualitative aspects of the generated reasoning chains with Rational Metareasoning.
- Experiments could include additional baselines and ablation studies (see below for details).


**Comments on clarity**

- Line 159: the statement: *“Initially, in the exploration phase, we approximate the optimal policy   by using rejection sampling on our student policy $\pi_\theta$”* is unclear.  First, the notion of an optimal policy $\hat{\pi}^*$ hasn’t been defined. Second, in rejection sampling, all reasoning chains with the reward above a certain threshold should be retained. Yet, the proposed algorithm only selects a single reasoning chain that maximises the reward. The motivation for this choice is unclear (see question 2).
- It seems that there are a few inconsistencies in the presentation of Algorithm 1. Assignment of $\pi_0$ is missing. In line 2, the entire dataset $\mathcal{D}$ is assigned to  $\mathcal{D}\_n$, but from the comment and the main text it follows that $\mathcal{D}\_n$ is subsampled from $\mathcal{D}$.  The reward function in the algorithm should should be made dependent on $\pi_{n-1}$, i.e.  $\mathcal{R}\_{\pi_{n-1}}$. Storing the rewards in $r_{i,k}$ (line 7 of Algorithm 1) seems redundant given the later $\arg\max$ again uses $\mathcal{R}$. The input to this argmax should be $\mathcal{R}\_{\pi_{n-1}}(x_i, z_{i,k}, y_i)$ rather than $\mathcal{R}(x, z_{i,k}, y)$, I believe. The quantifier $\forall i$ is also missing in this step.


**Experiments**

- It would be interesting to compare Metareasoning with Direct Few-Shot prompting, where the LLM is explicitly instructed to provide concise responses (e.g. “Keep your answer concise”.). It has been previously demonstrated that such a statement shortens the expected answer length, often without sacrificing the performance [1, 2].
- The batching technique with increasing $T$ is a design choice which should be tested in an additional ablation study.
- The proposed expert-iteration algorithm could be additionally compared to other fine-tunning algorithms, like PPO, to better motivate the particular choice of the training method.

**Minor writing suggestions**

- Line 032: “… many of these methods reduce costs at the expense of performance”—this statement would benefit by adding an example with an appropriate reference similar to the juxtaposed approach of chain-of-thought.
- Lines 071-075: the word “while” is repeated 4 times in 3 consecutive sentences, consider rephrasing.
- Line 097: There is no equation 3.1 in the paper

**References**

[1] https://arxiv.org/pdf/2401.05618 

[2] https://arxiv.org/pdf/2407.19825

questions:
1. Could the authors elaborate on how VOC translates to the setting of reasoning with LLMs? What would the computations $c$, beliefs $b$ and actions $a$ correspond to with respect to the terms presented in section 3: input $x$, reasoning chain $z$ and answer $y$?
2. Why is the LLM policy fine-tuned only with respect to the reasoning’s $\hat{z}$ that maximise the reward within the sampled batch of rasoning chains? In expert iteration, we would typically retain all reasoning chains $z_i$ for which the reward is above a certain threshold. Could the authors provide either a theoretical motivation behind their choice or run additional empirical studies comparing their training method to alternatives (e.g. expert iteration with threshold rejection, PPO) to justify the choice of the training method?
3. Is the $T$ in lines 165-169 referring to the number of fine-tunning steps used in the $\mathrm{train}(\pi, \mathcal{D}_n^*)$ subroutine at each iteration or is it the number of examples sampled from $\mathcal{D}$? 
4. How many training iterations $N$ of the algorithm are needed for convergence and what is the relationship between the compute time for policy training vs. final performance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730342814638

thread_type:
official_review

<thread_object>
creator: reviewer_6gdb

<content>
summary:
This paper presents a method for fine-tuning large language models to enhance the cost-performance trade-off in reasoning procedure. The authors introduce a reward function that evaluates the computational value of various reasoning chains, enabling the ranking of these chains. By integrating this reward function with expert Iteration, the method trains LLMs utilizing intermediate reasoning steps selectively, employing them according to the rewards. This approach significantly reduces inference costs on specific datasets while preserving task performance.

strengths:
1. The authors introduce an interesting problem in LLM reasoning, optimizing LLMs’ inference cost and performance at the same time. This is important issue in using LLMs especially considering the LLMs inference cost is becoming larger. 

2. The paper is well-written and easy to follow. Overall, I could follow the whole story that the authors want to present in this paper.

weaknesses:
1. Lack experiments in more realistic datasets. LLMs are not limited to tasks in text space; they are frequently utilized as agents that interact with external tools to perform complex tasks in various environments. Incorporating experiments on more realistic datasets, such as GAIA [1] and ToolBench [2], would provide valuable insights into the model's performance in more complex reasoning scenarios.

2. Currently, the method and experiments focus exclusively on CoT reasoning, generating trajectories for model fine-tuning and serving as baselines. However, CoT may not always produce the optimal trajectory and should not be considered a strong baseline. It would be beneficial to explore alternative reasoning methods for fine-tuning the models. At the very least, a discussion or justification regarding the choice of reasoning methods should be included in the paper.

[1] Mialon, G., Fourrier, C., Swift, C., Wolf, T., LeCun, Y., & Scialom, T. (2023). Gaia: a benchmark for general ai assistants. arXiv preprint arXiv:2311.12983.
[2] Guo, Z., Cheng, S., Wang, H., Liang, S., Qin, Y., Li, P., ... & Liu, Y. (2024). StableToolBench: Towards Stable Large-Scale Benchmarking on Tool Learning of Large Language Models. arXiv preprint arXiv:2403.07714.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730709522169

thread_type:
official_review

<thread_object>
creator: reviewer_eudp

<content>
summary:
This paper introduces a rational metareasoning approach for LLMs, using a novel Value of Computation (VOC)-based reward function to reduce inference costs without compromising performance. Tested across diverse benchmarks, the method achieves significant token savings (20-37%) compared to standard prompting techniques. The results highlight this approach as an effective solution for cost-efficient reasoning in LLMs, though further evaluation on task variety would enhance the findings’ robustness.

strengths:
* The paper introduces a unique rational metareasoning approach that balances inference cost with performance, addressing a crucial need in the efficient deployment of LLMs.
* The integration of the Value of Computation (VOC)-based reward function is well-designed and thoughtfully applied, showing careful consideration of LLM efficiency.
* The approach is tested across a diverse set of benchmarks, covering science knowledge, commonsense reasoning, math problem-solving, and logical deduction, as well as an out-of-domain generalization task (MMLU).
* Empirical results demonstrate notable reductions in token generation (20-37%), indicating that the method can achieve similar performance with fewer computational resources.
* The approach is valuable for scenarios where computational resources are limited, potentially benefiting applications needing cost-effective, high-quality language model outputs.

weaknesses:
* Limited Analysis of Time Complexity: While the paper focuses on token reduction, it lacks an analysis of the time complexity of the proposed method. A deeper investigation into time savings would provide a clearer picture of its practical efficiency.
* Narrow Range of LLMs and Tasks: The method was primarily tested on a limited selection of benchmarks and model architectures. Broader experimentation across different LLMs and a wider range of tasks would strengthen claims about the generalizability of the approach.
* Limited Robustness Testing: Although the results are promising, the paper lacks robustness checks to assess how performance holds under varied conditions, such as noisy inputs or more complex task requirements.
* Scalability Concerns: The feasibility of scaling this method to very large models or highly complex tasks is not fully addressed, leaving open questions about its applicability in more computationally intensive scenarios.

questions:
### Questions
* Section 3.2: What is the likelihood of not generating a correct answer? Will this happen frequently? Any comment or justification on this?
* Are the input tokens of failed cases counted as part of the input and output tokens?
* What's the time cost of the proposed method itself, rather than the inference costs? The complexity is largely from the online reinforcement learning used in the proposed methods, i.e., EI?
*  The reliance on a single, VOC-inspired reward function might limit flexibility across diverse task types. Is there any justification on using alternative reward structures could reveal greater adaptability and robustness?

### Minors
* Line 97: Eq. 3.1 --> Eq. (2)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730922734942

thread_type:
official_review

<thread_object>
creator: reviewer_1h6d

<content>
summary:
The paper proposes to fine-tune language models with an objective that maximizes task performance while minimizing the (logarithm of) the number of tokens used. This objective is optimized with expert iteration. Results suggest that this objective does indeed help models use the extra reasoning tokens only when they're actually needed, thus reducing compute cost on easier tasks.

strengths:
- Computational efficiency of LLM deployments is a timely topic relevant to the ICLR community.
- The proposed method is sufficiently simple that something like it might well get used in practice.
- The paper is well-written and clearly structured.

weaknesses:
### a) Insufficient baselines and ablations -- I don't feel like I get the "shape" of the proposed method and the potential alternatives all that well.

A few notes here:

1. How necessary is rationalization (described in the paragraph on line 131)? I assume the authors only use it because it's also used in STaR, the main baseline? Relatedly (and importantly), STaR should be described in more detail in the paper, and justified as a relevant baseline -- I needed to open up the original STaR paper to remind myself of what it does.

2. Re design of the computational cost function: what happens when the cost is linear instead of logarithmic in the number of tokens? (Authors mention they tried linear cost but don't provide details.) How exactly is $\gamma$ chosen, and how sensitive is the method to this choice?

3. Did the authors try a prompting baseline such as asking the model to think however long it needs, but answer as soon as it's ready? I'd be interested in seeing this kind of baseline with a few prompt variations, as this can plausibly also help avoid spending many tokens on easy tasks & spend more tokens on harder tasks when needed -- especially in larger and more capable models. I think models might well already be tracking whether further reasoning is useful: such tracking could plausibly arise just from self-supervised pretraining on large corpora that include texts written by humans who're tracking this implicitly. (Efficiency of such a prompting baseline could be improved further with prompt distillation, but I don't think this would be important for this paper).

4. I'm also interested in variations of the proposed method with RLHF-style algorithms other than expert iteration (e.g. iterated DPO).

### b) I did not find Figures 1 and 2 particularly useful, especially the subplots on the left side. Also, is information from subplots on the right side duplicated in the tables?

1. One idea for improving the figures could be plotting task performance on the y axis and the computational cost (e.g. input + output tokens) on the x axis. It might then be clearer whether your method is on the pareto frontier, and generally help understand the performance-efficiency tradeoff.

2. Regarding Figure 2: I think here it'd be helpful to include 5-shot and/or (0-shot, CoT) benchmarks for MMLU, to get a sense of how STaR and your method compare to these.

questions:
1. Re Section 4.3, how many expert iteration steps $n$ do you perform?
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a fine-tuning approach to improve the token cost of LLMs on reasoning tasks. The approach is motivated by metareasoning, and trains a reward model to capture the value of computation when decoding tokens for reasoning tasks. Compared to Chain-of-Thought prompting and STaR as baselines, the proposed fine-tuning achieves significant reduction in tokens used without harming overall reasoning performance on some benchmarks.

All of the reviewers agreed that the problem is well-motivated and timely, and the solution approach is sound. However all of the reviewers concurred that the paper is borderline at present, and identified the following avenues for improving the experiments and the contributions:
1. Does VoC finetuning help improve token efficiency on other approaches to reasoning beyond CoT? (e.g. Quiet-STaR, other prompting approaches like Tree-of-Thought).
2. Ablations to identify the contributions from the different components of the proposed approach (e.g. is RL with the VoC reward model necessary, or could we accomplish similar results by distilling successful reasoning chains directly).
3. Experiments on more domains (e.g. AQuA) to ensure that the proposed approach is indeed adaptive to the hardness of tasks and domains.