Reproduce “CODI: Compressing Chain-of-Thought into Continuous Space via Self-Distillation” on OpenResearch

About this paper

Zhenyi Shen, Hanqi Yan, Linhai Zhang, Zhanghao Hu, Yali Du, Yulan He

CODI trains a language model to replace written chain-of-thought reasoning with continuous hidden-state vectors, aligning a student’s answer-prompt activations with those of a written-reasoning teacher. On GSM8k, GPT-2 reaches 43.7% accuracy, close to CoT-SFT’s 44.1%, while six latent thoughts yield 3.1× compression and about 2.7× inference speedup; evidence is limited to tested models and tasks.

Key insights

• LLMs can effectively perform complex reasoning within their continuous latent space, matching or surpassing explicit natural language CoT performance, particularly at smaller model scales.
• A single-stage, self-distillation approach, where the model simultaneously learns explicit CoT as a reference and compresses it into latent representations, effectively mitigates "forgetting" issues seen in multi-stage implicit CoT methods.
• Reasoning in a continuous space, unconstrained by rigid linguistic imitation targets, promotes greater robustness and improved generalization to out-of-domain data compared to fine-tuning on explicit CoT.

Summary by alphaXiv.