When to Retrieve During Reasoning: Adaptive Retrieval for Large Reasoning Models

Guo, Wu, Yiu 의 논문 "When to Retrieve During Reasoning: Adaptive Retrieval for Large Reasoning Models"에 대한 상세한 기술 요약입니다.

1. 문제 제기

이 논문은 DeepSeek-R1 과 OpenAI o1 과 같은 대규모 추론 모델 (LRM) 과 전통적인 검색 증강 생성 (RAG) 시스템 간의 근본적인 시간적 불일치를 다룹니다.

• 불일치: 표준 RAG 는 추론이 시작되기 전에 컨텍스트를 주입하는 "검색 후 생성" 패러다임을 따릅니다. 그러나 LRM 은 다단계 추론, 자기 검증, 백트래킹을 포함하는 12,000~25,000 토큰에 달하는 확장된 사고의 사슬을 생성합니다. 지식의 공백은 종종 추론 중간에 발생하며, 쿼리 단계에서 예측할 수 없습니다.
• 기존 반복적 방법의 한계: 현재의 반복적 검색 방법 (예: IRCoT, FLARE) 은 LRM 에 적합하지 않습니다.
  1. 세분성 불일치: 이들은 토큰 또는 문장 수준에서 검색을 트리거하는 반면, LRM 은 추론 단계 수준 (여러 문장에 걸친 논리적 단위) 에서 작동합니다.
  2. 신호 부재: 토큰 확률 (FLARE) 이나 내부 주의 상태 (DRAGIN) 에 의존하는 방법은 이러한 내부 정보를 노출하지 않는 "완전 전용" 모델 (o1 등) 에서는 실패합니다.
  3. 효율성 붕괴: 고정 간격 검색 (예: 문장마다) 은 긴 사슬의 경우 과도하게 비용이 많이 들며, 상당한 지연 시간과 토큰 오버헤드를 추가합니다.

2. 방법론: ReaLM-Retrieve

저자들은 추론 사슬 동안 검색을 동적으로 통합하는 추론 인식 프레임워크인 ReaLM-Retrieve를 제안합니다. 이는 세 가지 핵심 구성 요소로 이루어져 있습니다.

A. 추론 단계 분할

문장 수준 분할 대신 시스템은 논리적 추론 단계를 식별합니다.

• 메커니즘: 경량 3 레이어 트랜스포머 분류기는 담화 마커 ("따라서", "그러나"), 논리적 연결어, 주제 전환 (임베딩 유사도를 통해), 그리고 구두점을 기반으로 단계 경계를 감지합니다.
• 학습: DeepSeek-R1 의 인간 주석 traces 로 학습되어 경계 식별에서 94.2% 의 F1 점수를 달성했습니다. 이를 통해 시스템은 토큰의 흐름이 아닌 이산적인 논리적 단위의 시퀀스로 추론 사슬을 처리할 수 있습니다.

B. 단계별 불확실성 감지 (RSUS)

시스템은 외부 지식이 필요한 단계를 결정하기 위해 **추론 단계 불확실성 점수 (RSUS)**를 계산합니다.

• 복합 점수: $RSUS(r_i) = \alpha \cdot U_{verb} + \beta \cdot U_{ent} + \gamma \cdot U_{cons}$
  • $U_{verb}$ (표현된 불확실성): 모델은 신뢰도 (0–100) 를 평가하도록 프롬프트됩니다. 완전 전용 모델의 경우, 학습된 MLP 프록시가 표면 특징과 개체 수를 기반으로 이를 추정합니다.
  • $U_{ent}$ (개체 엔트로피): 검색된 코퍼스 내의 명명된 개체들의 모호성/희소 커버리지를 측정합니다. 높은 엔트로피는 모델이 특정 개체에 대한 지식이 부족함을 나타냅니다.
  • $U_{cons}$ (일관성): 현재 단계의 여러 샘플링된 연속체 간의 합의를 측정합니다.
• 효율성: 이 계산은 시맨틱 엔트로피나 자기 일관성 샘플링의 막대한 비용 없이 약 8% 의 오버헤드만 추가합니다.

C. 학습된 개입 정책 및 효율적 통합

• 정책 ($\pi$): 컨텍스트 밴딧 모델이 언제 검색할지, 무엇을 쿼리로 구성할지, 그리고 증거를 어떻게 통합할지 결정합니다. 이는 REINFORCE 를 통해 훈련되어 검색 비용 (지연 시간 및 호출 수) 을 패널티로 부과하면서 답변 정확도를 극대화합니다.
• 쿼리 구성: 정책은 현재 단계를 반복하는 대신 정보 요구를 타겟팅하는 구체적인 쿼리를 생성합니다 (예: "[개체] 의 [개념] 과의 관계는 무엇입니까?").
• 효율성 최적화:
  • 암시적 압축: 중요도 임계값을 초과하는 검색된 문서의 문장만 유지하여 컨텍스트 확장을 73% 줄입니다.
  • 추측적 캐싱: 다음 단계 쿼리에 대한 문서를 병렬로 미리 가져옵니다. 추측이 맞으면 (37% 히트율) 검색 지연 시간이 제거됩니다.
  • KV 캐시 보존: 오픈 가중치 모델의 경우 변경되지 않은 컨텍스트 부분에 대해 캐시된 키 - 값 상태를 보존하여 첫 번째 토큰 도달 시간을 2.1 배 단축합니다.

3. 주요 기여

1. 단계별 세분성: 검색 트리거를 토큰/문장 수준에서 추론 단계 수준으로 전환하여 LRM 의 논리적 구조와 정렬합니다.
2. 블랙박스 호환 불확실성: 토큰 확률이나 내부 주의 상태에 대한 접근 없이 완전 전용 모델 (o1 등) 에서 작동하는 RSUS 를 도입합니다.
3. 효율성 최적화 통합: 추측적 캐싱과 컨텍스트 압축을 통해 단순 통합 대비 3.2 배의 검색당 오버헤드를 줄여 실용적인 배포를 가능하게 합니다.

4. 실험 결과

DeepSeek-R1 및 QwQ 모델을 사용하여 MuSiQue, HotpotQA, 2WikiMultiHopQA에서 평가되었습니다.

• 정확도: ReaLM-Retrieve 는 표준 RAG 대비 평균 10.1% 의 절대적 개선을, 가장 강력한 기존 방법 (ReARTeR) 대비 **4.6%**의 F1 점수 향상을 달성했습니다.
  • MuSiQue(2~4 홉 추론) 에서 71.2% F1을 달성했습니다 (IRCoT 의 65.4% 대비).
• 효율성:
  • 적은 호출: 고정 간격 접근법 대비 검색 호출을 47% 줄였습니다 (MuSiQue 에서 IRCoT 의 3.4 회 대비 1.8 회).
  • 지연 시간: IRCoT 대비 1.33 배 낮은 엔드투엔드 지연 시간과 단순 인터리빙 대비 1.50 배 낮은 지연 시간을 달성했습니다.
  • 비용: 증거가 불확실성을 더 빠르게 해결하여 추론 사슬이 짧아짐에 따라 출력 토큰 비용을 16% 절감했습니다.
• 검색 품질: **Recall@5 를 81.3%**로 향상시켜 (IRCoT 의 72.6% 대비), 시스템이 적절한 시점에 더 나은 증거를 검색함을 입증했습니다.
• 일반화: 정책은 다른 모델 아키텍처 (예: QwQ) 에 잘 일반화되며, HotpotQA 에서 훈련하고 MuSiQue 에서 테스트하는 등 벤치마크 간 전이 시 성능 저하가 최소화되었습니다.

5. 의의

• 패러다임 전환: 이 논문은 빈번하고 고정 간격의 검색보다 **"덜 많지만 더 적시에 이루어지는 검색"**이 우수함을 입증합니다. 이는 추론 집약적 작업에서 정확도 - 효율성 트레이드오프를 위한 새로운 파레토 프론티어를 확립합니다.
• 확장성: 추론 단계 수준에서 작동하고 완전 전용 모델을 지원함으로써, 이 프레임워크는 현재 많은 RAG 연구 노력에서 접근할 수 없는 가장 강력한 독점 LLM(예: o1) 에도 적용 가능합니다.
• 시스템 설계: 검색을 정적 전처리 단계에서 동적이고 적응적인 추론 파트너로 이동시켜, 장기 추론 작업에 검색을 통합하는 청사진을 제공합니다.

결론적으로, ReaLM-Retrieve 는 LRM-RAG 통합의 중요한 타이밍 불일치를 해결하여, 계산 비용을 크게 줄이면서 최첨단 정확도를 제공합니다.