RecaLLM: "Lost-in-Thought" 현상을 해결하기 위한 명시적 인-컨텍스트 검색 (Explicit In-Context Retrieval)

이 논문은 긴 문맥 (Long-context) 을 가진 대형 언어 모델 (LLM) 이 추론 (Reasoning) 과정을 거치는 동안 인-컨텍스트 정보 검색 능력이 급격히 저하되는 "Lost-in-Thought(사고 중 분실)" 현상을 발견하고, 이를 해결하기 위해 제안된 RecaLLM 모델에 대해 설명합니다.

1. 문제 정의: Lost-in-Thought (사고 중 분실)

기존의 긴 문맥 LLM 연구는 문맥의 길이를 늘리는 데 집중했으나, 모델이 긴 문맥에서 관련 정보를 효과적으로 활용하는 능력은 여전히 부족합니다. 특히, 추론 (Chain-of-Thought) 을 수행하는 모델에서 다음과 같은 새로운 병목 현상이 관찰되었습니다.

• 현상: 모델이 복잡한 추론 단계를 수행할 때, 그 과정에서 생성된 의미상 관련 있지만 주의력을 분산시키는 토큰들이 누적됩니다. 이로 인해 추론이 완료된 후, 원래 입력 문맥에 있는 정답을 찾아내는 인-컨텍스트 검색 (In-Context Retrieval) 성능이 심하게 저하됩니다.
• 발견: 저자들은 여러 오픈소스 LLM 을 대상으로 실험한 결과, 수학 문제를 풀고 그 해답을 딕셔너리에서 찾아야 하는 과제에서, 모델은 정답 키 (Key) 를 올바르게 추론해 내더라도, 해당 키에 대응하는 값 (Value) 을 문맥에서 정확히 복사해 내지 못하거나 (Hallucination), 아예 찾지 못하는 경우가 빈번함을 확인했습니다.
• 원인: 추론 과정에서 생성된 토큰들이 새로운 "주의 분산 요소 (Distractors)"로 작용하여, 모델이 문맥의 특정 부분을 정확하게 기억하고 복제하는 능력을 방해하는 것으로 분석되었습니다.

2. 방법론: RecaLLM

RecaLLM 은 추론과 검색을 분리하지 않고, 추론 과정 중에 명시적으로 인-컨텍스트 검색을 수행하도록 설계된 모델입니다. 핵심 구성 요소는 다음과 같습니다.

2.1. 명시적 회상 스페인 (Recall Spans) 및 토큰

모델은 추론이 필요한 정보를 찾았을 때, 특수한 구분자 토큰인 <|start recall|>과 <|end recall|>을 사용하여 해당 정보를 원문 (Verbatim) 그대로 복사하는 구간을 명시적으로 표시합니다. 이는 추론과 검색을 번갈아 수행하는 인터리빙 (Interleaving) 구조를 만듭니다.

2.2. 제약된 디코딩 (Constrained Decoding)

회상 스페인 구간 내에서는 모델이 임의의 텍스트를 생성하는 것이 아니라, 검색 가능한 문맥 (Searchable Context) 내에서만 유효한 다음 토큰을 선택하도록 제한합니다.

• 구현: 로짓 (Logit) 마스킹을 사용하여, 현재까지 생성된 접두사 (Prefix) 와 문맥에서 정확히 일치하는 다음 토큰만 허용합니다.
• 효과: 이 메커니즘은 모델이 정보를 "기억"하거나 "생성"하는 것이 아니라, **문맥에서 정확히 복사 (Copy)**하도록 강제하여 할루시네이션을 원천적으로 차단합니다.
• 오버헤드: 이 마스킹 연산은 GPU 프론트패스와 병렬로 CPU 스레드에서 수행되므로, 생성 지연 (Latency) 에 거의 영향을 주지 않습니다.

2.3. 훈련 파이프라인

RecaLLM 은 두 단계의 후속 훈련 (Post-training) 을 거칩니다.

1. 지도 미세조정 (SFT) 콜드 스타트: 교사가 주석한 추론 궤적 (Teacher-annotated rollouts) 을 사용하여 모델이 새로운 회상 토큰을 자연스럽게 사용하고, 문맥 정보를 <recall> 스페인으로 복사하는 법을 학습합니다.
2. 강화 학습 (RL - GRPO): 다양한 작업 (단일/다중 홉 QA, 검색, 추론 검색 등) 에 대해 강화 학습을 수행합니다.
   • 복합 보상 함수: 최종 정답의 품질뿐만 아니라, 검색된 증거 (Evidence) 의 정확성을 함께 평가합니다.
   • 보상 구성: 포맷팅 정확도, 정답 품질, 그리고 검색된 증거의 F1 점수를 결합하여, 모델이 단순히 정답만 맞추는 것이 아니라 올바른 근거를 찾아내도록 유도합니다.

3. 주요 기여

1. Lost-in-Thought 현상의 규명: 추론 과정이 인-컨텍스트 검색 능력을 저하시킨다는 새로운 실패 모드를 발견하고 이를 정량화했습니다.
2. RecaLLM 아키텍처 제안: 추론과 명시적 검색을 번갈아 수행하며, 제약된 디코딩을 통해 정보의 충실한 복제를 보장하는 새로운 모델 구조를 제시했습니다.
3. 검증 가능한 검색 메커니즘: 제약된 디코딩을 통해 검색된 내용이 원문과 일치함을 보장함으로써, 추론 과정의 신뢰성을 높이고 보상 설계 및 해석 가능성을 개선했습니다.
4. 효율적인 장문맥 일반화: 최대 10K 토큰의 짧은 문맥 데이터로 훈련되었음에도 불구하고, 128K 토큰까지의 긴 문맥 평가에서 뛰어난 성능을 보였습니다. 이는 기존 장문맥 모델들이 긴 훈련 데이터를 필요로 했던 것과 대조적입니다.

4. 실험 결과

RecaLLM 은 RULER와 HELMET라는 두 가지 주요 장문맥 벤치마크에서 기존 모델들을 압도하는 성능을 보였습니다.

• RULER 벤치마크:
  • RecaLLM-Qwen2.5-7B 는 7~8B 클래스 모델 중 평균 점수 (92.8) 에서 최상위권을 기록했으며, 14B 모델인 LoongRL-14B 보다도 높은 성능을 보였습니다.
  • 문맥 길이가 길어질수록 성능 향상 폭이 커졌습니다 (예: 4K 에서 +6.8, 128K 에서 +16.1).
• HELMET 벤치마크:
  • RecaLLM 은 베이스 모델 대비 평균 16~17 점 이상 향상되었습니다.
  • 특히 **Recall(회상), ICL(인-컨텍스트 학습), Re-rank(재순위화), Citation(인용)**과 같이 정확한 정보 검색이 필수적인 작업에서 비약적인 개선을 보였습니다 (예: Llama 모델의 ICL 점수가 3.0 에서 64.1 로 급상승).
• 효율성: RecaLLM 은 10K 토큰 이하의 데이터로 훈련되었으나, 512K 토큰 데이터로 훈련된 ProLong 이나 16K 토큰 데이터로 훈련된 LoongRL 보다 더 좋은 결과를 내며, 계산 비용과 데이터 수집 비용을 크게 절감할 수 있음을 입증했습니다.

5. 의의 및 결론

이 연구는 장문맥 LLM 의 성능 향상을 위해 단순히 문맥 창 (Context Window) 을 늘리는 것만으로는 부족하며, 추론과 검색의 상호작용을 최적화해야 함을 강조합니다.

• 패러다임 전환: "검색은 추론 전에 한 번만 수행한다"는 기존 접근법에서 벗어나, 추론 과정 중에 동적으로 검색을 수행하고 그 결과를 검증 가능하게 만드는 새로운 패러다임을 제시했습니다.
• 신뢰성 확보: 제약된 디코딩을 통해 모델이 문맥 정보를 왜곡 없이 정확히 복사하게 함으로써, 할루시네이션을 줄이고 모델의 판단 근거를 투명하게 만들었습니다.
• 미래 방향: RecaLLM 은 장문맥 처리를 위한 비용 효율적인 솔루션을 제공하며, 추론의 일관성을 유지하기 위한 자기 회상 (Self-recall) 및 더 복잡한 에이전트 워크플로우로 확장될 수 있는 가능성을 열었습니다.

결론적으로, RecaLLM 은 "Lost-in-Thought"라는 근본적인 문제를 해결함으로써, 적은 계산 비용과 짧은 훈련 데이터로도 고품질의 장문맥 추론이 가능함을 증명했습니다.