Maximizing Financial LLM RAG Performance: Knowledge Graph and GNN-based Context Retrieval Optimization Strategy
Traditional RAG (Retrieval Augmented Generation) shows limitations in grasping the subtle relationships within complex financial data. This article presents an in-depth strategy to innovatively enhance the context retrieval accuracy and explainability of LLMs by leveraging Knowledge Graphs and Graph Neural Networks (GNNs). This approach will be a game-changer, helping financial LLMs move beyond simple information retrieval to generate more reliable and insightful answers based on relational knowledge.
1. The Challenge / Context
The utilization of LLMs (Large Language Models) in the financial domain is rapidly increasing across various fields such as regulatory compliance, market analysis, and customer service. However, for LLMs to operate effectively, a RAG system that provides accurate and reliable information is essential. Existing RAG systems primarily rely on Vector Similarity Search. This method finds relevant information based on the semantic similarity between documents or chunks and provides it to the LLM.
However, financial data possesses complex relationships and hierarchical structures that are difficult to grasp solely through the semantic similarity of words. For example, the financial status of a specific company is intricately intertwined with numerous factors such as affiliate relationships, market regulations, macroeconomic events, and even statements made by specific individuals. Vector search cannot directly model these 'relationships', leading to the following problems:
- Low Accuracy and Hallucination: The lack of relational information increases the likelihood of LLMs making incorrect inferences or generating non-existent information (hallucination). In the financial sector, even minor errors can lead to critical consequences.
- Limited Explainability: It is difficult to clearly present why specific information was used in an answer and what relationships that information holds. This poses a serious obstacle to regulatory compliance and internal audits.
- Limitations in Processing Complex Queries: Traditional RAG is vulnerable to multi-entity, multi-relationship queries such as "When Company A acquired Company B, what impact did Regulation C have, and what were the ripple effects on Market D?"
These issues undermine the reliability of financial LLMs and make their application in real-world tasks difficult. We need a new approach that can provide 'structured knowledge' as context, going beyond mere 'information fragments'. Knowledge Graphs and GNN-based context retrieval are the answer.
2. Deep Dive: Knowledge Graph and GNN-based Context Retrieval
Knowledge Graphs and GNNs are a powerful combination to overcome the fundamental limitations of financial LLM RAG. Let's examine how they work in detail.
2.1. Knowledge Graph (KG)
A Knowledge Graph is a structured data model that represents real-world entities and their relationships as nodes and edges. For example, in a financial knowledge graph, 'Samsung Electronics (node)' and 'ARM (node)' can be connected by a 'takeover attempt (edge)' relationship, and the 'takeover attempt (node)' can have additional relationships such as 'occurred in March 2023 (edge)' and 'impacted the semiconductor market (edge)'.
Value in the Financial Domain:
- Explicit Relationship Modeling: Complex financial relationships such as corporate M&A, regulatory compliance, market trends, investment portfolios, and risk factors can be clearly defined and stored.
- Context Enrichment: Beyond simply extracting words within text, it provides rich context through 'hidden' or 'inferable' relationships between entities.
- Data Integration: Integrates information extracted from structured data (financial statements) and unstructured data (news articles, reports, regulatory documents) to build a consistent knowledge base.
2.2. Graph Neural Networks (GNNs)
GNNs are deep learning models specialized for graph-structured data. They learn embeddings for each node by using both node and edge information. These embeddings compress not only the features of the node but also the structural and semantic information of surrounding nodes and relationships.
Impact on Context Retrieval Optimization:
- Relationship-based Embedding: GNNs learn not only the attributes of nodes but also how nodes are connected. This is essential for understanding the context in which a specific entity is important in financial queries (e.g., which companies it partnered with, which regulations it is subject to).
- Subgraph Search: Using the embeddings learned by GNNs, it can efficiently search for a 'subgraph' or 'knowledge path' that best explains the intent of a query, going beyond finding a single node most similar to the user query.
- Ranking and Filtering: Information elements within the retrieved subgraph are ranked according to their relevance to the query, and the most important parts are prioritized as LLM context, reducing unnecessary information and highlighting key insights.
In conclusion, the Knowledge Graph structures complex relationships in the financial domain, and GNNs intelligently explore and extract the most suitable knowledge fragments for a query based on these structured relationships, thereby maximizing the performance of the RAG system.
3. Step-by-Step Guide / Implementation
Now, let's look at the specific implementation steps to optimize financial LLM RAG based on Knowledge Graphs and GNNs. This process covers everything from data preprocessing to GNN model training and RAG system integration.
Step 1: Financial Data Source Analysis and Entity-Relationship Extraction (NER/RE)
The first step is to identify and extract meaningful entities and their relationships from unstructured financial text. This forms the basis for building the Knowledge Graph.
Description:
- Define Entity Types: Define entity types specific to the financial domain, such as 'Company', 'Person', 'Financial Product', 'Regulation', 'Market', 'Date', 'Value'.
- Define Relationship Types: Define meaningful relationships between entities, such as 'Acquisition/Merger', 'Investment', 'Subject to Regulation', 'Impacts', 'Announces', 'Related to'.
- NER (Named Entity Recognition): Identify words or phrases of defined entity types in the text. Pre-trained models (e.g., Spacy, Stanza) can be used, or accuracy can be improved by leveraging BERT/RoBERTa-based models fine-tuned with financial domain data.
- RE (Relationship Extraction): Classify relationships between identified entity pairs. This can use rule-based, statistical, or deep learning-based methodologies (e.g., OpenNRE, GCNs for RE).
import spacy
from spacy.tokens import Span
from spacy.language import Language
# Step 1.1: 금융 도메인에 특화된 NER 모델 로드 또는 파인튜닝
# 여기서는 예시로 일반 한국어 모델을 사용하며, 실제로는 금융 특화 모델이 필요합니다.
try:
nlp = spacy.load("ko_core_news_sm")
except OSError:
print("Downloading 'ko_core_news_sm' model. This may take a moment...")
spacy.cli.download("ko_core_news_sm")
nlp = spacy.load("ko_core_news_sm")
text = "2023년 3월, 삼성전자는 ARM 인수를 검토했으나 협상 결렬로 이어졌다. 이는 반도체 시장에 큰 영향을 미칠 수 있다. 금융위원회는 해당 사안을 예의주시했다."
doc = nlp(text)
print("--- 엔티티 추출 예시 ---")
# SpaCy의 기본 엔티티와 함께, 금융 도메인에 특화된 엔티티 추가 정의가 필요
# 예: '금융위원회'를 'ORGANIZATION' 대신 'REGULATOR'로 분류
# 커스텀 NER 컴포넌트 추가 예시 (매우 간소화된 형태)
@Language.component("financial_ner_enhancer")
def financial_ner_enhancer(doc):
if "금융위원회" in doc.text:
start = doc.text.find("금융위원회")
end = start + len("금융위원회")
# 기존 엔티티와 겹치지 않는지 확인 후 추가
if not any(ent.start_char <= start < ent.end_char or
ent.start_char < end <= ent.end_char for ent in doc.ents):
financial_regulator_ent = Span(doc, doc.char_span(start, end).start, doc.char_span(start, end).end, label="REGULATOR")
doc.ents = list(doc.ents) + [financial_regulator_ent]
return doc
nlp.add_pipe("financial_ner_enhancer", after="ner")
doc = nlp(text) # 파이프라인 변경 후 다시 처리
extracted_entities = []
for ent in doc.ents:
extracted_entities.append((ent.text, ent.label_))
print("Detected Entities:", extracted_entities)
print("\n--- 관계 추출 개념 예시 ---")
# 실제 관계 추출은 텍스트 파싱, 의존성 구문 분석, 또는 딥러닝 모델이 필요합니다.
# 여기서는 추출된 엔티티를 기반으로 간단한 규칙을 가정합니다.
# (엔티티1, 관계_타입, 엔티티2) 형태로 관계를 표현
relationships = []
companies = [ent.text for ent in doc.ents if ent.label_ == "ORG"] # ORG는 회사로 간주
dates = [ent.text for ent in doc.ents if ent.label_ == "DATE"]
regulators = [ent.text for ent in doc.ents if ent.label_ == "REGULATOR"]
if "삼성전자" in companies and "ARM" in companies and "인수" in text:
relationships.append(("삼성전자", "인수_검토", "ARM"))
if "ARM 인수" in text and "반도체 시장" in text and "영향" in text:
relationships.append(("ARM 인수", "영향_미침", "반도체 시장"))
if "금융위원회" in regulators and ("삼성전자" in companies or "ARM" in companies):
relationships.append(("금융위원회", "주시함", "인수_사안"))
print("Extracted Relationships:", relationships)
Step 2: Knowledge Graph Construction and Storage (Graph Database)
The extracted entities and relationships are converted into an actual Knowledge Graph and stored in a graph database. Neo4j, ArangoDB, and Amazon Neptune are widely used.
Description:
- Schema Definition: Define node labels (e.g., Company, Regulation, Event) and relationship types (e.g., ATTEMPTED_ACQUISITION, REGULATED_BY, HAS_IMPACT). Each node has a unique ID and properties (e.g., name, date, value).
- Data Insertion: Convert extracted entities into nodes and relationships into edges, then insert them into the graph database. This uses graph query languages like Cypher (Neo4j) or Gremlin (Apache TinkerPop compatible DB).
// Neo4j Cypher 예시: 노드 및 관계 생성
// Step 1에서 추출된 정보들을 바탕으로 그래프를 구축합니다.
// 1. 엔티티 노드 생성
CREATE (c1:Company {name: '삼성전자', industry: 'Electronics'})
CREATE (c2:Company {name: 'ARM', industry: 'Semiconductor IP'})
CREATE (e1:Event {type: 'AcquisitionAttempt', description: 'ARM 인수 검토', date: '2023-03'})
CREATE (m1:Market {name: '반도체 시장', segment: 'Global Semiconductor'})
CREATE (r1:Regulator {name: '금융위원회', country: 'South Korea'})
CREATE (st1:Status {name: '협상 결렬', outcome: 'Failed'})
// 2. 관계 엣지 생성
// 삼성전자가 ARM 인수를 검토함
MERGE (c1)-[:ATTEMPTED_TO_ACQUIRE {date: '2023-03'}]->(c2)
// 인수 검토 이벤트
MERGE (c1)-[:INITIATED_EVENT]->(e1)
MERGE (e1)-[:RELATED_TO]->(c2)
// 협상 결렬로 이어짐
MERGE (e1)-[:RESULTED_IN]->(st1)
// 반도체 시장에 영향을 미침
MERGE (e1)-[:HAS_IMPACT {magnitude: 'High'}]->(m1)
// 금융위원회가 해당 사안을 주시함
MERGE (r1)-[:MONITORS {scope: 'Acquisition Activity'}]->(e1)
// 3. 지식 그래프 조회 예시
// 삼성전자와 관련된 2-hop 이내의 모든 정보 검색
MATCH (s:Company {name: '삼성전자'})-[r*1..2]-(o)
RETURN s, r, o
// 반도체 시장에 영향을 미친 이벤트 검색
MATCH (e:Event)-[:HAS_IMPACT]->(m:Market {name: '반도체 시장'})
RETURN e.description, e.date
Step 3: GNN Model Training and Embedding Generation
Based on the constructed Knowledge Graph, a GNN model is trained to generate embeddings for each node (entity) and edge (relationship). These embeddings map the structural and semantic information of the graph into a high-dimensional vector space.
Description:
- GNN Model Selection: Choose a model suitable for the data characteristics and objectives from various GNN architectures such as GraphSAGE, GCN, GAT. Knowledge Graph Embedding (KGE) models like TransE and DistMult are also effective for relational knowledge modeling.
- Feature Vector Initialization: The initial feature vector for each node can be the entity's name embedding (Word2Vec, FastText, Sentence-BERT) or property values.
- Learning Objective: GNNs can be trained through tasks such as node classification, link prediction, or graph classification. For RAG, the primary goal is to improve the quality of node embeddings. This is achieved by training to ensure adjacent nodes have similar embeddings or to predict specific relationships.
import torch
from torch_geometric.nn import SAGEConv # GraphSAGE 레이어 사용
from torch_geometric.data import Data
import numpy as np
# 예시 그래프 데이터 생성 (실제 KG 데이터는 훨씬 복잡)
# num_nodes: 그래프 내 노드 수
# num_node_features: 각 노드의 초기 특징 벡터 차원 (예: 엔티티 이름의 임베딩 차원)
# num_edges: 엣지 수
num_nodes = 100
num_node_features = 768 # Sentence-BERT 임베딩 차원 가정
num_edges = 400
# 각 노드의 초기 특징 벡터 (예: 엔티티 이름에 대한 Sentence-BERT 임베딩)
x = torch.randn(num_nodes, num_node_features)
# 엣지 리스트 (COO 형식: [ [source_nodes], [target_nodes] ])
edge_index = torch.randint(0, num_nodes, (2, num_edges))
data = Data(x=x, edge_index=edge_index)
# GraphSAGE 모델 정의
# in_channels: 입력 노드 피처 차원
# hidden_channels: 히든 레이어 차원
# out_channels: 출력 임베딩 차원 (GNN이 최종적으로 생성할 임베딩 차원)
class GNNEmbedder(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels):
super().__init__()
# GraphSAGE의 Convolution 레이어
self.conv1 = SAGEConv(in_channels, hidden_channels)
self.conv2 = SAGEConv(hidden_channels, out_channels)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index).relu() # 첫 번째 레이어 통과 후 ReLU 활성화 함수 적용
x = self.conv2(x, edge_index) # 두 번째 레이어 통과
return x
embedding_dim = 128 # 최종 노드 임베딩 차원
model = GNNEmbedder(num_node_features, 256, embedding_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# GNN 모델 학습 과정 (실제로는 손실 함수 정의 및 에폭 반복)
# 학습 목표는 노드 분류, 링크 예측 또는 Contrastive Learning 등이 될 수 있습니다.
# 여기서는 임베딩 생성만을 위한 개념 코드입니다.
# model.train()
# for epoch in range(num_epochs):
# optimizer.zero_grad()
# out = model(data.x, data.edge_index)
# loss = ... # 예를 들어, 노드 분류 손실 또는 링크 예측 손실
# loss.backward()
# optimizer.step()
# 학습된 모델을 사용하여 노드 임베딩 생성
model.eval()
with torch.no_grad():
node_embeddings = model(data.x, data.edge_index)
print("Generated GNN Node Embeddings shape:", node_embeddings.shape)
# 이 node_embeddings는 이제 지식 그래프의 각 노드(엔티티)의 풍부한 컨텍스트 정보를 담고 있습니다.
Step 4: GNN-based Context Retrieval and Ranking (RAG Integration)
When a user query comes in, the most relevant subgraph or knowledge path is retrieved from the Knowledge Graph based on this query, and then structured as context to be passed to the LLM.
Description:
- Query Embedding: Convert the user query into a vector using an embedding model like Sentence-BERT.
- Initial Node Search: Calculate the similarity (e.g., cosine similarity) between the query embedding and the node embeddings generated by the GNN to identify the most relevant initial node(s) in the Knowledge Graph.
- Graph Traversal and Subgraph Extraction: Starting from the identified initial node(s), traverse the graph using the relational information learned by the GNN. This traversal can involve finding adjacent nodes within a few hops or finding paths that follow specific relationship patterns. The relevance of the traversal path can be ranked using the GNN's node/edge embeddings.
- Context Construction: Serialize the extracted subgraph or knowledge path into a text format that the LLM can understand, forming the final context. Template-based explanation generation or natural language conversion techniques can be utilized in this process.
- LLM Integration: Pass the constructed context along with the user query to the LLM to generate an answer.
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
# 가정: Step 3에서 생성된 node_embeddings가 존재합니다.
# node_embeddings_np = node_embeddings.cpu().numpy()
# 예시를 위해 임의의 노드 임베딩과 노드 이름 생성
num_kg_nodes = 100
embedding_dim = 128 # GNN output dimension
node_embeddings_np = np.random.rand(num_kg_nodes, embedding_dim)
node_names = [f"node_{i}" for i in range(num_kg_nodes)]
# 실제로는 node_names는 '삼성전자', 'ARM', '인수_검토_이벤트' 등 실제 엔티티 이름이어야 함.
node_names[0] = '삼성전자'
node_names[1] = 'ARM'
node_names[2] = '인수_검토_이벤트'
node_names[3] = '반도체 시장'
node_names[4] = '금융위원회'
# 1. 쿼리 임베딩 생성
query_text = "삼성전자의 ARM 인수 시도가 반도체 시장에 미친 영향과 금융위원회의 입장은 무엇인가요?"
embedding_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
query_embedding = embedding_model.encode(query_text, convert_to_tensor=True).cpu().numpy().reshape(1, -1)
# 2. KG 노드 임베딩과 쿼리 임베딩 비교 (가장 유사한 시작 노드 찾기)
similarities = cosine_similarity(query_embedding, node_embeddings_np)
most_similar_node_idx = np.argmax(similarities)
starting_node_name = node_names[most_similar_node_idx]
print(f"Query: '{query_text}'")
print(f"Most similar starting node in KG: '{starting_node_name}' (Similarity: {similarities[0, most_similar_node_idx]:.4f})")
# 3. 지식 그래프 탐색 및 관련 서브그래프/경로 추출 (가상의 함수)
# 실제로는 Neo4j, DGL, PyG API를 사용하여 그래프 탐색 쿼리를 실행합니다.
# 예를 들어, Cypher 쿼리:
# MATCH (n {name: '삼성전자'})-[r*1..2]-(m) WHERE type(r) IN ['ATTEMPTED_TO_ACQUIRE', 'HAS_IMPACT', 'MONITORS'] RETURN n, r, m
# 여기서는 설명을 위해 간단한 텍스트 리스트를 반환하는 함수로 대체합니다.
def retrieve_relevant_subgraph_info(start_node_name, kg_data_mock, k_hops=2):
# 이 함수는 실제 Cypher 쿼리 결과 또는 DGL/PyG 그래프 탐색 결과를 모방합니다.
retrieved_info = []
if "삼성전자" in start_node_name or "ARM" in start_node_name or "인수_검토_이벤트" in start_node_name:
retrieved_info.extend([
"엔티티: 삼성전자 (회사)",
"엔티티: ARM (회사)",
"관계: 삼성전자는 ARM 인수를 검토함 (2023년 3월)",
"관계: 인수 검토는 '협상 결렬' 상태로 이어짐",
"관계: ARM 인수 시도는 '반도체 시장'에 '높은 영향'을 미침",
])
if "금융위원회" in start_node_name or "인수_사안" in start_node_name:
retrieved_info.extend([
"엔티티: 금융위원회 (규제기관)",
"관계: 금융위원회는 '인수_검토_이벤트'를 '주시함'",
"금융위원회는 주요 기업의 인수합병 동향을 면밀히 모니터링합니다."
])
# 중복 제거 및 정렬
return sorted(list(set(retrieved_info)))
# 가상의 KG 데이터 (실제로는 Graph DB 연결 객체)
mock_kg_data = {"nodes": node_names, "embeddings": node_embeddings_np}
retrieved_context_elements = retrieve_relevant_subgraph_info(starting_node_name, mock_kg_data, k_hops=2)
# 4. LLM에 전달할 최종 컨텍스트 구성
context_for_llm = "다음은 사용자 질의와 관련된 지식 그래프 정보입니다:\n" + "\n".join(retrieved_context_elements)
print("\n--- Context for LLM ---")
print(context_for_llm)
# 5. LLM 호출 (예시)
# 실제 LLM 호출은 OpenAI, Anthropic, 또는 로컬 LLM API를 사용합니다.
# from openai import OpenAI
# client = OpenAI()
# response = client.chat.completions.create(
# model="gpt-4",
# messages=[
# {"role": "system", "content": "You are a financial expert. Answer questions based on the provided context accurately and concisely."},
# {"role": "user", "content": f"Context: {context_for_llm}\n\nQuestion: {query_text}"}
# ]
# )
# print("\n--- LLM's Answer (simulated) ---")
# print(response.choices[0].message.content)
print("\n--- LLM's Answer (simulated for demonstration) ---")
print("시뮬레이션 답변: 삼성전자의 ARM 인수 시도는 2023년 3월에 검토되었으나, 협상 결렬로 이어졌습니다. 이 시도는 반도체 시장에 높은 영향을 미칠 것으로 예상되었으며, 금융위원회는 이러한 주요 인수합병 동향을 면밀히 주시하고 있었습니다. 이는 시장의 안정성과 공정성을 유지하기 위한 금융위원회의 일반적인 역할과 일치합니다.")
4. Real-world Use Case / Example
Over the past 10 years, I have participated in various financial technology projects and have witnessed multiple times how data complexity undermines the reliability of AI systems. Especially in the financial sector, where regulations are strict and accuracy is paramount, it's not enough to simply 'provide an answer'; one must be able to explain 'why that answer was given'. Knowledge Graph and GNN-based RAG offer overwhelming advantages precisely in this aspect of Explainability. While traditional vector search merely states "based on similar documents," KG+GNN can present clear knowledge paths like "because Company X is subject to Regulation Y, and a risk occurred due to Event Z." I am convinced that this is one of the ultimate goals of financial domain AI.
Use Case: Customized Investment Risk Analysis and Regulatory Compliance Review
Scenario: An asset management company is building a new AI-based system to automatically generate customized investment product risk analysis reports that comprehensively consider a client's existing portfolio, market trends, and the latest financial regulations.
Problems with Existing Methods: If a client's portfolio, investment product descriptions, and regulatory documents are fed into an LLM as text, the LLM attempts to process all information as a single context. In this process, it may miss key relationships (e.g., how a specific product correlates with a particular asset in the client's portfolio, or which regulations it violates) or focus on irrelevant information, generating inaccurate or generic reports. Explainability is also low, making additional expert review essential.
Application of KG+GNN based RAG:
- Knowledge Graph Construction:
- Client Node: A node with attributes such as client's investment propensity, existing assets, and risk tolerance.
- Investment Product Node: A node with attributes such as the product's constituent assets, operating strategy, fees, and past returns.
- Company Node: Asset management company, target investment companies, etc.
- Market Node: Stock market, bond market, specific industry sectors, etc.
- Regulation Node: Financial Investment Business Regulations, Capital Markets Act, etc.
- Relationship Edges: 'Holds (Client-Asset)', 'Includes (Product-Asset)', 'Manages (Asset Manager-Product)', 'Impacts (Market-Asset)', 'Regulates (Regulation-Product)', etc. Add recent news/event nodes and 'Related to' edges.
- GNN Model Training: Train a GNN on the constructed Knowledge Graph to generate high-quality embeddings for each node (client, product, asset, regulation, etc.).
- Query and GNN-RAG Operation:
- User Query: "If 'New Growth Engine Investment Fund' is added to Kim Cheol-su's portfolio, what are the expected risks and regulations to comply with?"
- GNN-RAG Process:
- Generate a query embedding and identify relevant initial nodes in the Knowledge Graph, such as 'Kim Cheol-su client' and 'New Growth Engine Investment Fund'.
- Through GNN-based graph traversal, extract the most relevant subgraph within 2-3 hops from these nodes. For example, it discovers a knowledge path indicating that a specific technology company invested by the 'New Growth Engine Investment Fund' is subject to 'export regulations' due to 'recent international disputes', and this regulation requires a 'high-risk investment warning' according to 'Article X of the Capital Markets Act'.
- The GNN can also identify a relationship where the fund's main investment sector overlaps with a sector already heavily weighted in 'Kim Cheol-su client's' existing portfolio, increasing portfolio concentration risk.
- LLM Answer and Report Generation:
- The LLM receives the extracted subgraph (in serialized text format) as context and generates a customized risk analysis report for Kim Cheol-su.
- Report Content: "The New Growth Engine Investment Fund has high exposure to Company A, a technology company, which faces export regulation risks due to recent international trade disputes. According to Article X of the Capital Markets Act, this matter requires a clear high-risk investment warning to the client. Furthermore, since your current portfolio is already concentrated in the technology sector, adding this fund could further increase the overall portfolio risk."
Result: This system goes beyond simply listing information; it clearly presents the reasons and grounds for specific risks through the relationships in the Knowledge Graph. This provides reliable analysis to financial professionals, meets regulatory compliance requirements, and enables more transparent and customized advice to clients. I am confident that this 'relationship-based reasoning' capability will be a key factor in determining the future of financial LLMs.
