GNN-RAG: Knowledge Graph 위에서 효율적인 LLM 추론을 위한 Graph Neural Network 기반 RAG

GNN-RAG: Graph Neural Retrieval for Efficient Large Language Model Reasoning on Knowledge Graphs

Jan 1, 2025•Costas Mavromatis, G. Karypis•View PDF

TL;DR Highlight

GNN으로 Knowledge Graph에서 관련 경로를 먼저 추려낸 뒤 LLM에 넘기면 KGQA 정확도와 속도가 모두 올라간다.

Who Should Read

Knowledge Graph 기반 QA 시스템을 구축 중이거나, RAG 파이프라인에서 그래프 구조 데이터를 활용하고 싶은 백엔드·ML 개발자. Wikidata, Freebase 같은 대규모 그래프 DB를 LLM과 연결하는 방법을 고민하는 사람.

Core Mechanics

Knowledge Graph(KG, 엔티티와 관계로 연결된 그래프 DB)에서 RAG를 할 때 모든 경로를 LLM에 던지면 토큰 폭발 → GNN이 먼저 관련 서브그래프만 추려서 LLM 입력을 압축
GNN(그래프 신경망)을 리트리버로 써서 질문과 관련 높은 엔티티·관계 경로를 점수화해 Top-K만 선택
LLM은 압축된 자연어 경로만 받아서 추론하므로, 전체 그래프를 직접 탐색하는 방식보다 훨씬 적은 컨텍스트로 정확한 답 생성 가능
GNN 리트리버와 LLM 리즈너를 분리해서 GNN은 가볍게 파인튜닝, LLM은 프롬프트만으로도 동작 → 비용 효율적
멀티홉 추론(여러 관계를 건너야 답이 나오는 질문)에서 기존 임베딩 기반 RAG 대비 명확한 강점

Evidence

WebQSP 벤치마크에서 Hits@1 기준 최신 KGQA 모델 대비 경쟁력 있는 수치 달성 (논문 내 Table 기준 상위권)
GNN 리트리버가 평균 KG 경로 후보를 수천 개 → 수십 개로 압축해 LLM 입력 토큰 수를 대폭 절감
CWQ(ComplexWebQuestions) 데이터셋에서도 멀티홉 질문 처리 성능이 기존 RAG 베이스라인 대비 유의미하게 향상
GPT-3.5/GPT-4 등 다양한 LLM 백엔드와 결합 시 일관된 성능 향상 확인

How to Apply

Freebase/Wikidata 연동 QA 서비스라면, 질문 입력 시 GNN 리트리버로 관련 서브그래프 경로만 추출한 뒤 '엔티티A → 관계 → 엔티티B' 형태 텍스트로 변환해 LLM 프롬프트에 삽입하는 파이프라인 구성
그래프 DB가 너무 커서 LLM에 직접 넣기 어려운 경우, 가벨한 GNN 모델(예: 2-3레이어 RGCN)을 리트리버로 파인튜닝해 Top-K 경로만 뽑는 전처리 단계를 RAG 체인 앞에 추가
멀티홉 질문(예: '영화 X의 감독이 태어난 도시의 시장은?')을 처리해야 한다면, 단순 벡터 유사도 검색 대신 GNN 기반 경로 스코어링으로 교체해볼 것

Code Example

snippet

# GNN-RAG 파이프라인 개념 코드 (PyG + LangChain 스타일)
import torch
from torch_geometric.nn import RGCNConv

# 1. GNN 리트리버: 질문 엔티티 주변 서브그래프에서 관련 경로 추출
class GNNRetriever(torch.nn.Module):
    def __init__(self, in_channels, hidden, num_relations):
        super().__init__()
        self.conv1 = RGCNConv(in_channels, hidden, num_relations)
        self.conv2 = RGCNConv(hidden, hidden, num_relations)
        self.score = torch.nn.Linear(hidden, 1)

    def forward(self, x, edge_index, edge_type):
        x = self.conv1(x, edge_index, edge_type).relu()
        x = self.conv2(x, edge_index, edge_type).relu()
        return self.score(x).squeeze(-1)  # 각 노드의 관련도 점수

# 2. 점수 높은 경로를 자연어로 변환
def paths_to_text(top_k_paths):
    """[(entity1, relation, entity2), ...] -> str"""
    lines = [f"{e1} --[{rel}]--> {e2}" for e1, rel, e2 in top_k_paths]
    return "\n".join(lines)

# 3. LLM 프롬프트에 경로 삽입
def build_prompt(question, kg_paths_text):
    return f"""다음 Knowledge Graph 경로를 참고해서 질문에 답하세요.

[KG 경로]
{kg_paths_text}

[질문]
{question}

[답변]"""

# 사용 예
# paths = gnn_retriever.get_top_k_paths(question_entity, k=20)
# prompt = build_prompt("영화 Inception의 감독 국적은?", paths_to_text(paths))
# answer = llm(prompt)

Terminology

Knowledge Graph엔티티(사람, 장소, 개념)와 그 사이의 관계를 노드-엣지 그래프로 표현한 DB. 예: '스티브 잡스 --[설립]--> 애플'처럼 사실을 구조화해서 저장.

GNN그래프 구조 데이터를 학습하는 신경망. 각 노드가 이웃 노드 정보를 모아서 자신의 표현을 업데이트하는 방식으로 동작.

KGQAKnowledge Graph Question Answering의 약자. KG에서 정보를 찾아 질문에 답하는 태스크.

RGCNRelational Graph Convolutional Network. 엣지(관계)의 종류가 여러 개인 그래프를 처리할 수 있는 GNN 변형.

멀티홉 추론한 번에 답을 찾는 게 아니라 여러 관계를 연쇄적으로 따라가야 답이 나오는 추론. '영화 감독의 출신 학교 설립자는?' 같은 질문.

서브그래프전체 거대한 그래프에서 특정 목적에 맞는 일부 노드와 엣지만 잘라낸 부분 그래프.

RAGRetrieval-Augmented Generation. 외부 DB에서 관련 정보를 먼저 검색한 뒤 그걸 컨텍스트로 넣어서 LLM이 답하게 하는 방식.