기술 요약: 확산 언어 모델을 위한 자기 증강 검색 (SARDI)

문제 정의

검색 증강 생성(RAG)은 외부 지식에 대규모 언어 모델(LLM)을 접지(grounding)시키기 위해 점점 더 많이 사용되고 있습니다. 그러나 지배적인 자기회귀(AR) 디코딩 패러다임은 RAG에 있어 두 가지 근본적인 한계를 가집니다:

1. 지연 시간(Latency): AR 생성은 토큰을 순차적으로 방출하며, 이는 지연 시간을 출력 길이에 직접적으로 종속시킵니다.
2. 검색 제약: AR 모델은 오직 "확정된(committed)" 좌측-우측 방향의 접두사(prefix)를 기반으로만 검색할 수 있습니다. 멀티홉 질의응답(QA)에서 이는 결정적인 병목 현상이 되는데, 왜냐하면 후속 추론 단계에 필요한 증거는 초기 질문에는 명시되지 않은 "브릿지 엔티티(bridge entities)"에 의존하는 경우가 많기 때문입니다. 정적 검색은 모델이 브릿지 엔티티를 이미 생성한 후에야 해당 증거를 찾을 수 있으며, 이는 오류를 수정하거나 생성을 효과적으로 유도하기에는 너무 늦을 수 있습니다.

확산 언어 모델(DLM)은 병렬 디코딩 능력을 제공하지만, 기존의 DLM용 RAG 접근 방식은 정적이고 단일 샷(single-shot)인 검색 패러다임에 국한되어 있어, 확산 궤적(diffusion trajectory)의 고유한 구조적 이점을 활용하지 못하고 있습니다.

방법론: SARDI

저자들은 이산형 DLM의 디노이징 과정과 검색을 교차시키는 학습 불필요(training-free) 동적 RAG 프레임워크인 **SARDI(Self-Augmenting Retrieval for Diffusion Language Models)**를 제안합니다.

핵심 메커니즘

SARDI는 DLM이 전체 응답을 병렬로 정제하며, 모든 토큰 위치가 최종 출력으로 확정되기 전에 일련의 중간 상태 $\{x_t\}$와 각 토큰의 예측값 및 신뢰도 점수를 생성한다는 사실을 활용합니다.

1. 룩어헤드 검색 (Lookahead Retrieval): 확정된 접두사를 기다려야 하는 AR 모델과 달리, SARDI는 현재 디노이징 상태의 추측적이고 낮은 신뢰도를 가진 토큰들을 사용하여 검색 쿼리를 구성합니다.

   • 각 디노이징 단계 $t$에서 모델은 모든 마스킹된 위치에 대해 토큰을 예측합니다.
   • 신뢰도 $c_i \geq \tau_q$ (쿼리 임계값)를 만족하는 토큰들을 추출하여 프록시 시퀀스 $\tilde{x}_t$를 형성합니다.
   • 검색 쿼리는 원래의 질문과 새롭게 나타나는 중간 응답을 결합한 $s_t = q \parallel \text{Detokenize}(\tilde{x}_t)$로 구성됩니다.
   • 이를 통해 시스템은 모델이 해당 엔티티를 최종 출력으로 확정할 만큼 충분히 확신하기 전에도 "후속 홉(later-hop)" 증거(예: 브릿지 엔티티의 위치)를 찾아낼 수 있습니다.

2. 신뢰도 기반 언마스킹 (Confidence-Based Unmasking):

   • 신뢰도 $c_i \geq \tau_c$ (커밋 임계값, 여기서 $\tau_c \geq \tau_q$)를 만족하는 토큰은 출력으로 확정됩니다.
   • 만약 어떤 토큰도 $\tau_c$를 만족하지 못하면, 진행을 보장하기 위해 가장 신뢰도가 높은 단일 마스크 토큰을 언마스킹합니다.
   • 이는 고신뢰도의 증거에 기반한 스팬(span)이 먼저 확정되는 반면, 불확실한 스팬은 다음 검색 라운드로부터 정제된 증거를 기다리게 하는 커리큘럼을 생성합니다.

3. 동적 증거 갱신 (Dynamic Evidence Refresh):

   • 매 디노이징 단계마다 SARDI는 업데이트된 쿼리 $s_t$를 기반으로 새로운 $K$개의 패시지를 검색합니다.
   • 새로운 증거 $D_{t-1}$은 이전 컨텍스트를 완전히 대체하여 다음 디노이징 단계의 조건이 됩니다.

주요 설계 선택 사항

• 임계값 분리: 쿼리 임계값($\tau_q$)과 커밋 임계값($\tau_c$)의 분리가 핵심입니다. $\tau_q \leq \tau_c$ (흔히 $\tau_q=0$으로 설정)로 설정함으로써, SARDI는 토큰이 안정되어 출력으로 확정되기 훨씬 전부터 추측적인 토큰들을 검색기에 노출합니다. 이는 검색은 노이즈가 있는 쿼리에 강건하지만, 잘못된 토큰을 확정하는 것은 출력을 오염시키기 때문에 안전한 전략입니다.
• 리트리버 불가지론 (Retriever Agnostic): 이 프레임워크는 모든 희소(sparse, 예: BM25) 또는 밀집(dense) 리트리버와 함께 작동합니다.
• 학습 불필요 (Training-Free): SARDI는 학습된 검색 컨트롤러나 쿼리 생성기를 요구하지 않으며, 추론 궤적을 생성할 수 있는 모든 이산형 DLM을 위한 플러그 앤 플레이 모듈로 작동합니다.

주요 기여

1. 룩어헤드 신호: 본 논문은 확산 궤적에서 버려지는 낮은 신뢰도의 토큰들이 검색을 위한 가치 있는 신호가 될 수 있음을 입증하며, 이것이 AR 생성보다 브릿지 엔티티를 더 일찍 드러낸다는 것을 보여줍니다.
2. 병렬 디코딩 호환성: 저자들은 증거에 기반하여 생성을 수행하는 것이 토큰 간 상호 정보량(특히 엔티티 스팬에 대해)을 급격히 감소시킨다는 것을 실증적으로 보여주며, RAG가 DLM의 병렬 디코딩 강점에 독특하게 적합한 체제임을 밝힙니다.
3. SARDI 프레임워크: 중간 확산 상태에 검색을 조건화하여, 추가 학습 없이도 디노이징 궤적 전반에 걸쳐 동적 증거 갱신을 가능하게 하는 최초의 프레임워크입니다.

실험 결과

저자들은 다섯 가지 멀티홉 QA 벤치마크(2WikiMultiHopQA, HotpotQA, MuSiQue, CofCA, SynthWorlds-RM)에서 SARDI를 평가했습니다.

• 성능: SARDI는 정적 확산 검색 베이스라인을 크게 상회합니다. 예를 들어, 2WikiMultiHopQA에서 $\tau_c=0.95$일 때 Exact Match(EM)를 43.7%에서 59.1%로 높였습니다.
• AR과의 비교: SARDI는 모든 학습 불필요 자기회귀 베이스라인(FLARE 및 AdaptiveRAG 포함)과 대등하거나 더 나은 정확도를 보이면서도, 최대 8배 높은 처리량(더 낮은 지연 시간)으로 작동합니다.
• RL 에이전트와의 비교: SARDI는 RL로 학습된 에이전트인 Search-R1과 대등하거나 더 나은 정확도를 달면서도, 강화 학습이나 보상 설계가 필요하지 않아 3~8배 더 빠르게 실행됩니다.
• 어블레이션 (Ablations):
  • Lookahead: 공격적인 룩어헤드($\tau_q=0$)가 가장 높은 정확도를 기록하여, 추측적 토큰이 검색을 개선함을 확인했습니다.
  • 질문 유형: 성능 향상은 브릿지 엔티티를 필요로 하는 멀티홉 추론 질문에 집중되었으며, 싱글홉 질문에서는 변화가 거의 없었습니다.
  • 리트리버: 결과는 BM25와 밀집 리트리버(E5-base-v2) 모두에서 유효했습니다.
  • 갱신 빈도: 매 단계마다 갱신하는 것이 최적이지만, 2단계마다 갱신해도 EM 손실이 1~2%에 불과하여 비용 절감을 위한 최적화 가능성을 시사했습니다.

의의 및 주장

본 논문은 SARDI가 검색 증강 생성의 "품질-지연 시간 경계(quality–latency frontier)"를 해제한다고 주장합니다. 자신의 생성물에 대한 미래를 "엿볼 수 있는" 능력과 증거에 기반했을 때 토큰 간 의존성이 줄어든다는 확산 모델의 구조적 이점을 활용함으로써, SARDI는 현재의 자기회귀 방식보다 더 정확하면서도 훨씬 빠른 동적 멀티홉 추론을 가능하게 합니다.

저자들은 현재 이 방법이 추론 궤적을 생성하는 모델에 의존하고 있으며, 이는 작은 규모의 DLM의 경우 가벼운 지도 학습 기반 미세 조정(supervised fine-tuning)이 필요할 수 있다고 언급했습니다(초기 AR 모델과 유사함). 그러나 DLM이 성숙해짐에 따라 이러한 요구사항은 줄어들 것입니다. 이 연구는 디노이징 궤적이 동적 검색을 위한 자연스러운 신호를 제공한다는 점을 시사하며, SARDI를 지식 집약적 과업을 위한 강력하고 학습이 필요 없는 솔루션으로 포지셔닝합니다.