Markovian Transformers for Informative Language Modeling

Thank you for your review!

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.

In the Author comment we talk about how we have extended our method to the Wikipedia dataset using Llama. Our perturbation and cross-model generalization results extend nicely to the Wiki-trained model.

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!

We are happy to change the line styles.

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

We have expanded our experiments significantly beyond mathematical reasoning. In particular, we've trained Llama-3.1-8B-Instruct to generate CoT for predicting Wikipedia text. These experiments demonstrate three key findings:

1. Training improves Wikipedia language modeling performance
2. The CoT becomes increasingly fragile to perturbations during training, suggesting genuine reliance on the reasoning
3. The generated CoTs generalize across architectures, successfully helping Mistral predict future text

These results are detailed in our new appendix on Wikipedia experiments, demonstrating the method's applicability beyond mathematical reasoning and across different model architectures.

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

We appreciate this organizational feedback. We've moved the GSM8K training results to the Experiments section for better visibility of these key outcomes. However, we maintain the detailed training methodology in the main paper because these specific implementation details are crucial to achieving stable reinforcement learning with language models. Without these carefully tuned training procedures (which we detail in Section 4.3), the approach does not converge successfully, even with larger models. These details represent core technical contributions that enable the practical application of our method.

We appreciate your detailed comments!

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.

We were originally thinking about how to define "faithfulness of an explanation" or a language model's "truthfulness" more formally. One very abstracted notion of a truthful message is one that counterfactually improves the recipient's future predictions of their observations. One immediately runs into the question "counterfactually with respect to what?", and we picked the counterfactual message to be that of a slightly younger version of the sender (the pre-trained model) so we can talk about how much more "truthful" a person/LM is getting over time. But we soon found that truthfulness is not a good word for this quantity. For instance, if I say "2+2=4" to you, that may be True in a Platonic sense, but it does badly according to our metric since you already knew "2+2=4" so it doesn't help you predict your environment better. So we opted for "informativeness" instead, which works for the "2+2=4" example. There may be some kind of mutual information connection that is more formal, but we didn't philosophize too much in that direction because at the end of the day it seems more practical to try minimizing KL than to make some kind of variational estimate of mutual information.

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.

We agree with this sentiment, and while this technique helps for GSM8K, which is written in natural language and therefore somewhat more open-ended, it is still not nearly as open-ended as, say, Wikipedia text. This is why we have trained models to produce CoTs which help predict Wikipedia text, which we mention in the author comment and we detail in Appendix A.

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.

To help address this point, we've added the Appendix D discussing baselines in detail.

Our process for picking an RL algorithm was less intuition based and more like trying a large array of different algorithms, techniques, and hyperparameters, and seeing which ones worked.

For instance in Figure 2, we plot the results of using policy gradient/REINFORCE versus expert iteration versus PPO, and we find that for arithmetic they are largely comparable, though PPO outperforms on average. By expert iteration here we really just mean only training on actions which did one standard deviation better than the historical average (See 4.2.1).

The detailed methods we give such as LoRA, gradient clipping, a frozen critic, and value function estimates using a weighted decay all stem from having tried the method without these additions and finding the performance to be lacking.

The main reason we haven't included more ablations is that each run with Mistral takes several H100 hours -- that said, we are adding some now.

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. Teaching Large Language Models to Reason with Reinforcement Learning. Havrilla et al. 2024 Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs. Ahmadian et al. 2024

We are happy to add these references, and we can add ablations in the appendix.

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.

The space of states is the set of length-n token sequences, where in practice we set n to be the smallest value we can find such that training converges. Indeed the action space of good CoTs is quite sparse, similarly to the sparsity of useful natural language strings in the space of all possible strings.

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"?

We have found success with the GSM8K (11%) reasoning task and some success with Wikipedia.

Thank you for the detailed response!

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.

We agree that checking whether the CoTs help Llama predict next tokens is a proxy. We discuss potential better evaluation approaches in our "Steganography" appendix, addressing three key concerns:

Threat Model 1: Information Accessibility

• Concern: CoTs may encode information that only high-capacity models can utilize
• Test: Evaluate using intentionally less capable models
• Rationale: Success with simpler models suggests accessible encoding

Threat Model 2: Model-Specific Encoding

• Concern: CoTs might only work for certain architectures
• Test: Evaluate across diverse models from Huggingface
• Rationale: Cross-architecture generalization indicates model-agnostic representations

Threat Model 3: Machine-Human Gap

• Concern: CoTs might be fundamentally inaccessible to humans
• Test: Measure humans' ability to predict subsequent text
• Rationale: Direct human evaluation provides strongest evidence of interpretability

For human trials, we envision two approaches: (i) Interpretability -- Have people predict which completion the LM assigns higher probability to, given the CoT (ii) Informativeness -- Test how well people can predict future observations given the CoT

The informativeness objective is most interesting when the domain being predicted is larger than the CoT size, as with Wikipedia examples, because then the CoT must meaningfully compress future information. This motivates learning to produce CoT such that people can succeed on multiple choice questions from a broad distribution unknown to the LM during training.

The scope of human trials is large enough that we feel it's best relegated to future work. We have added an appendix demonstrating generalization between Llama 8B and Mistral on Wikipedia text prediction.

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 arithmetic.

We hope that our results on Wikipedia help to address this concern! Also, when we train Mistral on the GSM8K reasoning task, we observe an increase from 24.64% n=1 to 35.71% n=1 accuracy on the test set.

There seem to be no baselines and ablation designed.

We address this in a new appendix "On Baselines for Faithful CoT" where we discuss three categories of potential baselines. We find that either no clearly applicable baseline exists, or we are already including it in our plots. We are currently running hyperparameter ablations.

The fragility analysis is insightful but lacks depth.

For the Wikipedia task, we've added 4 types of perturbation (character deletion, front/back truncation, random replacement) at 5 severity levels. Our training makes CoT fragile to each type.

The method section is presented before establishing the limitations of existing CoT techniques.

Since this is a novel setup, we cannot directly use existing techniques like QuietStar where models access history before the CoT. Our methods section builds up from simpler techniques to make the explanation more digestible.

There are few tables but quite a few definitions and equations.

We've moved the GSM8K training plot to the experiments section. However, the training method details are crucial as they enable stable RL training with LLMs.

Is there a way to combine the interpretability with actual human perception?

Measuring human predictions is challenging - even determining how surprised someone is by text is non-trivial. We could start with multiple choice predictions, where people guess which completion is more likely given a CoT.

Is the optimal CoT compression observed in your experiments?

We were specifically discussing question-answer tasks where the answer copies the question - in such cases, any CoT smaller than the question is definitionally a compression.

Were there more ablation study or comparison conducted?

We are running hyperparameter ablations for the appendix. For comparisons, please see Figure 2 and the baselines appendix.

• The main concern “it’s uncertain if this CoT is genuinely interpretable by humans” is not addressed after the rebuttal.

  • This manuscript motivates the method in both the abstract and introduction and argues that interpretability techniques are important. For example, the abstract highlight the problem of “interpretability”, and this paper is to “address this issue”. In the introduction, the author motivates with “high-stakes scenarios”. Given this motivation, it is not sufficient if how interpretability improves for human readers is studied. If the authors consider human interpretability and study should be another work, then major revision is needed for the motivation part of this paper. “Envisioned methods” are without results are not enough to prove the case, so they are not enough to serve as evidences.
  • I appreciate the authors trying to address the comments in the appendix, but a proposed method is not enough to support the claims. Further, like many has pointed out, the appendix will not be treated as important as the main body.
  • I further notice in the “steganography” appendix, L1011 and 1012 are undeleted comments of the authors, indicating the manuscript isn’t fully proofread.

• Baselines and insufficient analysis of section 5

  • The baselines designed for this paper should also be tie to the motivation (which is currently interpretability). A study needs to be setup to show how the method improves the interpretability. In my review I pointed out that 5.2 is lacking depth, the same applies to 5.3 (if not more), since if the key motivation of this paper is interpretability, I would expect a much richer version of the analysis section.

• Not enough tasks/datasets to prove the generalizability and effectiveness.

  • CoT is a technique generally applicable to many tasks, and if we want to show an improved version of that, more tasks and datasets need to be tested. The “Wikipedia” experiment is not convincing enough as an alternative challenging reasoning task, and one more next token prediction task still cannot convince me the “generalizability” of this work.

I do not see this paper passing the acceptance threshold without significant revision of text, more experiments and more insightful and convincing arguments.

Peer-review is a very noisy process. There is a decent chance that this work will be appreciated the next time it is submitted. Dear authors, please keep trying.