PaperlayerView exact revision on arXiv

Source-linked AI summary

CODI: Compressing Chain-of-Thought into Continuous Space via Self-Distillation

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

arXiv:2502.21074v3cs.CL

TL;DR

Existing implicit CoT methods have underperformed explicit natural-language CoT, motivating more efficient continuous-space reasoning. CODI jointly trains explicit-CoT teacher and implicit-CoT student tasks, aligning designated hidden states to compress reasoning. It matches CoT-SFT performance at GPT-2 scale and reports robustness, generalization, and interpretability, while its probing approach has scope limitations.

  • Problem

    Prior implicit CoT methods have underperformed standard explicit CoT, despite continuous representations offering potential efficiency and robustness benefits.

  • Method

    CODI jointly trains explicit-CoT teacher and implicit-CoT student tasks, transferring reasoning through hidden-state alignment at a designated token.

  • Results

    CODI is the first continuous CoT method to achieve performance comparable to CoT-SFT with GPT-2, reaching 99% of its accuracy on GSM8k.

  • Takeaways & Limitations

    CODI supports reasoning in compact continuous space while demonstrating robustness, generalization to complex datasets, and interpretability.

  • Takeaways & Limitations

    CODI’s token-level probing cannot fully reconstruct multi-token entities, and its answer-prompt and distillation-token choices may be suboptimal.

Abstract

from arXiv · show

Chain-of-Thought (CoT) reasoning enhances Large Language Models (LLMs) by encouraging step-by-step reasoning in natural language. However, leveraging a latent continuous space for reasoning may offer benefits in terms of both efficiency and robustness. Prior implicit CoT methods attempt to bypass language completely by reasoning in continuous space but have consistently underperformed compared to the standard explicit CoT approach. We introduce CODI (Continuous Chain-of-Thought via Self-Distillation), a novel training framework that effectively compresses natural language CoT into continuous space. CODI jointly trains a teacher task (Explicit CoT) and a student task (Implicit CoT), distilling the reasoning ability from language into continuous space by aligning the hidden states of a designated token. Our experiments show that CODI is the first implicit CoT approach to match the performance of explicit CoT on GSM8k at the GPT-2 scale, achieving a 3.1x compression rate and outperforming the previous state-of-the-art by 28.2% in accuracy. CODI also demonstrates robustness, generalizable to complex datasets, and interpretability. These results validate that LLMs can reason effectively not only in natural language, but also in a latent continuous space. Code is available at https://github.com/zhenyi4/codi.

1 Introduction

CODI addresses the limitations of natural-language and prior continuous-space CoT by compressing explicit reasoning into continuous representations through single-stage self-distillation. The framework targets efficient, robust, and interpretable implicit reasoning.

  • Natural-language CoT supports step-by-step reasoning but uses discrete tokens as its reasoning medium.
  • Continuous representations may improve efficiency by reducing communication-heavy reasoning tokens and avoid overfitting to superficial linguistic cues.
  • Prior implicit CoT methods progressively replace or pretrain additional continuous thinking tokens, yet remain substantially behind CoT-SFT.
  • CODI performs single-stage self-distillation by jointly training explicit-CoT teacher and implicit-CoT student tasks.
  • CODI transfers reasoning knowledge by aligning hidden activations at a designated answer-generation token without forcing continuous thoughts to match a fixed target.
  • Experiments report robustness, generalization to complex CoT datasets, and interpretability for the proposed continuous-space reasoning framework.

2 Related Work

Related work studies implicit CoT, knowledge distillation, and self-distillation as alternatives for improving reasoning efficiency and transferring CoT capabilities. CODI combines these directions while using distinct teacher and student tasks within one model.

  • Implicit CoT replaces verbalized intermediate steps with additional continuous computation, but earlier empirical methods showed limited scale and effect.
  • Prior finetuning methods distill explicit CoT reasoning into continuous representations, with Coconut using curriculum learning to replace CoT tokens gradually.
  • Traditional knowledge distillation commonly trains smaller students to imitate step-by-step outputs from larger teacher language models.
  • CODI uses self-distillation to transfer behavior within the same model while giving the teacher richer input contexts than the student.
  • Because its teacher and student tasks differ, CODI can also be viewed as multitask learning.

3 CODI: Continuous Chain-of-Thought via Self Distillation

CODI trains explicit and implicit reasoning jointly, using hidden-state alignment to supervise continuous thoughts while preserving answer prediction. Its design includes learnable thought boundaries, feature-space distillation, and inference-time continuous decoding.

  • 3 CODI: Continuous Chain-of-Thought via Self Distillation: CODI bypasses vocabulary-space autoregression by connecting the last hidden representation directly to the subsequent input.
  • 3 CODI: Continuous Chain-of-Thought via Self Distillation: The teacher learns explicit CoT, while the student learns implicit CoT; their hidden activations are aligned through a weighted combination of training losses.
  • Teacher Task: The teacher task models concatenated CoT reasoning and answer tokens with cross-entropy loss.
  • Student Task: The student autoregressively propagates continuous thoughts from a learnable <bot> token until a learnable <eot> token, then predicts the answer.
  • Student Task: A two-layer MLP and layer normalization transform continuous-thought representations before the next reasoning step.
  • Distillation in Feature Space: Feature-space distillation treats the CoT-induced hidden-state shift at a query token as information that can transfer explicit reasoning capability to implicit CoT.
  • The Distilled Token: The selected distillation token is the colon in “The answer is:”, positioned immediately before answer generation.
  • Self-Distillation: CODI aligns teacher and student hidden activations across all layers at a selected token, using stop-gradient so knowledge flows only from teacher to student.

4 Experiments

CODI is evaluated across mathematical and commonsense reasoning tasks, including in-domain, out-of-domain, efficiency, compression, and ablation settings. It consistently improves over prior implicit CoT methods while often matching or exceeding explicit CoT baselines.

  • Mathematical Reasoning: CODI surpasses Coconut by over 20% on GSM8k with both GPT-2 and LLaMA-1b, reaching 99% and 90% of CoT-SFT performance, respectively.These results establish CODI’s effectiveness on in-domain mathematical reasoning and its scalability to the larger model.
  • Commonsense Reasoning: CODI surpasses CoT-SFT on commonsense reasoning with GPT-2, while CoT data hurts LLaMA-1b performance when its reasoning pattern diverges from the training pattern.The authors attribute CODI’s GPT-2 advantage to reduced burden from generating coherent natural-language CoTs.
  • Efficiency: 2.7× and 5.9× inference speedups correspond to 3.1× and 8.2× CoT compression for compact and verbose CoTs, respectively.CODI uses six continuous thoughts enclosed by two special tokens, yielding eight reasoning tokens in total.
  • Compression Ratio: CODI outperforms Coconut across all tested compression ratios, with both methods peaking at six continuous thoughts.Fewer thoughts may lack sufficient expressiveness, whereas more than six introduce optimization challenges without further benefits.
  • OOD Robustness: CODI outperforms CoT-SFT across all three OOD benchmarks with GPT-2 and outperforms iCoT and Coconut with LLaMA-1b.The authors attribute this robustness to CODI’s reduced tendency to overfit exact natural-language CoT annotations.

5 Further Analysis

CODI’s continuous thoughts exhibit interpretable intermediate reasoning patterns, despite lacking explicit imitation targets. Decoding analyses show that these intermediate results often match reference solutions across varying reasoning lengths.

  • CODI’s continuous thoughts exhibit a degree of interpretability that standard token-by-token fine-tuning cannot trivially learn.
  • Projecting the last hidden state into vocabulary space exposes observable intermediate results within CODI’s continuous thoughts.Operands contributing to these results can often appear among the latent representation’s top-ranked attended tokens.
  • Table 3 compares CODI’s top-5 intermediate results with reference CoT across problems requiring different numbers of steps.
  • 97.1% of single-intermediate-result cases correctly match the reference intermediate solution.
  • For CoT sequences up to 3 steps, CODI decodes valid intermediate results with consistently over 75% accuracy.

6 Conclusion

The paper introduces CODI as a paradigm for reasoning in continuous space. Experiments characterize it as an effective, robust, generalizable, and interpretable implicit CoT approach, while motivating future evaluation on more diverse and challenging tasks.

  • CODI is introduced as a novel paradigm for reasoning in continuous space.
  • Experiments demonstrate CODI’s effectiveness as a state-of-the-art implicit CoT approach while achieving a high compression ratio.
  • CODI shows robustness, generalizability to complex datasets, and interpretability.
  • Future research should explore CODI on more diverse and challenging tasks.

7 Limitations

CODI’s continuous reasoning improves efficiency and offers probing-based interpretability, but current supervision, optimization, and scaling choices constrain its scope.

  • Interpretability: Token-level probing cannot fully reconstruct multi-token entities, such as rare numbers split by the tokenizer.The technique decodes only the first token, leaving remaining components unobserved.
  • Knowledge transfer: CODI transfers knowledge through the token generating the first answer token, but this choice may be suboptimal for some answer formats.Answers beginning with “-” are removed because they reduce performance; other probe locations and prompts may provide alternative supervision.
  • Optimization: Longer continuous reasoning chains may face optimization challenges because intermediate gradients are absent until the sequence end.With six continuous thought tokens, the first token’s gradient arrives from six or more steps away.
  • Scaling: The paper does not scale CODI training to larger models because the authors lack sufficient computation resources.A concurrent study demonstrated latent-reasoning scaling to 3.5B parameters and 800 billion tokens with 4096 GPUs.
  • Scope: Latent reasoning may cover reasoning types that are difficult to verbalize, including visual-spatial, emotional, social, and motor reasoning.The cited concurrent work focuses on pre-training, whereas CODI is presented as an efficient fine-tuning approach for existing pretrained LLMs.

A Implementation Details

The appendix specifies optimizer, scheduling, precision, LoRA, learning-rate, epoch, and continuous-token-removal settings for GPT-2 and LLaMA-1b baselines and CODI comparisons.

  • Shared settings: All GSM8K and Commonsense experiments use AdamW, cosine scheduling, 3% linear warm-up, effective batch size 128, and bfloat16 LoRA finetuning.LoRA uses rank 128 and alpha 32; α and β are set to 1.
  • CODI: GPT-2 CODI uses learning rate 3e-3, γ=1, and 40 training epochs.These settings are reported alongside the shared optimizer and LoRA configuration.
  • CODI: LLaMA-3.2-1b CODI uses learning rate 8e-4, γ=20, and 10 epochs, requiring approximately 48 hours on one A100 80GB.The larger γ compensates for the smaller distillation-loss magnitude observed for this model.
  • Baselines: GPT-2 iCoT training uses learning rate 5e-5 and 100 epochs, while LLaMA-1b uses learning rate 1e-5 and 50 epochs.Continuous-token removal differs by dataset and model, and LoRA is not used.
  • Baselines: GPT-2 Coconut training runs 25 epochs without continuous tokens followed by 25 epochs with them, for 50 epochs total.The LLaMA-1b comparison uses five epochs for each stage and no LoRA.

B Proof: CoTs Contribute a Shift in Hidden Activation

The proof models the last answer-prompt token’s attention and isolates the rationale’s contribution, concluding that CoT primarily shifts hidden activations in a learnable way.

  • Setup: The analysis considers a typical example composed of question Q, rationale R, answer prompt P, and final answer A.The last prompt token q is analyzed at transformer layer l.
  • Attention analysis: The attention output a_l is computed from concatenated inputs [Q; R; P] using key and value parameters, with d as a scaling factor.The derivation analyzes the output activation from the attention heads of q.
  • Simplification: The proof omits softmax and the scaling factor for simplicity because they do not affect the core conclusion.The resulting derivation focuses on the additive contribution of the rationale.
  • No-CoT comparison: Without CoT, W_no-CoT accounts for attention contributions from Q and P, excluding the rationale.The corresponding activation is defined as the attention activation excluding CoT.
  • CoT contribution: The rationale contributes the additional term W_V R(W_K R)^T q to hidden activation, acting primarily as an additive shift.The paper argues that a distance metric can capture and learn this shift.
  • Supporting material: The appendix provides training and evaluation benchmark statistics, examples, and decoded-thought comparisons supporting the empirical analysis.Examples include GSM8K-style arithmetic, commonsense, and corrupted or unusual inputs.

F Ablations on the Hyperparameter

Ablations show that β should match distillation-loss scale, while γ trades faster convergence for lower final implicit-CoT performance when increased.

  • β ablation: β weights the distillation loss; β=1 works well for GPT-2, while LLaMA models require larger values because their loss is about 10 times smaller.For LLaMA-1b, increasing β from 1 to 5 substantially improves accuracy, after which performance plateaus through β=30; β=20 is selected.
  • γ ablation: γ controls the relative weight of explicit teacher reasoning and implicit student reasoning during training.It directly determines the balance between the two objectives.
  • γ ablation: Higher γ accelerates convergence but lowers final performance because the implicitly evaluated student task receives less training emphasis.The explanation attributes faster early learning to greater emphasis on natural-language CoT.

G Ablations on the Choice of the Distillation Token.

CODI’s ablation tests varied answer-prompt language and the designated distillation token, finding no statistically significant accuracy differences from baseline.

  • Robustness: None of the alternative prompt designs differed significantly from the 39% baseline mean across five independent runs.Significance was assessed using an informal t-test criterion of ±2×std (1.8).
  • Ablation setup: Prompts 2–3 varied language, while prompts 4–7 varied the distillation token, using the prompt’s last token.
  • Evaluation setup: Table A6 reports the robustness test for CODI trained on GSM8k-Aug for 20 epochs.
  • Additional materials: Figure A1 illustrates CODI’s example Python implementation.
  • Additional materials: Figures A2 and A3 present CODI interpretability examples for two-step and one-step problems, respectively.
Loading 2502.21074v3…