능동적 잠재 계획을 위한 강화학습

초록

본 논문은 기존의 라틴 토큰(잠재 토큰) 기반 체인‑오브‑생각(Chain‑of‑Thought) 추론이 다수의 정답 CoT 라벨을 무작위로 모방함으로써 발생하는 표현 편향을 극복하고자, 조건부 변분 오토인코더(VAE)와 보조 일관성 보상을 결합한 강화학습 프레임워크인 ATP‑Latent을 제안한다. LLaMA‑1B 모델에 적용한 실험에서 정확도 4.1% 향상과 토큰 사용량 3.3% 감소를 달성하였다.

상세 분석

ATP‑Latent은 기존 잠재 CoT 방법이 “하나의 라벨을 그대로 모방”하는 한계에 주목한다. 동일 질문에 대해 여러 개의 올바른 언어 CoT가 존재함에도 불구하고, 기존 방식은 임의의 하나만을 지도학습 시켜 잠재 토큰 공간을 비정상적으로 왜곡시킨다. 이를 해결하기 위해 저자는 두 단계 접근법을 설계했다. 첫 번째 SFT 단계에서는 라벨 CoT를 직접 모방하는 대신, 조건부 VAE를 학습한다. 인코더는 질문‑답 쌍을 입력받아 연속적인 잠재 토큰 시퀀스를 생성하고, 디코더는 이 토큰을 다시 언어 CoT로 복원한다. VAE의 KL‑다이버전스 정규화는 잠재 공간을 매끄럽게 만들어, 이후 강화학습 단계에서 탐색이 의미 있는 연속 공간을 보장한다. 두 번째 RL 단계에서는 전통적인 정답 정확도 보상에 더해 “일관성(coherence)” 보상을 도입한다. 일관성은 동일 질문에 대해 VAE‑디코더가 복원한 여러 잠재 토큰 시퀀스 간의 내용적 일치도를 측정해 계산한다. 즉, 정책이 생성한 잠재 토큰이 VAE‑디코더를 통과했을 때 서로 다른 경로에서도 일관된 언어 설명을 산출하면 높은 보상을 받는다. 이 보상은 명시적 정답 신호가 없는 상황에서도 정책을 부드럽게 정렬시키는 역할을 한다.

구현 측면에서 저자는 LLaMA‑1B에 “latent head”와 “stop head”라는 두 개의 MLP를 추가해 토큰 생성과 종료 시점을 제어한다. 또한, 탐색을 위해 가우시안 노이즈를 잠재 토큰에 직접 더하는 대신, VAE가 제공하는 연속적 확률분포를 활용해 샘플링한다. 이는 기존 연구가 제시한 무작위 탐색보다 더 구조화된 탐색을 가능하게 한다. 실험에서는 수치 추론 벤치마크(예: GSM‑8K, MathQA 등) 네 가지에 대해 기존 SFT 기반 잠재 방법(Coconut, SIM‑CoT)과 최신 RL 기반 방법을 비교했으며, ATP‑Latent은 평균 정확도 4.1%p 상승과 토큰 사용량 3.3% 감소를 기록했다. 특히, “overthinking” 현상이 감소해 추론 속도가 크게 개선되었다.

이 논문의 핵심 기여는 (1) 잠재 토큰 공간을 매끄럽게 만드는 VAE 기반 학습 설계, (2) 디코더 일관성을 활용한 보조 보상으로 RL 탐색을 안정화, (3) 실험을 통해 토큰 효율성과 정확도 모두에서 기존 최첨단 방법을 능가함을 입증한 점이다. 다만, VAE 학습 비용과 디코더 품질에 크게 의존한다는 한계와, 현재는 1B 규모 모델에만 검증했으므로 대규모 모델에 대한 확장성 검증이 필요하다.