Latent Reasoning with Normalizing Flows

LLM의 chain-of-thought 추론을 이산 텍스트 토큰 대신 연속 잠재 공간에서 수행하면서도 자기회귀 생성, 확률적 샘플링, KV-cache 호환성, 가능도 평가를 모두 보존하는 NF-CoT 프레임워크를 제안한다.

NF-CoT는 normalizing flow를 LLM의 인과적 스트림 내부에 배치하여 연속 chain-of-thought에 텍스트 토큰과 동등한 모델링 지위를 부여한다. 기존 latent reasoning 방법이 자기회귀 생성, 확률적 샘플링, KV-cache 디코딩, 가능도 추정 중 일부를 포기하는 문제를 해결하며, NF head가 연속 사고 위치를 생성하고 LM head가 텍스트 위치를 생성하는 통합 아키텍처를 제안한다. 학습은 명시적 CoT 감독을 연속 사고로 증류하는 지도 가능도 학습과 잠재 추론 공간에서의 policy-gradient 최적화를 모두 지원한다. MBPP, HumanEval, LiveCodeBench 등 코드 생성 벤치마크에서 명시적 CoT 및 기존 latent reasoning 베이스라인 대비 패스율을 향상시키면서 중간 추론 비용을 크게 절감한다.

핵심 요약

• Normalizing flow를 LLM 백본 내부에 배치하여 연속 CoT에 대한 정확한 가능도를 제공하고, 좌-우 디코딩과 KV-cache를 그대로 지원한다.
• 연속 사고 위치는 NF head가, 텍스트 위치는 LM head가 담당하며, 동일한 인과적 스트림에서 공동으로 처리된다.
• 지도 가능도 학습과 policy-gradient 최적화를 통합된 인터페이스에서 모두 지원한다.
• 잠재 공간에 가우시안 섭동을 가해도 pass@1이 거의 유지되어, 학습된 연속 사고 공간이 국소적으로 매끄럽고 기능적으로 강인함을 보여준다.
• 코드 생성 벤치마크에서 명시적 CoT와 기존 latent reasoning(LaDiR 등) 대비 정확도 향상과 추론 비용 절감을 동시에 달성한다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.