Markovian Transformers for Informative Language Modeling

Authors: Scott Viteri, Max Lamparth, Peter Chatain, Clark Barrett
Primary area: reinforcement learning
Keywords: chain of thought reasoning, reinforcement learning, scalable oversight, language modeling, proximal policy optimization

Abstract:
Chain-of-Thought (CoT) reasoning holds great promise for explaining language model outputs, but recent studies have highlighted significant challenges in its practical application for interpretability. We propose to address this issue by making CoT causally essential to prediction through two key components: factoring next-token prediction through intermediate CoT text, and training CoT to predict future tokens independently of other context. This results in "Markovian" language models, where CoT serves as a fixed-size state for future token prediction. Our approach optimizes for "informativeness" – the improvement in next-token predictions using a trained CoT compared to a baseline. Using Proximal Policy Optimization (PPO) for arithmetic problems and policy gradient for GSM8K, we demonstrate effectiveness on both arithmetic problems with Mistral 7B and the GSM8K benchmark with Llama 3.1 8B, where the model learns to produce CoTs that are 33.20% more effective at predicting answers than the pre-trained baseline. The increased sensitivity of model performance to CoT perturbations provides strong evidence of CoT reliance. Furthermore, we show that CoTs trained for one model generalize to help other models predict answers, suggesting these CoTs capture reasoning patterns that transfer across different interpreters. This work advances the development of more interpretable language models, potentially enabling their extension to arbitrarily long contexts and enhancing AI reasoning capabilities across various domains.

Decision: reject
Reviewer scores: 5, 6, 6, 10

Reviews and discussion:
<thread>
creation date:
1730471748621

thread_type:
official_review

<thread_object>
creator: reviewer_ddh5

<content>
summary:
__Post-rebuttal update__:

After the rebuttal has concluded, I feel the need to express my strong support for this paper. I believe the proposed method has the potential to become an industry-defining standard, which ICLR should be proud to be the publisher off. The authors have done a lot to further improve the paper from a decent submission to an excellent submission that should be highlighted at a conference. While one can always conduct more experiments to support one's claims even more strongly, I think the remaining requests made by other reviewers are unrealistic. The paper should be accepted as is.

---------

The paper addresses an issue of Chain-of-Thought (CoT) reasoning in LLMs where the LM's final answer does not always depend on the CoT. The paper's idea is to enforce informativeness by conditioning the answer model on the generated CoT only without other context. To this end, the paper formally defines Markovian Language Models and (informative) update functions, from which the policy gradient procedure is derived. Applying the framework to the specific use case of CoT reasoning, the paper experiments with several RL techniques such as expert iteration, policy gradient, and PPO. The model is applied to a simple arithmetic task of adding 15 numbers as well as GSM8K, and shows that the model a) improves performance on the task b) is sensitive to perturbation in the CoT reasoning and c) produces CoTs that are sensible to a different language model such that its performance on the task is improved.

strengths:
* The paper addresses an important limitation in Chain-of-Thought reasoning, which is of relevance to the broader ICLR community.
* The core idea is intuitive and simple.
* The paper is well written.
* The results on the simple arithmetic task and the math task are promising.
* The claims that the proposed method improves the generated CoTs in terms of interpretability and informativeness are well supported.

weaknesses:
The method is evaluated only on few tasks and models, limiting how sure we can be that this is a useful method. Especially an application to language modeling would be very insightful and potentially extremely impactfull. However, while more is always better when it comes to experimental results, I think that this initial set of experiments support the ideas presented well and should suffice for publication.

questions:
Please use different line styles in your figures so colorblind people can make sense of them. Otherwise Figure 2 and 3 are really hard to parse!
</content>
</thread_object>

</thread>
<thread>
creation date:
1730653614457

thread_type:
official_review

<thread_object>
creator: reviewer_etns

<content>
summary:
This work introduces and explains the construction of a Markovian Language Model to study causality in chain-of-thought reasoning. With a limited state (the previous state and its observation) on which to condition, the model is trained (fine-tuned) to maximize an informativeness objective through PPO, and is empirically shown to improve performance on mathematics tasks such as GSM8K and toy addition problems.

strengths:
* The setup of this causally-guided model is pretty novel, and the finding that this improves performance by optimizing on an “information” metric is an impactful finding. 
* The selection of the RL training technique (PPO) is supported by highlighting the limitations of other considered methods (expert iteration and policy gradient).
* The work is written in a way that was simple to follow, which I appreciated.
* The gains on GSM8K (24.64% --> 35.71%) are meaningful (although, a bit tucked away in the paper's text).

weaknesses:
* I understand the general intuition surrounding the design of the informativeness function, but it would be good to add some discussion on why the expected reward over the trajectory actually constitutes / addresses “informativeness” under your construction. 
* While math tasks have a more well-defined structure (the order in which their steps may be pursued), this is less clear for other tasks without such a clear structure in natural language, for instance. It would be good to examine this approach on at least one such task to further support the method’s general efficacy.
* Despite the intuition-based process of selection for the RL training strategy, there are recent works that advocate in favor of expert iteration and REINFORCE / vanilla policy gradient for LLM reasoning and RLHF [1, 2]. To this effect, including such approaches for comparison in the results section (or in the appendix) would strengthen the defense of the PPO method chosen. It would be helpful to include some evidence supporting the limitations posed.
* While I appreciate documenting the design choices in Section 4.3, some justification behind them would be beneficial, either through ablations (it’s fine for these to be in the appendix) or relevant references. 

1. Teaching Large Language Models to Reason with Reinforcement Learning. Havrilla et al. 2024 
2. Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs. Ahmadian et al. 2024

questions:
* Is the space of states task-conditional? This isn’t apparent based on the formulation in Section 3.1 (and is unclear by the wording in line 162). If not, then it would seem that the set of relevant “CoT states” would be very sparse relative to the complete space. 
* As posed in the weaknesses section, does this method extend to other reasoning tasks (e.g. in natural language or code) whose structure is less “linear”?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702622431

thread_type:
official_review

<thread_object>
creator: reviewer_cnes

<content>
summary:
This paper proposes a metric to measure the informativeness of CoT tokens, and then uses RL to train the model to generate highly informative CoT tokens, in order to improve the correctness of the final answer. Experiments in random addition problems and GSM8K math problems demonstrate the effectiveness of the proposed metric and RL methods. The paper also shows that more informative tokens will bring gains in interpretability.

strengths:
The technical ideas, including the proposed metric and RL methods, are new, well-motivated, and technically reasonable.  

The experiments in random addition and GSM8K are positive.

weaknesses:
The experiments are limited to a synthetic math problem setting and GSM8K, and the only trained model is Mistral 7B. 

The presentation needs a better organization. E.g., some major results are placed in the appendices, but training details are in the main paper.

questions:
Have you tried other open-source models like llama? Not use CoT of Mistral in it, but use your method to finetune Llama.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730855858947

thread_type:
official_review

<thread_object>
creator: reviewer_vxem

<content>
summary:
This paper proposes a framework where the reasoning steps are used as fixed-size states, which limits the model’s context (text bottleneck), and force the model to use the reasoning steps as input. This method design is inspired by the fact that past CoT literature find that the final answer might not be sensitive to the CoT trace. In the experiment, the authors show that the model trained with this method is indeed more fragile against CoT pertubations.

strengths:
- The Markovian framework blocks the model from attending back to the original question and force it to use the CoT context for generation. This provides a new view and framework for analyzing CoT effects.
- The reinforcement learning-based approach demonstrates improved performance on tasks requiring multiple steps, 
- The CoT steps generated from this method seem to be more interpretable, from two dimensions: 1) pertubation of the CoT could lead to more model errors 2) the reasoning can be carried over to another model.

weaknesses:
-  While the approach improves the model’s reliance on CoT, it’s uncertain if this CoT is genuinely interpretable by humans. The transferrability between Mistral and Llama should only serve as an indirect proof.
-  It is actually fine to focus on QA task, but probably we'd like to see how this can generalize to more domains other than arithemtic. Would this paradigm also work for other reasoning task as well?
- There seem to be no baselines and ablation designed, so it is a bit hard to position the effectiveness of the method against other methods.
- The fragility analysis is insightful but lacks depth. A more detailed investigation into which types of perturbations impact CoT reliability most could provide valuable insights.
-  Writing-wise, the paper writing is clean, but probably some adjustment of the sections flow and emphasis would be nice.
   -  For instance, while the method section is quite detailed, it’s presented before establishing the limitations of existing CoT techniques clearly, which makes it harder to understand the innovation.
   - There are few tables but quite a few definitinos and equations. I'd suggest consider streamline the method descriptions and move some of them to the appendix, while adding more discussion and insights in the main body.

questions:
- Is there a way to combine the interpretability with actual human perception? Though informativeness here can be used to improve model quality, it is also very helpful from human level. This is probably mentioned in F. But it seems to me F is more about how to encode human interpretability in training.
- In F it is mentioned that "optimal CoT would be a compression of the question, which can potentially be difficult for humans ". Is this observed in your experiments?
- Were there more ablation study or comparison conducted?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732448265165

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Thank you for your comments and reviews!

We received **3 main pieces of feedback:** (1) not enough tasks (especially open-ended ones) or models, (2) lack of baselines, and (3) some methods in the main paper body and some results in the appendix should swap locations.

*In response to 1:*

**We trained a different model, Llama-3.1-8B-Instruct, on predicting Wikipedia text, and we've added all the relevant plots and details in Appendix A.**
We take 200 tokens of an article, and learn to produce 50 tokens of CoT text which serve as context for predicting the next 100 tokens of the article. In Figure 6, we show 4 separate training runs where the evaluator model improves its predictions given the CoTs, increasing the average probability assigned to the correct next token from 8.2% to 10.5%.
In Figures 7 and 8, we demonstrate that our perturbation and generalization results extend robustly to this new setting.

*In response to 2:*

**We've added an Appendix D carefully discussing what baselines would make sense in this context, finding no option that is clearly appropriate.**
At a high level, a baseline in this project could refer to: (i) a method for optimizing the informativeness objective, (ii) an optimization target for faithful language model reasoning, or (iii) a procedure to produce CoTs that are fragile to perturbation. We address (i) through explicit comparison in Figure 2, and we are not aware of any existing approaches for (ii). For (iii), we examine 3 potential approaches in the appendix, showing that they either (a) fail to generalize to natural language settings, (b) introduce experimental confounds, or (c) are already captured in our existing plots.

*In response to 3:*

**We moved the GSM8K results plot to the main experiments section of the paper.** Regarding the training details, we believe these are fundamental to the project's contribution, representing both the majority of our research effort and the insights we expect to be most valuable to other researchers implementing similar approaches.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732910562060

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Second Round

comment:
Thank you for the second round of feedback. We'd like to highlight several key updates:

1. Performance: Using Llama 3.1 8B instead of Mistral **improved our GSM8K performance gain from 11% to 33.2%** (35.94% n=1 exact-match accuracy given only CoT before training → 69.14% after training), strengthening evidence that the method enables meaningful reasoning improvements. In section 5.2, we've included the updated training figure, showing similar results across three separate runs to demonstrate robustness.
2. Interpretability & Generalization: In sections 5.5 and Appendix B, we show Llama’s GSM8K and Wikipedia CoTs transferring across a diverse set of models (**Mistral, Phi, and GPT2**), addressing two key interpretability concerns:
   - Transfer to GPT2 (a much smaller model) suggests the CoTs encode information accessibly
   - Success across model families (e.g., Phi)  indicates model-agnostic representations

Our experiments now comprehensively span:
- Multi-step arithmetic (Mistral)
- Mathematical reasoning (Llama on GSM8K)
- General language modeling (Llama on Wikipedia)

This demonstrates the method's viability across structured reasoning and open-ended text prediction tasks. The Wikipedia experiments specifically validate our approach for general language modeling - they weren't intended as an additional reasoning task.

_Regarding scope:_ While human trials could provide additional validation, we believe demonstrating (i) successful optimization, (ii) increased CoT reliance, and (iii) cross-model generalization establishes a meaningful technical contribution. We've devoted substantial compute resources ($30K) to validate these core claims.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733287609007

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Full Rebuttals Process

comment:
_Note: All section and appendix references have been updated to match the final submitted version of the paper from the second round of reviews._

The review process has helped us understand and address key areas of interest in evaluating our work.

**Initial Reviews Highlighted Six Main Suggestions:**
1. Model coverage beyond Mistral 7B
2. Increased task diversity, particularly for natural language
3. Additional baseline comparisons
4. Increased CoT perturbation analysis
5. Swapping training details in the main paper for GSM8K results in the appendix
6. More comprehensive interpretability evaluation

**Our First Response Added:**
1. Addressing #1 -- Llama 3.1 8B training
2. Addressing #2 -- Added Wikipedia experiments (8.2% → 10.5% improvement) [Section 5.3]
3. Addressing #3 -- Detailed analysis of potential baselines, concluding that we are already evaluating on the appropriate baselines [Appendix E]
4. Addressing #4 -- More detailed perturbation analysis on new model (Llama) and dataset (Wikipedia) [Appendix B]
5. Addressing #5 -- Swapped RL stability training details in the main paper with GSM8K training details in the appendix [GSM8K Section 5.2, Stability Appendix A].

**Second Round Suggestions:**
1. More difficult reasoning datasets

**Our Second Response Added:**
1. Addressing above -- _Demonstrated dramatic improvement on GSM8K using Llama (35.94% → 69.14%)_ [Section 5.2]
2. Addressing #6 -- Implemented additional interpretability metrics showing _successful transfer from Llama to Mistral, Gpt2, and Phi_ (chosen for diversity in size and model family) [Section 5.5]

This progression showcases how our work evolved from initial arithmetic and GSM8K experiments with Mistral to a **comprehensive evaluation spanning multiple architectures, tasks, and evaluation frameworks**. Our extensions - particularly the strong GSM8K results and successful cross-model transfer - demonstrate the robust capabilities and interpretability of our approach.

The resulting paper provides thorough technical validation while maintaining appropriate scope for a research contribution, with clear directions for future work in human evaluation and additional task domains.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes an approach to improving the reasoning capabilities of language models by using CoT as an information bottleneck. Using a reward based on the differences between the predictions of a model conditioned on context vs not, the authors use RL techniques to train their model, and find that on GSM8K, the produced CoT are more effective than the baseline. 

On the plus side, this paper proposes an interesting(ish) idea that can be practically operationalized based on metrics derived from LLM probabilities. The idea is coupled with preliminary experiments that show its promise. On the negative side, the experiments are limited in scope (even with the addition of Wikipedia results after the rebuttal), and some of the claims regarding the better interpretability of the generated CoT tokens are not well supported in my opinion.

I am therefore recommending that this paper be rejected.