Understanding Chain-of-Thought in LLMs through Information Theory

First of all, we would like to thank the reviewer for their time and feedback on our paper. Here below, we discuss the thought-provoking questions raised by the reviewer.

The proposed paradigm seems challenging when attempting to explain existing R1-like work, particularly in the context of handling exploratory nonlinear or branched inference paths. How does the model account for these complexities, and are there any strategies for maintaining accuracy and coherence in such cases?

Thank you for raising this important question, especially now that reasoning models like R1/O1 have become increasingly popular. Taking a step back, our method is built on an information-theoretic framework that evaluates each step of the chain-of-thought reasoning process. We measure the information gain (IG) at every step to determine whether that step is adding useful information toward predicting the correct final answer. The underlying assumption is that a correctly executed reasoning step should yield a positive IG, meaning it contributes meaningfully to the overall prediction. If a step is executed incorrectly or is unnecessary, the IG will be low or close to zero.

For settings like R1/O1-like reasoning traces, which often involve exploratory or branched inference steps, our framework remains applicable. More specifically, if a model initially follows an incorrect path, our method will indicate that early steps have low or negative information gain. When the model later self-corrects, our method is able to label subsequent steps with a positive IG, effectively indicating that we are on track towards the correct outcome. This demonstrates that our framework is still able to reliably work even in settings like R1 if the steps are clearly delineated. This is also evident from our PRM data experiments, where the steps deemed irrelevant by human annotators have low IG, while the correct steps have a high IG.

We will make sure to add this clarification in the final version of the paper.

I noticed that the robustness of the model isn't fully addressed. Given that the model's probability distribution appears similar to a reward function, it seems potentially vulnerable to adversarial manipulation or instability. Could the authors elaborate on any measures taken to improve the model's resilience to such issues?

Thank you again for pointing out this important issue of robustness. We have actually conducted experiments specifically designed to assess the robustness of our framework. In particular, we investigated the impact of spurious correlations on the evaluation of intermediate reasoning steps. In our experiments, standard methods such as Math-Shepherd and ORM were affected by spurious signals. They incorrectly inferred the usefulness of intermediate steps when spurious correlations were present.

To be concrete, in Figure 3, we designed an experiment where we injected spurious correlations into step 3, linking its correctness to a spurious feature of the previous step. In this case, we clearly demonstrated that existing methods are unable to detect the precise step where the error occurs, whereas our method is able to pinpoint it. In addition, Table 2 (Arithmetic dataset) further illustrates the lack of robustness of existing methods. When only the final step of the chain-of-thought is problematic, Math-Shepherd flags all steps wrong and ORM provides an uninformative score for each step. In contrast, our approach (IG) directly computes the useful information content in each step by measuring the predictive power added toward the final answer. This design choice makes our method robust to such spurious correlations.

Next, although our supervisor model can be interpreted as a reward model, it differs significantly from conventional ones that rely solely on correctness or preference signals. As mentioned above, our framework aims to predict the final answer tokens and not a binary label, contrary to existing methods.

Lastly, we acknowledge that computing information gain can become challenging when the chain of thought is adversarially long; in those cases, it may be necessary to employ stronger supervisor models to accurately capture the information gain at each step. We only tested with GPT-2 and LLama3. Our experiments, including those with reasonably long chains of thought as seen in the PRM evaluations (Table 2), confirm that our approach remains robust even under these long chain conditions. We plan to explore the analysis of very long chains of thought in future work.

References

We thank the reviewer for these references and will include them in the final version of our paper.

We thank the reviewer again for their insightful comments to improve our paper. If the above addresses all outstanding questions, we hope that the reviewer would consider raising their score. We remain happy to answer any further questions.

First of all, we would like to thank the reviewer for their time and constructive comments to improve our paper. Here below we clarify all the questions raised by the reviewer.

The empirical evaluations and adaptability to diverse CoT structures

Our framework evaluates each reasoning step using the information gain (IG) metric, measuring its contribution toward the final answer. Although we present our method with linear chains-of-thought for clarity, it applies equally well to non-linear structures. For example, in O1/R1-like reasoning where multiple paths are explored, early erroneous steps yield low or negative IG, while correct steps show positive IG. Our experiments on the MATH/PRM dataset confirm that steps deemed uninformative by human annotators consistently exhibit low IG, demonstrating that our method robustly adapts to varied chain-of-thought structures.

The proposed method is established under somewhat strong assumptions,

We appreciate the reviewer’s insightful comment. Firstly, we would like to clarify that we have used Assumptions 3.1 and 3.2 to rigorously motivate our use of information-gain for failure mode detection in LLM CoTs. In practice, however, our methodology remains applicable to real-world datasets (as shown by our PRM data experiments), where it outperforms the ORM baseline. Secondly, we emphasise that these assumptions are formalizations of phenomena which make intuitive sense in practice. Specifically:

Assumption 3.1: This assumption formalizes the idea that each correct reasoning step should contribute additional information for predicting the final answer. In our framework, a positive information-gain (IG) indicates that a step is informative, whereas a low or negative IG suggests that the step is either erroneous or superfluous. Importantly, this criterion is scalable and remains applicable to long chains-of-thought (CoTs). Whether the chain is short or long, if each step is clearly delineated, our method evaluates each step based on the IG metric. As demonstrated in our PRM experiments, even in cases where early steps may show low IG due to initial missteps, the framework still identifies the transition to correct reasoning later in the chain.

Assumption 3.2: While this assumption might seem strong, it is intuitively grounded in the operational characteristics of large language models. It essentially posits that the operations the model applies during reasoning are restricted to a set of primitives (or their compositions) learned during training. This is a reasonable expectation, as models are generally more effective when they perform tasks that are similar to those encountered during training.

Despite the theoretical rigor of these assumptions, our experiments on uncontrolled datasets (such as PRM800K and the Llama-3-8B arithmetic tasks) demonstrate that our framework remains effective in realistic settings where parts of these assumptions may not strictly hold.

In summary, while these assumptions are used to rigorously motivate our framework, our experiments indicate that the methodology is robust and practically applicable beyond the constrained formal setting.

Definition "Unidentifiability" for interpreting few-show generalization.

We thank the reviewer for this interesting suggestion. Although our focus isn’t on few-shot generalization, our method can be extended to few-shot settings. In our formulation, unidentifiability measures whether a task lies outside the span of learned primitives. In a few-shot setting, if adding examples increases the information gain for a reasoning step, it suggests that the examples help the model perform that step correctly; if not, the task remains unidentifiable. We see this as a promising direction for future research which is outside the scope of this paper.

It would be convincing if the authors could give more empirical evaluations on real datasets.

We would like to clarify that, in addition to our synthetic data and controlled GSM8K experiments, our submission also includes evaluations on real-world data:

• PRM800K Dataset: This dataset covers real mathematical problems from high school to post-graduate levels. Our experiments on PRM800K show that our methodology predicts the correctness of intermediate steps more accurately than the ORM baseline, making our IG a cost-effective proxy for human-annotated labels.
• Llama-3-8B Arithmetic Experiment: Although generated by us for a specific arithmetic task, this uncontrolled experiment demonstrates that our method correctly identifies errors, specifically, misapplications in the final addition step. In comparison, the baselines either provide uninformative results (ORM) or exhibit a high false-positive rate (Math-Shepherd). Together, these experiments underscore the practical utility of our approach in realistic, uncontrolled settings.

We hope that the above have addressed all the reviewer's questions and that the reviewer would consider raising their score.

We thank the reviewer for their thoughtful questions and feedback. Below, we address the concerns raised:

Scalability/Generalization: Additional training required of a supervisor model, training details and construction of new dataset

Training and Data Setup: Our method requires additional training of a supervisor model, primarily a small GPT-2 (117M parameters), fine-tuned on about 10,000 samples within 3–6 GPU hours on an A100 (learning rates: 5e-6 to 5e-7, batch size: 64). Although Llama performed slightly better on arithmetic tasks, GPT-2 is generally similarly effective and more cost-efficient on other datasets. Importantly, our approach does not require constructing a separate dataset; it only uses the final answers and the model’s CoT outputs e.g., just 3,000 samples in our PRM800K experiments.

Generalization: We acknowledge that this extra training may limit generalizability, a common challenge in reward modeling also faced by methods like ORM, and we appreciate the reviewer’s suggestion; we believe that future work, such as exploring in-context IG estimation techniques, will help address these limitations.

Empirical validation on non-mathematical tasks

In line with recent works on LLM reasoning [Wang et al., 2024b; Havrilla et al., 2024], we have focused on mathematical datasets since these provide a rigorous testbed for LLM reasoning, and determining the correctness of intermediate steps is unambiguous in these datasets. However, we agree that extending this to other reasoning datasets would be valuable in future.

Identifying LLM reasoning weaknesses and Llama-8b performance breakdown

Our experiments demonstrate that our framework effectively identifies specific reasoning weaknesses compared to existing methods. In fact, this is how we discovered that Llama is not proficient at adding large and small numbers together. Specifically, we observed a real and notable accuracy drop in Llama-3-8b for arithmetic tasks involving numbers greater than 10⁵, likely due to limited exposure in its training data. We believe investigating this inherent limitation further is an interesting avenue for future work and will emphasize this clearly in the final version of our paper.

Long CoT with self-correction (R1/O1 style generations)

Thanks for raising this important point. Taking a step back, our method relies on an information-theoretic framework that measures information gain (IG) at each chain-of-thought step. A correctly executed step yields a positive IG, while an erroneous or unnecessary step shows low or near-zero IG. As the reviewer rightly mentioned, in R1/O1-like traces with exploratory or branched inference, our framework actually remains robust: early missteps have low or negative IG, and when the model self-corrects, subsequent steps exhibit positive IG, signaling a return to the correct path.

It is evident that not all reasoning following an erroneous step is useless to the final result

We would like to clarify that our assumption 3.1 specifically considers the case where LLM's final answer is incorrect. More generally, in cases of self-correcting CoTs, our framework remains applicable as we explained above. For more details, please refer to our first response to Reviewer yLdQ.

Mapping each reasoning step to primitive tasks appears to require substantial manual summarization

We only use categorization to compute aggregate information gain, when the goal is to obtain a high-level view of an LLM's intermediate reasoning. This categorization is applied only on the evaluation split, not during training. Our PRM experiments show that our framework identifies errors without explicit categorization. In addition, to get annotations, for instance in GSM8K, we can efficiently prompt an LLM to classify each substep into categories like ['Addition', 'Subtraction', …, 'Other']. We appreciate this comment and will clarify it in the final version.

Theoretical proofs are based on several strong assumptions

We have used these assumptions to rigorously motivate the use of information-gain in our framework. In practice, however, our method remains applicable to real-world datasets (as shown by our PRM experiments) with potentially non-linear/branching CoTs. Additionally, we would like to emphasize that both of these assumptions are intuitively plausible.

Briefly, Assum. 3.1 states that each correct step should add information for predicting the final answer, while wrong/irrelevant steps should not. Likewise, Assum. 3.2 posits that the operations the model applies are restricted to those learned during training. Despite these assumptions, our experiments on uncontrolled datasets (such as PRM800K and the Llama-3-8B arithmetic tasks) show that our framework remains effective in realistic settings where parts of these assumptions may not strictly hold.

Lastly, we hope our clarifications above addressed the reviewer's concerns, and the reviewer would consider increasing their score.