[논문리뷰]TRACE the Evidence:Constructing Knowledge-Grounded Reasoning Chains for Retrieval-Augmented Generation (EMNLP findings, 2024)

Jinyuan Fang, Zaiqiao Meng, and Craig MacDonald. 2024. TRACE the Evidence:Constructing Knowledge-Grounded Reasoning Chains for Retrieval-Augmented Generation. In Findings of the Association for Computational Linguistics: EMNLP 2024, Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (Eds.). Association for Computational Linguistics, Miami, Florida, USA, 8472–8494.

1. Problem Statement

이 논문은 Retrieval-augmented Generation (RAG) 기반 Multi-hop Question Answering (MHQA) task를 다룬다. Question \(q\)와 retriever이 검색한 문서 집합 \(D_q = \{d_1, \ldots, d_N\}\)이 주어졌을 때, 이를 기반으로 LLM이 정답 \(a\)를 생성하는 것이다.

기존 RAG 방법론들은 \(D_q\) 전체를 reader LLM의 context로 넣지만, multi-hop QA에서는 여러 문서에 흩어진 supporting evidence를 정확히 연결해야 하므로 관련없는 문서 (irrelevant documents) 들이 추론을 방해한다. TRACE는 검색된 문서를 먼저 ⟨head entity,relation,tail entity⟩ 형태의 KG triple로 변환하고, 질문에 필요한 triple들을 autoregressive하게 선택해 knowledge-grounded reasoning chain을 구성한 뒤, 이 chain 또는 chain이 가리키는 documents만 사용해 답을 생성한다.

2. Limitations of Existing Works

• [Retriever-reader RAG의 irrelevant context 취약성] 기존 RAG는 retriever가 반환한 문서 전체를 reader LLM에 prepend하여 답을 생성한다. 그러나 retriever는 완벽하지 않기 때문에 검색된 문서 안에 irrelevant documents가 포함될 수 있고, 이러한 irrelevant context는 LLM을 혼동시켜 잘못된 추론을 유도한다. 이 문제는 step-by-step reasoning에 필요한 multi-hop questions에서 특히 두드러진다.
• [문서 단위 context의 supporting evidence 식별 한계] 기존 vanilla RAG는 문서 전체를 context로 사용하므로, 실제 정답에 필요한 supporting evidence와 불필요한 정보가 함께 LLM (=reader)에 전달된다. 문서나 문장은 여러 fact가 섞인 coarse-grained representation이기 때문에, 질문에 직접 필요한 개별 factual knowledge를 정밀하게 분리하기 어렵다. TRACE는 이를 해결하기 위해 문서를 더 fine-grained하고 concise한 KG triples로 변환한다.
• [CoT 기반 multi-hop retrieval의 hallucination 한계] IRCoT와 같은 chain-of-thought 기반 방식은 LLM이 생성한 CoT sentences를 이용해 retrieval과 reasoning을 수행한다. 하지만 LLM이 생성한 CoT sentences는 hallucination을 포함할 수 있고, 부정확한 intermediate reasoning은 이후 retrieval과 answer generation을 오염시킬 수 있다. TRACE는 reasoning step을 자유 생성된 CoT sentence가 아니라 retrieved documents에서 생성된 KG triples 위에서 선택하게 하므로, reasoning process가 문서에 grounded된다.
• [기존 KG 활용 RAG의 외부 KG 의존성] 기존 KG-enhanced RAG 또는 KGQA 연구는 Wikidata 같은 pre-existing KG를 사용하는 경우가 많고, 일부 setting에서는 answer가 KG entity로 제한된다. 반면 TRACE는 answer가 entity뿐 아니라 임의의 phrase가 될 수 있는 일반 multi-hop QA setting을 다루며, 외부 KG가 아니라 retrieved documents 자체에서 question-specific KG 를 생성한다.

3. Methodology

3.1. Overview

TRACE의 전체 방법론은 KG Generation → Reasoning Chain Construction → Answer Generation으로 구성된다. Figure 2는 전체 pipeline을 KG Generator, Triple Ranker, Triple Selector, Reader로 나누어 보여준다.

$$p(a \mid q, D_q) \sim p(a \mid q, z, D_q) \cdot p(z \mid q, G_q) \cdot p(G_q \mid D_q)$$

TRACE의 전체 확률적 구조는 위와 같다. 여기서 \(p(G_q \mid D_q)\)는 검색된 문서로부터 KG \(G_q\)를 생성하는 KG Generator이고, \(p(z \mid q, G_q)\)는 KG에서 reasoning chain(\(z\))를 구성하는 Reasoning Chain Constructor이며, \(p(a \mid q, z, D_q)\)는 reasoning chain 또는 관련 문서를 context로 사용해 정답(\(a\)) 생성하는 Reader이다.

3.2. KG Generator

KG Generator는 검색된 문서를 question-independent KG triple들로 변환한다. **입력은 각 문서 \(d_i \in D_q\)이고, 출력은 해당 문서에서 생성된 triple 집합 \(G_{q,d_i}\)이다. 각 Wikipedia 문서는 title과 text로 구성되며, TRACE는 title을 head entity로 두고 text 내부 phrase를 tail entity로 삼아 ⟨title, relation, tail entity⟩ 형태의 triple들을 생성한다.

모든 문서를 한 번에 LLM에 넣지 않고 문서별로 독립적으로 KG triple들을 생성하는 이유는 긴 입력에서 중간 정보가 무시되는 lost-in-the-middle 문제를 완화하고, 각 문서의 KG triples를 offline precomputation할 수 있기 때문이다. 최종 KG는 각 문서에서 생성된 triple들을 합쳐 \(G_q\)로 구성된다.

3.3. Reasoning Chain Constructor

Reasoning Chain Constructor는 질문에 필요한 supporting evidence를 KG triple들 위에서 autoregressive하게 선택한다. 질문 \(q\), 생성된 KG \(G_q\), 이전까지 선택된 triple들 \(z_{<i}\)를 입력받아, reasoning chain \(z\)를 생성한다.

$$p(z \mid q, G_q) = \prod_{i=1}^{L} p(z_i \mid q, z_{<i}, \hat{G}_i) \\ \hat{G}_i = f(q, z_{<i}, G_q), \quad \forall i = 1, \ldots, L$$

논문은 reasoning chain construction을 위와 같이 수식으로 정의한다. \(L\)은 reasoning chain의 최대 길이이고, \(\hat G_i\)는 \(i\)번째 스텝에서 선택 후보가 되는 candidate triple들을 의미한다. 즉 TRACE는 한 번에 모든 evidence를 선택하지 않고, 현재까지 선택된 evidence를 조건으로 다음 evidence를 선택한다. 이 방식은 multi-hop reasoning에서 이전 hop이 다음 hop을 찾는 clue로 작동한다는 점을 반영한다.

Triple Reranker는 전체 KG에서 현재 reasoning state와 관련 있는 candidate triples를 top-K로 좁히는 모듈이다. 질문과 현재까지 선택된 트리플, question-grounded KG \(G_q\)를 입력받아 최종 후보 트리플 \(\hat G_i\)를 정한다. Triple Reranker는 사전학습된 bi-encoder \(\text{Enc}(\cdot)\)이며 질문과 이전 triple들을 하나의 쿼리로 구성해 인코딩하고, 각 KG의 triple은 독립적으로 인코딩한다.

$$q_i = \text{Enc}(q, z_{<i}), \quad t_j = \text{Enc}(t_j), \quad \forall t_j \in G_q$$

이후 내적을 통해 각 question-triple pair별 relevance score를 구하고 이를 기반으로 상위 \(K\)개의 트리플을 선택하여 candidate triple set \(\hat G_i\)를 만든다.

Triple Selector는 candidate triple들 중 reasoning chain을 가장 잘 확장하는 다음 triple \(z_i\)를 선택하는 모듈이다. Triple Selector의 선택 확률은 option token의 logit으로 정의된다. Option token이란, 입력으로 들어가는 triple들을 indexing 했을 때 그 index를 의미한다. 예를 들어,

• A. <h1, r1, t1>
• B. <h2, r2, t2>

이 두 가지 triple이 있을 때, option token은 A 와 B 가 된다.

$$p(z_i \mid q, z_{<i}, \hat{G}_i) = \text{Softmax}(l(c_1), \ldots, l(c_K))$$

수식은 위와 같다. 여기서 \(c_1, \ldots, c_K\)는 option token이고, \(l(\cdot)\)는 해당 option token의 logit이다. 이 방식은 LLM이 새로운 triple을 자유롭게 생성하지 않고, KG에 grounded된 후보 중 하나만 선택하도록 제한한다. 따라서 hallucination 가능성을 줄이고 reasoning chain의 근거성을 높인다.

Chain Construction은 beam search를 이용해 여러 개의 reasoning chains를 생성하는 과정이다. TRACE는 먼저 모든 \(d_i \in D_q\)에 대해 \(G_{q, d_i} = \text{KG Generator}(d_i)\)를 생성하고, \(G_q\)에 추가한다. 이후 각 스텝마다 Triple Reranker를 통해 candidate tirple들을 얻고, Triple Selector가 triple distribution을 계산한다. 여기서 beam search 기반 decoding을 통해 \(R\)개의 reasoning chain을 생성한다.

3.4. Answer Generation

Answer Generation은 constructed reasoning chains를 직접 쓰거나, reasoning chains가 지목한 documents를 사용해 답을 생성하는 단계이다. TRACE는 두 가지 variant를 제안한다.

• TRACE-Triple은 reasoning chains 자체를 context로 reader에 넣는다. 이 경우 triples는 chain 순서대로 이미 정렬되어 있으므로 그대로 concatenate한다.
• TRACE-Doc은 reasoning chain 안의 각 triple이 생성된 원래 document에 vote를 부여하고, vote 수에 따라 \(D_q\)의 문서들을 순위화한다. Vote를 받지 못한 문서는 제거되며, 선택된 문서들만 reader context로 사용된다.

따라서 TRACE-Triple은 concise evidence representation을 직접 사용하는 방식이고, TRACE-Doc은 reasoning chain을 document selector로 사용하는 방식이다.

4. Experiments

RQ 1. How does TRACE perform against the baselines?

TRACE는 세 dataset 모두에서 vanilla RAG와 기존 baselines보다 높은 성능을 보인다. Table 1에서 TRACE-Triple과 TRACE-Doc은 w. all documents보다 평균 EM 기준 각각 14.03, 13.46 향상된다. 이는 retrieved documents 전체를 reader에 넣는 것보다, KG triples로 구성된 reasoning chains를 통해 supporting evidence만 선별하는 것이 multi-hop QA에 더 효과적임을 보여준다.

IRCoT와 비교해도 TRACE가 더 우수하다. IRCoT는 LLM이 생성한 CoT sentences를 기반으로 retrieval을 수행하므로 hallucination의 영향을 받을 수 있지만, TRACE의 reasoning chains는 retrieved documents에서 생성된 KG triples에 grounded되어 있다. 따라서 TRACE는 reasoning 과정의 신뢰성과 answer generation 성능을 동시에 개선한다.

Figure 3은 reader model을 LLaMA3, Mistral, Gemma로 바꾸어도 TRACE의 성능 향상이 유지됨을 보여준다. 즉 TRACE의 효과는 특정 reader에 종속된 것이 아니라, reasoning chain 자체가 일반적으로 유효한 evidence representation임을 의미한다.

RQ 2. What is the advantage of generating KGs from documents?

KG triples는 sentences나 documents보다 supporting evidence를 더 정밀하게 표현한다. Table 2에서 TRACE-Triple은 KG Generator를 제거한 w. sentences와 w. documents보다 일관되게 높은 성능을 보인다. 특히 documents를 그대로 reasoning unit으로 사용하는 경우 성능 저하가 더 크다. 이 결과는 KG triple이 하나의 factual relation만 담는 fine-grained representation이기 때문이다. 문장이나 문서는 여러 정보가 섞여 있어 irrelevant information을 포함하기 쉽지만, KG triple은 필요한 evidence를 더 명확하게 분리한다.

RQ 3. Can the reasoning chain constructor effectively identify supporting evidence?

Reasoning Chain Constructor는 단순히 top-ranked triples를 가져오는 방식보다 효과적이다. Table 2에서 constructor를 제거하고 E5-Mistral로 top-T triples만 선택하면 TRACE-Triple보다 성능이 크게 낮아진다. 이는 multi-hop QA에서 개별 triple의 relevance만 보는 것으로는 충분하지 않다는 것을 보여준다. TRACE는 이전에 선택된 \(z_{<i}\)를 조건으로 다음 triple \(z_i\)를 선택한다. 따라서 첫 번째 evidence가 다음 evidence를 찾는 clue가 되고, 여러 문서에 흩어진 supporting evidence가 논리적으로 연결된다.

5. Conclusion

Limitations

• [Wikipedia 중심 KG generation] TRACE는 Wikipedia documents의 title과 text 구조를 활용해 triples를 생성한다. 따라서 모든 entity pair 사이의 relation을 일반적으로 추출하는 방식은 다루지 않는다.
• [KG triples와 reasoning chains 자체의 직접 평가는 부족] 저자는 annotated data 부족으로 인해 생성된 KG와 reasoning chains를 직접 정량 평가하지 못하고, 최종 QA 성능과 qualitative examples로 간접 평가한다.
• [Triple Selector logits 접근이 필요] TRACE는 option token logits를 사용해 triple distribution을 정의한다. 따라서 black-box LLM을 사용할 경우 beam search 기반 chain construction이 제한될 수 있다.

Contributions

• TRACE는 RAG-MHQA를 위한 knowledge-grounded reasoning chain framework를 제안한다.
• Reasoning chains가 compact하고 효과적인 context임을 보인다. TRACE-Triple은 전체 문서보다 훨씬 짧은 context를 사용하면서도 높은 성능을 달성한다.
• TRACE는 vanilla RAG와 IRCoT보다 높은 성능을 보이며, 여러 reader model에서도 효과가 유지된다.

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