기술 요약: 서사적 과업을 위한 경량화된 잠재 추론 (Lightweight Latent Reasoning for Narrative Tasks)

문제 정의

대규모 언어 모델(LLM)은 강화 학습(RL)을 통해 긴 사고 사슬(reasoning traces)을 생성함으로써 복잡한 작업에서 성능을 향কে 개선해 왔습니다. 그러나 이러한 접근 방식은 특히 수천 개의 토큰이 포함되는 서사적 과업(narrative tasks)에서 훈련 및 추론 시 높은 계산 비용을 초래합니다. 단계(step)를 건너뛰기 위해 이산적인 토큰 대신 연속적인 임베딩을 생성하는 "잠재 추론(latent reasoning)" 방식이 효율성을 높이기 위해 제안되었으나, 기존 방식들은 서사 영역에서 다음과 같은 중대한 한계에 직면해 있습니다:

1. 데이터 부족: 서사적 과업(예: 이야기 생성, 플롯 구멍 탐지)은 잠재 헤드(latent heads)를 효과적으로 훈련하는 데 필요한 고품질의 추론 사슬 데이터셋이 부족합니다.
2. 적응성: 서사적 입력은 매우 다양하여, 고정된 잠재 추론 전략을 적용하기 어렵습니다.
3. 성능 격차: 기존의 잠재 추론 방법들은 기존의 추론 사슬에 의존하거나, 스토리 플래닝과 같이 검증 불가능한 도메인에서 표준 비잠재(non-latent) RL 훈련의 성능을 따라잡지 못하는 경우가 많습니다.

방법론: LiteReason

저자들은 서사 이해 및 생성을 위해 설계된 잠재 추론 프레임워크인 LiteReason을 제안합니다. 이는 경량화된 **추론 프로젝터(Reasoning Projector)**를 표준 RL 훈련과 통합하여, 모델이 이산적인 토큰 샘플링과 연속적인 잠재 추론을 교차하여 사용할 수 있도록 합니다.

핵심 아키텍처

• 추론 프로젝터 (Reasoning Projector): 기본 LLM에 추가된 작은 다층 퍼셉트론(MLP) 스택입니다. 이는 모델의 마지막 은닉 상태(last hidden state)를 입력받아 연속적인 토큰 임베딩 벡터를 예측합니다.
• 교차 모드 (Interleaved Modes): 모델은 두 가지 모드로 작동합니다:
  • 이산 모드 (Discrete Mode): LM 헤드를 통한 표준 토큰 샘플링입니다.
  • 잠재 모드 (Latent Mode): 모델이 생성한 특수 태그인 <implicit_thought>에 의해 트리거됩니다. 이 모드에서 추론 프로젝터는 연속적인 임베딩을 예측하여 모델에 다시 주입하며, 이를 통해 명시적인 텍스트 토큰을 생성하지 않고도 여러 추론 단계를 "건너뜁니다". 이 태그에는 수행할 복잡도(잠재 단계 수)를 나타내는 숫자가 포함됩니다.

훈련 파이프라인

LiteReason 훈련은 세 단계로 구성됩니다:

1. 데이터 수집: 저자들은 거절 샘플링(rejection sampling)을 통해 초기 데이터셋을 구축합니다. 각 프롬프트에 대해 여러 추론 사슬을 샘플링하고, 양의 보상을 받은 것들(플롯 구멍 탐지의 경우 정확도, 장(chapter) 예측의 경우 대조적 개선도)을 필터링합니다.
2. 지도 미세 조정 (SFT): 추론 프로젝터가 "건너뛴" 단계들을 위한 유용한 임베딩을 예측하도록 훈련됩니다. 추론 사슬 내의 일정 비율의 문장들이 <implicit_thought> 태그로 대체됩니다. 프로젝터는 남은 명시적 추론 단계들 사이의 간극을 메우는 법을 배웁니다. 이 단계 동안 기본 LLM은 동결됩니다.
3. 강화 학습 (RL): 전체 모델은 그룹 상대 정책 최적화(GRPO)를 사용하여 미세 조정됩니다.
   • 정책 모델은 프로젝터를 언제 활성화할지(즉, 언제 잠재 추론을 사용할지) 결정하는 법을 배웁니다.
   • 그래디언트는 이산적 토큰 샘플링 단계에 대해서만 계산되며, 프로젝터는 정책 그래디언트를 통해 직접 업데이트되지 않습니다.
   • 주기적 SFT: 모델의 추론 능력이 진화함에 따라 프로젝터가 낙후되는 것을 방지하기 위해, 매 RL 에포크마다 새로 생성된 추 traces를 사용하여 작은 SFT 단계를 수행합니다.

서사 과업을 위한 보상 설계

• Flawed Fictions (검증 가능): 이진 보상을 사용합니다 (플롯 구멍 탐지가 정확하면 1, 그렇지 않으면 0).
• Next Chapter Prediction (비검증 가능): 기존의 VR-CLI(Completion Likelihood Improvement를 통한 검증 가능한 보상) 목적 함수를 확장합니다. 저자들은 다른 책에서 무작위로 선택된 장(chapter)의 가능도 개선도를 빼줌으로써 일반적인 조언에 페널티를 주는 대조적 개선(Contrastive Improvement) 보상을 도입합니다. 이는 더 다양하고 이야기 특화된 추론을 장려합니다.

주요 기여

1. 새로운 프레임워크: 서사 과업을 위해 잠재 추론을 이산적 토큰 샘플링과 교차하여 RL 훈련 중에 사용할 수 있는 최초의 방법인 LiteReason을 소개합니다. 이를 통해 기존의 추론 데이터셋 없이도 모델이 잠재 및 이산 추론을 동적으로 선택할 수 있게 합니다.
2. 효율성-성능 트레이드오프: 이 방법은 표준 비잠재 RL 훈련에 필적하는 성능 향상을 달성하면서도 토큰 생성을 획기적으로 줄입니다.
3. 과업 적응성: 수학/논리 벤치마크(잠재 추론이 이전에 지배적이었던 분야)와 크게 다른 긴 컨텍스트 서사 과업(Flawed Fictions 및 Next Chapter Prediction)에 잠재 추론을 성공적으로 적용했습니다.

실험 결과

저자들은 Qwen 2.5-7B-Instruct를 사용하여 Flawed Fictions(플롯 구멍 탐지)와 Next Chapter Prediction(NCP)에 대해 LiteReason을 평가했습니다.

• 베이스라인 대비 성능:
  • LiteReason은 모든 훈련 없는(Mixture of Inputs, Soft Thinking) 방식 및 훈련된 잠재 추론 베이스라인(COCONUT, CoLaR)보다 성능이 뛰어납니다.
  • Flawed Fictions에서 LiteReason은 **87.42%**의 정확도를 달성하여, 표준 비잠재 RL의 성능 향상(88.71%)의 **96%**에 도달하는 동시에, 91.5% 적은 토큰(34개 vs 114개)을 생성했습니다.
  • NCP에서 LiteReason은 대조적 개선도 측면에서 베이스라인 대비 613%의 향상을 보였으며, 비잠재 RL 이득의 **69%**에 도달하는 동시에 77.5% 적은 토큰을 생성했습니다.
• 훈련 효율성: LiteReason은 표준 RL에 비해 RL 훈련 중 약 50% 적은 토큰을 필요로 하여 계산 비용을 크게 절감합니다.
• 추론 속도: LiteReason 모델은 실제 환경에서 RL 훈련된 베이스라인보다 약 3배 빠르게 답변을 생성합니다.
• 일반화: 다른 잠재 방식들(CO-NUT, CoLaR)이 도메인 외 작업(예: GSM-Hard 수학)에서 성능 붕괴를 겪은 것과 달리, LiteReason은 일반적인 모델 능력을 유지했으며 다양한 모델 크기(Qwen 3 4B 및 Gemma 3 1B로 테스트)에 걸쳐 견고함을 보여주었습니다.

의의 및 주장

본 논문은 LiteReason이 RL 훈련을 성능-계산 트레이드오프 곡선의 더 효율적인 부분으로 성공적으로 유도한다고 주장합니다. 고정된 잠재 단계나 고품질의 기존 추론 데이터셋에 의존하는 이전 방식들의 엄격한 제약을 피하면서, Lite-Reason은 모델이 추론을 잠재 공간으로 압축하는 방법을 학습할 수 있게 합니다.

저자들은 LiteReason을 자원이 제한된 서사 과업을 위한 실질적인 솔루션으로 제시하며, 연속적인 잠재 추론이 최적에 가까운 성능을 달성하면서도 계산 오버헤드를 크게 줄이는 방식으로 효과적으로 결합될 수 있음을 입증했습니다. 또한, 이것이 서사 도메인에 적용된 최초의 연속 잠재 추론 사례이며, RL 중에 잠재 및 이산 추론을 동적으로 교차시킨 첫 번째 사례임을 밝히고 있습니다.