💬 NLP

DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

DiffRetriever 는 확산 언어 모델을 위해 프롬프트에 여러 개의 마스킹된 위치를 추가하고 단일 양방향 패스로 이를 읽는 병렬 대표 토큰 검색 방법을 도입하여, 순차적 자기회귀 다중 토큰 접근법 및 기존 대비적 미세 조정 검색기보다 우수한 성능과 효율성을 달성합니다.

원저자: Shuai Wang, Yin Yu, Shengyao Zhuang, Bevan Koopman, Guido Zuccon

게시일 2026-05-11

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Shuai Wang, Yin Yu, Shengyao Zhuang, Bevan Koopman, Guido Zuccon

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대한 도서관에서 냅킨에 쓴 한 문장을 바탕으로 완벽한 책을 찾아보려 한다고 상상해 보세요. 이것이 컴퓨터의 '검색기(retriever)'가 하는 일입니다: 사용자의 질의를 받아 가장 잘 맞는 문서들을 찾아냅니다.

오랫동안 이를 수행하는 최선의 방법은 사용자의 문장을 읽고 전체 아이디어를 나타내는 단 하나의 단어만 적어내는 '똑똑한 사서'(AI 모델)를 활용하는 것이었습니다. 이는 빠르지만, 때로는 한 단어가 복잡한 생각을 포착하기에 충분하지 않습니다.

연구자들은 사서가 더 정밀하게 여러 개의 단어를 쓰도록 시도했습니다. 그러나 기존 유형의 AI(자기회귀형)를 사용할 경우, 사서는 이 단어들을 한 줄로 하나씩 차례대로 써야 했습니다. 20 단어를 요청하면, 첫 번째를 쓰고 기다렸다가 두 번째를 쓰고, 다시 기다리는 식으로 이어졌습니다. 이는 너무 느려서 추가된 단어들이 기다린 시간만큼의 가치가 없어 보였으며, 결과도 크게 나아지지 않았습니다.

DiffRetriever 등장: '집단 사고' 사서

이 논문은 **확산 언어 모델(Diffusion Language Model)**이라는 다른 유형의 AI를 사용하는 DiffRetriever라는 새로운 사서를 소개합니다.

여기서 마술 같은 비법이 있습니다:
단어를 한 줄로 하나씩 쓰는 대신, 확산 방식의 사서는 20 개의 빈 자리(마스크된 위치)가 있는 빈 페이지를 보고 20 개의 자리를 정확히 동시에 채웁니다.

이렇게 생각해보세요:

기존 방식 (자기회귀): 친구에게 영화에 대해 설명해 달라고 요청합니다. 친구는 "그 영화는..." (잠시 멈춤) "..." (잠시 멈춤) "훌륭한..." (잠시 멈춤) "액션..." (잠시 멈춤) "영화였습니다."라고 말합니다. 다음 단어를 시작하기 전에 각 단어가 끝날 때까지 기다려야 합니다. 20 단어를 원한다면 시간이 무한히 걸립니다.
새로운 방식 (확산): 친구에게 20 개의 빈 칸이 있는 시트를 건네줍니다. 그리고 "이 칸들을 설명으로 채워줘"라고 말합니다. 친구는 시트 전체를 보고 영화에 대해 생각한 뒤, 동시에 모든 20 개의 칸을 한 번의 시선으로 채웁니다.

논문이 발견한 점

속도 대 품질: 연구자들은 기존 '하나씩' 작성자와 새로운 '한 번에 모두' 작성자 두 가지 유형의 AI 를 테스트했습니다.
- 기존 작성자에게 여러 단어를 생성하도록 요청하면 속도가 점점 느려졌고, 품질은 실제로 향상되지 않았습니다.
- 반면, 새로운 작성자에게 여러 단어를 생성하도록 요청하면 한 단어를 작성하는 것과 똑같이 빠르았습니다(모든 것을 한 번에 수행했기 때문) 하지만 품질은 크게 향상되었습니다. 같은 시간 동안 훨씬 더 나은 설명을 얻은 것과 같습니다.
최고의 수행자: 새로운 '한 번에 모두' 작성 사서 (구체적으로 'Dream'이라 불리는 모델) 를 학습시킨 후, 이 모델은 테스트에서 관련 문서를 찾는 데 있어 가장 뛰어난 성능을 보였습니다. 기존 다중 단어 시도와 다른 최신 AI 모델을 모두 제치고 1 위를 차지했습니다.
'완벽한' 예산: 연구자들은 또한 다른 질문들이 서로 다른 수의 단어가 필요하다는 점도 발견했습니다. "날씨는 어때요?"와 같은 간단한 질문은 2 단어만 필요할 수 있지만, 복잡한 질문은 16 단어가 필요할 수 있습니다.
- 현재는 하나의 숫자 (예: 4 또는 16) 를 선택하여 모든 질문에 동일하게 적용합니다.
- 그들은 만약 작성하기 전에 각 특정 질문에 정확히 몇 단어가 필요한지 마법처럼 알 수 있다면, 시스템은 학습된 모델들보다도 더 나아질 것이라고 발견했습니다.
- 또한 질문의 길이와 같은 간단한 단서로도 필요한 단어 수를 예측할 수 있음을 발견했습니다. 이는 미래에 각 질문에 대해 작성할 단어 수를 자동으로 조정하여 양쪽 세계의 장점을 모두 얻을 수 있는 시스템을 구축할 수 있음을 시사합니다.

요약하자면

이 논문은 검색에 여러 단어를 사용하는 문제의 핵심이 여러 단어를 사용한다는 아이디어가 아니라, 그들을 하나씩 쓰는 구식 방식에 있었다는 것을 증명합니다. 모든 단어를 한 번에 작성하는 새로운 AI 로 전환함으로써, 그들은 다중 단어 검색을 빠르고, 저렴하며, 훨씬 더 정확하게 만들었습니다. 마치 느린 기차가 물건을 배달하는 것을 기다릴 필요가 없다는 것을 깨닫고, 드론이 모든 것을 한 번에 내려주는 방식이 필요하다는 것과 같습니다.

기술 요약: DiffRetriever

문제 정의

최근 연구, 특히 PromptReps(Zhuang 등, 2024) 은 자동회귀 언어 모델 (AR LM) 이 쿼리와 패시지에 대한 대표 토큰 (밀집 벡터 및 희소 로지트) 생성을 위한 프롬프트를 통해 제로샷 검색기로 기능할 수 있음을 입증했습니다. 그러나 이 방식을 멀티토큰 검색 (하나 대신 $K$ 개의 대표 토큰 생성) 으로 확장하는 것은 비효율적이었으며, AR 모델에서는 종종 효과가 없었습니다. AR 모델은 토큰을 순차적으로 생성하므로, $K$ 개의 대표 토큰을 생성하려면 $K$ 번의 순전파 (forward pass) 가 필요하여 인코딩 지연 시간이 $K$ 에 비례하여 선형적으로 증가합니다. 이전의 멀티토큰 PromptReps 시도들은 이러한 상당한 지연 시간 비용에도 불구하고 단일 토큰 디코딩 대비 일관된 성능 향상을 보여주지 못해, 한계가 멀티토큰 개념 자체인지 순차적 생성 병목 현상인지 불분명하게 남았습니다.

방법론

저자들은 순차적 생성 병목 현상을 극복하기 위해 Diffusion Language Models(DiffLM) 에 특화된 검색 프레임워크인 DiffRetriever를 제안합니다.

핵심 메커니즘

토큰을 하나씩 생성하는 AR 모델과 달리 DiffLM 은 양방향 어텐션 하에서 마스킹된 위치를 공동으로 채웁니다. DiffRetriever 는 이를 다음과 같이 활용합니다:

프롬프트 구성: 대표 토큰 프롬프트 (예: "The words are [MASK]...") 에 $K$ 개의 마스킹된 위치 ([MASK]) 를 추가합니다.
병렬 디코딩: 전체 시퀀스 (프롬프트 + $K$ 개의 마스크) 를 단일 양방향 순전파로 DiffLM 에 입력합니다.
표현 추출: $K$ 개의 마스킹된 위치에서 $K$ 개의 밀집 은닉 상태와 $K$ 개의 희소 로지트 벡터를 동시에 추출합니다.

점수 매김 및 학습

점수 매김: 프레임워크는 PromptReps의 하이브리드 점수 매김 레시피를 채택합니다:
- 밀집 점수 (Dense Score): 쿼리와 패시지 은닉 상태 간의 ColBERT 스타일 후기 상호작용 (MaxSim) 을 통해 계산됩니다.
- 희소 점수 (Sparse Score): 최대 풀링 및 ReLU 활성화된 로지트를 통해 계산됩니다.
- 하이브리드 점수: 정규화된 밀집 점수와 희소 점수의 선형 보간입니다.
파인튜닝: 저자들은 밀집 및 희소 점수 모두에 대해 대조적 파인튜닝 (InfoNCE 손실) 을 적용하며, 이는 AR 및 DiffLM 백본 모두에 적용 가능합니다.

실험 설정

본 연구는 7B~8B 파라미터의 네 가지 백본을 비교합니다:

자동회귀 (AR): LLaMA3-8B, Qwen2.5-7B.
확산 (Diff): Dream-7B(Qwen2.5 에서 초기화), LLaDA-8B(처음부터 학습).
베이스라인: BM25, PromptReps(AR), DiffEmbed(DiffLM 을 평균 풀링 인코더로 사용), 그리고 대조적으로 파인튜닝된 RepLLaMA.
데이터셋: MS MARCO, TREC DL 2019/2020, 그리고 BEIR-7(일곱 개의 다양한 도메인 외 데이터셋).

주요 결과

1. 멀티토큰 효율성 및 효과성

확산 대 자동회귀: 테스트된 모든 확산 백본에서 멀티토큰 DiffRetriever 는 단일 토큰 DiffRetriever 보다 크게 향상되었습니다. 반면, 멀티토큰 AR 모델은 단일 토큰 AR 모델 대비 평탄하거나 부정적인 향상을 보였습니다.
지연 시간: 멀티토큰 DiffRetriever 의 인코딩 비용은 $K$ 와 관계없이 거의 일정하게 유지됩니다 (단일 순전파). 반면 AR 멀티토큰 비용은 선형적으로 증가합니다. 제로샷 설정 ( $K \le 20$ ) 에서 AR 멀티토큰은 일관된 성능 향상 없이 확산 모델 대비 약 15 배의 지연 시간 패널티를 부과합니다.
성능 역전: 제로샷 설정에서 AR 모델은 초기에 $K=1$ 에서 DiffLM 보다 성능이 우수합니다. 그러나 $K>1$ 에서 DiffLM(특히 Dream 과 LLaDA) 은 AR 대응 모델을 능가합니다. 예를 들어, Dream 멀티토큰은 단일 토큰 버전 대비 MS MARCO 성능을 거의 두 배로 늘려 Qwen2.5 를 추월합니다.

2. 최첨단 성능

제로샷: LLaDA에서의 DiffRetriever 는 비교 대상 제로샷 시스템 중 가장 높은 BEIR-7 평균 점수를 달성하여, AR 백본에서의 PromptReps 와 DiffEmbed 베이스라인을 능가했습니다.
파인튜닝: 지도 학습 파인튜닝 후 Dream에서의 DiffRetriever 는 전체적으로 가장 강력한 BEIR-7 검색기가 되어, 파인튜닝된 PromptReps, DiffEmbed, 그리고 대조적으로 파인튜닝된 RepLLaMA 를 능가했습니다.
백본 시너지: Qwen2.5 에서 초기화된 Dream 백본은 LLaDA 에 비해 제로샷에서 파인튜닝으로 전환될 때 더 큰 성능 변동을 보여, AR 사전 학습 사전 지식과 양방향 확산 디코딩의 결합이 검색에 매우 효과적임을 시사합니다.

3. 예산 분석 및 오라클 여유분

고정 대 적응형: 본 연구는 모든 쿼리에 고정된 예산 $(K_q, K_p)$ 을 사용하는 것은 상당한 성능을 놓치고 있음을 확인했습니다.
오라클 분석: 고정된 예산에서 대조적 파인튜닝을 능가하는 "쿼리별 오라클"(정답을 사용하여 각 쿼리에 최적의 $K$ 를 선택하는 것) 이 고정된 예산의 모든 조합에서 고정된 베이스 모델에서 더 나은 성능을 보였습니다.
예측 가능성: 최적의 예산은 저렴한 쿼리 특징 (길이 및 섀넌 엔트로피) 과 양의 상관관계를 보이므로, 적응형 예산 선택을 위한 학습된 라우터가 실현 가능한 미래 방향임을 시사합니다.

중요성 및 주장

본 논문은 멀티토큰 검색을 둘러싼 모호함을 해소한다고 주장합니다:

병목 현상은 순차적 생성입니다: 이전 멀티토큰 AR 검색기의 실패는 멀티토큰 개념 자체 때문이 아니라 순차적 디코딩 비용 때문이었습니다. DiffRetriever 는 확산 모델을 통해 순차적 생성과 분리될 때 멀티토큰 검색이 효과적이고 저렴함을 입증합니다.
학습과 정렬된 목적 함수: 사전 학습된 형태로 확산 모델에 쿼리하는 것 (병렬로 마스킹된 위치 예측) 은 모델을 DiffEmbed 에서 수행된 것처럼 표준 BERT 스타일 인코더로 재사용하는 것보다 검색 작업으로 더 잘 전이됩니다.
적응형 잠재력: 고정 예산 성능과 쿼리별 오라클 성능 사이의 큰 격차는 적응형 예산 선택이 검색 효율성과 효과성을 개선할 중요한 미개척 분야임을 강조합니다.

저자들은 DiffRetriever 가 확산 기반 검색에 대한 새로운 최첨단 성능을 확립했다고 결론지으며, 병렬 대표 토큰이 단일 토큰 확산과 순차적 자동회귀 접근법 모두를 능가하면서도 낮은 지연 시간을 유지함을 입증했습니다.