Ko, Youmin, Sung Jong Seo, and Hyunjoon Kim. "Cooperative Retrieval-Augmented Generation for Question Answering: Mutual Information Exchange and Ranking by Contrasting Layers." The Thirty-ninth Annual Conference on Neural Information Processing Systems.
https://arxiv.org/abs/2512.10422
Abstract
RAG 기반의 QA는 여전히 부정확한 검색 및 환각이 존재한다. 이 문제를 다루기 위해 본 논문에서 제안하는 CoopRAG는, retriever와 LLM이 협력해서 정보를 교환하며, retriever의 여러 레이어가 협력해서 문서 검색 및 재정렬을 진행한다. 다음 방식으로 프레임워크가 구성된다.
- 질문을 부분 질문과 추론 체인으로 나누고(unroll), 명확하지 않은 부분은 마스킹한다.
- 부분 질문과 추론 체인으로 문서를 검색한다.
- retriever의 여러 레이어로 문서를 리랭킹한다.
- LLM으로 마스킹을 채우면서 추론 체인을 재구축한다.
1. Introduction
LLM은 틀린 답변을 만들 수 있는 단점이 있으며, RAG는 이를 완화하는 방법론으로 주목을 받았다. 활용분야로는 Knowledge Base Question Answering(KBQA), Multi-Hop Question Answering(QA), Knowledge Graph Completion(KGC), Recommender System 등이다. 최근의 연구들은 LLM의 단점을 개선하고자 다양한 시도를 했다.
- 미리 인코딩된 쿼리와 문서를 토큰 레벨에서 매칭해서 검색( fine-granular late interaction scoring)
- LLM으로 생성된 가상의 답변으로 쿼리를 증강해서 사용
- LLM이 생성한 요약 및 KG 기반으로 새로운 구조를 만들고, 새롭게 연결되는 여러 문서를 탐색
이런 방법들은 특히 QA에서 여전히 검색/환각 문제가 있다. 이 문제는 3가지 문제에서 발생한다.
- 질문은 보통 짧고 제한적이다. 따라서 가장 관련 있는 문서를 찾고 고품질의 추론을 하기 어렵다. 이에 대한 방법론인 query rewriting은 유망한 방식이다. 하지만 이 방법의 성능은 외부 리소스와 지도 학습의 품질에 의존한다. 또한 재작성된 쿼리가 검증되지 않는 문제도 있다.
- RAG의 검색 성능은 대조 학습 목적함수에 의존한다. 이는 retriever가 중요한 정보가 아닌 쿼리 패턴에 맞는 문서를 찾게 한다. 또한 경험적으로 in-batch negative(배치 내 다른 문서), heuristic positive(규칙으로 증강한 문서)를 사용한 경우 유사도의 품질이 떨어진다. 표현 학습 관점에서 transformer 기반 retriever는 하위 레이어에 있는 문법 및 표면적 정보에 의존한다. retriever의 상위 레이어는 더 복잡한 의미적 관계를 인코딩한다.
- LLM에게 필요한 지식을 제공할 기회를 주지 않는다. 여러 전근법들은 LLM이 초기 답변을 정제하게 했으나, LLM이 모르는 것을 기반으로 지식을 채우지는 못한다.
이 문제를 다루기 위해 새로운 RAG 프레임워크인 Cooperative RAG를 제안한다. retriever와 LLM이 지식을 교환하는데 협력하고, retriever의 하위 레이어와 상위 레이어가 정확한 검색을 위해 협력한다.
CoopRAG는 retriever, LLM간의 양방향성 시너지를 만들며, 효과적으로 LLM의 내부 지식을 활용해서 retriever가 쿼리와 유사한 문서를 찾고 재정렬하게 하며, 신뢰성 높은 추론이 가능하게 한다.
2. Related Work
Query-Augmentation for RAG
질문에 충분한 정보가 없으면 LLM, Retriever의 성능이 떨어진다. 다양한 쿼리 증강 기법이 문서 검색 성능을 높이기 위해 도입되었다. LLM이 질문에서 키워드를 추출해서 상위 수준의 개념, 가상의 정답, 가상의 문서, 하위 질문 등으로 질문을 증강하는 방식이다. LLM을 질문을 의역하는데 사용해서 검색 시 사용하기도 한다. 반대로, 검색된 문서를 요약해서 질문을 증강하여 다시 검색하는 방법도 있다. 이러한 방법들은 LLM이 환각을 생성하는 경우 틀린 정보를 포함하게 되므로 오히려 검색 및 추론 성능이 저하될 수 있다.
Dense Retrieval with Pre-trained Language Models
기존 단어 횟수 기반 방법론들은 의미적 관계를 포착하지 못하는 한계가 있었다. Dense Retrieval은 쿼리와 문서를 밀집된 벡터로 인코딩하고 벡터 간의 유사도를 계산한다. 쿼리와 문서는 효율적 유사도 연산을 위해 단일 벡터로 임베딩된다. 이러한 압축은 표현이 소실되는 문제가 있어서, CoIBERT는 토큰-레벨 임베딩을 제안했다. 이런 방법은 다른 Transformer 레이어에 있는 지식 타입을 구분하지 못하므로, 왜곡된 유사도 점수와 검색 성능의 한계로 이어진다.
Structure-Augmented RAG
문서와 지식의 구조화된 특징을 활용하는 검색 기법이 관심을 얻었다. RAPTOR와 Proposition은 문서를 작은 문서로 나눠서 임베딩, 클러스터, 요약을 통해 긴 문맥 정보를 포착하는 계층적 표현을 생성한다. SiReRAG와 HopRAG는 복잡한 쿼리를 LLM으로 간단하게 변경하고, 이를 통해 KG 내 문서 간의 의미적 연결성을 찾는다. GraphRAG와 LightRAG는 LLM을 활용해서 텍스트에서 트리플을 추출하고 그래프 내에서 의미적 유사도를 찾을 수 있게 계층적으로 KG를 구성한다. HippoRAG와 HippoRAG2는 쿼리에서 명사는 노드로, 관계는 엣지로 사용해서 KG를 만든다. 하지만 이런 방식은 만드는데 높은 비용이 들고, 종종 수많은 관련 없는 추론 경로를 가지는 과밀집 그래프를 만든다.
Reaoning-Enhanced Approaches for Complex QA
LLM이 복잡한 질문을 해결하는데 여러 스텝의 추론이 이득을 제공함은 증명되었다. 선형 추론과 여러 브랜치 경로, 그래프 구조 등의 방법론 등이다. CoK는 LLM이 중간 지식을 생성하게 하는 프롬프트 방식이며, 외부 모델을 사용해서 지식을 검증한다. CoQ와 question decomposition은 질문을 답변 가능한 작은 질문으로 만들어서 최종 답변을 추론한다. 하지만 이 접근법은 LLM이 낮은 신뢰성을 가지는 항목에 대해 RAG를 하지 않는다. IRCoT는 각 단계에서 여러번 RAG를 해서 각 추론 내용을 활용하는데, 이 논문에 적용될 수 있다.
3. Method
3.1 Overview
검색, 재정렬을 위해 사전학습된 인코더를 파인튜닝한다. 그리고 학습된 인코더에 모든 문서를 입력해 각 문서를 대표하는 [CLS] 토큰 임베딩을 얻고, 전처리 단계에서 vectorDB에 저장한다.
- Question Unrolling에서 LLM으로 원본 질문을 하위 질문들과 (모르는 부분이 마스킹된) 추론 체인으로 나눈다.
- 질문은 하위 질문들과 추론체인으로 확장되며, 확장된 질문으로 top-n 문서를 검색한다.
- 이 문서들을 검색기의 멀티 레이어를 활용해서 재정렬해 top-k (k<n)의 문서를 검색한다.
- LLM으로 top-k 문서를 활용해 마스킹을 채운다.
- LLM이 재구축된 추론 체인, 원본 질문, 하위 질문, top-k 문서로 질문에 대한 답변을 생성한다.
3.2 Question Unrolling
질문으로 관련 문서를 검색하거나 추론하기에 자주 정보가 부족하다. Question Unrolling 과정에서는 LLM이 입력을 하위 질문과 추론 체인으로 나눈다. 추론 체인은 몇몇 엔티티가 마스킹된 트리플 시퀀스로 구성된다. LLM은 단계별 추론을 위한 근거 트리플을 생성하며, 이 때 생성한 엔티티 중 신뢰도가 낮은 것도 존재한다. 이러한 엔티티로 텍스트를 생성하거나 질문을 증강하는 것은 잘못된 검색, 증강으로 이어진다. 따라서 이런 엔티티를 마스킹한다.(<UNCERTAIN>) 최종적으로 Unrolled Question $U$는 질문, 하위 질문, 추론 경로가 합쳐진 형태이다.
$U=Q||S||R$
$\mathcal{S}=\{s_1,s_2,...,s_{|\mathcal{S}|}\}$
$\mathcal{R}=\{(e_1,r_1,e'_1),(e_2,r_2,e'2),\dots,(e{|\mathcal{R}|},r_t,\langle\text{FILL}\rangle)\}$
마지막 트리플의 tail 엔티티는 <FILL>로 마스킹되는데, 질문의 최종 정답이 생성된다. 위 과정에서 LLM은 추론 체인의 엔티티가 신뢰도가 낮은지 확인하는데만 사용된다. 이는 다른 방법론들과 상이한데 Question Decomposition은 LLM으로 하위 질문을 반복적으로 수정하고, CoQ는 최종 답변을 생성하기 위해 하위 질문에 반복적으로 답변한다.
3.3 Unrolling-Augmented Retrieval
unrolling-augmented retrieval(UAR) 단계에서 unrolled question U에 가장 가까운 top-n 문서를 검색한다. 이를 위해 파인튜닝된 인코더로 U의 토큰-레벨 임베딩과 모든 문서 D의 임베딩을 미리 계산해놓는다.
$\mathbf{q}0, \mathbf{q}1, \dots, \mathbf{q}{|U|} := \text{Normalize}(\text{Encoder}([\text{CLS}] \ q_1 \ q_2 \ \dots \ q{|U|}))$ … 1개의 unrolled question
$\mathbf{d}0, \mathbf{d}1, \dots, \mathbf{d}{|D|}:= \text{Normalize}(\text{Encoder}([\text{CLS}] \ d_1 \ d_2 \ \dots \ d{|D|}))$ … 1개의 문서
- Normalize : L2 정규화
- $\mathbf{q}_i, \mathbf{d}_i$ : $U$, $D$의 i번째 토큰 임베딩
- $\mathbf{q}_0, \mathbf{d}_0$ : $U$, $D$의 CLS 토큰 임베딩
$\mathbf{q}_0$와 각 문서의 $\mathbf{d}_0$로 유사도 계산을 해서 top-n 문서 $D_U=\{D_1,D_2,...,D_n\}$를 검색한다. Faiss를 사용해서 유사도는 빠르게 연산된다.
3.4 Ranking by Contrasting Layers
1차 검색으로 얻은 $D_U$의 문서를 재정렬해서 top-k 문서를 얻는다. 지식은 트랜스포머 모델의 특정 레이어에 존재하게 되므로, 멀티 레이어를 활용한 재정렬이 가능하다. 본 논문은 DoLA에서 영감을 얻은 새로운 정렬 방법 RaLa(ranking by contrasting layers)를 제안한다. RaLa는 트랜스포머 인코더의 하위/상위 레이어에 있는 정보 표현의 차이를 활용한다. 하위 레이어는 표면적 정보에 집중하고, 상위 레이어는 더 추상적인 의미 관계에 집중하는 특징이 있다. RaLa는 질문에 대한 각 레이어의 유사도를 비교해서, 차이가 많이나는 문서에 높은 가중치를 부여한다.
인코더의 전체 레이어를 2~4개 버킷으로 나누고, 각 버킷별로 랜덤하게 하나의 레이어를 선택한다. 선택된 중간 레이어들의 버킷은 $\mathcal{C}=\{1,2,...,|\mathcal{C}|\}$이며, $L$은 인코더의 마지막 레이어이다. early exit을 통해 중간 레이어에서 연산중인 값을 가져온다. 각 문서에 대해 $l \in \mathcal{C}$번째 은닉층은 $d_0^{(l)}, d_1^{(l)},...,d_D^{(l)}$ (토큰별로 나뉨)으로 표현된다.
재정렬 시 모든 문서 $D$와 unrolled question $U$ 사이의 점수를 계산한다. 각 쿼리 토큰 임베딩 $\mathbf{q_i}$과 문서 토큰들의 임베딩($\mathbf{d_0}$~$\mathbf{d_D}$)의 유사도를 계산하고, MaxSim 연산자로 최댓값(Maximum similarity)를 얻는다. 이 과정을 전체 쿼리 토큰에 대해 반복해서 각 쿼리 토큰별로 얻은 유사도의 평균을 계산한다. 각 쿼리 토큰 임베딩 $\mathbf{q_i}$과 문서 토큰들의 임베딩($\mathbf{d_0}$~$\mathbf{d_D}$)의 유사도 계산에는 $g(q_i,d_j)$ 연산을 사용한다. 이는 쿼리 토큰과 마지막 레이어 유사도 $\langle \mathbf{q_i,d_j}^{(L)}\rangle$와 중간 레이어 유사도 $\langle \mathbf{q_i,d_j}^{(l)}\rangle$의 최대 점수 차이이다.
$score(U,D)=avg_{i=0}^{|U|}(\max_{j\in\{0,1,...,|D|\}}g(q_i,d_j))$
$g(q_i,d_j)=\max_{l\in\mathcal{C}}(\langle \mathbf{q_i,d_j}^{(L)}\rangle-\langle \mathbf{q_i,d_j}^{(l)}\rangle)$
하지만 모든 후보 레이어에 대해 유사도 계산을 진행하는 것은 계산 오버헤드가 너무 크다. 따라서 average maximum similarity를 활용해서, 문서-쿼리 유사도 가중치를 곱하는 식으로 로직을 단순화한다.
$\text{score}o(U,D) = \omega{U,D} \cdot \text{avg}{i=0}^{|U|} \left( \max{j\in\{0,\dots,|D|\}} \langle \mathbf{q}_i, \mathbf{d}_j^{(L)} \rangle \right)$
- $\omega_{U,D}$ : $g(q_0,d_0)$
$\omega_{U,D}$에서 MaxSim으로 가장 유사도 차이가 큰 레이어 값을 선택하는 이유는, 쿼리가 문서와 관련이 있는 경우에 초기 레이어에 비해 마지막 레이어에서 출력이 변할 확률이 높기 때문이다.
*버킷 사이즈의 경우 증가할수록 재정렬 성능이 개선되나, 이는 연산량과 트레이드오프 관계이다. 따라서 본 논문에서는 연산량을 줄이기 위해 버킷사이즈를 4로 지정한다.
3.5 Reasoning Chain Completion and Reasoning
검색 결과가 충분한 정보를 제공하지 않을 때, LLM은 추론 과정에서 낮은 신뢰도를 가진다. 본 논문에서는 LLM 입력으로 신뢰도를 제공하는 경우 환각을 줄일 수 있다고 주장한다. LLM은 query unrolling으로 얻은 추론 체인을 업데이트한다. unrolled question과 top-k 문서로 LLM이 <UNCERTAIN> 마스크를 채우며, 최종적으로 추론체인의 <FILL> 엔티티를 채운다. 결과적으로 질문, 하위 질문, 재구축된 추론 체인, top-k 문서로 LLM이 최종 답변을 생성한다.
3.6 Difficulty-Aware Training of Retriever
쉬운 질문에 대한 학습은 오버피팅으로 이어질 수 있다. 어려운 질문은 주로 모호하고 흔하지 않은 시나리오이다. 검색용 사전학습 인코더를 파인튜닝하고 sample-wise loss reweighting(배치 내 샘플별 가중치 부여)로 어려운 질문에 더 많은 비중을 둔다.
파인튜닝 과정에서 학습 데이터셋의 질문을 unroll해서 CoIBERT에 따라 하나의 positive, 하나의 negative 문서와 연관시킨다. 미니배치 $\mathcal{B}$의 크기는 b이고, unrolled question은 $\mathcal{U}=\{U_1, U_2, ..., U_b\}$이며, 관련된 모든 문서는 $\mathcal{D}=\{D_1,D_2,...,D_{2b}\}$이다. $U_i$에 대해 $D_i$는 positive 문서이고 $D_{b+i}$는 negative 문서이다. 각 미니배치에 대해 다음과 같이 loss를 연산한다.
$\mathcal{L}{\mathcal{B}} = \sum{U_i \in \mathcal{U}} \alpha_{U_i} \mathcal{L}(U_i, D_i) \quad \text{where} \quad \mathcal{L}(U_i, D_i) = -\log \frac{\exp(\text{score}o(U_i, D_i)/\tau)}{\sum{D \in \mathcal{D}} \exp(\text{score}_o(U_i, D)/\tau)}$
가중치 $\alpha_{U_i}$는 $U_i$의 어려움을 의미하며 어려울수록 높은 가중치를 부여한다. infoNCE loss 함수인 $\mathcal{L}(U_i,D_i)$에서 $\tau$는 negative 샘플의 중요성을 제어하는 하이퍼 파라미터이다.
실험을 통해 question unrolling 없이 멀티-홉 QA를 했을 때, 정답 문서의 수가 증가할수록 (멀티-홉 질문일수록) 검색 성능이 저하된다는 것을 확인했다. 또한 LLM이 복잡한 멀티-홉 질문을 더 많은 하위 질문으로 나누는 경향이 있으므로, 정답 문서 개수에 따라 하위 질문이 늘어났다. 따라서 $\alpha_{U_i}$의 경우 하위질문의 개수를 활용해서 정의한다.
$\alpha_{U_i}=\text{log}(1+|S_{U_i}|)$
- $S_{U_i}$ : $U_i$의 하위 질문 집합
sample-wise loss 가중치를 적용해서 인코더가 더 도전적인 케이스를 학습하게 유도한다. 어려운 문제는 때로는 많은 정보를 얻을 수 있으며, 인코더 일반화에 중요한 역할을 한다.
4. Experiments
데이터셋
- single-hopQA : NaturalQuestions
- Multi-hopQA : HotpotQA, MuSiQue, 2WikiMultihopQA
- Factual QA : SimpleQA, FreshQA
MPNET을 인코더로 사용한다.
Retrieval 성능
- 베이스라인 : HippoRAG 계열, HopRAG, SiReRAG
- 평가지표 : R@2, R@5
검색에서 SOTA를 달성한다. 이 결과는 unrolling-augmented retrieval(UAR)의 마스킹 추론 경로에 의한 성능 개선으로 예상된다.
QA 성능
- 베이스라인 : HippoRAG 계열, HopRAG, SiReRAG
- 평가지표 : EM, F1
베이스라인에 대해 SOTA를 달성한다. 여러 레이어를 활용해 재정렬하고 추론 경로를 채우는 것이 QA 성능을 높인 것으로 예상된다.
SimpleQA, FreshQA에 대해서는 Correct/Incorrect/Not Attempted로 구분되는 ChatGPT grader 평가지표도 사용한다. (https://arxiv.org/abs/2411.04368) 이 지표에서도 HippoRAG2, HopRAG보다 높은 성능을 보인다.
Ablation (Gap-Aware, Difficulty-Aware 가중치)
쿼리-문서 유사도를 계산할 때 언어모델의 레이어별 지식을 사용하는 Gap-Aware 방식을 없애는 경우 가장 크게 점수가 떨어졌다. Difficulty-Aware Weight를 제거한 모델의 경우에도 점수가 떨어진다.
RaLa의 Gap-Aware 가중치 기반의 유사도 점수 비교
정답 문서, 오답문서, 방해문서(정답과 헷갈림)의 점수 차이를 비교한다. Gap-Aware 가중치 추가 시 각 집단 간의 차이가 더 커진다. RaLa를 적용했을 때 더 문서를 더 잘구별하게 됨을 알 수 있다.
Ablation (Uncertainly 마스크)
Uncertain 마스크를 제거했을 때 검색(R@2), QA(EM) 성능 모두 성능이 떨어진다. 이는 신뢰도 점수가 낮은 정보를 활용하는 것보다, 마스킹하는 것이 더 나음을 증명한다.
또한 마스킹은 모델의 혼란스러움에 대한 지표인 평균 엔트로피를 감소시킨다. 이는 마스킹이 모델의 환각을 줄일 수 있음을 의미한다.
파인튜닝 Retrieval 성능 비교
- 베이스라인 : Contriever, CoIBERTv2, ReSCORE
파인튜닝 이전 성능과 각 데이터셋에 대해 파인튜닝한 성능을 비교한다. RaLa는 비교 대상에 비해 높은 검색률을 보인다.
Contribution & Limitation
Contribution
- 질문을 하위 질문과 마스킹된 추론 체인으로 나누는 Question Unrolling을 제시했다.
- 신뢰성이 낮은 엔티티는 마스킹이 되어 추후 검색된 문서와 LLM으로 마스킹을 채운다.
Limitation
- 실제 사용시 배치사이즈, 토큰 길이에 따라 RaLa 연산 비용이 매우 크다.