yuntian-deng/implicit-cot-math-mistral7b

Mistral-7B math with implicit chain of thought

This Mistral-7B checkpoint answers grade-school math word problems without generating written intermediate reasoning steps. It uses stepwise internalization from From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step: explicit reasoning steps are gradually removed during fine-tuning. At inference, the answer tokens are still generated autoregressively.

• Interactive demo
• Training code and GSM8K setup
• Paper record and citation downloads
• Implicit chain-of-thought reading list · Internalized reasoning

Input and output

Use a plain-text question followed by a space and the tokenizer's EOS token, </s>. This token is the trained separator between the question and the generated answer. Do not replace it with a chat template.

Asumi's bookshelf has 120 books. She has 10 books on history, twice that many books on literature, and the rest are science books. How many science books does Asumi have? </s>

The demo returns the generated continuation, without the input question or special tokens. It uses greedy decoding with a limit of 10 new tokens, so longer outputs can be truncated.

Example

Install PyTorch and a current Transformers release. This example mirrors the demo and requires a CUDA GPU with bfloat16 support and enough memory for a 7B model.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "yuntian-deng/implicit-cot-math-mistral7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, dtype=torch.bfloat16
).to("cuda").eval()

question = (
    "Asumi's bookshelf has 120 books. She has 10 books on history, "
    "twice that many books on literature, and the rest are science books. "
    "How many science books does Asumi have?"
)
question = " ".join(question.split())
prompt = question + " " + tokenizer.eos_token
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=10,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )
answer_ids = outputs[0, inputs["input_ids"].shape[-1]:]
print(tokenizer.decode(answer_ids, skip_special_tokens=True))

Training context and limitations

The checkpoint's configuration identifies mistralai/Mistral-7B-v0.1 as its base. The authors' repository documents GSM8K training and a demo-specific format change: the separator between the reasoning and answer is internalized, while the separator after the input question remains.

This is a research demonstration, not a general-purpose chat assistant or a reliable math solver. Answers can be incorrect, and no explanation is provided to help check them. Verify results independently. Performance on a few examples does not establish GSM8K accuracy or performance on other kinds of math; the original paper and the demo have different output-format details.

Citation

The method used by this checkpoint is stepwise internalization:

@misc{deng2024internalize,
  title = {{From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step}},
  author = {Yuntian Deng and Yejin Choi and Stuart Shieber},
  year = {2024},
  eprint = {2405.14838},
  archivePrefix = {arXiv},
  primaryClass = {cs.CL},
  url = {https://arxiv.org/abs/2405.14838}
}

Related earlier work, Implicit Chain of Thought Reasoning via Knowledge Distillation, studies a different training approach based on hidden-state distillation. See its paper record for citation downloads.

@misc{deng2023implicit,
  title = {{Implicit Chain of Thought Reasoning via Knowledge Distillation}},
  author = {Yuntian Deng and Kiran Prasad and Roland Fernandez and Paul Smolensky and Vishrav Chaudhary and Stuart Shieber},
  year = {2023},
  eprint = {2311.01460},
  archivePrefix = {arXiv},
  primaryClass = {cs.CL},
  url = {https://arxiv.org/abs/2311.01460}
}