확산 언어 모델을 위한 자기 증강 검색 (Self-Augmenting Retrieval for Diffusion Language Models)

이산 확산 언어 모델 (Discrete diffusion language models)은 전체 응답을 병렬적으로 반복적인 디노이징 (denoising) 과정을 통해 생성합니다. 각 단계에서 모델은 모든 마스크된 (masked) 위치에 대해 잠정적인 토큰 (tentative tokens)을 예측하며, 확신이 있는 예측은 출력으로 확정하고 확신이 낮은 예측은 폐기합니다. 본 논문에서는 폐기된 토큰들이 사실 검색 증강 생성 (Retrieval-Augmented Generation, RAG)을 위한 유용한 룩어헤드 신호 (lookahead signal)임을 보여줍니다. 즉, 확신도가 낮은 토큰이라 할지라도 디노이징 궤적 (denoising trajectory)의 초기 단계에서 중요한 개체 (salient entities)를 드러내는 경우가 많으며, 이를 통해 출력이 최종 확정되기 전에 더 강력한 증거를 검색할 수 있습니다. 우리는 이러한 특성을 활용하여 확산 언어 모델을 위한 자기 증강 검색 (Self-Augmenting Retrieval for Diffusion Language Models, SARDI)을 제안합니다. SARDI는 이러한 룩어헤드 토큰을 사용하여 디노이징 과정 동안 검색을 안내하는 동적인 RAG 프레임워크입니다. SARDI는 별도의 학습이 필요하지 않고 (training-free), 검색기 모델에 구애받지 않으며 (retriever-agnostic), 추론 능력을 갖춘 모든 이산 확산 언어 모델에 적용 가능합니다. 5개의 멀티홉 질의응답 (multi-hop QA) 벤치마크 전반에서 SARDI는 기존의 학습이 필요 없는 확산 모델 및 자기회귀 (autoregressive) 검색 베이스라인보다 최대 8배 높은 처리량 (throughput)을 기록하며 성능을 앞질렀습니다.