[논문리뷰]Descriptive and Discriminative Document Identifiers for Generative Retrieval(AAAI, 2025)

Jiehan Cheng, Zhicheng Dou, Yutao Zhu, and Xiaoxi Li. 2025. Descriptive and discriminative document identifiers for generative retrieval. In Proceedings of the Thirty-Ninth AAAI Conference on Artificial Intelligence and Thirty-Seventh Conference on Innovative Applications of Artificial Intelligence and Fifteenth Symposium on Educational Advances in Artificial Intelligence (AAAI’25/IAAI’25/EAAI’25). AAAI Press, Article 1280, 9 pages. https://doi.org/10.1609/aaai.v39i11.33253

1. Problem Statement

이 논문은 Generative Document Retrieval (생성 기반 문서 검색) 문제를 다룬다. 즉, 주어진 질문에 대해 문서 식별자(Document Identifier, DocID)를 직접 생성함으로써 관련 문서를 검색하는 end-to-end generative trieval 방식을 제안한다. 기존 생성 기반 검색(Generative Retrieval) 연구들은 주로 DocID의 문서 표현력(”Decriptiveness”)에만 집중하였고, 문서 구별력(”Discriminativeness”)는 고려하지 않았다. 구체적으로 기존 연구들은 다음과 같은 한계점을 보인다.

• [DocID의 의미적 유사성 부족] 기존 Generative Retrieval 기법은 주로 문서의 텍스트 표현(예: 제목, URL, synthetic query 등)에만 기반하여 DocID를 생성하므로, 의미적으로 유사한 문서들이 표현 차이(동의어, 서술 구조 등) 때문에 서로 다른 DocID를 가지게 됨. 이로 인해 동일 질의에 대해 유사 문서가 검색되지 않는 문제가 발생할 수 있다.
• [Corpus 인식 부족한 DocID 설계] 기존 연구들은 각 문서 단위로만 DocID를 설계하였으며, 동일한 문서가 어떤 코퍼스 내에 존재하더라도 유사 문서들과 구분(discrimination) 되지 못함. 즉, DocID가 문서의 의미를 잘 나타내더라도 다른 문서들과의 식별성(discriminativity)이 부족하여 검색 혼동을 야기한다.
• [정적(Static) DocID 사용의 한계] 대부분의 기존 연구(예: Ultron, SE-DSI, NOVO, TSGen)는 사전에 정의된 정적 DocID(static DocID)를 사용함. 이 방식은 학습 중 DocID가 문서 의미에 맞게 조정되지 않아, 모델이 문서 구조를 충분히 학습할 수 없다.

2. Methodology

2.1. Overview

D2Gen은 크게 세 단계로 구성된다. 핵심 아이디어는 “문서 이해 → 서술적·변별적 DocID 선택 → 그 DocID로 재학습”의 순환 최적화 구조이다.

• i) 먼저 문서 이해 단계에서 인코더-디코더 기반 생성형 리트리버(generative retrieval)에 대조학습을 보조로 결합하여 쿼리-문서 표현을 학습하고, 문서의 토큰을 n그램을 통해 분할하고, n그램에 대한 마지막 레이어의 어텐션을 수집해 “문서 & n그램 관련도 행렬”을 만든다.
• ii) 다음으로 DocID 선택 단계에서 전 코퍼스를 가로지르는 이 행렬을 분석하여 TF-IDF를 확장한 NR-IDR(ngram Relevance Inverse Document Relevance) 점수를 계산하고, 각 문서에 대해 의미를 가장 잘 설명하면서도 유사 문서와 혼동되지 않게 하는 상위 n그램을 중복 제거하여 DocID로 확정한다.
• iii) 마지막으로 리트리버 학습 단계에서 실제 쿼리와 품질 및 다양성 필터링을 거친 합성 쿼리로 모델을 추가 학습하고, 새로 선택된 D2-DocID로 라벨을 갱신해 위 세 단계를 반복함으로써, 생성되는 DocID가 문서의 핵심 의미를 ‘서술’하고 코퍼스 내 유사 문서와 “변별”되도록 동시 최적화한다.

2.2. Document Understanding

입력 / 출력

• 입력: 쿼리-문서 쌍 (\(Q, D\))
• 출력: 문서의 의미적 표현(embedding)과 DocID를 생성하기 위한 확률 분포

생성기반 검색의 인코더-디코더 구조에서 질의 \(Q\)를 인코딩하고 DocID를 자가회귀적으로 생성하며, 이에 더해 인코더 표현을 대조학습으로 정렬하여 문서 의미를 “정확히” 추출하도록 훈련하는 단계이다. 이 단계의 산출물(문서-n그램 어텐션 스코어)은 이후 DocID 선택 단계의 기저 행렬을 형성한다. 이를 위해 두 가지 손실 함수를 정의하여 사용한다.

[생성(Generation)을 위한 손실]: Next Token Prediction

$$\mathbf{e}_i = \text{Decoder}(\text{Encoder}(Q), \text{id}_{<i})$$ $$P(id_i \mid q, \text{id}_{<i}, \Theta_{e, d}) = \text{SoftMax}(\text{Lmhead}(\mathbf{e}_i))$$ $$\mathcal{L}_r(\Theta_{e,d}) = \displaystyle\sum_{i}^l \log P(\text{id}_i \mid q, \text{id}_{<i}; \Theta_{e, d})$$

• \(\Theta_{e, d}\): 인코더-디코더의 학습 가능 파라미터
• \(\text{Lmhead}\): 디코더 출력을 어휘(vocabulary) 분포로 변형하는 선형 레이어

[대조학습 (인코더 정렬) 손실]: InfoNCE

$$\mathbf{h}_q = \text{MeanPooling}(\text{Encoder}(Q))$$ $$\mathbf{h}_d = \text{MeanPooling}(\text{Encoder}(D))$$ $$\mathcal{L}_c = -\log \frac{e^{s(\mathbf{h}_q, \mathbf{h}_d)/\tau}}{e^{s(\mathbf{h}_q, \mathbf{h}_d)/\tau} + \sum_{(q, d^-) \notin \mathcal{C}}e^{s(\mathbf{h}_q, \mathbf{h}_d)/\tau}}$$

최종적으로 문서 이해를 위한 손실함수는 다음과 같이 정의된다.

$$\mathcal{L}_p = \mathcal{L}_r + \lambda \mathcal{L}_c$$

2.3. DocID Selection

문서에서 길이 1−3의 모든 n그램을 후보로 추출하고, “문서 인코딩 시 마지막 층 어텐션”을 평균해 각 n그램의 문서 관련성을 점수화한다. 전체 코퍼스 차원에서 문서×n그램 희소행렬 \(\mathcal{M}\)을 만들고, TF-IDF에서 착안한 “ngram Relevance-Inverse Document Relevance (NR-IDR)”를 정의하여 서술성(문서 대표성, Descriptive)과 판별성(다른 문서 대비 구별력, Discriminative)을 동시에 반영한다. 이후 중복 제거 규칙으로 상위 \(n_g\)개의 n그램을 최종 DocID로 선택한다. 문서 \(D = (d_1, d_2, \cdots, d_n)\), n그램 \(g = (d_i, \cdots, d_j)\)에 대해 문서-n그램 사이의 관련도와 행렬은 다음과 같이 표현할 수 있다.

$$\text{Rel}(D, g) = \text{MP}\Big(\text{Att}\big(\text{Encoder}(D)\big)[d_i:d_j] \Big)$$ $$\mathcal{M}[D, g] = \text{Rel}(D, g), \quad \text{for every } D \in \mathcal C, g \in \mathcal{G}_D$$

이 때, MP는 Mean Pooling이고, \(\mathcal{G}_D\)는 문서 \(D\)의 n그램 집합, \(\mathcal C\)는 전체 코퍼스 집합을 의미한다. 중요한 점은 행렬 \(\mathcal{M}\) 에서 행이 문서이고, 열은 n그램이며, 각 성분은 마지막 레이어로부터 얻은 어텐션 스코어를 의미한다는 것이다. 한 문서가 전체 n그램 중 극히 일부에만 값이 존재하므로, 할당되지 않은 원소는 None으로 두어 희소 행렬이 된다 (<span style=”font-size:80%”>※ 그림에서는 행이 n그램으로 되어 있는데, 이는 오류이다.</span>).

이렇게 얻은 희소행렬은 다시 역문서 관련도(Inverse Document Relevance, IDR)을 계산하는데 사용된다. 코퍼스 전체 문서들에서 형성된 열 기준 통계로 판별력을 부여하기 위해, 각 n그램 열 \(j\)에 대해 해당 열이 값을 가진 행(문서) 인덱스 집합 \(I_j\)를 정의하고, 역문서 관련도를 다음과 같이 계산한다.

$$\text{IDR}[j] = \log \Big(\frac{1}{\text{Mean}_{k \in I_j}(\mathcal{M}[k, j])} \Big) = \log\Big( \frac{\vert I_j \vert}{\sum_{k \in I_j} \mathcal{M}[k,j]} \Big)$$ $$I_j = \{ k \in [0, \vert\mathcal C\vert ] \mid \mathcal{M}[k, j] \neq \text{None} \}$$

마지막으로, 희소 행렬(\(\mathcal {M}\))과 역문서 관련도 (\(\text{IDR}\))를 통해 본문에서 제안한 NR-IDR 점수를 계산한다. 이는 문서 \(D_i\)와 해당 문서의 n그램 \(g_j\)에 대해, 문서 내부의 중요도(관련도, 빈도)와 코퍼스 차원의 판별력(IDR)을 결합한 점수를 의미한다. 수식으로 표현하면 다음과 같다.

$$\text{NR-IDR}[i, j] = \mathcal{M}[i, j] * \sqrt{\text{TF}[i, j]} * \text{IDR}[i]$$

• \(\mathcal{M}[i, j]\) : 문서 \(D_i\)와 n그램 \(g_j\)의 의미적 관련도(어텐션 스코어)
• \(\text{TF}[i, j]\): 문서 \(D_i\) 내 n그램 \(g_j\)의 등장 빈도
• \(\text{IDR}[\cdot]\) : 역문서 관련도, 코퍼스 차원의 구별력을 반영함

각 문서 \(D_i\)에 대해, 해당 문서에서 추출한 n그램 집합 \(G_{D_i}\)에 대해 \(\text{NR-IDR}[i, :]\)를 기준으로 내림차순으로 정렬한다. 이후 단어 프로토타입 중복 규칙으로 순차적 중복 제거를 수행하여, 이미 선택된 DocID들의 각 단어 프로토타입과 모두 중복되는 n그램은 건너뛰고, 중복되지 않는 n그램만 채택한다. 이렇게 중복 제거 후 상위 \(n_g\)개의 n그램을 해당 문서의 DocID로 확정한다.

2.4. Retrieval Learning: 증강 질의 선별과 재학습

DocID를 갱신한 뒤, 검색 성능을 더 끌어올리기 위해 “실제 질의 + 선별된 합성 질의”를 혼합해 재학습한다. 합성 질의(pseudo query)는 원문과 문서 분절에서 생성하며, “다양성·품질” 필터를 통과한 질의만 사용한다. 이 단계의 손실은 Section 2.2의 Next Token Prediction 손실과 동일하며, DocID 라벨을 갱신하여 iteration 가능하다.

[데이터 구성]

• 문서 분절(dividing): 각 문서 \(D\)를 문장 단위로 이어붙인 패세지 집합 \(P = \{p_1, p_2, \cdots \}\)로 만들고, 윈도우 크기를 3, 중첩 1문장으로 슬라이딩한다 (CNN에서 Stride와 유사).
• 합성 질의(pseudo-query): 문서 \(D\)에서 \(n_d=10\)개(\(Q_d\)) + 각 패세지 당 3개의 \(Q_p\)를 생성하고, 둘을 합쳐 사용 (\(Q = Q_d \cup Q_p\))

$$\text{sim}_q = \max_{q' \in Q_s} [M_\beta(q) \cdot M_\beta(q')]$$

• 필터링 기준: 외부 Dense Retrieval \(M_\beta\) 로 \(q\in Q\)의 MRR@10을 산출하여 능력을 측정하고, 이미 선택된 집합 \(Q_s\)와의 최대 유사도를 계산한다. 이 유사도가 임계치를 넘어가면 \(q\)를 사용한다.

최종적으로 i)Document Understanding - ii)DocID Selection - iii)Retrieval Learning을 반복해 점진적으로 성능을 향상시킨다고 한다.

3. Experiments

3.1. Main Results

MS MARCO처럼 비정형 다출처 코퍼스에서는 문서 의미 이해와 구분력이 중요한데, D2-DocID가 이를 보완함으로써 큰 개선을 보였다. 반면 NQ는 Wikipedia 기반이라 구조적 일관성이 높아 개선폭이 작다.

3.2. Ablation Study

• NR-IDR 방식이 TF-IDF, 단순 Relevance Ranking보다 높은 성능을 보임 (M@10 = 0.607 vs 0.586).
• Contrastive Loss를 제거하면 성능 저하 (M@10: 0.592 → 0.607).
• n-gram 개수는 3개일 때 최적, 너무 많거나 적으면 성능 저하.
• Iteration 1→2회 시 성능 소폭 상승 (0.603 → 0.610).
• 즉, NR-IDR의 성능 gain이 가장 크다.

3.3. Generality

일반화 실험(Table 3)은 기존 생성형 리트리버의 본체와 학습법은 그대로 둔 채 DocID만 교체하여, 제안한 D2-DocID의 독립적 효과를 검증한 설정이다. Ultron의 경우 기존 Title+URL 식별자에서 M@10이 0.400→0.538로 +0.138p(상대 +34.5%), R@1이 0.296→0.431, R@10이 0.678→0.757로 모두 크게 상승하였다. SE-DSI에서도 기존 pseudo-query 식별자를 D2-DocID로 바꾸면 큰 폭으로 성능이 개선되었다. 즉, 모델 구조·학습데이터·최적화는 동일해도 DocID만 교체하면 성능이 유의하게 향상되며, 이는 D2-DocID가 모델에 종속되지 않고 독립적으로 검색력을 끌어올린다는 점을 직접적으로 보여준다이다.

또한 동일한 본 논문의 D2Gen 본체 안에서도 DocID 종류만 바꿔 비교하면, Title+URL(0.579)이나 pseudo-query(0.585)를 사용할 때보다 D2-DocID가 M@10=0.607, R@1=0.504, R@10=0.810으로 가장 높다는 결과가 확인된다. 이는 NR-IDR 기반 n-gram 선택이 문서 의미를 기술적으로 반영하면서 코퍼스 내 유사 문서와의 식별성을 동시에 확보해, 생성 디코딩 단계에서 혼동을 줄이고 올바른 DocID 생성을 유도하기 때문임을 시사한다이다.

4. Conclusion

Limitations

• 문서 내 n-gram 기반 DocID는 문서가 길거나 구조적으로 복잡한 경우에는 일부 의미 누락 가능성이 있음.
• pseudo-query 생성 품질이 DocID 학습 성능에 직접 영향을 미치므로, query generation 품질 관리가 필요함.
• 본 연구는 주로 영어 중심 코퍼스(MS MARCO, NQ)에 대해 검증되었으며, 다국어/도메인 특화 상황에서는 일반화가 제한될 수 있음.

Contribution

• [D2-DocID 설계]: 문서의 의미를 기술(descriptive)하면서 동시에 코퍼스 내 유사 문서와 구분(discriminative)되는 학습 가능한 DocID 제안.
• [NR-IDR 알고리즘]:n-gram 수준의 semantic relevance와 inverse document relevance를 통합하여 문서 식별자 선택의 효율성과 정확도를 높임.
• [D2Gen 모델]: Contrastive Learning 기반 Document Understanding + Iterative DocID 학습 구조를 결합한 최초의 paired generative retrieval 모델 제안.
• [검증된 일반화 능력]: D2-DocID는 다른 Generative Retriever에도 적용 가능하며, MS MARCO와 NQ 모두에서 SOTA 성능을 달성함.

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