CoTP: Expanding Reasoning Potential in Foundation Models by Learning Diverse Chains of Thought Patterns

Official implementation of "Expanding Reasoning Potential in Foundation Model by Learning Diverse Chains of Thought Patterns", accepted by ICLR 2026.

📖 Introduction

Recent progress in large reasoning models (LRMs) has largely been driven by reinforcement learning (RL). However, current approaches often utilize CoT data indiscriminately. In this paper, we:

1. Define reasoning potential for the first time as the inverse of the number of independent attempts required to correctly answer a question
2. Abstract atomic reasoning patterns from CoT sequences with commonality and inductive capabilities
3. Propose CoTP framework using dual-granularity algorithm (reasoning patterns + token entropy) to efficiently select high-value CoT data

Key Results

• 9.58% improvement on AIME 2024 and 2025 with only 10B-token CoTP data
• 7.81% boost in downstream RL performance upper bound
• Enables 85A6B MoE model to achieve SOTA results on challenging mathematical reasoning benchmarks

🏗️ Framework Overview

The CoTP framework consists of three main stages:

1. Core Set Construction: Build reference set enriched with valuable reasoning patterns
2. Pattern Chain Annotation: Extract atomic reasoning patterns from CoT sequences
3. Data Selection: Use weighted DTW to select high-value CoT data from pool

📦 Installation

Requirements

• Python >= 3.8
• PyTorch >= 2.0.0
• CUDA >= 11.8 (for GPU training)

Setup

# Clone the repository
git clone https://github.com/YOUR_USERNAME/CoTP.git
cd CoTP

# Create conda environment
conda create -n cotp python=3.10
conda activate cotp

# Install dependencies
pip install -r requirements.txt

🚀 Quick Start

1. Pattern Annotation

Annotate reasoning patterns in your CoT data using a strong reasoning model (e.g., DeepSeek-R1):

python source_code/annotation/annotate_cot_patterns.py \
  --input_path data/input_cot.jsonl \
  --output_path data/annotated_cot.jsonl \
  --model_path deepseek-ai/DeepSeek-R1-Distill-Qwen-7B \
  --batch_size 64 \
  --max_new_tokens 2048

2. Compute Token Entropy

Calculate token-level entropy for fine-grained reasoning analysis:

python source_code/annotation/compute_cot_entropy.py \
  --input_path data/annotated_cot.jsonl \
  --output_path data/cot_with_entropy.jsonl \
  --model_path your_reference_model \
  --batch_size 32

3. Compute Pattern TF-IDF Weights

Calculate importance scores for reasoning patterns:

python source_code/data_selection/compute_pattern_tfidf_importance.py \
  --input_path data/cot_with_entropy.jsonl \
  --output_path data/pattern_weights.jsonl

4. Data Selection with Hungarian Algorithm

Select high-value CoT data using dual-granularity weighted DTW:

python source_code/data_selection/data_selection.py \
  --core_path data/core_set.jsonl \
  --source_path data/source_pool.jsonl \
  --output_path data/selected_cotp.jsonl \
  --lambda_weight 0.8 \
  --ngram_n 2

Key Parameters:

• lambda_weight: Balance between pattern distance and entropy distance (0.8 works well)
• ngram_n: Character n-gram size for pattern similarity (1 or 2 recommended)

📊 Data Format

Input CoT Data Format

{
  "identity": "unique_id",
  "question": "Your math problem here",
  "question_type": "math",
  "answer": "Final answer",
  "cot": "Step-by-step reasoning process..."
}

Annotated Pattern Format

{
  "identity": "unique_id",
  "question": "...",
  "answer": "...",
  "cot": "...",
  "pattern": {
    "pattern_list": [
      {
        "id": 1,
        "name": "Pattern Name",
        "description": "...",
        "features": "...",
        "role_in_this_case": "..."
      }
    ],
    "how_CoT_utilizes_patterns_in_this_case": {
      "process_description": "...",
      "pattern_chain": [1, 2, 3, 4]
    }
  },
  "pattern_with_weight": {
    "Pattern Name": 0.85
  },
  "output_token_entropies": [
    {"token": "Let", "entropy_nats": 0.123},
    {"token": "'s", "entropy_nats": 0.456}
  ]
}

🧪 Experiments

Mid-Training

# Using the selected CoTP data for mid-training
bash source_code/pretrain/pretrain_gpt.sh

Supervised Fine-Tuning (SFT)

# Prepare SFT data (see sft/settings.json for configuration)
python source_code/sft/prepare_sft_data.py

Reinforcement Learning (GSPO)

# Run GSPO training
bash source_code/rl/run_gspo.sh

📈 Results

Performance on AIME 2024 & 2025

Dataset	AIME 2025	AIME 2024	HMMT 2025	BeyondAIME	MATH500	AVG
KnowEdu	0.33	1.22	5.10	0.00	45.80	10.49
LongCoTPool	21.89	24.90	15.63	7.90	85.40	31.14
CoTP (Ours)	28.02	37.92	20.73	10.20	90.80	37.53

Scaling to 60B Tokens

CoTP demonstrates exceptional scalability, achieving 4.72% average improvement on AIME 2024 & 2025 when scaled from 30B to 60B tokens.

📚 Citation

If you find this work helpful, please consider citing:

@inproceedings{zhang2026cotp,
  title={Expanding Reasoning Potential in Foundation Model by Learning Diverse Chains of Thought Patterns},
  author={Zhang, Xuemiao and Ren, Can and Tu, Chengying and Weng, Rongxiang and Wang, Shuo and Yan, Hongfei and Wang, Jingang and Cai, Xunliang},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026}
}

📄 License

This project is licensed under the Apache License 2.0 - see the LICENSE file for details.

🙏 Acknowledgments

• Thanks to the open-source community for various tools and datasets
• Special thanks to DeepSeek for the R1 model series
• Thanks to the reviewers and area chairs at ICLR 2026

📧 Contact

For questions and feedback, please contact:

• Xuemiao Zhang: zhangxuemiao@pku.edu.cn
• Can Ren: 2401210098@stu.pku.edu.cn
• Chengying Tu: tuchengying@stu.pku.edu.cn

🔗 Links

• Paper: arXiv:2509.21124
• Dataset: [Coming Soon]
• Checkpoints: [Coming Soon]

---

Note: Core datasets and model checkpoints will be released soon. Stay tuned!