12 분 소요

0. Introduction

Paper link

Code link

Model collection

한 줄 요약: DREAM은 labeled query-positive-negative pair의 contrastive loss 대신, candidate document가 frozen LLM의 target next-token prediction을 얼마나 돕는지를 relevance signal로 사용하며, retriever score를 selected query-focused attention head에 주입해 autoregressive loss의 gradient를 standalone dense retriever까지 전달한다.

이 논문을 지금 볼 가치가 있는 이유는 다음과 같음.

  • Dense retriever의 relevance를 semantic similarity가 아니라 downstream generation utility로 정의한다.
  • Frozen LLM의 next-token prediction loss를 retriever supervision으로 사용하는 differentiable interface를 만든다.
  • Explicit negative label 없이 candidate probability normalization만으로 document competition을 만든다.
  • 아무 attention head에 score를 넣는 것이 아니라 query-focused retrieval head를 사전 탐지해 routing한다.
  • 0.5B, 1B, 3B embedding backbone과 8B scaling check에서 method consistency를 확인한다.

Dense retrieval model은 보통 query $q$와 relevant document $d^+$를 가깝게, negative document $d^-$를 멀게 만든다. InfoNCE류 objective는 강력하지만 다음 비용이 있다.

  • Positive pair label이 필요하다.
  • Hard negative mining이 필요하다.
  • False negative가 training signal을 오염시킬 수 있다.
  • Relevance definition이 annotation policy에 고정된다.
  • Generator가 실제로 유용하게 쓰는 document와 retriever similarity가 다를 수 있다.

DREAM은 질문을 바꾼다.

어떤 document가 query와 비슷한가가 아니라, 어떤 document를 보여주면 LLM이 target output을 더 잘 예측하는가?

Relevant document를 LLM context에 넣으면 target passage의 next-token prediction loss가 낮아져야 한다. 이 loss를 retriever가 직접 받을 수 있다면, human relevance label 없이도 generation utility를 학습할 수 있다.

문제는 retriever와 LLM이 separate model이라는 점이다. Retriever가 candidate score를 내도 standard LLM attention과 연결되어 있지 않으면 next-token loss의 gradient가 retriever로 흐르지 않는다.

DREAM은 retriever score를 frozen LLM의 selected attention head 안에 주입한다. Candidate document별 probability가 query token에서 document token으로 가는 attention mass를 조절한다. 그러면 target cross-entropy loss가 retriever score에 미분 가능하게 연결된다.

이 method의 핵심은 LLM을 embedding model로 distill하는 것이 아니다. LLM은 training-time judge이자 differentiable relevance evaluator이고, inference에서는 standalone retriever만 사용한다.

1. Problem Setting

1-1. Problem definition

Trainable retriever를 $f_\phi$라고 하자. Query와 candidate document를 shared embedding space로 encode하고 L2-normalize한다.

\[e_q = \frac{f_\phi(q)}{\|f_\phi(q)\|_2}\] \[e_{d_j} = \frac{f_\phi(d_j)}{\|f_\phi(d_j)\|_2}\]

Similarity는 cosine score다.

\[s_\phi(q,d_j) = e_q^\top e_{d_j}\]

$K$개 candidate document에 대해 temperature-scaled distribution을 만든다.

\[p_\phi(d_j \mid q) = \frac{\exp(s_\phi(q,d_j)/\tau)} {\sum_{k=1}^{K}\exp(s_\phi(q,d_k)/\tau)}\]

Standard contrastive learning은 어떤 $d_j$가 positive인지 label을 사용한다. DREAM은 frozen LLM $g$가 target $y$를 예측하는 loss를 사용한다.

\[\mathcal{L}_{\mathrm{NTP}} = - \sum_{t=1}^{|y|} \log p_g(y_t \mid d_{1:K}, q, y_{<t}; p_\phi)\]

여기서 $p_\phi$가 selected attention head의 document-level attention mass에 들어간다. 따라서 다음 gradient가 생긴다.

\[\nabla_\phi \mathcal{L}_{\mathrm{NTP}} \neq 0\]

Frozen LLM parameter는 update하지 않고 retriever parameter $\phi$만 학습한다.

1-2. Why previous approaches are insufficient

1) Contrastive relevance label은 expensive and brittle하다

Dense retriever는 query-positive pair와 in-batch or mined negative에 크게 의존한다.

  • Human label은 비싸다.
  • Synthetic positive는 generator bias를 가진다.
  • BM25 or teacher-mined negative는 false negative를 포함할 수 있다.
  • Dataset마다 relevance granularity가 다르다.

특히 RAG에서는 answer generation에 유용한 document가 lexical or topical similarity만으로 정의되지 않는다.

2) Generator-aware retriever는 보통 black-box reward가 된다

Retriever output을 generator에 넣고 answer accuracy나 likelihood를 reward로 사용할 수 있다. 하지만 discrete top-k selection은 gradient를 끊는다.

Policy gradient or reinforcement learning을 사용할 수 있지만 variance와 compute cost가 크다. DREAM은 LLM attention에 score를 직접 주입해 differentiable path를 만든다.

3) All-head intervention은 supervision을 희석한다

Transformer의 모든 attention head가 retrieval을 담당하지 않는다. Syntax, position, copying, induction 등 다양한 head가 있다.

Retriever score를 random or irrelevant head에 넣으면 LLM prediction loss가 document relevance에 민감하게 반응하지 않는다. DREAM에서 head selection은 optional trick이 아니라 supervision channel의 핵심이다.

4) LLM fine-tuning은 moving target을 만들 수 있다

Judge LLM과 retriever를 함께 update하면 loss를 낮추는 두 경로가 생긴다.

  1. Retriever가 better document weighting을 학습한다.
  2. LLM이 current weighting에 적응한다.

두 번째가 너무 쉬우면 retriever supervision이 약해진다. DREAM은 judge를 freeze해 relevance criterion을 고정한다.

2. Core Idea

2-1. Main contribution

1) Relevance as autoregressive utility

DREAM의 relevance definition은 다음과 같다.

Query-conditioned target prediction을 개선하는 document가 relevant하다.

이는 query-document semantic similarity보다 generator-centric하다. Target $y$가 query answer, answer-supporting passage, downstream text일 수 있으며, relevant candidate에 더 높은 probability를 주면 frozen LLM loss가 낮아진다.

2) Query-focused retrieval head selection

Frozen LLM에서 query token이 candidate document token을 어떻게 attention하는지 분석해 retrieval-relevant head를 고른다.

Content-free query baseline도 사용한다. 예를 들어 meaningful query 대신 N/A를 넣었을 때의 document attention을 빼, generic positional or document-length bias를 줄인다.

Head별 document ranking을 target relevance와 비교하고 NDCG@10으로 rank한다. Top 16 head가 main setting이다.

3) Document-level score injection with token-level preservation

Retriever는 document별 probability $p_\phi(d_j q)$를 만든다. 그러나 document 내부에서 어떤 token이 중요한지는 frozen LLM이 이미 알고 있을 수 있다.

DREAM은 두 level을 분리한다.

  • Inter-document allocation: Retriever probability가 결정
  • Intra-document allocation: Original LLM attention pattern을 유지

즉 document $d_j$에 배정된 total attention mass는 retriever가 정하고, 그 mass를 document token 안에서 나누는 방식은 LLM attention을 따른다.

4) Learnable gate

Original attention과 DREAM-guided attention을 바로 완전히 교체하지 않는다. Learnable gate가 두 attention을 혼합한다.

개념적으로 selected head $h$의 output attention은 다음처럼 볼 수 있다.

\[\alpha_h' = (1-\gamma_h)\alpha_h + \gamma_h \tilde{\alpha}_h(p_\phi)\]

$\gamma_h$는 learned gate다. Frozen LLM의 native computation을 보존하면서 retriever signal을 넣는다.

5) Implicit candidate competition

Candidate probability는 합이 1이다.

\[\sum_{j=1}^{K} p_\phi(d_j|q) = 1\]

한 document의 mass를 높이면 다른 document mass는 낮아진다. Explicit negative label이 없어도 candidate pool 안에서 competition이 생긴다.

Target prediction에 도움이 되지 않는 document에 높은 mass를 주면 loss가 올라가므로 negative pressure를 받는다.

2-2. Design intuition

DREAM은 LLM attention을 training API처럼 사용한다.

Retriever와 generator는 원래 modular하다. Retriever는 score를 내고 generator는 selected text를 읽는다. 이 interface는 inference에는 편하지만 gradient가 끊긴다.

DREAM은 training time에만 internal attention을 열어 두 model을 연결한다.

  1. Retriever가 candidate distribution을 만든다.
  2. Distribution이 selected head의 document mass를 바꾼다.
  3. LLM이 target을 prediction한다.
  4. NTP loss가 document weighting quality를 평가한다.
  5. Gradient가 retriever로 돌아간다.
  6. Training 후 frozen judge는 제거한다.

이 design은 teacher score distillation과 다르다. Teacher가 explicit relevance label을 출력하지 않는다. Teacher의 generative loss가 indirect but task-grounded signal을 준다.

또 하나의 intuition은 frozen judge가 stable environment라는 점이다. Retriever가 학습하는 동안 utility landscape가 움직이지 않아 optimization이 더 명확하다.

3. Architecture / Method

3-1. Overview

Item Description
Goal Contrastive label 없이 autoregressive loss로 dense retriever 학습
Trainable module Dual-encoder embedding model with LoRA
Frozen judge Llama-3.1-8B-Instruct
Input order Candidate documents, query, target passage
Candidate set One Wikipedia document split into 16 chunks
Retriever score L2-normalized last-token embedding cosine
Routing target Top 16 query-focused retrieval heads
Attention injection Document mass replaced, intra-document pattern preserved
Objective Frozen judge target next-token cross-entropy
Inference Standalone retriever only

3-2. Module breakdown

1) Retriever encoder

DREAM은 causal LLM backbone을 embedding encoder로 사용한다.

  • Query and document를 separately encode한다.
  • Last non-padding token hidden state를 pooling한다.
  • L2 normalization한다.
  • Cosine similarity를 계산한다.

Public release는 PEFT LoRA adapter 형태다.

DREAM model Base backbone
DREAM-0.5B Qwen2.5-0.5B
DREAM-1B Llama-3.2-1B
DREAM-3B Llama-3.2-3B

2) Training sample layout

One Wikipedia document를 16 chunks로 나눈다. 그중 하나를 target passage $y$로 정하고, target에서 answerable한 query를 Qwen3-14B로 생성한다.

Candidate documents는 full 16 chunks다. Target information이 포함된 chunk와 같은 document의 distractor chunk가 함께 들어간다.

Frozen decoder input order는 causal constraint를 고려한다.

\[[d_1; d_2; \ldots; d_K; q; y]\]

Candidate와 query가 target보다 앞에 있어야 target next-token prediction에 condition될 수 있다.

3) QRHead detection

Query-focused retrieval head detection은 별도 stage다.

  1. Candidate documents와 query를 frozen LLM에 넣는다.
  2. Query token에서 each document로 가는 attention mass를 모은다.
  3. Content-free query attention을 subtract한다.
  4. Head별 document ranking을 만든다.
  5. Known target chunk ranking의 NDCG@10을 계산한다.
  6. Top heads를 training routing target으로 저장한다.

Public repository에는 qrhead_repo/가 이 역할을 한다.

4) DREAM routing

Selected head and query row에서 original attention matrix를 document block별로 분해한다.

Original document $d_j$ 내부 token distribution을 $a_{j,m}$이라고 하자.

\[\sum_{m \in d_j} a_{j,m} = 1\]
DREAM-guided attention은 document mass $p_\phi(d_j q)$와 intra-document distribution을 곱한다.
\[\tilde{\alpha}_{j,m} = p_\phi(d_j|q) a_{j,m}\]

이렇게 하면 retriever가 document selection을 담당하고 frozen LLM이 local token selection을 유지한다.

5) Gate and frozen judge

Selected head마다 guided attention을 original attention과 gate로 섞는다. Judge weight는 freeze한다.

Loss gradient path는 다음과 같다.

\[\mathcal{L}_{\mathrm{NTP}} \rightarrow \alpha_h' \rightarrow p_\phi(d_j|q) \rightarrow s_\phi(q,d_j) \rightarrow f_\phi\]

6) Standalone retrieval

Training 이후 inference에서는 frozen judge, candidate concatenation, attention injection이 필요 없다.

Query and corpus document를 DREAM encoder로 embed하고 ANN search를 수행한다. 따라서 training은 무겁지만 serving topology는 standard dense retriever와 같다.

4. Training / Data / Recipe

4-1. Synthetic training data

Data source는 Wikipedia and Natural Questions corpus다.

Construction은 다음과 같다.

  1. Wikipedia document를 선택한다.
  2. 16 chunks로 split한다.
  3. One chunk를 target으로 정한다.
  4. Qwen3-14B가 target에서 answerable query를 생성한다.
  5. 16 chunks 전체를 candidate pool로 사용한다.
  6. Target chunk text를 autoregressive prediction target으로 둔다.

이 setting은 human relevance label을 쓰지 않지만 completely supervision-free는 아니다.

  • Target chunk identity가 construction으로 알려져 있다.
  • Query는 large LLM이 synthetic하게 만든다.
  • Candidate가 same source document에 묶여 있다.

따라서 “label-free”보다 “explicit contrastive pair-free”라고 이해하는 편이 정확하다.

4-2. Judge and retriever setup

Component Setting
Judge LLM Llama-3.1-8B-Instruct
Judge update Frozen
Selected heads Top 16 QRHeads
Candidate count 16
Retriever adaptation LoRA
LoRA rank 32
LoRA alpha 64
Training steps 1,500
Batch size 1 sample
Gradient accumulation 32

Learning rate와 일부 optimizer detail은 final paper and config에서 확인해야 한다. 원문에 명시되지 않은 값은 추정하지 않는다.

4-3. Compared objectives

Baselines는 다음 범주를 포함한다.

  • BM25
  • InfoNCE trained on same constructed data
  • RePlug-style LLM supervision
  • Revela

InfoNCE baseline을 같은 candidate pool and data로 학습해, gain이 data advantage가 아니라 objective and routing에서 오는지 비교한다.

4-4. Engineering notes

1) QRHead detection을 먼저 실행해야 한다

Public repo는 QRHead detection result를 자동 포함하지 않는다. Judge model and training data로 head ranking JSON을 생성한 뒤 DREAM config를 실행한다.

2) Full artifact는 모두 공개되지 않았다

Repository README 기준으로 다음 large artifact는 포함되지 않는다.

  • Full training data
  • Full checkpoints outside released adapters
  • Processed datasets
  • Evaluation outputs
  • QRHead detection results

Reproduction에는 data construction과 head detection을 다시 수행해야 한다.

3) Adapter serving

Released model은 LoRA adapter다. Base causal model에 PEFT adapter를 load하고 final hidden state를 last-token pooling해 embedding을 얻는다.

4) Judge memory cost

Training에서는 8B judge가 16 documents, query, target을 함께 처리한다. Retriever 0.5B or 1B를 학습하더라도 judge context가 memory and compute bottleneck이 될 수 있다.

5. Evaluation

5-1. Main results

논문은 BEIR 9 tasks and RTEB 14 tasks에서 NDCG@10을 평가한다.

1) Main scale sweep

Retriever scale BEIR average RTEB average
DREAM-0.5B 0.4163 0.4788
DREAM-1B 0.4888 0.5514
DREAM-3B 0.5074 0.5892

Scale이 커질수록 두 benchmark average가 올라간다.

2) Revela 대비 improvement

Revela result는 다음과 같다.

Scale BEIR Revela RTEB Revela
0.5B 0.4011 0.4107
1B 0.4075 0.4499
3B 0.4315 0.4945

DREAM의 gain은 BEIR에서 약 0.015-0.081, RTEB에서 약 0.068-0.102다.

RTEB에서 gain이 더 크다는 점은 autoregressive utility signal이 table/entity or structured retrieval에서 특히 도움이 될 가능성을 보여준다. 다만 task별 variance를 함께 봐야 한다.

3) Head selection이 결정적이다

Llama-3.2-1B backbone ablation은 다음과 같다.

Head routing BEIR RTEB
Selected query-focused heads 0.4888 0.5514
Fully random heads 0.0637 0.0320

Random head injection은 사실상 training signal을 파괴한다. Frozen LLM loss를 사용한다는 것만으로는 충분하지 않고, loss가 document relevance에 민감한 circuit에 연결되어야 한다.

4) Top 16 heads가 best trade-off다

Head 수를 늘리면 initially signal coverage가 좋아진다. Top 16에서 peak를 보이고, 더 약한 head까지 포함하면 noise가 늘어 performance가 떨어진다.

5) Frozen judge가 LoRA-updated judge보다 낫다

Judge를 LoRA로 함께 update하면 NTP loss는 더 쉽게 낮아질 수 있지만 retriever supervision은 약해진다. Frozen judge가 stronger retrieval result를 낸다.

이는 teacher stability가 중요한 evidence다.

6) Candidate count가 늘면 performance가 좋아진다

Candidate 4에서 16으로 늘릴수록 document competition이 강해지고 retrieval result가 개선된다.

그러나 candidate를 무한히 늘리면 judge context and compute가 커진다. Paper의 tested range 밖은 추가 확인이 필요하다.

7) 8B scaling check

DREAM-8B는 다음 average를 보고한다.

Model BEIR RTEB
DREAM-8B 0.5531 0.6417
E5-Mistral 0.5526 0.6031
Qwen3-Embedding-8B 0.6348 0.7383

DREAM-8B는 E5-Mistral과 competitive하지만 Qwen3-Embedding-8B보다 낮다. 논문도 absolute SOTA claim보다 autoregressive supervision의 scalability를 강조한다.

5-2. What really matters in the experiments

1) Same-data InfoNCE comparison

DREAM이 better data를 사용해서 이기는지 확인하려면 same constructed query and candidate pool의 InfoNCE baseline이 중요하다.

Method gain은 next-token utility signal and attention routing에서 나와야 한다.

2) Random-head collapse가 mechanism을 지지한다

Random head result가 매우 낮다는 것은 QRHead selection이 paper의 load-bearing component임을 보여준다.

동시에 practical fragility도 의미한다. New judge model마다 retrieval head detection을 다시 해야 할 수 있다.

3) NTP objective가 embedding geometry를 만든다

논문은 DREAM embedding이 strong alignment를 유지하면서 RePlug and Revela보다 better uniformity를 보인다고 분석한다.

Explicit contrastive repulsion 없이도 normalized candidate competition이 embedding collapse를 줄일 수 있다는 해석이다.

4) Benchmark average만으로 deployment readiness를 말하기 어렵다

BEIR and RTEB average는 broad transfer를 보여주지만 다음은 별도다.

  • Multilingual retrieval
  • Long-document retrieval
  • Domain-specific legal or medical retrieval
  • ANN index latency
  • Quantization robustness
  • Query/document instruction format
  • Reranking interaction

6. Limitations

  1. Completely label-free training은 아니다.
    • Synthetic query와 known target chunk가 필요하다.
    • Human relevance pair를 없앴지만 construction supervision은 남는다.
  2. Candidate pool이 same Wikipedia document의 16 chunks다.
    • Corpus-scale hard negative와 domain-shift negative를 충분히 반영하지 않을 수 있다.
  3. Frozen judge dependence가 크다.
    • Judge가 어떤 document를 useful하다고 보는지가 retriever relevance notion이 된다.
    • Judge bias or knowledge gap이 retriever에 전달될 수 있다.
  4. White-box attention access가 필요하다.
    • Closed API model에서는 QRHead detection and score injection을 수행하기 어렵다.
  5. QRHead selection이 brittle할 수 있다.
    • Model architecture, layer count, tokenizer, context format이 바뀌면 detection을 다시 해야 한다.
  6. Training compute가 크다.
    • Small retriever라도 8B judge and 16-document context를 매 step forward한다.
  7. Released artifact가 partial하다.
    • Full data, processed outputs, QRHead detection result가 없어 exact reproduction cost가 있다.
  8. Synthetic query distribution이 제한적이다.
    • Qwen3-14B query style이 downstream benchmark query와 다를 수 있다.
  9. Evaluation language and domain coverage가 제한된다.
    • Multilingual and enterprise proprietary corpus transfer는 별도 확인이 필요하다.
  10. Generator utility와 retrieval faithfulness가 항상 일치하지 않는다.
    • LLM이 parametric knowledge로 target을 예측하거나 plausible continuation을 만들면 truly evidential document보다 stylistically helpful document를 선호할 수 있다.

7. My Take

7-1. Why this matters for my work

DREAM의 가장 흥미로운 점은 retriever supervision을 “semantic pair label”에서 “downstream model이 실제로 쓸 수 있는 information”으로 옮긴 데 있다.

RAG system에서 embedding similarity가 높아도 generator가 answer를 만들 때 필요한 fact, relation, table row를 못 찾을 수 있다. 반대로 lexical overlap이 낮아도 target reasoning을 크게 돕는 document가 있다.

DREAM은 generator utility를 embedding training에 직접 연결하는 방법을 보여준다. 특히 document AI and evidence retrieval에서는 다음 extension을 생각할 수 있다.

  • Target answer generation loss
  • Evidence quote reconstruction loss
  • OCR-grounded span prediction loss
  • Citation verification loss
  • Structured JSON field completion loss

Retriever가 downstream task loss를 줄이는 document를 학습하도록 할 수 있다.

7-2. Reuse potential

1) Task-specific retriever pretraining

Enterprise QA에서 target을 answer text가 아니라 validated evidence span으로 둔다.

  • Candidate: Document chunks
  • Query: User question
  • Target: Gold evidence sentence or structured field
  • Judge: Frozen domain LLM

이렇게 하면 relevance가 answer utility and evidence support에 더 가까워진다.

2) Reranker training

Standalone dual encoder뿐 아니라 lightweight reranker score를 attention mass에 넣을 수 있다. Candidate set이 작을 때 cross-encoder alternative로 연구할 수 있다.

3) Multi-target supervision

Target을 하나만 쓰지 않고 다음을 조합한다.

  • Answer
  • Rationale
  • Citation quote
  • Structured output

각 target loss가 다른 relevance dimension을 줄 수 있다.

4) Frozen multi-judge ensemble

One frozen judge의 bias를 줄이기 위해 different LLM family or task head의 loss를 합칠 수 있다.

\[\mathcal{L} = \sum_m w_m \mathcal{L}_{\mathrm{NTP}}^{(m)}\]

다만 compute가 크게 늘어난다.

5) Query-focused circuit audit

DREAM은 head selection이 decisive함을 보여준다. New architecture에서 retrieval circuit을 먼저 찾고, 그 circuit을 intervention point로 사용하는 general recipe가 가능하다.

6) Candidate curriculum

Training early stage에는 same-document chunks를 사용하고, later stage에는 corpus-level hard negative와 cross-domain distractor를 넣을 수 있다.

Implicit competition이 false negative에 얼마나 robust한지도 확인해야 한다.

7-3. Follow-up papers

  • RePlug: Retrieval-Augmented Black-Box Language Models
  • Revela and LLM-supervised retriever training
  • Query-Focused Retrieval Heads Improve Long-Context Reasoning
  • E5 and E5-Mistral
  • Qwen3 Embedding Technical Report
  • Dense Passage Retrieval
  • Unsupervised and generative retrieval objectives
  • RAG end-to-end retriever-generator training

8. Summary

  • DREAM은 dense retriever를 contrastive pair label 대신 frozen LLM next-token loss로 학습한다.
  • Retriever score를 document probability로 바꾸고 selected query-focused attention head에 주입한다.
  • Document-level mass는 retriever가, document 내부 token pattern은 frozen LLM이 담당한다.
  • Candidate probability normalization은 explicit negative 없이 implicit competition을 만든다.
  • QRHead selection, frozen judge, Top 16 heads, 16 candidates가 주요 load-bearing design이다.
  • Inference에서는 standard standalone embedding model로 사용할 수 있지만 training compute와 judge dependence가 크다.

댓글남기기

참고