[논문리뷰]Multi-level Relevance Document Identifier Learning for Generative Retrieval

Fuwei Zhang, Xiaoyu Liu, Xinyu Jia, Yingfei Zhang, Shuai Zhang, Xiang Li, Fuzhen Zhuang, Wei Lin, and Zhao Zhang. 2025. Multi-level Relevance Document Identifier Learning for Generative Retrieval. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar (Eds.). Association for Computational Linguistics, Vienna, Austria, 10066–10080. https://doi.org/10.18653/v1/2025.acl-long.497

Problem Statement

MERGE는 Generative Retrieval (GR) 환경에서 고품질의 DocID를 생성하는 방법을 다루고 있다. 기존 GR에서는 문서의 텍스트 표현만을 기반으로 DocID를 생성하기 때문에, 유사한 의미를 가지는 문서라도 표현 방식 차이로 인해 DocID 간의 의미적 유사성이 약해질 수 있다. 이를 해결하기 위해, 저자들은 쿼리를 문서 간 의미 연결의 매개체로 활용하여 다단계(relevance level) 의미 정보를 반영한 DocID 학습 방식을 제안한다. MERGE에서 말한 기존 연구의 한계는 다음과 같다.

[문서 내용 기반 DocID 생성의 한계] 기존 Generative retrieval 기법은 주로 문서의 텍스트 표현만을 기반으로 DocID를 생성하기 때문에, 의미적으로 유사한 문서라도 표현 방식의 차이(동의어, 서술 구조 변화 등)로 인해 생성된 DocID 간의 의미적 유사성이 약해질 수 있다. 이로 인해 동일한 질의에 대해 유사 문서들이 서로 다른 식별자를 가지게 되어 검색 일관성이 저하된다.

[이진 수준 Relevance 학습의 한계] 많은 Generative retrieval 연구들은 쿼리–문서 간 관련성을 이진(관련/비관련) 수준에서만 반영하여 학습한다. 이러한 접근은 관련 문서들 간의 세부적인 유사도 차이를 반영하지 못하고, 결과적으로 동일 relevance 그룹 내의 문서 구분 능력이 떨어진다.

[쿼리–문서 간 계층적 의미 정보 미활용] 기존 방법들은 쿼리와 문서 간의 다단계 relevance 관계를 효과적으로 학습 과정에 반영하지 않는다. 특히, 유사 문서 간 관계를 중재하는 ‘쿼리’의 연결 역할이 활용되지 않아, 계층적 의미 정보(e.g., Exact–Substitute–Complement 구분)가 DocID에 반영되지 못한다.

• 계층적 의미 정보는 ESCI 데이터셋에 내제된 relevance label 구조를 의미한다.
  • E (Exact): 쿼리 조건과 속성(사이즈, 색상, 기능 등)을 모두 만족하는 완전 일치 상품
  • S (Substutute): 쿼리와 대체 가능하지만, 일부 속성이 다르거나 부가 조건을 만족하지 않는 상품
  • C (Complement): 쿼리 제품과 함께 사용할 수 있는 보완 제품
  • I (Irrelevant): 쿼리와 관련 없는 상품
• 즉, 계층적 의미 정보는 단순히 ‘관련 vs 비관련’이 아니라, 관련(Exact) → 덜 관련(Substitute) → 보완적(Complement) → 비관련(Irrelevant) 이렇게 다단계의 의미적 거리를 표현하는 라벨 구조를 의미한다.
• 일반적인 QA 데이터셋 (e.g., SimpleQA, NaturalQuestions, TriviaQA, HotpotQA, 2WikimultihopQA)들은 모두 binary relevance (정답 문서 vs 오답 문서)로 label이 주어져있고, 어떤 문서가 정답에 더 가까운지에 대한 정보가 없으므로, “정답 문서에 가까운 정도”를 label로 만들어야지 relevance관련된 아이디어를 사용할 수 있다.

Methodology

Overview

MERGE는 먼저 쿼리–문서 쌍을 입력으로 받아, 문서를 PLM을 통해 임베딩한 뒤 RQ-VAE를 사용해 다단계 코드북 기반의 계층적 DocID를 생성한다. 이 과정에서 Multi-relevance Query–Document Alignment로 서로 다른 relevance level 문서 임베딩을 쿼리에 정렬하고, Outer-level Relevance Contrastive Learning으로 관련/비관련 문서를 구분하며, Inner-level Multi-level Relevance Learning으로 관련 문서들 간의 세부 구분을 학습한다.

학습된 RQ-VAE는 각 문서에 고유하고 의미 계층이 반영된 DocID를 부여하며, 이를 GR 학습 데이터로 변환한다. 마지막으로, Seq2Seq 기반 GR 모델이 쿼리를 입력받아 해당 문서의 DocID를 생성하도록 학습하여, 질의 시 생성된 DocID를 통해 문서를 검색한다.

RQ-VAE Training

Step 1. DocID Learning via Multi-level Relevance

• 입력: PLM(BERT/T5/mT5)으로 인코딩된 문서 임베딩
• 출력: 계층적구조를 가진 DocID (\(c_0, c_1, \cdots, c_{m-1}\)), 각 \(c_l\)은 레벨 \(l\)의 코드북에서 선택된 인덱스

먼저 문서 임베딩을 입력으로 받아 다층 코드북 기반의 이산 DocID 시퀀스를 생성해야 한다. PLM을 통해 문서를 임베딩하여 \(d\) 벡터를 생성하고, 이를 RQ-VAE 모듈에 입력시킨다. 그리고 \(m\)개의 코드북에서 단계별로 고른 인덱스 \((c_0, c_1, \cdots, c_{m-1})\)로 구성된 계측적 DocID를 출력한다. 이때 DocID는 coarse-to-fine의 위계를 가지며, 이후 GR에서 쿼리를 입력받아 해당 DocID를 생성·검색하는 용도로 활용된다.

RQ-VAE 절차를 세분화하면 다음과 같다. 먼저 DNN 인코더 \(E\)가 문서 임베딩 \(d\)를 잠재 표현 \(z = E(d)\)으로 변환하고, 이를 초기 잔차 \(r_0 = z\)로 둔다. 각 레벨 \(l\)마다 코드북 \(C_l =\{e_k^l\}_{k=1}^K\) 에서 잔차 \(r_l\)에 최근접인 코드워드를 고른 뒤 \(c_l = \text{argmin}_k \vert \vert r_l - e_k^l \vert \vert\) 로 표기하고, 잔차를 \(r_{l+1} = r_l - e_{c_l}^l\)로 갱신한다.

이 과정을 \(m\)번 반복하여 인덱스 튜플 (\(c_0, c_1, \cdots, c_{m-1}\))을 얻는다. 코드북 충돌을 막기 위해 각 코드북은 잠재표현에 대한 k-means 초기화를 사용한다. 최종 양자화 표현은 \(\hat{z} = \sum_{l=0}^{m-1}e_{c_l}^l\) 이며, 디코더 \(D\)를 통해 입력 임베딩을 복원하는 손실 함수 \(\mathcal{L}_{\text{recon}}\)과, 코드북과 잔차의 상호 근접을 유도하는 commitment 손실 \(\mathcal{L}_{\text{rq}}\)을 정의해서 사용한다. 구체적인 수식은 아래와 같다.

$$\mathcal{L}_{\text{recon}} = \vert \vert d - D(\hat{z}) \vert \vert_2^2 $$ $$\mathcal{L}_{\text{rq}} = \sum_{l=0}^{m-1} \Big ( \vert \vert \text{sg}[r_l] - \mathbf{e}_{c_l}^l \vert \vert_2^2 + \alpha \vert \vert \mathbf{r}_l - \text{sg}[e_{c_l}^l] \vert \vert_2^2 \Big)$$

이 때, \(\text{sg}[\cdot]\)은 stop-gradient로 해당 항의 역전파를 차단하여 코드북과 잔차를 안정적으로 맞물리게 한다. 최종적으로 RA-VAE loss는 다음과 같다.

$$\mathcal{L}_{\text{RQ-VAE}} = \mathcal{L}_{\text{recon}} + \mathcal{L}_{\text{rq}}$$

Step 2. Multi-level Relevance DocID Learning

다음으로, MERGE는 단순 문서 임베딩 기반 ID 학습의 한계를 보완하기 위해 다단계 관련성을 DocID 학습에 주입한다. 첫 단계는 Multi-relevance Query–Document Alignment로, 실제 검색 환경에서 짧은 쿼리와 긴 문서 간 분포 불일치를 완화하고 쿼리를 중심축으로 문서 표현을 정렬하는 장치이다. 쿼리 \(q\)는 해당 쿼리에서 가장 높은 관련성 레벨 \(L\)에 속하는 문서 잠재표현의 평균으로 정의되며

$$\mathbf{q} = \frac{1}{\vert \mathcal{D}_q^L \vert} \sum_{i=1}^{\vert \mathcal{D}_q^L \vert} \mathbf{z}_{i}^{q, L}$$

로 표현된다. 레벨 \(j\)의 관련 문서 집합 \(\mathcal{D}_q^j\)를 쿼리 \(q\)에 정렬하는 계층형 손실 함수 \(\mathcal{L}_{\text{align}}\)을 적용하며, 수식은 다음과 같다.

$$\mathcal{L}_{\text{align}} = \frac{1}{|Q|} \sum_{q \in Q} \sum_{j=1}^L \frac{w_j}{\vert \mathcal{D}_q^j \vert} \sum_{k=1}^{\vert \mathcal{D}_q^j \vert} \text{Dist}(\mathbf{d}_{q, k}^j - \mathbf{q})$$

이 때, \(w_j = \frac{1}{L - j+1}\)는 높은 레벨에 더 큰 정렬 압력을 가하는 가중치이고, \(\text{Dist}(\cdot)\)은 distance를 의미한다. 즉, \(d_{q, k}^j - q\) 벡터의 거리를 나타낸다. 거리 척도는 코사인 기반을 사용하여 레벨별 문서가 쿼리 표현에 질서 있게 수렴하도록 유도한다.

Step 3. Outer-level Relevance Contrastive Learning.

정렬만으로는 관련 없는 문서와의 경계가 충분히 뚜렷하지 않으므로, 두 번째 단계로 Outer-level Relevance Contrastive Learning을 도입한다. 이는 각 양자화 레벨 \(l\)에서 InfoNCE를 적용하여 쿼리 \(q\)에 관련된 문서 의 잔차 \(r_l\)끼리는 가깝게, 관련 업는 문서(In-batch negative)와는 멀어지게 만든다. 이를 수식으로 표현하면 다음과 같다. 이 손실은 특히 하위 코드북 레벨에서 거친 분리를 형성하여 이후의 미세 구분이 효과적으로 이루어지도록 발판을 제공한다

$$\mathcal{L}_{\text{outer}}^l = \sum_{d_i^q, d_j^q \in D_q^{\text{rel}}} \log \frac{\exp(\text{sim}(\mathbf{r}_l^{d_i^q}, \mathbf{r}_l^{d_j^q}) / \tau)}{\sum_{d \in \mathcal{D}_{\text{batch}}} \exp(\text{sim}(\mathbf{r}_l^{d_i^q}, \mathbf{r}_{l}^d)/\tau)}$$

Step 4. Inner-level Multi-level Relevance Learning

세 번째 단계는 Inner-level Multi-level Relevance Learning으로, 관련 문서들 사이의 등급 차이(예: Exact, Substitute, Complement)를 표현공간에 각인시키는 과정이다. 레벨 \(l\)에서 삼중항 (\(d, d^+, d^-)\)을 구성하되, 앵커 \(d\)와 양성 \(d^+\)는 더 높은 관련성 레벨, 음성 \(d^-\)는 더 낮은 레벨에서 선택한다. 이 과정에서 손실 함수는 triplet loss를 사용한다. \(\mathcal{T}\)는 학습하는 triplet의 집합이다.

$$\mathcal{L}_{\text{inner}}^l =\displaystyle\sum_{(d, d^+, d^-) \in \mathcal{T}} \max(0,\gamma + \text{sim}(\mathbf{r}_l^d, \mathbf{r}_l^{d^-}) - \text{sim}(\mathbf{r}_l^d, \mathbf{r}_l^{d^+}))$$

마진 \(\gamma\)를 통해 “상위-레벨은 더 가깝고 하위-레벨은 더 멀다”는 질서를 학습한다. 이를 통해 동일 ‘관련’군 내부에서도 등급별 미세 구분이 가능한 DocID가 형성된다.

Step 5. Full DocID Learning

이 세 손실을 레벨별로 어떻게 결합할지는 위계적 최적화 전략으로 요약된다. 저레벨 코드북일수록 비관련 분리를 강하게, 고레벨로 갈수록 등급 구분을 강하게 만드는 것이며 이를 위해 \(\beta_l\)을 도입하여 \(\mathcal{L}_{\text{rel}}\)을 다음과 같이 정의한다.

$$\mathcal{L}_{\text{rel}} = \frac{1}{m}\sum_{l=0}^{m-1} \Big( \beta_l \mathcal{L}_{\text{outer}}^l + (1- \beta_l)\mathcal{L}_{\text{inner}}^l \Big)$$

최종적으로 DocID 학습의 전체 목적함수는 \(\mathcal{L}_{\text{IDgen}} = \mathcal{L}_{\text{RQ-VAE}} + \lambda_1\mathcal{L}_{\text{align}} + \lambda_2\mathcal{L}_{\text{rel}}\)로 통합된다.

Fine-tuning After DocID Pre-training

먼저 학습이 끝난 RQ-VAE로 모든 문서에 고유한 ID를 부여한다. 예를 들어 문서의 ID 튜플이 (1, 0, 2)이면 이를 <a_1><b_0><c_2> 와 같은 토큰 시퀀스로 구성하고, <a_1> 과 같은 각 요소는 학습용 vocabulary의 개별 토큰으로 추가한다. 이 고유한 ID들이 정답 시퀀스가 된다.

학습은 Seq2Seq (T5)로 수행하며, DocID의 앞자리(low-level codebook) 토큰이 더 중요하다는 가정하에 위치 가중 손실을 적용한다.

$$\mathcal{L} = - \sum_{t=1}^T w_t \cdot \log P(y_t \vert y_{ \ <t}, q), \quad w_t = \frac{1}{\sqrt{t}}$$

손실 함수는 위와 같으며, \(q\)는 입력된 질문, \(T\)는 토큰 길이, \(y_t\)는 DocID의 \(t\)번째 토큰이다. 이 weighted scheme은 모델이 초기 토큰의 정확도에 더 많은 주의를 기울이게하여, 생성 기반 검색에서 초기 단계 수렴과 후보 압축을 돕도록 설계되었다.

Experiments

Main Results

Table 1은 Sparse(BM25), Dense(DPR, Sentence-T5, mMPNet), Generative(DSI variants, NCI, LTRGR, RIPOR, vanilla RQ-VAE)와의 직접 비교 결과를 제시한다. MERGE는 English/Spanish/Japanese 공통으로 R@10, NDCG@100에서 가장 높은 수치를 기록했고, 단순 코드북 양자화만 사용한 vanilla RQ-VAE를 크게 앞섰다. 특히 GR 강자 RIPOR를 능가하면서도 학습 파이프라인은 더 단순하다는 점이 강조된다. 이는 MERGE의 정렬(Alignment)–바깥쪽 대조(Outer)–안쪽 대조(Inner) 손실 설계가 DocID의 계층적 의미와 판별성을 동시에 강화했기 때문임을 뒷받침한다.

Ablation Study

• 특히 \(\mathcal{L}_{\text{outer}}\)(바깥쪽 대조 손실) 제거 시 NDCG@100이 12.58→10.81로 가장 크게 떨어졌다.
• GR에서 비관련 문서를 초기에 강하게 분리하는 메커니즘(\(\mathcal{L}_{\text{outer}}\))이 가장 결정적이며, 쿼리 중심 정렬(\(\mathcal{L}_{\text{align}}\))과 레벨 내 등급 구분(\(\mathcal{L}_{\text{inner}}\)), 그리고 DocID 초반 토큰 가중(\(w_t\))이 누적적으로 기여함을 보여준다.

Table 2는 English ESCI에서의 제거 실험을 통해 각 모듈의 기여도를 정량화한다. Outer 대조 손실은 coarse 레벨 분리를 만들어 이후 미세 구분의 토대를 제공하므로, 제거 시 성능 낙폭이 가장 크다. Alignment는 쿼리–문서 분포 불일치를 보정하여 관련성 신호를 안정화하고, Inner는 관련 내 등급 차이를 표현 공간에 각인시켜 DocID의 세밀한 판별력을 높인다. 위치 가중은 DocID 앞자리에 더 큰 학습 신호를 줘 초기 토큰 정확도를 향상시키며, 이는 GR에서의 검색 성공률과 직결된다.

Visualization of Document Represenation (PCA)

• vanilla RQ-VAE는 서로 다른 쿼리(예: #1 “car wash cannon”, #2 “sasquatch cookie cutter”, #3 “logitech mx master 3”)에 대한 문서 표현이 상호 겹침을 보이며, 특히 쿼리 #2 일부가 쿼리 #1 근처에 위치하는 등 경계가 모호해졌다.
• MERGE는 동일한 설정에서 고관련 문서가 쿼리 근처로 더 조밀하게 모이고, 서로 다른 쿼리 간 표현 공간 경계가 더 명확해졌다(도형 상 점선 박스 강조 구간).
• 다단계 관련성 손실(Alignment/Outer/Inner)이 결합되면 표현 공간에서 쿼리 중심의 계층적 구조가 형성되고, 이는 DocID 유사도–비유사도 관계에 직접 반영되어 생성 기반 검색의 정밀도를 높이는 효과가 있다.

Figure 3은 RQ-VAE의 DNN 인코더가 만든 문서 표현을 PCA로 2D 투영해 쿼리별로 시각화한 결과이다. MERGE의 경우, 쿼리-문서 정렬과 이진/등급 구분 학습의 누적 효과로 클러스터 간 간극이 확대되고 클러스터 내부 밀집도가 증가하였다. 이는 단순 재구성 기반 RQ-VAE가 캡처하지 못한 질의 중심의 위계적 시맨틱이 MERGE에서 성공적으로 학습되었음을 보여준다.

Distribution of DocID Layers

• MERGE는 하위 레벨(layer 1–2)에서 동일 쿼리의 관련 문서들이 동일/유사 ID 토큰으로 더 많이 수렴하며, 반대로 상위 레벨(layer 4)에서는 분산도가 vanilla와 유사하게 유지된다. Table 3의 Average Unique ID Count에서도 layer1(4.65 vs 6.55), layer2(12.00 vs 12.90)로 MERGE가 더 집중되어 있고, layer4는 유사한 수치로 개별성을 유지한다.
• 시각화 그래프에서 MERGE는 낮은 레벨 토큰 노드의 연결/중첩(공유)이 더 두드러지며, 높은 레벨 토큰으로 갈수록 분기가 늘어 세분화가 강화된다.
• MERGE의 DocID는 “coarse(공유 의미) → fine(개별 식별)”의 이상적인 위계를 실현한다. 이는 검색 시 초기 토큰의 회수/정확을 높이고, 후행 토큰에서 정밀 판별을 보장하여, 효율(짧은 생성으로도 후보 압축)과 정확(후속 토큰으로 미세 구분)을 동시에 달성하게 한다.

Figure 4는 여섯 개 쿼리에 대해 레벨별 DocID 토큰 노드(각 레벨 256개)와 연결을 네트워크로 시각화한다. MERGE는 초기 레벨에서 공통 토큰을 더 폭넓게 공유하여 유사 문서 동류화를 이루고, 이후 레벨에서 토큰 다양성을 확보해 개체 식별성을 유지한다. Table 3의 수치가 이를 정량적으로 뒷받침한다. 이러한 계층적 토큰 배분은 GR의 생성 과정(왼쪽에서 오른쪽으로 토큰 생성)과 정합적이어서, 앞자리 토큰 집중으로 빠른 수렴과 높은 초기 재현을, 뒷자리 토큰 분화로 최종 판별력을 확보한다.

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