[논문 리뷰] KGARevion: An AI Agent for Knowledge-Intensive Biomedical QA

Xiaorui Su, Yibo Wang외 5명|arXiv (Cornell University)|2024. 10. 07.

Quality and Safety in Healthcare인용 수 4

한 줄 요약

KGARevion은 지식 그래프(KG) 기반의 LLM 에이전트로, 기반 지식 그래프에 대응하여 의료 삼중항을 생성하고 검증하며 수정함으로써 지식 집약적인 생물의학적 질문 응답을 향상시킨다. 표준 데이터셋에서 정확도를 5.2% 이상 향상시키며, 새로운 복잡한 의료 QA 벤치마크에서는 10.4% 향상되어 다중 소스 의료 추론에서 강건성과 적응성을 입증한다.

ABSTRACT

Biomedical reasoning integrates structured, codified knowledge with tacit, experience-driven insights. Depending on the context, quantity, and nature of available evidence, researchers and clinicians use diverse strategies, including rule-based, prototype-based, and case-based reasoning. Effective medical AI models must handle this complexity while ensuring reliability and adaptability. We introduce KGARevion, a knowledge graph-based agent that answers knowledge-intensive questions. Upon receiving a query, KGARevion generates relevant triplets by leveraging the latent knowledge embedded in a large language model. It then verifies these triplets against a grounded knowledge graph, filtering out errors and retaining only accurate, contextually relevant information for the final answer. This multi-step process strengthens reasoning, adapts to different models of medical inference, and outperforms retrieval-augmented generation-based approaches that lack effective verification mechanisms. Evaluations on medical QA benchmarks show that KGARevion improves accuracy by over 5.2% over 15 models in handling complex medical queries. To further assess its effectiveness, we curated three new medical QA datasets with varying levels of semantic complexity, where KGARevion improved accuracy by 10.4%. The agent integrates with different LLMs and biomedical knowledge graphs for broad applicability across knowledge-intensive tasks. We evaluated KGARevion on AfriMed-QA, a newly introduced dataset focused on African healthcare, demonstrating its strong zero-shot generalization to underrepresented medical contexts.

연구 동기 및 목표

• 다양한 전략(유사 추론, 규칙 기반 추론, 사례 기반 추론 등)을 필요로 하는 지식 집약적인 생물의학적 추론의 고유한 과제를 해결한다.
• LLM의 생물의학 QA에서의 한계, 즉 환각 현상, 열악한 검색 성능, 구조화된 의료 지식에 대한 기반 부족 문제를 극복한다.
• 비코딩된 LLM 지식과 의료 KG에서 유래한 구조화된 지식을 통합하여 추론 정확도와 신뢰성을 향상시키는 시스템을 개발한다.
• 다중 선택 QA 환경에서 흔히 발생하는 답변 순서 및 인덱싱 편향에 대비한 강건성을 확보한다.
• 일반화 능력과 난이도 증가 조건 하에서의 성능 평가를 위해 표준 및 새로 총집된 의미론적으로 복잡한 의료 QA 벤치마크를 모두 평가한다.

제안 방법

• 의료 질문을 체계적으로 처리하기 위해 4단계 에이전트 프레임워크(생성, 검색, 검토, 수정)를 사용한다.
• 입력 질문과 KG 임beddings를 프롬프트로 사용하여 미세조정된 LLM을 통해 후보 의료 삼중항(주어-서술어-목적어)을 생성한다.
• 사전에 훈련된 구조적 임베딩(예: OGB-biokg 또는 PrimeKG에서 유래)을 기반으로 벡터 검색을 통해 관련 KG 삼중항을 검색하여 검색을 지원한다.
• 기반 의료 지식 그래프와의 교차 확인을 통해 생성된 삼중항을 검토하고 검증하여 환각되거나 잘못된 사실을 걸러낸다.
• KG 검증 단계에서의 피드백을 바탕으로 삼중항을 반복적으로 정밀 조정함으로써 추론 경로를 수정한다.
• KG를 검색 소스로 사용하는 것이 아니라 사실 검증 도구로 활용하여, 최종 답변에 기여하는 것은 오직 의학적으로 정확하고 관련성이 있는 삼중항만을 보장한다.

Figure 1: a) Performance of LLMs on three new datasets introduced in this paper with questions of varying difficulty. b) Sample questions from new MedDDx-Basic, MedDDx-Intermediate, and MedDDx-Expert datasets.

실험 결과

연구 질문

• RQ1 표준 RAG 또는 순수 LLM 기반 베이스라인 대비 KG 기반 에이전트가 복잡한 지식 집약적인 생물의학적 질문 응답에서 LLM의 정확도를 향상시킬 수 있는가?
• RQ2 구조화된 KG 검증 통합이 LLM 환각 현상과 검색 오류가 존재하는 상황에서 추론의 강건성을 어떻게 향상시키는가?
• RQ3 KGARevion이 의미론적으로 복잡한 의료 QA, 특히 전문가 수준의 질문에서도 성능을 유지하는 정도는 어느 정도인가?
• RQ4 표준 LLM에 비해 KGARevion은 다중 선택 QA 환경에서 답변 순서 및 레이블링에 얼마나 민감한가?
• RQ5 다양한 KG와 LLM 백본을 통해 일반화할 수 있는가? 이 경우에도 높은 정확도와 신뢰성을 유지하는가?

주요 결과

• KGARevion은 15개의 기준 모델 대비 네 개의 골드 표준 의료 QA 데이터셋에서 정확도를 5.2% 이상 향상시켜 복잡한 의료 질문에서 뛰어난 성능을 입증한다.
• 새로 총집된 고복잡도 의료 QA 데이터셋 세 개(MedDDx-Basic, MedDDx-Intermediate, MedDDx-Expert)에서 KGARevion은 기준 LLM 대비 10.4%의 절대 정확도 향상을 달성한다.
• KGARevion은 다중 선택 설정에서 답변 순서 및 인덱싱에 대해 강건하며, LLaMA3-8B와 LLaMA3.1-8B의 경우 각각 2.59%와 3.86%의 정확도 손실만을 보이며, 순수 LLM은 각각 8.4%와 16.0%의 손실을 보인다.
• OGB-biokg와 같은 더 작은 KG를 사용할 때도 KGARevion의 성능은 안정적이고 높은 수준을 유지하며, 더 큰 KG인 PrimeKG에 의존하는 모델보다 뛰어난 성능을 보여, 검증 메커니즘의 효과성을 입증한다.
• KGARevion이 생성한 추론 경로는 설명 가능하고 기반 지식에 기반한 것으로, 검증된 삼중항들이 체계적이고 추적 가능한 의학적 논리 체인을 형성한다. 예를 들어 염색체 결손과 특정 증후군 간의 연결을 구축할 수 있다.
• KGARevion의 설계는 사실 기반(구조화된) 지식과 암묵적(비코딩된) 지식의 체계적 통합을 가능하게 하여, 미세한 의료적 차이와 의미적으로 유사하지만 임상적으로 다른 개념을 처리할 수 있다.

Figure 2: a) The overview of KGARevion . b ) The architecture of fine-tuning stage in the Review action, where embeddings get from KGs are structural embeddings, while concept embeddings from LLMs.