Skip to content
InternLMPublic

About

[ICLR 2026] An official implementation of "SIM-CoT: Supervised Implicit Chain-of-Thought"

Topics

Resources

Stars

219 stars

Watchers

3 watching

Forks

Latest commit

 

History

18 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

SIM-CoT: Supervised Implicit Chain-of-Thought

ICLR 2026

Paper · Project page · Models · Apache 2.0

Demo · Method · Results · Quickstart · Citation

Original SIM-CoT method diagram from Figure 2: each latent receives step-level supervision through an auxiliary decoder during training

SIM-CoT adds step-level supervision to implicit reasoning methods such as Coconut and CODI. An auxiliary decoder learns to recover each reasoning step from its corresponding latent during training. Remove the decoder at inference to keep reasoning compact; reuse it when inspecting the learned latents.

Demo

Watermelon example from Figure 4 and Appendix I.1: four latent steps produce answer 21, and an optional decoder reveals the intermediate arithmetic

Paper case visualization, based on Figure 4 and Appendix I.1. The animation illustrates latent reasoning and optional step decoding; it is not a live model run. The four steps are shown separately here; Figure 4 combines the two middle steps. Static view · Visual sources.

Method

  1. Identify the instability. More implicit reasoning steps can make latent states homogeneous and cause training collapse.
  2. Supervise each step during training. An auxiliary decoder predicts the k-th explicit reasoning step from the k-th latent.
  3. Keep inference compact. Remove the decoder, reason in latents, then generate the answer. Reattach the decoder for optional step inspection.
Original method and instability figures

Latent instability analysis: scaling implicit tokens, operator information, and latent diversity

Method comparison: answer-level Coconut, trajectory-level CODI, and step-level SIM-CoT

In the GPT-2 scaling study, Figure 3 reports up to 8 latents / 16 implicit tokens. Axis labels such as 8–16 denote a latent count and a token count.

Results

Accuracy and token counts from paper Tables 1–3, comparing Coconut and CODI with and without SIM-CoT

GSM8k-Aug accuracy (%) and average token counts from Tables 1–3. Token counts are unchanged within each baseline / SIM-CoT pair. Gains are percentage points (pp):

Base method Model Baseline → + SIM-CoT Gain Avg. tokens
Coconut GPT-2 36.6 → 44.8 +8.2 pp 12.2
Coconut LLaMA 3.2 1B 33.2 → 42.2 +9.0 pp 13.2
CODI LLaMA 3.2 3B 60.8 → 62.3 +1.5 pp 7.2
CODI LLaMA 3.1 8B 61.1 → 64.1 +3.0 pp 7.2

On GPT-2 / GSM8k-Aug, Coconut + SIM-CoT also exceeds explicit SFT-CoT by 2.1 pp (44.8% vs. 42.7%) with 2.3× token efficiency (12.2 vs. 27.6; ratio rounded). This is a token-count comparison. Larger-model in-domain accuracy remains below explicit SFT-CoT.

Full results: in-domain, out-of-domain, and token counts

The paper evaluates in-domain GSM8k-Aug and out-of-domain GSM-Hard, MultiArith, and SVAMP. These images retain the full results from Tables 1–3:

Table 1: GPT-2 results

Table 2: LLaMA 3.2 1B results

Table 3: LLaMA 3.2 3B and LLaMA 3.1 8B results

For GPT-2, CODI + SIM-CoT improves GSM8k-Aug by 0.6 pp (42.0% → 42.6%). The 4.3 pp result is the OOD average gain of Coconut + SIM-CoT (42.6% → 46.9%).

Quickstart

git clone https://github.com/InternLM/SIM-CoT.git
cd SIM-CoT
# Install the root requirements for the Coconut implementation.
python -m pip install -r requirements.txt

Choose Coconut or CODI below. Prepare the model, dataset, and checkpoint paths before launching; the repository contains training code and configurations, with pretrained checkpoints linked in the model collection. The implementations have separate dependency sets; use a dedicated environment for CODI with its requirements. The Coconut examples use eight GPU processes in a CUDA environment.

Coconut + SIM-CoT: prepare, train, and evaluate

All commands in this section run from Coconut/:

cd Coconut
bash preprocessing/gsm_icot.bash

Prepare the local GPT-2 model at pretrained/gpt2/. The preprocessing script creates data/gsm_train.json and data/gsm_test.json; each record contains a question string, steps list of strings, and answer string.

Set literal model, data, and checkpoint paths in the selected YAML configuration to your actual files. The supplied configurations refer to these checkpoints:

Configuration Purpose Default checkpoint
gsm_coconut.yaml Train the Coconut baseline from an explicit-CoT checkpoint. ckpts/gsm_cot/checkpoint_20
gsm_simcot.yaml Continue a Coconut checkpoint with step supervision. ckpts/gsm-coconut/checkpoint_24
gsm_simcot_eval.yaml Evaluate a trained SIM-CoT checkpoint. ckpts/gsm_coconut_simcot/checkpoint_30

Adjust these references to match your downloaded or trained checkpoints. The repository's ckpts/gsm_cot entry is a placeholder, not usable model weights.

# Train the Coconut baseline from your prepared CoT checkpoint.
torchrun --nnodes 1 --nproc_per_node 8 run.py args/gsm_coconut.yaml

# Continue from your selected Coconut checkpoint.
torchrun --nnodes 1 --nproc_per_node 8 run.py args/gsm_simcot.yaml

# Evaluate your selected SIM-CoT checkpoint.
torchrun --nnodes 1 --nproc_per_node 8 run.py args/gsm_simcot_eval.yaml

To create an explicit-CoT warm start first, configure args/gsm_cot.yaml. Set load_model_path: "None" (the quoted string) to skip checkpoint loading; the runner expects this sentinel rather than YAML null. Run that configuration before the Coconut stage:

torchrun --nnodes 1 --nproc_per_node 8 run.py args/gsm_cot.yaml
CODI + SIM-CoT: prepare, train, and evaluate

Start from the repository root, then use CODI/ as the working directory:

# In a separate CUDA environment, install the CODI dependency set.
python -m pip install -r CODI/requirements.txt
cd CODI

The bundled launch scripts are configuration examples. Update the model and checkpoint paths and SAVE_DIR for your environment. Dataset and result paths also appear in train.py and test.py; replace those author-specific absolute paths before running.

# Train CODI + SIM-CoT on GSM8k-Aug with LLaMA 3.2 3B.
bash scripts/train_llama3b_gsm8k-aug-decoder-2.sh

# In a CODI/ terminal with your output path configured, evaluate the checkpoint.
export SAVE_DIR="/path/to/your/results"
bash scripts/test_llama3b-copy.sh

The training and evaluation scripts each need the correct local model and checkpoint paths. Changing a path in the shell script does not update the separate dataset paths in Python.

Authors and affiliations

Xilin Wei, Xiaoran Liu, Yuhang Zang, Xiaoyi Dong, Yuhang Cao, Jiaqi Wang, Xipeng Qiu, Dahua Lin

Fudan University; Shanghai AI Laboratory; The Chinese University of Hong Kong; Shanghai Innovation Institute;

News

  • 2026-01-26: Accepted to ICLR 2026.
  • 2025-09-24: Paper and code released.

Citation

BibTeX — ICLR 2026
@inproceedings{wei2025simcot,
  title={{SIM-COT}: Supervised Implicit Chain-of-Thought},
  author={Wei, Xilin and Liu, Xiaoran and Zang, Yuhang and Dong, Xiaoyi and Cao, Yuhang and Wang, Jiaqi and Qiu, Xipeng and Lin, Dahua},
  booktitle={International Conference on Learning Representations},
  year={2026}
}

Acknowledgments

  • Coconut: The codebase we built upon. Thanks for their wonderful work.
  • CODI: Our work is based on this codebase; we are grateful for their valuable contribution.
  • LLaMA series: The amazing open-sourced large language model!
  • GPT2: An impressive open-source large language model!

License

This repository is released under the Apache License 2.0.

About

[ICLR 2026] An official implementation of "SIM-CoT: Supervised Implicit Chain-of-Thought"

Topics

Resources

Stars

219 stars

Watchers

3 watching

Forks

Releases

Packages

Contributors

Languages