ADu2021/chain-of-thought-distribution-lens

Chain-of-Thought: A Distribution Lens Analysis

This framework explains CoT reasoning effectiveness through a data distribution lens: CoT succeeds when test data matches training distribution, fails under distribution shift. Using a fully controllable synthetic environment (DataAlchemy), it demonstrates CoT is not reasoning but learned pattern matching, highly brittle to even moderate shifts.

Core Concept

Chain-of-Thought prompting dramatically improves LLM performance—but the mechanism remains unclear. This work hypothesizes: CoT encodes distributional assumptions from training data; it works brilliantly in-distribution but fails under shifts. By building a synthetic environment where distribution can be precisely controlled, the paper demonstrates this empirically and provides a theoretical bound on generalization.

Architecture Overview

• DataAlchemy Environment: Synthetic task environment (tokens, sequences, transformations) with full distributional control
• Three-Dimensional Distribution Analysis: Task (unseen transformations), length (longer sequences), format (perturbations)
• Theoretical Framework: Generalization bound via total variation distance
• Empirical Finding: CoT is a "brittle mirage"—effective in-distribution, fails under shifts

Implementation Steps

Step 1: Build DataAlchemy Synthetic Environment

import numpy as np
from typing import List, Tuple, Callable
from enum import Enum

class Token(Enum):
    """Basic tokens in DataAlchemy."""
    A = 'A'
    B = 'B'
    C = 'C'
    D = 'D'

class Element:
    """Sequence of tokens (ordered element)."""
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens

    def __repr__(self):
        return ''.join([t.value for t in self.tokens])

class Transformation:
    """Operation on elements (e.g., reverse, sort, deduplicate)."""
    def __init__(self, name: str, fn: Callable[[Element], Element]):
        self.name = name
        self.fn = fn

    def apply(self, element: Element) -> Element:
        return self.fn(element)

class DataAlchemy:
    """Fully controllable synthetic task environment."""

    def __init__(self):
        # Define basic transformations
        self.transformations = {
            'reverse': Transformation('reverse', lambda e: Element(list(reversed(e.tokens)))),
            'sort': Transformation('sort', lambda e: Element(sorted(e.tokens, key=lambda t: t.value))),
            'length_filter': Transformation('length_filter', lambda e: Element([t for t in e.tokens if t != Token.D])),
            'duplicate': Transformation('duplicate', lambda e: Element(e.tokens + e.tokens)),
        }

    def generate_element(self, length: int = 4, alphabet: List[Token] = None) -> Element:
        """Generate random element."""
        if alphabet is None:
            alphabet = list(Token)

        tokens = [np.random.choice(alphabet) for _ in range(length)]
        return Element(tokens)

    def generate_task_dataset(self, num_examples: int = 100,
                             element_length: int = 4,
                             transformation_names: List[str] = None) -> List[Tuple[Element, Element]]:
        """
        Generate dataset of (input, output) pairs for specific transformations.
        """
        if transformation_names is None:
            transformation_names = list(self.transformations.keys())

        dataset = []

        for _ in range(num_examples):
            element = self.generate_element(element_length)

            # Apply random transformation
            transformation = self.transformations[np.random.choice(transformation_names)]
            result = transformation.apply(element)

            dataset.append((element, result))

        return dataset

    def apply_format_perturbation(self, element: Element, perturbation_type: str = 'insert') -> Element:
        """Apply format-level perturbations."""
        if perturbation_type == 'insert':
            # Insert random token at random position
            insert_pos = np.random.randint(0, len(element.tokens) + 1)
            new_tokens = element.tokens[:insert_pos] + [Token.A] + element.tokens[insert_pos:]
            return Element(new_tokens)

        elif perturbation_type == 'delete':
            # Delete random token
            if len(element.tokens) > 1:
                delete_pos = np.random.randint(0, len(element.tokens))
                new_tokens = element.tokens[:delete_pos] + element.tokens[delete_pos + 1:]
                return Element(new_tokens)

        elif perturbation_type == 'modify':
            # Change random token
            modify_pos = np.random.randint(0, len(element.tokens))
            new_tokens = element.tokens.copy()
            new_tokens[modify_pos] = np.random.choice([t for t in Token if t != element.tokens[modify_pos]])
            return Element(new_tokens)

        return element

Step 2: Implement Distribution Shift Analysis

from scipy.spatial.distance import jensenshannon

class DistributionAnalyzer:
    """Analyze distribution discrepancy and CoT effectiveness."""

    def __init__(self, cot_model):
        self.model = cot_model

    def measure_total_variation_distance(self, train_data: List, test_data: List) -> float:
        """
        Compute total variation distance between training and test distributions.
        Higher distance = more severe distribution shift.
        """
        # Extract element statistics
        train_lengths = [len(e.tokens) for e, _ in train_data]
        test_lengths = [len(e.tokens) for e, _ in test_data]

        # Create histograms
        train_hist, bins = np.histogram(train_lengths, bins=range(1, 11), density=True)
        test_hist, _ = np.histogram(test_lengths, bins=bins, density=True)

        # Total variation distance = 0.5 * sum(|p - q|)
        tv_distance = 0.5 * np.sum(np.abs(train_hist - test_hist))

        return tv_distance

    def measure_cot_performance(self, model, test_data: List[Tuple[Element, Element]],
                              use_cot: bool = True) -> float:
        """
        Measure model accuracy with or without CoT.
        """
        correct = 0

        for element, expected_output in test_data:
            if use_cot:
                # Generate CoT step-by-step reasoning
                prompt = f"Element: {element}\nStep-by-step reasoning:\n"
                reasoning = model.generate(prompt, max_tokens=100)

                # Extract final answer from reasoning
                prompt_with_reasoning = f"{prompt}{reasoning}\nFinal output: "
                prediction = model.generate(prompt_with_reasoning, max_tokens=20)
            else:
                # Direct prediction
                prompt = f"Element: {element}\nOutput: "
                prediction = model.generate(prompt, max_tokens=20)

            # Check correctness
            if self._parse_element(prediction) == expected_output:
                correct += 1

        accuracy = correct / len(test_data)
        return accuracy

    def analyze_three_dimensions(self, algebra: DataAlchemy, model, base_train_data: List):
        """
        Analyze CoT performance across three distribution dimensions:
        1. Task generalization (unseen transformations)
        2. Length generalization (longer sequences)
        3. Format generalization (perturbations)
        """
        results = {
            'task_dim': {},
            'length_dim': {},
            'format_dim': {}
        }

        # Dimension 1: Task generalization
        print("Testing task generalization...")
        for unseen_transform in ['reverse', 'sort', 'length_filter']:
            # Train on other tasks
            train_tasks = [t for t in algebra.transformations.keys() if t != unseen_transform]
            train_data = algebra.generate_task_dataset(100, transformation_names=train_tasks)

            # Test on unseen task
            test_data = algebra.generate_task_dataset(50, transformation_names=[unseen_transform])

            cot_acc = self.measure_cot_performance(model, test_data, use_cot=True)
            direct_acc = self.measure_cot_performance(model, test_data, use_cot=False)

            tv_dist = self.measure_total_variation_distance(train_data, test_data)

            results['task_dim'][unseen_transform] = {
                'cot_accuracy': cot_acc,
                'direct_accuracy': direct_acc,
                'tv_distance': tv_dist
            }

        # Dimension 2: Length generalization
        print("Testing length generalization...")
        for test_length in [8, 10, 12]:
            train_data = algebra.generate_task_dataset(100, element_length=4)  # Short
            test_data = algebra.generate_task_dataset(50, element_length=test_length)  # Longer

            cot_acc = self.measure_cot_performance(model, test_data, use_cot=True)
            direct_acc = self.measure_cot_performance(model, test_data, use_cot=False)

            tv_dist = self.measure_total_variation_distance(train_data, test_data)

            results['length_dim'][f'length_{test_length}'] = {
                'cot_accuracy': cot_acc,
                'tv_distance': tv_dist
            }

        # Dimension 3: Format generalization
        print("Testing format generalization...")
        for perturbation in ['insert', 'delete', 'modify']:
            train_data = algebra.generate_task_dataset(100)

            # Apply perturbations to test data
            test_data = algebra.generate_task_dataset(50)
            perturbed_test = [
                (algebra.apply_format_perturbation(e, perturbation), out)
                for e, out in test_data
            ]

            cot_acc = self.measure_cot_performance(model, perturbed_test, use_cot=True)
            tv_dist = self.measure_total_variation_distance(train_data, perturbed_test)

            results['format_dim'][perturbation] = {
                'cot_accuracy': cot_acc,
                'tv_distance': tv_dist
            }

        return results

    def compute_generalization_bound(self, tv_distance: float, model_capacity: float = 1.0) -> float:
        """
        Theoretical bound: test_risk ≤ train_risk + C * TV_distance(P_train, P_test)

        C relates to model capacity. Higher TV distance = worse generalization.
        """
        # Simplified bound
        bound = 0.5 * tv_distance * model_capacity
        return bound

    def _parse_element(self, text: str) -> Element:
        """Extract Element from generated text."""
        tokens = [Token[c] for c in text if c in 'ABCD']
        return Element(tokens)

Step 3: Measure CoT Brittleness

def measure_cot_brittleness(analyzer: DistributionAnalyzer, results: dict) -> dict:
    """
    Quantify how fragile CoT is to distribution shifts.
    """
    brittleness = {
        'task': [],
        'length': [],
        'format': []
    }

    # Task brittleness
    for task, metrics in results['task_dim'].items():
        cot_acc = metrics['cot_accuracy']
        tv_dist = metrics['tv_distance']

        # Brittleness = how much CoT degrades with TV distance
        degradation = 1.0 - cot_acc  # Inverse accuracy
        brittleness_score = degradation / (tv_dist + 0.01)
        brittleness['task'].append(brittleness_score)

    # Length brittleness
    for length, metrics in results['length_dim'].items():
        cot_acc = metrics['cot_accuracy']
        brittleness['length'].append(1.0 - cot_acc)

    # Format brittleness
    for format_type, metrics in results['format_dim'].items():
        cot_acc = metrics['cot_accuracy']
        brittleness['format'].append(1.0 - cot_acc)

    return brittleness

def summarize_findings(results: dict, brittleness: dict):
    """Print summary of CoT analysis."""
    print("\n=== CoT Analysis Results ===\n")

    print("Task Generalization (Unseen Transformations):")
    for task, metrics in results['task_dim'].items():
        print(f"  {task}: CoT={metrics['cot_accuracy']:.1%}, TV={metrics['tv_distance']:.3f}")

    print("\nLength Generalization:")
    for length, metrics in results['length_dim'].items():
        print(f"  {length}: CoT={metrics['cot_accuracy']:.1%}")

    print("\nFormat Generalization (Perturbations):")
    for fmt, metrics in results['format_dim'].items():
        print(f"  {fmt}: CoT={metrics['cot_accuracy']:.1%}")

    print("\n=== Conclusion ===")
    print("CoT reasoning is highly brittle to distribution shifts.")
    print("It reflects learned patterns, not robust reasoning.")

Practical Guidance

When to Use:

• Analyzing CoT failure modes
• Understanding distribution sensitivity
• Designing more robust reasoning systems
• Academic research on LLM capabilities

When NOT to Use:

• Production systems (this is analytical, not predictive)
• Real-world tasks (synthetic environment is limited)
• Scenarios requiring immediate practical insights

Key Findings:

• CoT works brilliantly in-distribution but fails under shifts
• Generalization bound correlates with total variation distance
• Format perturbations more harmful than length changes
• CoT is pattern matching, not reasoning

Reference

Paper: Is Chain-of-Thought Reasoning of LLMs a Mirage (2508.01191)

• DataAlchemy synthetic environment with full distributional control
• Demonstrates CoT brittleness across three dimensions
• Theoretical bound on generalization via TV distance