Reasoning Under Compression

Tree-based modelling of LLM reasoning-trace compressibility using Microsoft OpenMementos.

Overview

This repository contains the code and materials for a PhD course project in Statistics for Data Mining and Machine Learning.

The project studies whether the compressibility of long-form reasoning traces generated by large language models can be predicted using tree-based supervised learning methods. The empirical object is the structure of the reasoning process: prompts, reasoning blocks, block summaries, and the way compression evolves along the chain of thought.

The main idea is that compressibility may not be only a prompt-level property. It may emerge dynamically as the reasoning trace unfolds.

References

Core project references:

• Kontonis et al., MEMENTO: Teaching LLMs to Manage Their Own Context.
• Wei et al., Chain-of-Thought Prompting Elicits Reasoning in Large Language Models.
• Kang et al., C3oT: Generating Shorter Chain-of-Thought without Compromising Effectiveness.

See references.bib for citation metadata.

Research question

Can tree-based models predict how compressible an LLM reasoning trace becomes as the reasoning chain unfolds?

More specifically:

• What can be predicted from the problem statement alone?
• How much does prediction improve after observing the first reasoning block?
• Does previous reasoning history help predict the compressibility of later blocks?
• Do math, code, and science traces display different compression dynamics?

Dataset

The project uses microsoft/OpenMementos, a public Hugging Face dataset associated with the MEMENTO framework.

The dataset contains long reasoning traces annotated with:

• problem: original prompt or task;
• response: full reasoning trace;
• domain: broad domain, such as math, code, or science;
• source: source dataset;
• difficulty: task-level difficulty metadata. Missing or non-code values are represented as not_applicable in engineered feature tables;
• sentences: sentence-level decomposition of the reasoning trace;
• blocks: block segmentation of the reasoning trace;
• block_summaries: compressed summaries, or mementos, for each block.

sentences, blocks and block_summaries are variables derived from response.

Conceptual framing

Compression is treated as a transformation from a richer representation to a shorter representation under an information-retention constraint.

In this project:

• original representation = reasoning block or reasoning trace;
• compressed representation = block summary / memento;
• compression measure = summary length divided by original block length;
• modelling objective = predict or explain compression behavior from prompt-level and trace-level features.

The project does not directly estimate a continuous compression ratio as its main task. Continuous compression measures are used to derive classification labels and may be explored in extensions.

The analysis distinguishes between:

1. Prompt-level compressibility Compressibility predicted from the problem, domain, source, and difficulty.

2. Block-level compressibility Compressibility predicted from the current reasoning block.

3. Sequential compressibility Compressibility predicted from the problem, current block, and previous reasoning history.

Classification target variables

The basic compression ratio is computed as:

block_compression_ratio = summary_words_t / block_words_t

or, using tokens:

block_compression_ratio_tokens = summary_tokens_t / block_tokens_t

The main classification target is then derived from this ratio.

Primary target:

high_compression_t = 1 if block_compression_ratio_t <= 25th percentile
high_compression_t = 0 otherwise

Lower compression ratios indicate stronger compression. Therefore, a high-compression block is one that can be represented by a relatively short memento compared with the original block.

The project interprets this as quality-constrained compression because OpenMementos summaries are produced through a summarization and refinement pipeline rather than arbitrary shortening.

Predictive designs

The project compares three information sets.

Model A — Prompt-only model

Inputs:

• problem;
• domain;
• source;
• difficulty;
• problem-level linguistic and symbolic features (symbolic density, code markers, reasoning/discourse markers).

Target:

• summary or trace high-compression class.

Interpretation:

What can be predicted before the LLM reasoning trace is observed?

Model B — Current-block model

Inputs:

• problem;
• domain;
• source;
• difficulty;
• problem-level linguistic and symbolic features (symbolic density, code markers, reasoning/discourse markers).
• current response block length;

Target:

• current summary or trace high-compression class.

Interpretation:

What becomes predictable after observing the current reasoning block?

Model C — History-aware sequential model

Inputs:

• problem;
• domain;
• source;
• difficulty;
• problem-level linguistic and symbolic features (symbolic density, code markers, reasoning/discourse markers).
• current response block length;
• block index;
• cumulative previous length;
• previous average compression ratio;
• previous block-length statistics.

Target:

• current or next-block compression behavior.

Interpretation:

Does compressibility emerge dynamically as the reasoning chain unfolds?

Methods

The project focuses on tree-based supervised learning methods.

Main models:

• Random Forest Classifier;
• Gradient Boosting Classifier.

The methodological comparison moves from a simple interpretable tree to ensemble methods that capture more complex non-linear interactions.

Repository Structure

• notebooks/01_data_loading_engineering.ipynb
• notebooks/02_prompt_level_compressibility.ipynb
• notebooks/03_block_level_compressibility.ipynb
• notebooks/04_sequential_compressibility.ipynb
• notebooks/05_reasoning_under_compression.ipynb
• notebooks/eda01_openmementos_probe.ipynb
• notebooks/eda02_feature_engineering.ipynb
• notebooks/eda03_baseline_modelling.ipynb
• src/reasoning_compression/features.py
• tests/

Limitations

This project models compression behavior, not reasoning correctness.

Compression ratio measures length reduction, not information retention by itself.

The summaries are produced by a specific MEMENTO-style pipeline, so the observed compression behavior should not be interpreted as a universal property of all LLM reasoning traces.

Tree-based models reveal predictive associations, not causal mechanisms.

Future work

Possible extensions include:

• estimating continuous compression ratios as a regression task;
• adding semantic similarity or information-retention measures;
• linking compression dynamics to final-answer correctness;
• comparing compression patterns across different LLM families;
• using probabilistic sequential models for explicit uncertainty updating;
• studying whether early reasoning blocks predict later reasoning reliability.

Project status

Proposal-stage research project for a PhD course in Statistics for Data Mining and Machine Learning.

Documentation

This project uses MkDocs with mkdocstrings for local documentation.

To preview the docs locally:

conda env create -f reasoning_env.yml
conda activate reasoning-compression
mkdocs serve

Then open http://127.0.0.1:8000/.