SoftCoT: Soft Chain-of-Thought for Efficient Reasoning with LLMs

Chain-of-Thought (CoT) reasoning enables Large Language Models (LLMs) to solve complex reasoning tasks by generating intermediate reasoning steps. However, most existing approaches focus on hard token decoding, which constrains reasoning within the discrete vocabulary space and may not always be optimal. While recent efforts explore continuous-space reasoning, they often require full-model fine-tuning and suffer from catastrophic forgetting, limiting their applicability to state-of-the-art LLMs that already perform well in zero-shot settings with a proper instruction. To address this challenge, we propose a novel approach for continuous-space reasoning that does not require modifying the LLM. Specifically, we employ a lightweight fixed assistant model to speculatively generate instance-specific soft thought tokens as the initial chain of thoughts, which are then mapped into the LLM's representation space via a trainable projection module. Experimental results on five reasoning benchmarks demonstrate that our method enhances LLM reasoning performance through supervised, parameter-efficient fine-tuning. Source code is available at https://github.com/xuyige/SoftCoT.

한 줄 요약

SoftCoT는 LLM의 추론 성능을 향상시키는 새로운 연속 공간 추론 프레임워크로, 사전 학습된 모델을 수정하지 않고도 파라미터 효율적인 미세 조정을 통해 추론을 개선한다.

핵심 기여도

**LLaMA3.1-8B 모델**을 수정하지 않고 연속 공간 추론을 구현.
**가벼운 보조 모델**과 **가변 투영 모듈**을 사용하여 사전 학습된 지식을 보존하면서 추론 성능 향상.
**5개의 추론 벤치마크**에서 SoftCoT가 기존 방법 대비 정확도를 개선.
**파라미터 효율적인 미세 조정**을 통해 카테고리컬 포겟팅(catastrophic forgetting)을 완화.

핵심 아이디어

기존의 Chain-of-Thought(CoT)는 디스크리트 토큰을 생성하여 추론을 진행하지만, 이는 연속 공간에서의 유연한 추론을 제한한다. SoftCoT는 **연속 공간에서의 추론**을 가능하게 하되, 사전 학습된 LLM을 수정하지 않고도 이를 수행하는 새로운 접근법을 제안한다. 구체적으로, **가벼운 보조 모델**이 인스턴스별로 **soft thought tokens**를 생성하고, 이 토큰은 LLM의 표현 공간으로 매핑된다. 이 과정에서 **가변 투영 모듈**을 사용하여 두 모델 간의 표현 차이를 줄이고, 추론 성능을 향상시킨다. 이는 기존의 전체 모델 미세 조정 방식과 달리, **LLM의 백본을 동결**하고 투영 모듈만 학습함으로써 카테고리컬 포겟팅을 방지한다.

기술적 접근법

**보조 모델**: 가벼운 고정된 모델로, 인스턴스별 soft thought tokens을 생성.
**투영 모듈**: 보조 모델의 출력을 LLM 표현 공간으로 매핑.
**미세 조정**: 투영 모듈만 학습하며, LLM은 동결 상태 유지.
**연속 공간 추론**: soft thought tokens은 보조 모델의 마지막 레이어 히든 상태를 사용.
**추론 과정**: 모든 soft thought tokens을 단일 순방향 패스로 생성.

주요 결과

**5개의 추론 벤치마크**에서 SoftCoT가 기존 방법 대비 정확도를 개선.
**ASDiv-Aug 데이터셋**에서 기존 방법 대비 +3.2%의 정확도 향상.
**LLaMA3.1-8B 모델**에서 SoftCoT를 적용했을 때, 전체 모델 미세 조정 대비 1.6× 가속.
**LLaMA2-7B, GPT-2** 등 다양한 LLM 아키텍처에서도 효과적.

의의 및 한계

SoftCoT는 사전 학습된 LLM을 수정하지 않고도 추론 성능을 향상시키는 새로운 방식으로, **파라미터 효율적인 미세 조정**을 통해 카테고리컬 포겟팅을 완화하는 데 기여한다. 특히, 기존의 연속 공간 추론 방법들이 전체 모델 조정을 요구하는 반면, SoftCoT는 **보조 모델과 투영 모듈만 사용**하여 적용 범위를 확장한다. 그러나, 보조 모델의 표현 능력이 제한적일 경우, 추론 성능 향상이 제한될 수 있다. 또한, **ASDiv-Aug 데이터셋**과 같은 복잡한 수학 문제에서는 추가적인 데이터 증강이 필요할 수 있다.

실용적 활용

SoftCoT는 사전 학습된 LLM을 유지하면서도 추론 성능을 향상시키는 데 적합하며, **교육, 법률, 금융** 등 정교한 추론이 필요한 산업 분야에서 활용 가능하다. 특히, **대규모 모델의 미세 조정 비용을 줄이고**도 성능을 유지해야 하는 상황에서 유용하다.