DREAM: Autoregressive Modeling을 이용한 Dense Retrieval 임베딩

Dense retriever 학습은 보통 긍정/부정 문서 쌍 라벨이 필요해 데이터 확보 비용이 높다. DREAM은 frozen LLM의 next-token prediction 손실을 활용해 라벨 없이도 retriever를 학습시켜 라벨 의존도를 낮춘다. 이 방식은 BEIR·RTEB 실험에서 기존 LLM-감독 방법들보다 일관된 성능 향상을 보였다.

문제 출발점과 기존 한계: Dense retriever는 일반적으로 contrastive objective를 사용해 쿼리-긍정 문서를 가깝게, 부정 문서를 멀게 학습한다. 이 방식은 고품질 긍정 라벨과 하드 네거티브 샘플링에 크게 의존하며, 라벨 수집과 네거티브 마이닝이 비용·복잡도의 병목이 된다.

해결 원리: Autoregressive LLM의 next-token prediction(NTP) 손실은 문서가 쿼리에 대해 얼마나 '유용한지'를 정량화할 수 있는 신호다. 입력으로 candidate documents, query, target을 연결한 뒤 retriever가 각 문서에 할당한 유사도 점수들을 특정 query-focused attention head에 주입하면, LLM의 NTP 손실 변화가 어느 문서가 예측을 돕는지에 대한 기울기를 생성한다. 이 기울기는 retriever의 파라미터로 역전파되어 embedding 모델을 직접 갱신한다.

구체적 기작: retriever는 query와 각 문서의 L2 정규화된 표현을 계산하고 내적으로 유사도 sφ(q,dj)를 얻는다. 이들 점수를 온도 τ로 나눠 softmax로 정규화하면 pφ(dj|q)가 된다. 선택된 attention head의 query-token 행에 대해 각 문서 내부 토큰에 대한 기존 attention 분포를 문서 내 분포로 보존한 채, 문서별 총 질량을 pφ(dj|q)로 재분배한다. 최종적으로 혼합 게이트 g를 통해 원래 attention과 score-guided attention을 선형 혼합한다. 이렇게 하면 LLM의 NTP 손실 변화가 retriever 점수에 민감해져 retriever가 유용한 문서를 더 높이 평가하도록 학습된다.

차이와 이점: 이 접근은 labeled positive/negative 쌍을 직접 요구하지 않으며, 문서 가중치의 정규화가 자연스러운 경쟁을 만들어 음성 예시(negatives) 없이도 불필요한 문서를 밀어내는 효과를 제공한다. 실험에서 Top 16 query-focused heads 주입, 후보 문서 K=16 사용, judge LLM 고정(frozen)이 가장 안정적인 신호를 제공함이 관찰되었다.

전체 접근과 핵심 아이디어: 각 학습 예시는 query q, 후보 문서들 d1…dK, 목표 패시지 y로 구성된다. retriever fφ가 각 쿼리-문서 쌍에 대해 L2 정규화된 임베딩 e_q, e_dj를 출력하고, sφ(q,dj)=e_q^T e_dj로 유사도를 계산한다. 이 유사도들을 온도 τ로 나누어 softmax로 정규화하면 pφ(dj|q)가 된다. pφ는 training시에만 frozen LLM 내부의 일부 attention head에 주입되어 LLM의 next-token prediction 손실이 retriever로 역전파된다.

핵심 메커니즘 상세: (L2 정규화) 임베딩 fφ(x)를 입력으로 → L2 노름으로 나눔 → 단위 벡터 e_x를 얻음 → 의미: 임베딩 길이 차이를 제거해 코사인 유사도(내적)로 비교 가능하게 함. (유사도 정규화) sφ(q,dj)를 입력으로 → exp(s/τ) 연산 후 정규화 → 출력 pφ(dj|q): 후보 문서들에 대한 확률 분포가 생성됨 → 의미: 문서들 간 경쟁을 만들어 한 문서 가중치의 증가는 다른 문서 가중치 감소를 유발함. (attention 재분배) 선택된 head의 query-token 행에서 각 문서 내부 토큰에 대한 기존 softmax attention α_h(a,b)를 문서별로 정규화하여 내부 분포 hat{α}_{h,j}(a,b)를 구함 → 이 내부 분포에 pφ(dj|q)를 곱해 문서 간 총 질량을 재할당함 → 의미: retriever가 문서 선택을 제어하고, LLM은 문서 내 토큰 우선순위를 보존함. (혼합) 최종 attention은 (1-g)·원래 + g·재분배 attention으로 혼합되며 g=σ(γ)로 학습됨.

학습 전략과 구현 세부사항: LLM은 학습 중 고정(frozen)시키는 것이 기본 설정이다. 논문 구현에서는 judge LLM으로 Llama-3.1-8B-Instruct를 사용하고, embedding backbones으로 Qwen2.5-0.5B, Llama-3.2-1B, Llama-3.2-3B, 8B 스케일 실험도 포함했다. retriever 쪽 파라미터는 embedding 모델의 q/k/v/o projection에 대해 LoRA(rank=32, alpha=64)를 적용해 적은 추가 파라미터로 학습했다. 학습 하이퍼파라미터는 learning rate=1e-4, gradient accumulation=32, batch size=1, 총 1500 steps, 후보 문서 K=16이 기본이다.

head 선별 절차: 학습 전 probe 집합(쿼리-후보-타깃 포맷)을 사용해 각 head의 query→문서 attention이 타깃을 지지하는지 측정한다. 각 head에 대해 query-토큰에서 각 후보 문서로 흘러가는 평균 attention을 계산하고, null query 기준값을 빼서 보정한 후 NDCG@10으로 정렬한다. 이 순위에서 상위 M(heads)을 선택해 H로 사용한다. 논문 실험에서는 M=16이 기본으로 채택되었다.

메인 벤치마크 결과: Table 1에서 BEIR과 RTEB의 평균 NDCG@10을 보고한다. 동일한 embedding backbone 기준으로 DREAM이 RePlug·Revela·InfoNCE를 일관되게 상회했다. 예: Llama-3.2-1B backbone에서 BEIR: Revela 0.4075 vs DREAM 0.4888, RTEB: Revela 0.4499 vs DREAM 0.5514. BEIR에서 DREAM의 Revela 대비 이득 범위는 0.015–0.081, RTEB에서는 0.068–0.102로 보고되었다.

ablation 결과: Head selection 영향: 선택된 query-focused heads를 사용한 경우 평균 BEIR NDCG@10은 0.4888, RTEB는 0.5514였으나 완전 무작위 heads는 각각 0.0637, 0.0320에 머물렀다. 선택된 헤드 수는 Top 1→Top 16까지 성능이 개선되고 Top 32/64에서 악화되어 Top 16이 경험적으로 최적이었다. Judge LLM 업데이트 영향: judge LLM을 frozen으로 유지하는 것이 LoRA로 judge도 업데이트하는 것보다 모든 비교에서 더 좋았다.

효과의 공간적 해석: 5,000 쿼리-긍정 쌍에서 alignment(쿼리-긍정 거리 제곱 평균)와 uniformity(임베딩 분포의 펼쳐짐)를 측정했다. RePlug는 좋은 alignment를 보이나 uniformity가 낮았고, Revela는 uniformity가 개선되었지만 alignment가 약했다. DREAM은 alignment를 Revela 수준으로 유지하면서 uniformity를 최상위로 만들었다.

스케일 점검: DREAM-8B(학습된 Llama-3.1-8B 기반)은 BEIR 0.5531, RTEB 0.6417을 기록했다. E5-mistral-7b-instruct와 유사한 BEIR 성능을 보였고, RTEB에서는 0.0386 포인트 우위였다. Qwen3-Embedding-8B에는 미치지 못했으나 비교 대상의 데이터·백본 차이를 명시적으로 주의해야 한다.

전체 아키텍처: 학습 중 입력은 candidate passages d1…dK, Question q, Target passage y를 문서-우선 순서로 연결한 단일 decoder-only LLM 시퀀스다. Retriever fφ는 query와 각 후보의 마지막 토큰 표현을 L2 정규화해 e_q, e_dj를 만들고, sφ(q,dj)=e_q^T e_dj를 계산한다. 이 sφ들을 온도 τ로 나누어 softmax 정규화하면 pφ(dj|q)가 되고, 선택된 query-token 행의 attention에 주입되어 LLM의 NTP 손실이 retriever로 역전파된다.

수학적/알고리즘적 기반: (유사도→가중치) 입력 sφ(q,dj) → exp(s/τ) 계산 → Σ로 나눠 pφ(dj|q) 산출 → 의미: 문서 간 경쟁적 확률 분포. (문서 내부 attention 보존) 원래 head의 attention α_h(a,b)를 문서별로 정규화해 hat{α}{h,j}(a,b)를 구함 → 각 문서 토큰들에 대해 합이 1이 되는 내부 분포 확보. (문서 간 재배분) α^R_h(a,b)=pφ(dj|q)·hat{α}{h,j}(a,b)로 문서별 총 질량 할당. (혼합 게이트) α'_h=(1-g)·α_h + g·α^R_h, g=σ(γ)로 학습되어 retriever 영향도 조절.

Prior work 대비 차별점: RePlug는 frozen-LLM likelihood를 distill하는 방식으로 retriever를 학습하지만 retriever 점수가 LLM의 forward pass에 직접 영향을 주지 않아 손실이 retriever에 직접 기울기를 주지 않는다. Revela는 LM 기반 시퀀스 언어 모델링 신호로 학습하나 LLM과 retriever를 공동 학습하고 query-target 명시적 설정을 사용하지 않는다. DREAM은 frozen LLM 내부의 query-focused heads를 직접 조작해 next-token loss가 retriever 파라미터로 직접 전달되도록 설계했다.

구현 및 학습 세부사항: judge LLM으로 Llama-3.1-8B-Instruct 사용(기본 frozen). embedding backbones: Qwen2.5-0.5B, Llama-3.2-1B, Llama-3.2-3B, 확장 검사로 Llama-3.1-8B 기반 DREAM-8B. retriever의 embedding projection 모듈들에 LoRA(rank=32, alpha=64)를 적용. 학습 하이퍼파라미터: lr=1e-4, grad accumulation=32, batch size=1, steps=1500, 후보 K=16. Head 선택은 Zhang et al. (2025a) 방법으로 probe 세트에서 NDCG@10 기준으로 상위 M(기본 16)을 취함.

논문이 명시한 한계: (1) 방법은 query-focused retrieval heads가 존재하는 frozen LLM에 의존한다; 모든 LLM에서 동일한 헤드 구조를 보장하지 못할 수 있다. (2) Judge LLM을 업데이트하면 손실 하강이 LLM 변화에 기인할 수 있어 retriever에 대한 기울기가 약해지므로 judge를 frozen으로 유지하는 설정이 더 안정적이다. (3) DREAM-8B가 경쟁 모델(Qwen3-Embedding-8B)을 넘어서지 못했으며, 모델·데이터 차이로 직접 비교 해석에 주의가 필요하다고 논문이 명시했다.