Meaningful

Meaningful

Greeting :)

[논문리뷰]TRACE: Constructing Knowledge-Grounded Reasoning Chains for Retrieval-Augmented Generation

Date:     Updated:

카테고리:

Jinyuan Fang, Zaiqiao Meng, and Craig MacDonald. “TRACE the Evidence: Constructing Knowledge-Grounded Reasoning Chains for Retrieval-Augmented Generation.” In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 8472–8494, Miami, Florida, USA.

Problem Statement

1

이 논문은 다중 문서에서 정보를 통합해 복잡한 질문에 답해야 하는 multi-hop question answering(MHQA) 문제를 해결함에 있어, 기존 RAG 모델들이 다음과 같은 한계점을 지적한다.

  • irrelevant한 문서를 검색하여 노이즈를 유입시키고, 이로 인해 추론 성능이 저하됨.
  • 특히 단순히 retriever가 가져온 문서를 LLM에 연결(concatenation)하는 기존 방식은, 다중 홉 추론에 필요한 근거(evidence)를 효과적으로 구성하지 못함.

이를 해결하기 위해 저자들은 문서들을 Knowledge Graph(KG)로 변환하고, 이 KG로부터 논리적으로 연결된 추론 체인(Reasoning Chain)을 구성하여, 보다 정제된 supporting evidence를 RAG 모델에 제공하는 방법인 TRACE를 제안한다.



Methodology

Overview

1

TRACE는 multi-hop 질문에 대한 정답을 생성하기 위해, 검색된 문서로부터 Knowledge Graph(KG)를 생성하고, 이 KG로부터 논리적으로 연결된 reasoning chain을 autoregressive하게 구성하며, 최종적으로 이 reasoning chain을 기반으로 정답을 생성하는 RAG 프레임워크이다. 전체 확률 모델은 다음과 같이 정의된다:

  • \(q\): 질문
  • \(D_q\): 질문에 대해 검색된 문서 집합
  • \(G_q\): \(D_q\)로부터 생성된 KG
  • \(z\): KG에서 생성된 추론 체인
  • \(a\): 생성된 정답

KG Generator

  • 입력: 검색된 문서들 (by Original Question, \(D_q\))
  • 출력: Knowledge Graph
    • Head: 문서의 title
    • Tail: Body (text에서 추출된 entity나 문장)

KG Generator는 검색된 문서 집합 \(D_q\)를 입력으로 받아, 각 문서에서 ⟨head, relation, tail⟩ 형태의 knowledge triple을 추출함으로써 KG \(G_q\)를 생성한다. 이때, 각 문서의 title을 head entity로 간주하고, 본문에서 출현하는 tail entity들과의 관계를 추출하여 triple을 구성한다. KG 생성은 LLM 기반 in-context prompting으로 수행되며, KG generator는 다음 확률식을 따르는 모듈이다:

\[p(G_q \vert D_q)\]

TRACE는 lost-in-the-middle 문제를 방지하기 위해, 모든 문서 \(d_i \in D_q\)에 대해 triple을 독립적으로 생성하고, 문서 간 triple 연결은 공통 entity를 통해 암묵적으로 연결한다. 즉, 긴 문서를 하나로 입력하는 방식 대신, 문서 단위로 triple을 추출하고 이를 합쳐 전체 KG를 구성한다:

\[G_q = \bigcup_{i=1}^{N} \text{KG Generator}(d_i)\]

문서가 Wikipedia 기반이기 때문에, title과 본문 간 관계를 추출하기 용이하며, 추출된 triple은 reasoning chain 구성의 후보로 사용된다. 이때의 prompt는 “title과 관련된 tail entity와 그 관계를 문장에서 추출하라”는 방식으로 구성된다.

Reasoning Chain Constructor

  • 입력: 질문 \(q\), KG \(G_q\)
  • 출력: Reasoning Chain \(z = [z_1, \dots, z_L]\)
  • reasoning chain을 autoregressive하게 구성하며, 각 단계마다 다음 triple을 선택한다.
  • 필요 없는 triple 생성을 방지하기 위해 adaptive chain termination 전략 도입 (선택지 A: “no need for additional triples”)

Reasoning Chain Constructor는 KG \(G_q\)를 기반으로 질문 \(q\)에 대한 reasoning chain \(z = [z_1, z_2, \dots, z_L]\)을 autoregressive하게 구성한다. 이 모듈은 \(p(z \vert q, G_q)\)의 근사로 작동하며, chain factorization은 다음과 같이 정의된다:

\[p(z \mid q, G_q) = \prod_{i=1}^{L} p(z_i \mid q, z_{<i}, \hat{G}_i)\]

여기서 \(\hat{G}_i\)는 i번째 step에서의 후보 triple 집합이다. 이 reasoning chain 생성을 위해 TRACE는 두 개의 핵심 서브모듈을 포함한다.

  • Triple Ranker
  • Triple Selector

Triple Ranker는 현재 질문과 이전 triple들을 합친 context를 bi-encoder 모델(E5-Mistral)을 통해 임베딩한 후, KG 내 모든 triple들과의 inner product를 계산하여 상위 K개의 triple을 후보군으로 선택한다.

이어서 Triple Selector는 LLaMA3-8B-Instruct 기반 LLM을 사용하여 multiple-choice 형식의 prompt로부터 다음 triple을 선택한다. 이때 “A. no need for additional triples”라는 특수 선택지를 포함시켜 reasoning chain의 종료 시점을 LLM이 스스로 판단하도록 하며, 이를 adaptive chain termination 전략이라 한다. 이 선택 과정은 softmax(logit)를 기반으로 다음과 같은 분포를 정의한다:

\[p(z_i \vert q, z_{<i}, \hat{G}_i) = \text{Softmax}(l(c_1), \dots, l(c_K))\]

전체 reasoning chain은 beam search를 통해 top-\(R\)개의 후보 chain을 구성하며, chain의 길이는 최대 \(L\)로 제한하거나 LLM의 종료 선택에 따라 조절된다. 이러한 autoregressive chain 구성 방식은 단순한 triple 수집이 아니라 논리적으로 연결된 근거를 단계적으로 추출하는 데에 핵심적인 역할을 수행한다.

Answer Generator

  • 입력: reasoning chain 혹은 이를 통해 선별된 문서
  • 출력: 정답

이 모듈은 구성된 reasoning chains를 활용하여 최종 답변을 생성한다. TRACE는 이 과정을 두 가지 방식으로 수행할 수 있다.

첫 번째 방식인 TRACE-Triple은 reasoning chain 내의 triple들을 그대로 context로 사용하여 정답을 생성하며, triple의 순서 자체가 추론의 논리 구조를 반영하므로 별도의 재정렬 없이 LLM에 그대로 입력된다.

두 번째 방식인 TRACE-Doc은 reasoning chain 내 각 triple이 생성된 원본 문서를 추적하고, 각 triple이 자신이 유래한 문서에 투표함으로써 득표 수 기준으로 문서를 정렬하여 context로 활용한다. 이 과정을 통해 noise가 포함된 전체 문서 대신 핵심 문서만을 사용하여 보다 정제된 context를 구성한다.

두 방식 모두 in-context learning 기반 LLM을 사용하며, 다음과 같은 조건부 확률 모델을 따른다

\[p(a \mid q, z, D_q)\]

실험 결과에 따르면, TRACE-Triple은 token 수가 상대적으로 적음에도 불구하고 기존 RAG 방식보다 높은 정확도를 기록하며, reasoning chain만으로도 충분한 정보를 제공할 수 있음을 입증한다.



Experiments

Main Results

1

HotPotQA, 2WikiMultiHopQA, MuSiQue 세 개 데이터셋에서 TRACE-Triple과 TRACE-Doc은 모든 데이터셋에서 최고 성능을 달성했다. Vanilla RAG 모델(w. all documents) 대비 TRACE-Triple과 TRACE-Doc은 평균적으로 EM 기준 각각 14.03%, 13.46%의 성능 향상을 보였다. 최고 성능 baseline인 IRCoT 대비해서도 TRACE-Triple과 TRACE-Doc은 평균적으로 EM 기준 각각 5.90%, 5.32%의 성능 향상을 달성했다. TRACE-Triple은 reasoning chains만을 context로 사용함에도 불구하고 2WikiMultiHopQA와 MuSiQue에서 최고 성능을, HotPotQA에서 두 번째 성능을 기록했다.

Ablation Study

1

KG Generator의 효과성 검증: TRACE-Triple 대비 sentences나 documents를 사용한 variants는 모든 데이터셋에서 성능이 저하되었다. 이는 KG triples이 더 fine-grained하고 concise한 지식 표현을 제공하여 irrelevant information의 포함을 최소화하기 때문이다.

Reasoning Chain Constructor의 필요성: Top-T triples만 사용하는 variant는 TRACE-Triple 대비 성능이 현저히 떨어졌다. 이는 autoregressive manner로 reasoning chains를 구성하는 것의 중요성을 보여주며, 이전에 식별된 supporting evidence가 후속 evidence 식별에 중요한 단서를 제공함을 시사한다.

각 컴포넌트의 기여도: Triple ranker 제거 시 모든 데이터셋에서 성능이 저하되었으며, 다른 ranker (DRAGON+, E5) 사용 시에도 E5-Mistral이 최고 성능을 보였다. Triple selector 제거 시에도 성능이 크게 저하되었으며, 다른 모델(Mistral, Gemma) 사용 시에도 LLaMA3이 최고 성능을 달성했다.

Analysis. Chain Length의 영향 분석

1

Maximum chain length L이 1에서 6으로 증가할 때 평균 chain length는 선형적으로 증가하지 않고 점진적으로 증가폭이 감소했다. 이는 adaptive chain termination strategy의 효과를 보여준다. 성능은 특정 threshold(HotPotQA에서 4) 이후 감소하는 경향을 보였는데, 이는 더 긴 reasoning chains에서 irrelevant하거나 redundant한 정보가 도입되어 reader를 혼란시키기 때문이다.

Case Study

1

TRACE가 구성한 reasoning chains는 multi-hop 질문에 대해 효과적인 추론 경로를 제공한다. 예를 들어 “When was the father of Albert Einstein born?” 질문에 대해 ⟨Albert Einstein, father, Hermann Einstein⟩, ⟨Hermann Einstein, date of birth, 3 July 1814⟩ 형태의 논리적 연결을 구성하여 정확한 답변을 도출한다. 이러한 reasoning chains는 모델이 답변을 생성하는 방식에 대한 해석 가능성을 제공한다.



Conclusion

Contributions

TRACE는 RAG 모델의 multi-hop reasoning 능력을 향상시키기 위해 knowledge-grounded reasoning chains를 구축하는 새로운 방법을 제안했다. Autoregressive 방식으로 reasoning chains를 구성하여 supporting evidence를 식별하고 통합하는 혁신적인 방법을 개발했다. 세 개의 multi-hop QA 데이터셋에서 vanilla RAG 대비 평균 14.03%의 EM 성능 향상을 달성했으며, reasoning chains가 다양한 reader 모델에서 효과적으로 일반화됨을 입증했다.

Limitations

이 연구는 Wikipedia에서 검색된 문서에 초점을 맞추어 KG generator가 주로 title과 text 내 entities 간의 관계만 생성한다는 제약이 있다. Annotated data 부족으로 생성된 KG triples와 reasoning chains의 품질을 직접적으로 정량평가할 수 없어 최종 QA 성능을 통해서만 평가했다는 한계가 있다. Reasoning chain constructor에서 triple selector의 logits에 접근이 필요한데, black-box LLM 사용 시에는 이것이 불가능할 수 있다는 실용적 제약도 존재한다.

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

댓글 남기기

최근 글 10 개 :)