[논문리뷰]DRAMA: Diverse Augmentation from Large Language Models to Smaller Dense Retrievers

Xueguang Ma, Xi Victoria Lin, Barlas Oguz, Jimmy Lin, Wen-tau Yih, and Xilun Chen. 2025. DRAMA: Diverse Augmentation from Large Language Models to Smaller Dense Retrievers. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Wanxi- ang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar(Eds.). Association for Computational Linguistics, Vienna, Austria, 30170–30186.https://doi.org/10.18653/v1/2025.acl-long.1457

Problem Statement

[Dense Retrieval의 한계] 제한된 데이터로 fine-tuning하기 때문에, multilingual task 등에서의 일반화 능력이 떨어진다.

[LLM-based Retrieval의 한계] LLM을 retriever의 용도로 fine-tuning한 모델은 그 효과와 robustness에 대해서는 입증되었지만, computational cost가 매우 크다.

이를 해결하기 위해, 이 논문에서는 i) retriever를 학습하기 위한 데이터를 LLM을 통해 augment하는 방법과, ii) LLM을 pruning하여 retriever의 backbone으로 활용하는 방법을 제시한다.

Methodology

Overview

Data Augmentation

Data Augmentation 단계에서는 쿼리(query, \(q\)), 긍정 문서(positive document, \(D^+\)), 음성 문서(negative document, \(D_{HN}\))로 구성된 triplet을 생성하여 retriever를 학습하는데 초점을 맞춘다. 이 때, 음성 문서는 긍정 문서와 유사하지만 정답이 아닌 distractor document (=hard negative document)이다.

1. Random Sentence Cropping (Cropped Sentence Query)

첫 번째 증강 방식은 Random Sentence Cropping이다. 각 문서에 대해 랜덤하게 하나의 문장을 뽑고, 이를 pseudo-query로 사용하여 사전에 retrieval로 학습된 Llama에 입력으로 주어 top-\(k\)개의 문서를 뽑는다. 이 중 top-\(1\)부터 \(m\)을 positive로, 그 하위를 distractor로 지정한다. 증강 방식 중 가장 cost가 작다.

2. Synthetic Query Generation

Synthetic query generation은 LLM의 능력을 더 적극적으로 활용하는 증강 방법이다. 여기에서는 각 문서를 Llama-70B 모델에 입력으로 주고, 각각에 대해 문서를 검색하는데 도움이 될만한 쿼리를 생성하도록 합니다. 이후 random sentence cropping과 마찬가지로 이 query를 입력으로하여 top-\(k\)개의 문서를 검색하고, 이를 positive와 distractor로 split한다. 증강 방식 중 cost는 중간 정도이다.

3. LLM Listwise Reranking

다음은 list-wise reranking이다. 이 단계는 앞에서 수행한 random sentence cropping과 synthetic query generation에서 얻은 합성 쿼리와 각 쿼리에 대한 top-\(k\)개의 문서를 70B 모델에 입력으로 주고, LLM이 query와의 연관성을 직접 판단하여 reranking을 수행하게 한다. 리랭킹의 결과로 정렬된 문서들에서 top-\(1\)만을 positive로 사용하고, 나머지 문서를 distractor로 사용한다. 이 방식은 LLM의 리랭킹을 요구하므로 증강 방식 중 가장 큰 cost를 요구하지만, 이 단계에서 얻은 데이터 쌍들을 활용하여 8b retriever 모델을 retrain한 것이 성능향상에 큰 도움이 되었다고 한다.

4. Synthetic Triplet Generation

가장 고비용의 증강 방식으로 LLM이 직접 query–positive–negative triplet을 생성. 하지만 실험에서는 성능 개선 효과가 거의 없었다.

Retreival Training

$$L(q, D^+, \{D_N\}) = -\log \frac{\exp(\text{Sim}(q, D^+)/\tau)}{\sum_{D_i \in \{D^+\} \cup \{D_N\}}\exp(\text{Sim}(q, D_i)/\tau)}$$

Retrieval 학습에는 위와 같이 InfoNCE loss가 사용된다. \(\{D_N\}\)은 in-batch 내에 존재하는 distractor의 집합이다.

Pruning

DRAMA는 기존의 encoder-only 구조 대신, decoder-only LLM을 prunining하여 retrieval backbone으로 사용한다. 즉, 파라미터 수를 전통적인 dense retriever의 수준으로 맞춰주기 위해 진행하는 단계이다. Llama 3.1의 pruning 버전인 3.2로 초기화를 하고, layer, hidden dimension, attention head, intermediate dimension에 대해 각각 parameter mask를 학습하는 방식으로 진행된다.

$$\tilde{\mathcal{L}}^{\text{head}}(\lambda, \phi, z) = \lambda^{\text{head}} \cdot \big(\displaystyle\sum_{i} z_i^{\text{head}} - H_{\mathcal T} \big) + \phi^{\text{head}} \cdot \big( \displaystyle\sum_{i} z_i^{\text{head}} - H_{\mathcal T} \big)^2$$

• \(z_i^{\text{head}}\): \(i\)번째 Attention Head가 활성화되었는지 여부를 나타내는 마스크 변수 (0 또는 1에 가까운 확률적 값).
• \(H_{\mathcal T}\): 목표 Attention Head 수 (Target head number).
• \(\lambda^{\text{head}}\): 1차 항에 대한 라그랑주 승수(Lagrange multiplier).
• \(\phi^{\text{head}}\): 2차 항에 대한 라그랑주 승수.이때, lagrange multiplier를 활용한다. 첫 번째 항 (first-order)은 실제 남겨진 Head 수와 목표값의 차이(sign 포함)를 반영. 선형 패널티이므로, 방향성(부족/과잉)을 조정하는 역할을 한다. 두 번째 항(second-order)은 차이의 제곱을 사용하여 편차 크기 자체를 강하게 벌주는 역할. 이는 제약 위반 정도를 제곱 패널티로 보정한다.

즉, 이 두 항의 조합은 제약 최적화(Constrained Optimization)에서 흔히 쓰이는 선형 + 제곱 벌점(augmented Lagrangian) 형태이다. 단순히 제약을 강제하는 것이 아니라, 모델 학습 과정에서 목표 Head 수에 맞춰가도록 유도한다.

$$\mathcal{L}_{\text{prune}}(\theta, z, \lambda, \phi) = \mathcal{L}(\theta, z) + \sum_{j=1}^{L_S} \tilde{\mathcal{L}}^{\text{head}}_j+ \sum_{j=1}^{L_S} \tilde{\mathcal{L}}^{\text{int}}_j+ \tilde{\mathcal{L}}^{\text{layer}}+ \tilde{\mathcal{L}}^{\text{hidden}}$$

최종적으로 lagrange multiplier기반의 loss와 mask별 loss를 더해 joint loss로 학습을 진행한다. 학습 이후에는 프루닝 과정에서 생기는 성능 저하를 회복하기 위해 통상적인 LLM 학습 object인 next-token prediction loss를 활용한다. 이 과정을 통해 0.1B, 0.3B의 backbone 모델을 얻어내고, 원본 데이터셋에 앞서 증강을 통해 얻은 data를 합쳐 infonce loss를 통한 retriever학습을 진행하게 된다.

Experiments

Main Results

DRAMA 모델은 다양한 크기에서 강력하고 일반화 가능한 성능을 보인다. Table 1은 DRAMA가 기존의 인코더 기반 검색기를 능가함을 보여준다. DRAMA-0.1B는 BEIR 벤치마크에서 56.9 nDCG@10을 달성하여 ArcticEmb-v2-M과 동등한 성능을 보이며, 다른 소형 모델들을 능가한다. DRAMA-0.3B는 BEIR에서 58.0 nDCG@10을 달성하여, ArcticEmb-v2-L보다 0.8점 높은 성능을 보이며, 1B 파라미터 모델인 Gecko와 동등한 성능을 보인다. DRAMA 모델들은 MIRACL과 같은 다국어 벤치마크에서도 기존의 최고 성능 모델인 M3-BGE-Dense를 능가하며, 영어 성능도 함께 유지한다.

Abalation

• FT: Original Data
• Sent: Random Sentence Cropping으로만 합성 데이터 구축
• QGen: Synthetic Query Generation으로만 합성 데이터 구축
• Rerank: LLM Rerank로만 합성 데이터 구축
• Triplet: LLM이 직접 생성한 triplet으로 학습

이 표에서 asterisk(*)가 붙은 옵션은 reranking의 output으로 retriever를 재학습하는 과정을 뺀 것이다. 세가지 augmentation 방법 중 한가지만 추가하여 실험 했을 때는 random cropping을 통한 augmentation을 활용했을 때가 가장 좋았고, 한가지만 제거하여 실험했을 때에는 reranking을 제거하는 것이 성능하락이 가장 컸다.

Effectiveness 1: Multilingual Long-context Retrieval Task

이 실험에서는 pruning된 LLM을 retriever의 backbone으로 활용하였을때의 장점을 설명한다. 왼쪽 표는 multilingual long-context retrieval 실험에 대한 성능 표이다. 이 실험에서 저자들은, 대부분의 0.1B 모델은 multilingual을 지원하지 않지만, DRAMA는 LLM을 pruning 한 backbone을 사용함으로써, 영어 뿐 아니라 multilingual task에 대해서도 준수한 성능을 보였다고 한다. 또한, long-context data로 학습이 되지 않은 상황에서도 long-context retrieval에 대해 우수한 성능을 보임을 알 수 있다.

Effectiveness 2: Using Pruned Llama3.2

이 실험도 pruning 된 LLM을 활용한 것의 장점을 보여주는 실험이다. 적은 parameter 수를 갖는 기존의 retriever보다 우수한 성능을 보임을 확인할 수 있다. 특히 저자들은 Modern Bert와 비교하며, 이 방식은 LLM을 활용한 것이기 때문에 multilingual을 지원하고, pre-train된 LLM을 활용하기 때문에, backbone부터 scratch로 학습할 필요가 없다는 장점이 있다고 주장한다.

Analysis: Cross-lingual Generalization

이 실험은 retriever 학습 시 영어 데이터로만 학습을 진행한 상태로 수행한 실험이다. 이때, de와 yo는 pruning에서 language model로써의 성능을 복원하는 과정의 데이터셋에 포함된 언어이고, pl은 포함되지 않은 언어이다. 대체적으로 parameter 수가 많을 수록 성능이 좋음을 알 수 있는데, yo의 1B에서 성능이 낮게 나오는것에 대해 저자들은,기존 1B의 모델이 yo에 대해 학습이 잘 되지 않았을 것이라고 추측하였다. 또, pl의 1B와 0.3B의 성능차이가 두드러짐을 확인할 수 있는데, 저자들은 이 결과를 통해 pruning단계에서 해당 언어에 노출이 되는지 여부가 매우 중요함을 알 수 있다.

Conclusion

Contribution

• LLM을 활용한 다양한 데이터 증강 방법을 제안하였다
• 프루닝된 decoder-only LLM을 retrieval backbone으로 사용함으로써 다양한 검새에서 robust하고, 장문 검색에서 더 나은 성능을 보였다.

Limitations

• 데이터 증강 과정에서 추가적인 연산 비용이 발생한다
• 일반화를 위해 프루닝 단계에서 다중 언어 데이터를 충분히 포함시켜야 한다.

NR 카테고리 내 다른 글 보러가기