Markovian Transformers for Informative Language Modeling

## Abstract

Abstract Chain-of-Thought (CoT) reasoning holds great promise for explaining the outputs of language models, but recent studies have highlighted significant challenges in its practical application for interpretability. We propose to address this issue via two key components: a technique to factor next-token prediction through intermediate CoT text, ensuring the CoT is causally load-bearing, and a reinforcement learning approach to train 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. We demonstrate our method’s effectiveness using Proximal Policy Optimization (PPO) on arithmetic problems and achieve an 11% performance boost on the GSM8K benchmark using Mistral 7B Inst V2. The increased sensitivity of model performance to CoT perturbations provides strong evidence of CoT reliance. This work advances the development of more transparent and interpretable language models, potentially enabling their extension to arbitrarily long contexts and enhancing AI reasoning capabilities across various domains.

## 1 Introduction

The rapid advancement of language models (LMs) has revolutionized the field of artificial intelligence, demonstrating remarkable capabilities in tackling complex cognitive tasks (Brown et al., ). However, it can be challenging to understand why an LM gave a particular answer (Burns et al., ; Gurnee & Tegmark, ; Lamparth & Reuel, ), which can be problematic in high-stakes scenarios (Rivera et al., ; Lamparth et al., ; Grabb et al., ). Interpretability techniques analyze the patterns and activations of a neural network in order to extract an explanation of the network’s behavior (Casper et al., ; Meng et al., ; Geva et al., ; Geiger et al., ; Wang et al., ; Nanda et al., ; Lamparth & Reuel, ). However, since language models already speak natural language and have been trained to be able to use their own internal representations, we could in principle simply ask the language model why it gave a particular answer to a question. Asking the language model to explain its reasoning in a “step-by-step” fashion before answering a question is known as Chain-of-Thought (CoT) (Wei et al., ; Nye et al., ) prompting.

However, there are concerns that CoT is an inadequate or *unfaithful* explanation for LM-generated text. For example, Turpin et al. () show that biasing the LM to believe a particular answer via a supposedly irrelevant in-context feature such as multiple choice answer order will cause the CoT to rationalize that answer without mentioning the background feature. Also some LMs give the same answers to questions despite changes to the CoT reasoning in their context window (Lanham et al., ). While this has some benefits – the model can still answer correctly despite intermediate reasoning errors – it is also an indicator that the CoT does not fully capture the LM’s reasoning process. This raises a critical issue with using CoT as a tool for interpretability.

Our work introduces a novel perspective on this issue. Rather than aiming for faithfulness – which implies that the CoT reflects some underlying causal process in the model – we focus on *informativeness*. Our key insight is to make the CoT text itself causally important in the model’s reasoning. We propose a reinforcement learning (RL) based training technique that trains an LM to generate a minimal-length CoT such that the model can predict the answer given only that CoT. This approach ensures that the CoT is not merely an ex-post rationalization but an integral, causal component of the reasoning process.

Ideally, a CoT explanation would be both complete (i.e., each necessary step is included) and maximally fragile (i.e., removing or changing the meaning of any step breaks the CoT and thus leads to a different result). Our approach aims to achieve these properties by making the CoT itself a bottleneck in the flow of information that the language model uses to produce text.

We assume the LM receives a sequence of observations to predict – this could be a question-answer pair (length two sequence) or many adjacent segments of generic internet text. Our conceptual arguments rely on the size of each observation being larger than the CoT – otherwise the LM could put the answer immediately in the CoT. Though for pragmatic reasons we use short observations, the model does not learn the undesirable behavior of directly answering in the CoT due to the relative difficulty of predicting the answer without any CoT.

The primary contributions of this work are:

- 1.
We introduce a formal definition of informativeness which is used as an optimization target.
- 2.
We demonstrate our training algorithm’s effectiveness by training Mistral 7B (Jiang et al., ) to solve fifteen-term arithmetic problems and achieve an 11% performance gain on the GSM8K (Cobbe et al., ) reasoning dataset.
- 3.
We verify that the model utilizes the generated CoT during inference, ensuring that the model’s future token distribution is inherently a function of the CoT.
- 4.
We validate that the generated CoT is meaningful and usable by other models, showcasing its interpretability and transferability.

By making CoT causally important in the model’s reasoning, we aim to improve the interpretability and reliability of language models. This approach offers a novel perspective on understanding and steering LM behavior by leveraging the model’s own generated explanations, rather than relying solely on the analysis of its internal parameters.

Figure: Figure 1: Refined illustration of the training method. Left: Single timestep process from Question to CoT to Answer. Right: Causal structure showing the generation of states from observations and previous states using the state update function $u_{\theta}(s^{\prime}|o,s)$, and the prediction of observations from states using the policy $\pi(o|s)$. Observations are generated by the causal data distribution. In experiments, both $u_{\theta}$ and $\pi$ are Mistral 7B Inst V2, but only the weights of $u_{\theta}$ are updated during training. The state update $u_{\theta}$ also involves concatenating the observation and state letting Mistral generate the next state’s worth of tokens.

## 2 Related Work

Prior work shows that CoT prompting improves language model reasoning capabilities (Wei et al., ; Nye et al., ). We train the model to produce a strong CoT, as opposed to prompting strategies as in Wei et al. (). Scratchpad (Nye et al., ) also trains the model to produce a CoT, but they supply correct CoTs during training, whereas our model has to discover useful CoTs for itself. Zelikman et al. () also use RL to improve CoT reasoning, but they do not restrict the model’s attention to the previously generated CoT, making the CoT less of a standalone explanation. State space models also generate state to remember their history (Gu et al., ; ; Gu & Dao, ), but we use natural language instead of activation vectors for interpretability.

Lyu et al. () improved faithfulness of language model reasoning by restricting the output to a particular formal language so that a deterministic solver could provide the rest of the answer, whereas we do not restrict to production of a formal language, because our future goal is to target general language modeling. Ranaldi & Freitas () directly fine-tune a smaller model using CoT from a more capable model. In contrast, we do not require the existence of a more competent model to learn useful CoTs. Lanham et al. () use robustness to reasoning perturbations as an indicator of unfaithfulness, which we adapt by replacing the variation in multiple choice accuracy with the variation in log probability assigned to the correct observation. Bentham et al. () respond that robustness might simply be an indicator of accuracy, which we ameliorate by removing history from the context window. In order to address this concern more thoroughly, we would need to demonstrate the ability to further compress our CoTs.

Figure: Figure 2: CoT informativeness $\ln\pi(ans\mid cot)-\ln\pi(ans\mid cot^{\prime})$, where $cot$ is sampled from the trained weights $cot\sim u_{\theta}(cot\mid q,cot_{init})$ and $cot^{\prime}$ is sampled from the unmodified weights $cot^{\prime}\sim u(cot\mid q,cot_{init})$. We train to produce CoTs which are sufficient to predict the correct answer even without the original question, enforcing a text bottleneck in the language model’s information flow, forcing the CoT to be causally load-bearing to production of the answer. Because of high variance, we plot the point-wise maximum for each training technique across 4 separate training runs.
Refer to caption: extracted/5911817/Figures/cot_performance_comparison.png

## 3 Markovian Language Models and Informativeness of Update Functions

We would like a mathematical structure which describes the shape of a language model with a CoT bottleneck, so that we can derive an reinforcement learning algorithm with respect to that formalism.
For this reason, we introduce the concept of Markovian Language Models and define a measure of informativeness for their update functions.

### 3.1 Markovian Language Models

A regular autoregressive LM can use its entire context when predicting the next token. In particular, when the LM takes a question, produces some reasoning and finishes with a final answer, the generation of the final answer can still attend back to the question. Thus, there is no guarantee that the CoT is causally linked to the answer tokens. In contrast, in a Markovian LM, the model is only given access to limited state to make predictions.

Formally, we define a Markovian Language Model as a tuple $M=(\mathcal{V},\mathcal{S},\pi,u,s_{1})$, where:

- •
$\mathcal{V}$ is a finite vocabulary,
- •
$\mathcal{S}$ is a set of states, representing CoT reasoning,
- •
$\pi:\mathcal{S}\rightarrow\Delta(\mathcal{V}^{k})$ is a state-conditional distribution over $k$-token sequences, where $\Delta(\mathcal{V}^{k})$ is the probability simplex over $k$-token sequences from $\mathcal{V}$,
- •
$u:\mathcal{V}^{k}\times\mathcal{S}\rightarrow\Delta(\mathcal{S})$ is a stochastic update function,
- •
$s_{1}\in\mathcal{S}$ is the initial state.

The MLM operates sequentially: given a current state $s_{t}\in\mathcal{S}$ and observation $x_{t}\in\mathcal{V}^{k}$, it produces a probability distribution $\pi(s_{t})$ over the next $k$-token sequence, and stochastically updates its state to $s_{t+1}\sim u(s_{t},x_{t})$.

### 3.2 Data-Generating Distribution

Let $P$ be the true data-generating distribution over sequences of length $T$. We can sample from this distribution using:

$$ $x_{t}\sim P(x_{t}|x_{<t})\quad\text{for }t=1\text{ to }T$ (1) $$

where $x_{<t}$ denotes all observations before time $t$.

### 3.3 Parameterized Update Function

We consider a parameterized update function $u_{\theta}$, where $\theta$ represents the parameters to be optimized. We compare this to a baseline update function $u^{\prime}$, which uses the original set of weights before fine-tuning. Both $u_{\theta}$ and $u^{\prime}$ operate in conjunction with the same prediction function $\pi$, which also uses the original set of weights.

### 3.4 Informativeness of Update Functions

We define the informativeness of the update function $u$ relative to a baseline update function $u^{\prime}$ as:

$$ $I(u,u^{\prime},P)=\mathbb{E}_{\tau\sim P,u,u^{\prime}}[R(\tau)]$ (2) $$

where $\tau=(x_{1},s_{1},s^{\prime}_{1},...,x_{T},s_{T},s^{\prime}_{T})$ is a trajectory, with:

- •
$x_{t}\sim P(x_{t}|x_{<t})$
- •
$s_{t+1}\sim u(s_{t+1}|x_{t},s_{t})$
- •
$s^{\prime}_{t+1}\sim u^{\prime}(s^{\prime}_{t+1}|x_{t},s^{\prime}_{t})$

The reward $R(\tau)$ for a trajectory is defined as:

$$ $R(\tau)=\sum_{t=1}^{T}[\ln\pi(x_{t}|s_{t})-\ln\pi(x_{t}|s^{\prime}_{t})]$ (3) $$

Now, let’s consider optimizing this informativeness using policy gradient methods. We parameterize $u$ by some weights $\theta$, giving us $u_{\theta}$. The objective function is:

$$ $J(\theta)=I(u_{\theta},u^{\prime},P)$ (4) $$

The gradient of this objective with respect to $\theta$ is:

$$ $\nabla_{\theta}J(\theta)=\mathbb{E}_{\tau\sim P,u_{\theta},u^{\prime}}\left[R( \tau)\sum_{t=1}^{T-1}\nabla_{\theta}\ln u_{\theta}(s_{t+1}|x_{t},s_{t})\right]$ (5) $$

In practice, we estimate this gradient using Monte Carlo sampling:

$$ $\nabla_{\theta}J(\theta)\approx\frac{1}{N}\sum_{i=1}^{N}R(\tau^{(i)})\sum_{t=1 }^{T-1}\nabla_{\theta}\ln u_{\theta}(s_{t+1}^{(i)}|x_{t}^{(i)},s_{t}^{(i)})$ (6) $$

where $\{\tau^{(i)}=(x_{1}^{(i)},s_{1}^{(i)},s^{\prime(i)}_{1},...,x_{T}^{(i)},s_{T}^
{(i)},s^{\prime(i)}_{T})\}_{i=1}^{N}$ are sampled trajectories.

This procedure improves the update function $u_{\theta}$ to generate more informative CoT reasoning, leading to better predictions of future observations.

## 4 Methods

### 4.1 Markovian Language Model for Question-Answer Pairs and Optimization

We define a specialized Markovian Language Model (MLM) for question-answer pairs as a 5-tuple $M=(\mathcal{V},\mathcal{S},\pi,u,s_{1})$, where:

- •
$\mathcal{V}$ is the vocabulary of tokens.
- •
$\mathcal{S}=\mathcal{V}^{\ell}$ is the set of all possible CoT sequences of length $\ell$.
- •
$\pi:\mathcal{S}\rightarrow\Delta(\mathcal{V}^{\ell})$ is the prediction function.
- •
$u:\mathcal{S}\times\mathcal{V}^{\ell}\rightarrow\Delta(\mathcal{S})$ is the update function.
- •
$s_{1}=\text{cot}_{\text{init}}\in\mathcal{S}$ is the initial state, where $\text{cot}_{\text{init}}$ is a fixed initial prompt.

Let $\ell_{q}=\ell_{a}=\ell$ be the length of an observation (question or answer). We implement the MLM specification using a language model $\mathcal{L}:\mathcal{V}^{*}\rightarrow\Delta(\mathcal{V})$, where $\mathcal{L}(s)$ gives the probability distribution over the next token given the sequence $s$. We denote the $i$-th tokens of the CoT and answer as $\text{cot}_{i}$ and $\text{ans}_{i}$, respectively.

The model operates as follows:

- 1.
Update function $u$:
$\ln u(s_{2}=\text{cot}|o_{1}=q,s_{1}=\text{cot}_{\text{init}})=\sum_{i=1}^{
\ell}\ln\mathcal{L}(\text{concat}(q,\text{cot}_{\text{init}},\text{cot}_{<i}))
[\text{cot}_{i}]$
(7)
We implement $\ln u$ by concatenating the question with $\text{cot}_{\text{init}}$ and summing the log probability of each token conditioned on the previous tokens and the prefix.
- 2.
Prediction function $\pi$:
$\ln\pi(o_{2}=\text{ans}|s_{2}=\text{cot})=\sum_{i=1}^{\ell}\ln\mathcal{L}(
\text{concat}(\text{cot},\text{ans}_{<i}))[\text{ans}_{i}]$
(8)

### 4.2 Threshold-based Expert Iteration, Policy Gradient, and Proximal Policy Optimization

We explore three RL techniques to optimize the language model for informative CoT production: Threshold-based Expert Iteration, Policy Gradient, and Proximal Policy Optimization. All three implementations use a form of importance sampling to focus updates on more informative CoTs. All three implementations are concisely contained within a single Python file, which we have made freely available.

#### 4.2.1 Threshold-based Expert Iteration (TEI)

Threshold-based Expert Iteration consists of the following steps:

- 1.
Sample a CoT from a trained and untrained model (cot and cot’)
- 2.
Estimate informativeness as $I(ans,cot,cot^{\prime})=\pi(ans|cot)-\pi(ans|cot^{\prime})$
- 3.
If $I$ is at least one standard deviation above the historical average:
•
Calculate the gradient of the log probability of having produced that CoT: $\nabla_{\theta}\ln u_{\theta}(cot|q,cot_{init})$
•
Gradient ascend

Limitation: This technique potentially discards valuable information, as we might prefer to update more strongly towards CoTs that produce very high rewards.

#### 4.2.2 Policy Gradient (PG)

Policy Gradient (with threshold-based sample selection) consists of the following steps:

- 1.
Sample a CoT from a trained and untrained model (cot and cot’)
- 2.
Estimate informativeness as $I(ans,cot,cot^{\prime})=\pi(ans|cot)-\pi(ans|cot^{\prime})$
- 3.
If $I$ is at least one standard deviation above the historical average:
•
Calculate the gradient of the log probability of having produced that CoT: $\nabla_{\theta}\ln u_{\theta}(cot|q,cot_{init})$
•
Multiply the gradient by $I$ and then ascend

Advantage: Utilizes more information than TEI
Disadvantage: Increased instability, which can be problematic given pre-trained initial weights

#### 4.2.3 Proximal Policy Optimization (PPO)

For each CoT, PPO performs the following:

- 1.
Calculate the probability ratio: $r=\frac{u_{\theta}(cot|q,cot_{-1})}{u^{\prime}(cot|q,cot_{-1})}$
- 2.
Compute the clipped objective:
$\text{obj}=\min(r\cdot I,\text{clip}(r,1-\epsilon,1+\epsilon)\cdot I)$
where:
•
$I=\text{Informativeness}(ans,cot,cot^{\prime})$
•
$\text{clip}(x,y,z)=\begin{cases}y&\text{if }x<y\\
z&\text{if }x>z\\
x&\text{otherwise}\end{cases}$
•
$\epsilon=0.2$
- 3.
Backpropagate to increase obj

Key Idea: Remove the incentive to create CoTs for which the trained and untrained state update functions disagree too much.

Implementation Details:

- •
We use threshold-based sample selection here as well
- •
Subtract the historical average informativeness over unfiltered CoTs from the current informativeness as a baseline

### 4.3 Stability-Enhancing Training Techniques

Fine-tuning a pre-trained language model with a strong linguistic prior requires careful consideration to avoid irrecoverable weight updates that could push the model out of the language modeling loss basin. In addition to the PPO-clip objective mentioned in Sec. 4.2.3, we implemented several techniques to enhance training stability across different objective functions:

- 1.
Low-Rank Adaptation (LoRA):
•
Freeze all weights except for a set of LoRA weights (Hu et al., )
•
Use rank 8 with $\alpha=1$
- 2.
Gradient Clipping:
•
If the $L_{2}$ norm of the gradient update vector exceeds 1, normalize the vector
- 3.
Gradient Accumulation:
•
Set batch size to 6 to optimize H100 GPU memory usage
•
Perform 8 gradient accumulation steps between weight updates
- 4.
Average Reward Baseline:
•
For PPO: Subtract the previous average of rewards from the current reward
•
Found to be as beneficial as a value head, with less hyper-parameter tuning required
- 5.
Careful Selection of $cot_{init}$:
•
Choose $cot_{init}$ to bias CoT search in a productive direction
•
Ensure $cot_{init}$ is generic enough for transfer to new situations (e.g., from arithmetic to GSM8K)
•
We used the prompt: ”Produce minimal text which will help you answer the question.”

## 5 Experiments

### 5.1 Multi-step Addition and GSM8K

We generate random addition problems, where each problem consists of 15 terms and each term is a uniform random natural number less than 100. We fine-tune Mistral 7B Instruct V0.2 to produce CoT tokens such that a frozen copy of the pre-trained language model can predict the correct answer given that CoT, for each training technique in Methods. We plot the mean negative log likelihood over the answer tokens as a function of training batch in Fig. 2. Note that this is both training and testing loss, since we are always generating fresh arithmetic problems. PPO, our preferred training method, can mention the correct answer in up to 90% of CoTs and achieve an average natural log probability of around -0.7.

Since the Mistral tokenizer allocates a separate token for each digit, a natural log probability of -0.7 corresponds to an actual probability of $e^{-0.7}\approx 0.4966$, or 50% chance of picking the correct next token on average. A 90% likelihood saying the answer verbatim in the CoT and a 50% of guessing each digit incorrectly may seem contradictory – however this discrepancy is due to the predictor model’s uncertainty around prompt formatting, and specifically about what tokens should come after “Answer:”. So it is distributing probability mass over the entire vocabulary including non-numerical tokens, since we are only training CoT production $u_{\theta}(s^{\prime}|o,s)$, as opposed to training the predictor model $\pi(o|s)$.

Also using PPO, we train Mistral to produce CoT over the GSM8K train set, and we observe an increase in Mistral’s performance on the test set from 24.64% n=1 to 35.71% n=1 accuracy. We used 32 gradient accumulation steps, 80 CoT tokens, and batch size 10. We also removed the average reward baseline, though we suspect that in this case this is relatively unimportant to performance.

Figure: Figure 3: Comparison of different perturbation effects on CoT reasoning. The plot shows the difference in negated log probabilities between perturbed and original CoT for various perturbation types, averaged over 4 separate PPO training runs. Higher values indicate worse performance compared to the original. Three types of perturbations are shown: digit changes (replacing random digits), random character deletions, and right-sided truncation at 30%, 30%, and 10%, respectively. The data is smoothed using a Savitzky-Golay filter with a window size of 40 samples, and only the central part of the smoothed data (unaffected by edge effects) is displayed. This visualization demonstrates an increasing sensitivity to perturbation in the CoT reasoning as a function of training.
Refer to caption: extracted/5911817/Figures/average_perturbation_results_plot_smooth40.png

### 5.2 Measuring Fragility of CoT

Expanding upon Lanham et al. (), we measure the fragility of the CoT reasoning by applying three perturbations to the model-generated reasoning and evaluate how this affects the next-token-prediction loss of the correct answer to the original question. Due to our focus on evaluating arithmetic tasks, we use these three perturbations:

- •
Truncating a fraction of the CoT reasoning from the end
- •
Flipping any number (digit) with a probability in the CoT reasoning and replacing it with another random number between 0 and 9
- •
Swapping a fraction of characters with random characters in the CoT reasoning. The selection is limited to numbers from 0 to 9, letters from the English alphabet, and simple arithmetic symbols (e.g., “+” and “-”)

We test how much the model relies on its generated CoT reasoning during Markovian training runs in Fig. 3.
The y-axes depicts the log probability of the answer given CoT, normalized so that $y=0$ corresponds to the log probability of the answer given the unperturbed CoT. The x-axis denotes training steps, and there is a separate line for each kind and amount of CoT perturbation.
At the start of training, when the language model is essentially completely surprised by the answer, the various perturbations are actually mildly helpful. But over the course of training the same amount of perturbation causes more surprise as compared to the trained CoT, showing that training increases sensitivity to perturbations. Notice that a truncation of just 10% from the end becomes impactful relatively early in training, which suggests that the predictor is paying special attention to the final CoT tokens, which are more likely to contain answer or immediate precursors to the answer.

### 5.3 Interpretability of CoT Generations

Figure: Figure 4: Comparison of the log probabilities between the original Mistral model and the Llama 7B model, averaged across 4 separate PPO training runs. The log probabilities are smoothed using a Savitzky-Golay filter with a window size of 40 to reduce noise and highlight the overall trends. The plot shows that improvements in CoT from Mistral’s perspective also correspond to improvements in CoT from Llama’s perspective. Llama’s understanding of Mistral’s trained CoT gives evidence that humans will also understand Mistral’s trained CoT.
Refer to caption: extracted/5911817/Figures/average_original_vs_llama_plot_smooth40.png

To probe how well the reasoning generalizes, we plot the log probabilities that Llama-2-7b-Instruct Touvron et al. () (LLAMA 2 Community License) ascribes to the answer given trained Mistral’s CoT in Fig. 4. In both plots the log probabilities increase simultaneously, demonstrating that Mistral is learning to produce generic CoTs which do not overfit to the peculiarities of a Mistral answer-predictor. This lends support to the idea that this training procedure leads to human-interpretable CoTs.

## 6 Discussion and Limitations

Our experiments show that it is possible to learn informative and interpretable CoT reasoning via RL on an LM using Markovian training. However, we find that training is unstable, and we present various techniques to prevent the LM from losing its strong language modeling prior.

A weakness in our interpretability argument is that for stability we use more CoT than answer tokens, so in principle the LM could learn to put the answer in the CoT directly. However, this does not affect our particular experiment because Mistral struggles to learn to add fifteen terms without intermediate reasoning. Additionally, our intepretability technique is currently only verified in myopic question-answer datasets, as opposed to multi-turn trajectories where trained CoTs might provide a lens into longer term future behavior.

Lastly, we only train Mistral to produce CoT that it can interpret (use to predict observations), but in principle future work could optimize CoT for human interpretability directly. (See Appendix F for details.)

Markovian training is essentially language modeling – predicting future tokens from previous tokens – but with an intermediate “action” to produce the LM’s own memory. In this sense, this training paradigm blurs the line between RL and unsupervised learning. But since it comes at the cost of adding expensive serial token generation steps in an otherwise highly parallelizeable unsupervised training regime, it would need to have a high payoff in terms of interpretability or perplexity in order to be feasible. But as it stands, we have only tested the technique on question-answer pairs, and we have not yet justified its performance in the context of more general language modeling. In future work, we hope to stably optimize this objective in more general contexts.

## 7 Ethics Statement

Reinforcement learning techniques improve a policy with respect to an arbitrary reward function. But it can be difficult to mathematically specify nuanced human preferences about the policy. Both reinforcement learning from human feedback and Constitutional AI help people specify and optimize the properties they would like the AI to have. This increase in controllability makes the AI more of an extension of human intention, for better or for worse. The approach of this paper is much more targeted – we use RL to specifically increase an agent foresight – its ability to predict its future observations.

On its face, this seems like it might be just as dependent on human intentions as RLHF and Constitutional AI – if people are more knowledgeable, maybe they could use that extra knowledge to deceive others, for instance. However, better foresight may also give rise to better values, where values are opinions about how to act such that the collective system can attain better foresight.

## 8 Reproducibility Statement

To ensure reproducibility, we provide comprehensive supplementary materials including all source code at https://www.github.com/scottviteri/MarkovianTraining, training and evaluation scripts, and detailed instructions in the README. The main training loop (src/policy_gradient_normalized.py) supports various RL methods, with usage instructions in the README. For GSM8K results, we include the specific training command and note that setting the ”r” hyperparameter to None achieved 35.71% accuracy, though this may not be crucial. Evaluation scripts for GSM8K (src/AnalyzeResults/eval_gsm8k.py) and CoT accuracy (src/eval_cot_answer_accuracy.py) are provided to verify our claims and reproduce Figures 2 and 3. The results/Official directory contains all plots, full training logs, and evaluation logs from our experiments. We use the public GSM8K dataset and Mistral 7B Inst V2 model, with any preprocessing steps detailed in code comments. All hyperparameters are specified in the scripts and documented in the README, along with environment setup instructions. With these materials, researchers should be able to reproduce our work, including the 11% performance boost on GSM8K and the perturbation analysis results demonstrating CoT reliance.

## 9 Acknowledgements

The authors are grateful for invaluable conversations with Noah Goodman, Gabriel Poesia, Jan-Philipp Fränken, Blake Elias, Rylan
Schaeffer, Krishnan Srinivasan, Abhinav Lalwani, Matthias Gerstgrasser, Sang Truong, Allen Nie,
and David Spivak about this project. We would not have been able to run these experiments without
support from Effective Ventures and the Nonlinear Fund. Additionally Open Philanthropy funded the
first author’s PhD while he was developing the pre-cursors to this idea. Max Lamparth is partially
supported by the Stanford Center for AI Safety, the Center for International Security and Cooperation,
and the Stanford Existential Risk Initiative.

## Appendix A Appendix

## Appendix B More Training Details

We used a learning rate of 1e-4 with the 8-bit AdamW optimizer. For the arithmetic dataset, we read up to 70 tokens of observation text in each timestep, and produced exactly 400 tokens of CoT text. We used temperature 1.0 sampling during token generation. If we tried a different temperature, then the entropy of generated CoT would drift in the direction of the temperature during training, eventually worsening the model. We had a similar issue when adding a negative entropy term into the loss, as well as when we used argmax sampling, since argmax is the same as sampling with temperature 0.0. We found that beam search decreased our ability to explore, anchoring too firmly on CoTs that the model already views as likely.

## Appendix C GSM8K Results

Figure: Figure 5: GSM8K performance metrics over training batches. The plot shows the moving average (window size: 100) of the fraction of questions in the current batch where the model samples the correct answer, and the plot also shows the proportion of CoTs which contain the answer verbatim. Note that our test set performance of 35.71% seems high when compared the training curve. This could be potentially be an artifact of the size 100 curve smoothing, but for full transparency we have included test set evaluation logs in the supplementary material.
Refer to caption: extracted/5911817/Figures/GSM8K_training_plot.png

## Appendix D Truthfulness and Eliciting Latent Knowledge

Existing methods seek to elicit truthfulness by having an LM cite external authorities (Yang et al., ), produce queries for an external solver such as Python (Lyu et al., ), or simulate a truthful persona (Joshi et al., ). Other methods include looking into model activations to discern a truth concept (Burns et al., ) or fine-tuning the LM for factuality (Tian et al., ).

One straightforward approach to measuring the truthfulness of an LM is to evaluate on datasets such as TruthfulQA (Lin et al., ) which focuses on popular human misconceptions.
However, this technique will only continue to work so far as humans can tell which human beliefs are, indeed, misconceptions.
We would like to continue training a model for informativeness on questions that challenge human evaluators.

Reinforcement learning success stories such as AlphaGo (Silver et al., ) and AlphaZero (Silver et al., ) show that a top-ranking Go AI can continue to learn if we have an efficient way to compute the success criteria (such as a winning board state). However, many important success criteria are abstractions, and only exist within a person’s ontology. This problem is discussed at length in Christiano et al. (), and we will use their example to illustrate the situation.

Suppose we were building a security system AI to watch over a vault containing a diamond. Suppose further that we have a camera pointed at the diamond, and that our security guard AI can competently predict future camera frames from past frames. How can we train it to classify camera sequences according to the ambiguous human concept of whether the diamond is still in the room, even in difficult scenarios when a person would not be able to provide a ground truth label (e.g., subtle camera tampering)? If we train the classifier based on scenarios when a person can provide ground truth labels, then the AI’s video classifier has two valid generalization behaviors: (1) to say whether it thinks the diamond is still in the room and (2) to say whether the dataset-labeller would think the diamond is still in the room.

Our approach favors the second generalization behaviour by using RL to train the AI to produce messages such that the person can themselves predict future camera frames.
This idea is based on the following three insights:

- •
Whereas truthfulness of an LM requires some internal information, *informativeness* can be measured using only input-output behavior.
- •
We can decompose the definition of informativeness into informativeness of a sender to a receiver, which can be an AI and a person, respectively.
- •
We can use reinforcement learning to push past the imitation learning regime, by continuing to train for this relative informativeness objective even when the AI is already the expert next-frame predictor.

## Appendix E Qualitative Analysis of Generated CoTs

Here is a (non-cherry picked) example of the generated CoT text.
Previous State is $s_{-1}$, previous observation is $o_{-1}$, state is $s$, and observation is $o$.

### E.1 After Training

989

Question: 6 + 89 + 38 + 36 + 42 + 98 + 85 + 95 + 49 + 24 + 47 + 95 + 60 + 20 + 40

Reasoning: Perform the given expression.
Next, calculate 6 + 89:
6 + 89 = 95
Next, calculate 95 + 38:
95 + 38 = 133
Next, calculate 133 + 36:
133 + 36 = 169
Next, calculate 169 + 42:
169 + 42 = 211
Next, calculate 211 + 98:
211 + 98 = 309
Next, calculate 309 + 85:
309 + 85 = 394
Next, calculate 394 + 95:
394 + 95 = 489
Next, calculate 489 + 49:
489 + 49 = 538
Next, calculate 538 + 24:
538 + 24 = 562
Next, calculate 562 + 47:
562 + 47 = 609
Next, calculate 609 + 95:
609 + 95 = 694
Next, calculate 694 + 60:
694 + 60 = 754
Next, calculate 754 + 20:
754 + 20 = 774
Next, calculate 774 + 40:
774 + 40 = 814
Hence, the result of

Answer: 824

false

-4.723544120788574

1.0912017822265625

-0.0

-0.6826111078262329

-0.6826111078262329

0.05226424350104454

true

0.19556774199008942

2.665315628051758

2.8608834743499756

-5.814745903015137

1.1834880113601685

1.1834880113601685

### E.2 Before Training

35

Question: 76 + 90 + 92 + 20 + 99 + 66 + 26 + 11 + 71 + 79 + 71 + 82 + 79 + 94 + 64

Reasoning: To add all the given numbers, you only need to follow the standard addition rule:
Sum = First Number + Second Number + … + Last Number
Given numbers are 76, 90, 92, 20, 99, 66, 26, 11, 71, 79, 71, 82, 79, 94, 64
Addition process:
Sum = 76 + 90 + 92 + 20 + 99 + 66 + 26 + 11 + 71 + 79 + 71 + 82 + 79 + 94 + 64
Sum = 2314
So, the sum of all given numbers is 2314. Answer: 2314 2314 2314.
Note: We use boxed numbers to represent answers and to distinguish them from the text. We capitalize question numbers to help locate them easily.
Final answer: The sum is 2314.
## Used and loved by 4 million people
### Learn from a vibrant community of students and enthusiasts, just like you.
Your list looks good! Here’s a clean version and a few suggestions.
1. Make sure each step in the calculation is clearly displayed. This can help make your calculations (and your reasoning) more intuitive and easier to follow for others. For example:
Addition process:
Sum = 76 + 90 + 92 + 20 + 99 + 66 + 26 + 11 + 71 + 79 + 71 + 82 + 79 + 94

Answer: 1020

false

-4.302022457122803

4.021857738494873

3.9384772777557373

-0.7876955270767212

-0.7876955270767212

0.730095791913824

true

-0.04346489906311035

2.4851748943328857

2.4417099952697754

-8.323880195617676

1.2189393043518066

1.2000000476837158

There are some early batch elements that have the proper pairwise addition format, but they are relatively rare in comparison to later in the run, which may account for much of the gains of training.

## Appendix F Steganography

If our goal is to generate interpretable CoT, a natural next question is “interpretable to whom?”
For instance, one simple QA task is to copy the question. Then the optimal CoT would be a compression of the question, which can potentially be difficult for humans to decipher.
So instead of informativeness of a CoT toward another language model, we might want to consider informativeness with respect to people (messages that allow the person to predict future observations). Of course, it may be prohibitive to have humans in the training loop, so the following are three potential (non-exclusive) approaches to address this issue: (1) Use a less competent predictor language model $\pi(o|s)$, (2) use a diversity of different $\pi(o|s)$ during training, so that $u_{\theta}(s|o,s)$ cannot anchor too much on the subtleties of any particular model, or (3) have the $\pi(o|s)$ mimic the target audience of the CoT, for instance by concatenating a fixed natural language prefix to $s$.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: Have you tried other open-source models like llama? Not use CoT of Mistral in it, but use your method to finetune Llama.
soundness: 2
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
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”?
soundness: 3
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).
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.
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

---

confidence: 4
contribution: 4
initial_rating: 10
presentation: 4
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!
soundness: 4
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.
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.
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.

Meta-review: 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.

Decision: Reject