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

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.

한 줄 요약

CODI는 자연어 사고 과정을 연속 공간으로 압축해 GPT-2 규모에서 3.1배 압축률과 28.2% 정확도 향상을 달성한 새로운 사고 방식 학습 프레임워크이다.

핵심 기여도

CODI는 자연어 기반 사고 과정(Explicit CoT)을 연속 공간(Implicit CoT)으로 압축하는 새로운 자가 증류(Self-Distillation) 프레임워크를 제안.
GPT-2 규모에서 GSM8k 데이터셋에서 기존 최고 성능 대비 28.2% 정확도 향상 및 3.1배 압축률 달성.
사고 과정의 핵심 토큰의 은닉 상태(activation) 정렬을 통해 지식 전달.
복잡한 CoT 데이터셋에서도 일반화 가능하며 해석 가능성(interpretability)을 보여준다.

핵심 아이디어

기존 CoT는 자연어 토큰을 사용해 단계별 사고를 유도하지만, 이는 계산 효율성과 과적합 위험을 초래한다. CODI는 자연어 대신 연속 공간에서 사고를 표현함으로써 효율성과 안정성을 동시에 확보하려는 접근법이다. 핵심 아이디어는 **자가 증류(Self-Distillation)**을 통해, **Explicit CoT(교사)**에서 **Implicit CoT(학생)**로 지식을 전달하는 방식이다. 교사는 교차 엔트로피 손실로 CoT 토큰을 학습하고, 학생은 연속 공간에서 사고를 생성한 후 최종 답변을 도출한다. 학생의 연속 사고는 특정 타겟과 일치하지 않아도 되며, 대신 **최종 답변 생성 시점의 은닉 상태 정렬**을 통해 교사의 사고 패턴을 학습한다.

기술적 접근법

**CODI 프레임워크**: 교사(Explicit CoT)와 학생(Implicit CoT)를 동시에 학습.
**자가 증류(Self-Distillation)**: 교사의 CoT 토큰을 학습한 후, 학생의 연속 사고는 교사의 **최종 답변 생성 시점의 은닉 상태**와 정렬.
**하이퍼파라미터**: 압축률 3.1x, 정확도 28.2% 향상.
**데이터셋**: GSM8k, SVAMP, GSM-HARD, MultiArith 등 수학적 추론 데이터셋 사용.
**모델 규모**: GPT-2 수준의 LLM에서 실험.

주요 결과

**GSM8k 데이터셋**: CODI는 GPT-2 규모에서 기존 최고 성능 대비 **정확도 28.2% 향상**.
**압축률**: 3.1배.
**SVAMP, GSM-HARD, MultiArith**: 복잡한 CoT 데이터셋에서도 일반화 성능 보여.
**해석 가능성**: 연속 공간에서도 사고 과정이 해석 가능함을 입증.

의의 및 한계

CODI는 자연어를 사용하지 않고도 LLM이 효과적으로 사고할 수 있음을 입증하며, 추론 효율성과 안정성을 동시에 달성하는 새로운 패러다임을 제시한다. 특히, **자가 증류를 통한 단일 학습 단계**로 커리큘럼 학습의 **forgetting 문제**를 회피한 점이 학술적 의의이다. 그러나 본 연구는 **수학적 추론 데이터셋**에 초점을 맞추었으며, **일상적 추론**이나 **다중 모달리티** 상황에서의 성능은 아직 검증되지 않았다. 또한, **연속 공간 내 사고의 해석 가능성**은 부분적으로만 입증되었으며, 보다 체계적인 분석이 필요하다.

실용적 활용

CODI는 대규모 LLM의 추론 과정을 **효율화**하고, **과적합 방지** 및 **해석 가능성 향상**에 기여할 수 있어, **AI 추론 시스템**, **교육용 AI**, **자동화된 문제 해결 플랫폼** 등에 적용 가능하다. 특히, **자연어 대신 연속 공간을 활용한 추론**은 **실시간 처리** 및 **자원 제한 환경**에서 유리하다.