Papers·어제
DREAM: LLM의 next-token prediction 손실로 dense retriever 학습 — BEIR에서 3B 기준 2.3점↑
DREAM은 LLM의 autoregressive next-token prediction 손실을 이용해 dense retriever를 학습하는 프레임워크입니다. retriever가 생성한 query-document 유사도를 frozen LLM의 attention head에 주입해, LLM이 target을 예측할 때 각 문서가 받는 attention을 조절하고, 그 예측 손실이 retriever를 학습시킵니다. BEIR과 RTEB 벤치마크에서 0.5B~3B 규모의 embedding backbone에 대해 기존 대비 일관된 성능 향상을 보였습니다. 단, 학습 시 LLM이 고정되어 있어 retriever만 업데이트되므로, LLM 자체를 fine-tuning하는 방식과의 비교는 추가 연구가 필요합니다.
DREAM은 별도의 contrastive pair 없이 LLM의 next-token prediction 손실로 dense retriever를 학습하는 방법입니다.
핵심 결론
- 태스크 — dense retrieval embedding 학습 — BEIR 및 RTEB 벤치마크 평가.
- 성능 — 0.5B~3B backbone에서 기존 contrastive baseline 대비 BEIR 평균 1.5~2.3점 향상.
- 저자 — Yixuan Tang 단독 — 소속은 공개되지 않았습니다.
방법
- 핵심 아이디어 — retriever가 계산한 query-document 유사도를 frozen LLM의 일부 attention head에 scalar bias로 주입합니다.
- 학습 루프 — LLM이 target token을 예측할 때, 각 문서의 유사도가 attention logit에 더해져 문서가 받는 attend 양을 결정합니다.
- 그래디언트 — next-token prediction 손실이 attention head를 통해 retriever로 역전파되어 embedding을 업데이트합니다.
- LLM — frozen 상태로 유지되며, retriever만 학습됩니다.
한계·조건
- 비교 대상 — contrastive baseline(SimCSE, Contriever)과 비교했으나, LLM 자체를 fine-tuning하는 방식(예: REPLUG)과의 직접 비교는 없습니다.
- 코드 — GitHub 공개 여부는 논문에 명시되지 않았습니다.
- 계산량 — frozen LLM의 forward pass가 필요하므로, contrastive 방식보다 학습 비용이 높을 수 있습니다.
편집자 한 줄
contrastive pair 수집이 어려운 도메인에서 특히 쓸모 있어 보이네요. 다만 LLM이 고정되어 있어 retriever가 LLM의 사전 지식을 넘어서기 어려울 수 있다는 점은 후속 연구가 필요합니다.
- #dense-retrieval
- #llm
- #dream
- #beir
- #contrastive-learning
Yixuan Tang