deep-think-tokens 0.1.0

Track Deep-Thinking Tokens in Transformer Models

Deep Think Tokens

Track deep thinking tokens and logit lenses in transformer models.

About

This is an implementation of the paper "Think Deep, Not Just Long: Measuring LLM Reasoning Effort via Deep-Thinking Tokens" by Chen et al.

Deep-thinking tokens are tokens where internal predictions undergo significant revisions in deeper model layers prior to convergence. The deep-thinking ratio (proportion of deep-thinking tokens in a generated sequence) exhibits a robust and consistently positive correlation with accuracy on reasoning tasks, substantially outperforming length-based and confidence-based baselines.

Key applications include:

• Reasoning quality prediction: Better than raw token counts for measuring reasoning effort
• Think@n: Test-time scaling strategy that prioritizes samples with high deep-thinking ratios (DTR)
• Early rejection: Enable early rejection of unpromising generations based on low DTR, reducing inference costs

Installation

pip install deep-think-tokens

Or install directly from GitHub:

pip install git+https://github.com/maxzuo/deep-think-tokens.git

Example Usage

from deep_think_tokens import (
    add_deep_thinking_tokens_hooks,
    plot_divergences,
    deep_thinking_ratio,
)
from transformers import AutoModelForCausalLM

# Load a model
model = AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-1.5B-Instruct')

# Or add hooks for deep thinking tokens
tracker = add_deep_thinking_tokens_hooks(model)

# Generate text and collect divergences
# ...

# Collect results
results = tracker.collect()

# Plot divergences
plot_divergences(results, tokens=['token1', 'token2', ...])

# Calculate deep thinking ratio
dtr = deep_thinking_ratio(results, g=0.5, p=0.9)