Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Understanding Chain-of-Thought in LLMs through Information Theory

Jean-Francois Ton, Muhammad Faaiz Taufiq, Yang Liu

TL;DR

This paper tackles the challenge of evaluating chain-of-thought (CoT) reasoning in LLMs without relying on annotated intermediate steps. It introduces an information-theoretic framework that quantifies information-gain at each reasoning step via a supervisor model to estimate p(Y|X^M_j), enabling detection of unidentifiable or non-contributory steps. The approach is validated on toy data, GSM8K, and PRM800K, where it outperforms outcome-based baselines like ORM and Math-Shepherd in identifying faulty CoT steps and providing granular, task-specific diagnostics. The results suggest that information-gain-based evaluation offers more reliable, scalable insights for CoT quality and generalizes across nonlinear or exploratory reasoning patterns, with potential extensions to non-mathematical domains.

Abstract

Large Language Models (LLMs) have shown impressive performance in complex reasoning tasks through the use of Chain-of-Thought (CoT) reasoning, allowing models to break down problems into manageable sub-tasks. However, existing CoT evaluation techniques either require annotated CoT data or fall short in accurately assessing intermediate reasoning steps, leading to high rates of false positives. In this paper, we formalize CoT reasoning in LLMs through an information-theoretic lens. Specifically, our framework quantifies the `information-gain' at each reasoning step, enabling the identification of failure modes in LLMs without the need for expensive annotated datasets. We demonstrate the efficacy of our approach through extensive experiments on toy arithmetic, GSM8K and PRM800k datasets, where it significantly outperforms existing outcome-based methods by providing more accurate insights into model performance on individual subtasks.

Understanding Chain-of-Thought in LLMs through Information Theory

TL;DR

This paper tackles the challenge of evaluating chain-of-thought (CoT) reasoning in LLMs without relying on annotated intermediate steps. It introduces an information-theoretic framework that quantifies information-gain at each reasoning step via a supervisor model to estimate p(Y|X^M_j), enabling detection of unidentifiable or non-contributory steps. The approach is validated on toy data, GSM8K, and PRM800K, where it outperforms outcome-based baselines like ORM and Math-Shepherd in identifying faulty CoT steps and providing granular, task-specific diagnostics. The results suggest that information-gain-based evaluation offers more reliable, scalable insights for CoT quality and generalizes across nonlinear or exploratory reasoning patterns, with potential extensions to non-mathematical domains.

Abstract

Large Language Models (LLMs) have shown impressive performance in complex reasoning tasks through the use of Chain-of-Thought (CoT) reasoning, allowing models to break down problems into manageable sub-tasks. However, existing CoT evaluation techniques either require annotated CoT data or fall short in accurately assessing intermediate reasoning steps, leading to high rates of false positives. In this paper, we formalize CoT reasoning in LLMs through an information-theoretic lens. Specifically, our framework quantifies the `information-gain' at each reasoning step, enabling the identification of failure modes in LLMs without the need for expensive annotated datasets. We demonstrate the efficacy of our approach through extensive experiments on toy arithmetic, GSM8K and PRM800k datasets, where it significantly outperforms existing outcome-based methods by providing more accurate insights into model performance on individual subtasks.

Paper Structure

This paper contains 65 sections, 2 theorems, 22 equations, 10 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Proposed Framework: Setup and Notation
  3. Compositionality
  4. Primitive Tasks
  5. Unidentifiability
  6. Operationalising Our Framework
  7. An information-theoretic perspective
  8. Intuition
  9. Task execution in LLMs
  10. Intuition
  11. Testing for unidentifiability using information-gain
  12. Related Works
  13. Experiments
  14. Outcome Reward Model (ORM) havrilla2024glore
  15. Math-Shepherd (MS) wang2024mathshepherdverifyreinforcellms
  16. ...and 50 more sections

Key Result

Theorem 3.3

Let $\Gamma^M \subseteq \Gamma$ denote the primitive tasks available in the training data. Let $\lambda$ be a task with decomposition $\lambda = \lambda_1 \circ \ldots \circ \lambda_T$. If $\lambda_k$ is the first task in the decomposition of $\lambda$ which is unidentifiable in $\Gamma^M$ (i.e. $k

Figures (10)

  • Figure 1: Prompt \ref{['eq:composition_prompt']} requires compositional application of tasks.
  • Figure 2: Bayesian network corresponding to Assumption \ref{['ass:bayenet']}.
  • Figure 3: Heatmaps (left) quantifying the correctness of different sub-tasks for the 5 LLMs using the different baselines, and the associated classification metrics (right). Red color in the heatmaps indicates a significant drop in the plotted metrics (an incorrectly executed sub-task).
  • Figure 4: The distribution of $(x, y)$ for incorrect samples: Llama-3-8B struggles to add large and small numbers (represented by the top-left and bottom-right shaded regions).
  • Figure 5: Toy data results: Sample-wise information-gain trajectories for 15 randomly chosen prompts with wrong final answers.
  • ...and 5 more figures

Theorems & Definitions (8)

  • Definition 2.1
  • Definition 2.2: Primitive tasks
  • Definition 2.3: Span of tasks
  • Definition 2.4: Unidentifiability
  • Theorem 3.3
  • Proposition 3.4
  • proof : Proof of Theorem \ref{['prop:conditional_indep']}
  • proof : Proof of Proposition \ref{['thm:conditional_mi']}