Adaptive-masking(AGE)을 이용한 GraphRAG용 그래프 임베딩

왜 중요한가

GraphRAG는 그래프 구조의 관계 정보를 LLM에 제공해 도메인 지식 활용을 가능하게 하나, 그래프 임베딩과 LLM의 텍스트 임베딩 간 잠재 공간 불일치가 성능 저해 원인으로 작용했다. AGE는 이 불일치를 해결하기 위해 그래프 임베딩 학습 과정을 LLM의 마스크 기반 자기지도학습 패턴에 맞춰 재구성함으로써 고정된 LLM 환경에서도 검색된 그래프 정보를 효과적으로 재활용할 수 있게 했다. 이 접근은 전체 LLM을 파인튜닝하지 않고도 비매개변수 리트리버 기반 파이프라인에서 성능을 크게 높여 실무적 비용 제약이 있는 상황에서 유용성이 높다.

핵심 기여

적응형 마스킹 기반 그래프 임베딩 학습 메커니즘

AGE는 노드 샘플러를 도입해 그래프에서 예측이 어려운 핵심 노드를 샘플링으로 보존하고, 보조 노드만 마스킹해 예측하도록 구성되었다. 이 구조는 무작위 마스킹이 비판적 정보 손실로 이어지는 문제를 완화하며 마스크 기반 SSL 신호를 그래프의 관계 표현 학습에 집중시켰다. 노드 샘플러는 MHA와 선형층으로 확률 분포를 만들고 샘플링율 rho에 따라 실제 마스킹을 수행하도록 학습되었다.

JEPA 형태의 개념 인코더-디코더와 타깃 인코더 병렬 학습

본 연구는 재구성 대신 타깃 인코더의 상위 임베딩을 예측하는 JEPA 방식을 채택해 의미적 수준의 그래프 표현을 학습했다. 개념 인코더-디코더는 핵심 노드로부터 보조 노드 임베딩을 유추하는 역할을 맡고 타깃 인코더는 하드웨어 친화적인 LLM 입력 차원으로 최적화된 교사 특성을 제공했다. 이 분리 학습은 복원성에 의존하는 기존 GA 대비 보조 정보의 압축과 관계성 캡처에 유리했다.

RL 영감의 샘플링 손실로 샘플러 최적화

노드 샘플러의 비미분 샘플링 연산을 직접 최적화하기 위해 L_NS 손실을 도입해 샘플 확률을 보정했다. L_NS는 보조 노드 예측 오차가 큰 노드일수록 샘플 확률을 높이는 구조로 설계되어 예측이 어려운 노드를 핵심으로 이동시키는 신호를 제공한다. 이 방식은 REINFORCE 스타일의 정책 그라디언트 아이디어를 차용해 샘플링 분포를 안정적으로 개선했다.

다양한 벤치마크와 규격에 걸친 성능 검증

ExplaGraphs, SceneGraphs, WebQSP, CWQ 등 네 개 벤치마크에서 AGE를 평가해 여러 설정에서 일관된 성능 향상을 확인했다. 예를 들어 G-Retriever 기준 Llama3.2 1B에서 ExplaGraphs 정확도가 0.5595에서 0.8267로 26.72 포인트 상승했다. 또한 JEPA와 노드 샘플러의 결합이 무작위 마스크나 GA 대비 가장 큰 개선폭을 보이는 것이 실험으로 확인되었다.

핵심 아이디어 이해하기

그래프 데이터를 LLM에 입력하기 위해서는 그래프의 구조적 관계가 LLM의 텍스트 임베딩 공간과 '정렬'되어야 한다. 기존 방법은 그래프 토큰을 텍스트화하거나 그래프 인코더 출력을 직접 LLM에 주입했으나 그래프의 간결성 때문에 무작위 마스킹 기반 SSL은 중요 노드의 정보를 손실하거나 재구성 난이도가 커 표현 학습이 부진했다. 이러한 불일치는 고정된 LLM을 사용하는 환경에서 특히 성능 저하로 이어졌다. AGE는 핵심 아이디어로서 중요한 노드는 마스킹에서 제외하고 보조 노드만 예측하게 함으로써 SSL 신호를 관계성 학습에 집중시킨다. 구현상으로는 그래프 인코더가 생성한 노드 특징에서 노드 샘플러가 확률 분포를 계산하고 샘플링율 rho로 핵심 노드 집합을 결정한다. 핵심 노드 정보만으로 보조 노드를 재구성하는 과제를 부과하면, 개념 인코더-디코더가 핵심과 보조의 관계를 내재화한 잠재 표현 z_key를 학습하게 된다. 학습 목표는 JEPA 방식으로 재구성 대신 타깃 인코더의 임베딩을 예측하는 것이며, 이로 인해 하위 수준 노이즈를 제거하고 의미적 관계를 더욱 집중적으로 캡처한다. 결과적으로 AGE는 그래프 임베딩을 LLM의 텍스트 임베딩 공간에 더 잘 맞추어 고정 LLM 환경에서도 검색된 그래프 정보를 효과적으로 활용할 수 있게 만든다.

방법론

전체 파이프라인은 입력 처리, 서브그래프 구성, 임베딩 학습, 추론 네 단계로 구성되며 임베딩 학습 단계에서 AGE 모듈이 작동한다. AGE는 그래프 인코더, 노드 샘플러, 개념 인코더-디코더, 타깃 인코더, 그래프 구조 기반 집계기(GNN GSA)와 프로젝터로 이루어져 있다. 그래프 인코더는 서브그래프의 각 노드에 대해 d_g 차원의 특징을 출력하고 이 특징이 노드 샘플러와 타깃 인코더의 입력이 된다. 노드 샘플러는 MHA_NS, 선형층, softmax로 노드별 확률 p_NS를 계산하고 샘플링율 rho에 따라 N_key를 결정해 핵심 노드 집합 I_key와 보조 노드 집합 I_aux를 만든다. 핵심 노드 특징 h_key는 개념 인코더 MHA_CE로 잠재 z_key로 인코딩되고, z_key는 보조 노드를 위한 플레이스홀더 z_aux와 결합되어 전체 노드 순서를 유지하는 z를 만든다. 개념 디코더 MHA_CD는 z와 위치인코딩을 입력으로 받아 h_out을 출력하며 h_out은 최종적으로 GNN GSA와 MLP Proj를 통해 LLM 입력 차원 h_bar_g로 집계된다. 학습은 세 가지 손실로 분리되어 모듈별로 최적화된다. 타깃 인코더와 하위 모듈 일부는 프롬프트 튜닝 손실 L_PT로 직접 downstream 성능에 맞춰 학습되고, 개념 인코더-디코더와 그래프 인코더 일부는 타깃 임베딩을 예측하는 L_target으로 학습된다. 노드 샘플러는 L_NS로만 갱신돼 샘플 확률이 예측 에러와 연동되도록 한다. 이 손실 분리는 모듈 간 파라미터 중복을 피하고 각 파트에 명확한 학습 신호를 제공한다.

주요 결과

주요 벤치마크에서 AGE는 G-Retriever 기반 설정에 일관된 향상을 보였다. 대표적으로 Llama3.2 1B 백본에서 G-Retriever는 ExplaGraphs 정확도 0.5595를 보였고 AGE를 적용한 동일한 설정은 0.8267을 기록해 26.72 포인트의 상승을 나타냈다. Llama3.2 3B에서도 G-Retriever 0.7761에서 AGE 0.9260으로 유의미한 향상이 관찰되었다. 추가로 PEFT(LoRA)를 결합한 설정에서도 AGE는 성능을 더 끌어올려 여러 스케일에서 우수한 결과를 냈다. Ablation 실험에서는 GA 대비 JEPA가 더 우수했고 무작위 마스크 대비 학습 가능한 노드 샘플러가 추가 이득을 제공해 JEPA + 노드 샘플러 조합이 가장 큰 개선폭을 보였다. 샘플링율 rho 민감도 실험에서 rho = 0.3이 ExplaGraphs와 WebQSP에서 전반적으로 최적에 가까운 성능을 보여 본 실험군의 기본값으로 채택되었다. LLM 기반 리트리버 방식과의 직접 비교에서는 일부 프로프라이어터리 방법에 비해 분명한 우위가 없었으나 AGE를 AMAR와 결합한 경우 ReKnoS 같은 방법을 상회하는 결과가 관찰되었고, 데이터셋 특성에 따라 WebQSP와 CWQ에서 성능 우열이 달라지는 양상이 나타났다.

기술 상세

전체 아키텍처는 그래프 인코더 GNN_GE가 먼저 서브그래프 S*의 노드별 특징 h_in을 생성하고, 이 특징이 노드 샘플러와 타깃 인코더의 입력으로 사용되는 구조다. 노드 샘플러 내부는 MHA_NS로 특징을 처리하고 Linear와 Softmax로 각 노드의 선택 확률 p_NS를 만들어 rho로 결정된 수만큼 핵심 노드를 샘플링한다. 샘플링된 핵심 노드 특징 h_key는 위치 인코딩을 더해 MHA_CE로 인코딩되어 z_key가 되고, z_key와 보조 노드 플레이스홀더 z_aux를 결합한 z를 MHA_CD가 디코딩해 h_out을 생성한다. 타깃 인코더 MHA_TE는 h_in 전체에 대해 타깃 임베딩 h_target을 생성하고 이 하위 표현은 downstream 프롬프트 튜닝 손실 L_PT로 최적화된다. h_out은 L_target으로 h_target을 모방하도록 학습되며 GNN_GSA와 MLP_Proj를 거쳐 LLM 입력 차원으로 투영된 후 추론 시 LLM에 연결된다. 손실 분배 설계로 L_PT, L_target, L_NS는 서로 다른 파라미터 집합을 최적화하도록 구성되어 손실 가중치 조정이 불필요하다고 실험적으로 확인되었다. 학습 세부사항으로는 샘플링율 rho = 0.3이 ExplaGraphs와 WebQSP에서 전반적으로 최적 성능을 보였고 실험은 Llama3.2(1B,3B), Llama2(7B,13B), Llama3.1(8B) 등 여러 고정 LLM 백본과 함께 수행되었다. L_NS는 로그 확률에 예측 에러를 곱하는 형태로 구성되어 비미분 샘플링 연산에 대한 RL 영감의 보상을 제공하며, 이로 인해 예측이 어려운 노드가 핵심으로 재분류되는 경향이 발생해 JEPA의 보조 예측 과제가 안정화되었다.

한계점

논문은 몇 가지 한계를 명확히 밝혔다. 첫째, 모든 그래프에 대해 동일한 고정 샘플링율 rho를 사용했기 때문에 그래프별 핵심 노드 밀도 차이를 자동으로 반영하지 못한다. 둘째, AGE는 본 연구에서 GraphRAG 과제에만 실험적으로 적용되어 다른 모달리티나 완전히 다른 태스크에 대한 보편성은 아직 검증되지 않았다. 셋째, 대규모 LLM에 대한 효용성은 계산 자원 한계로 실험되지 않아 그 범위에서의 성능은 미확인 상태다. 넷째, 본 접근은 그래프 표현을 LLM 입력에 맞추는 데 초점을 두어 전통적 그래프 학습 과제인 노드 분류나 링크 예측의 직접적 성능 향상에 대한 이론적 보장은 제공하지 않는다.

실무 활용

AGE는 고정된 LLM을 사용하는 프로덕션 환경에서 그래프 검색 결과를 LLM 입력으로 정렬해 성능을 개선할 수 있는 실용적 방법을 제공한다. 비매개변수 리트리버를 유지하면서도 그래프 임베딩 품질을 높여 추가적인 LLM 파인튜닝 비용을 줄이는 데 유리하다. 구현 관점에서는 그래프 인코더와 개념 인코더-디코더, 샘플러 모듈을 기존 파이프라인에 추가하는 형태로 통합 가능하다.

지식 그래프 기반 질의응답 시스템에서 검색된 서브그래프의 구조 정보를 LLM에 정합하여 응답 정확도를 개선하는 경우
도메인 특화 QA 서비스에서 고정된 상용 LLM을 그대로 사용하면서 그래프 임베딩만 업데이트해 비용과 시간 제약을 완화하려는 경우
복잡한 관계 정보를 가진 소규모 그래프에서 핵심 정보 보존을 통해 추론 안정성을 높여야 하는 추천 또는 시멘틱 검색 파이프라인

코드 공개 여부: 미확인

키워드

GraphRAG(그래프 기반 RAG)JEPA(Joint-Embedding Predictive Architecture)Adaptive-masking(적응형 마스킹)Node Sampler(노드 샘플러)Masked SSL(마스크 기반 자기지도학습)

용어 해설

Joint-Embedding Predictive Architecture(JEPA): — 입력 일부를 재구성하는 대신 타깃 인코더의 상위 수준 임베딩을 예측하는 자기지도학습 아키텍처로, 불필요한 세부 복원을 배제하여 의미적 수준의 표현을 학습하도록 설계되었다. 본 논문에서는 그래프 임베딩의 디스크리미너빌리티 문제를 완화하기 위해 JEPA 형태로 개념 인코더-디코더와 타깃 인코더를 구성했다. JEPA는 그래프에서 핵심노드를 보존하면서 보조노드의 관계성을 압축해 LLM 입력 공간에 정렬하는 역할을 한다.
적응형 마스킹(Adaptive-masking): — 무작위 마스킹 대신 학습 가능한 샘플러가 마스킹 대상 노드를 선택하는 기법으로, 그래프의 핵심 노드는 마스킹에서 제외해 보조 노드를 예측하는 자기지도학습 신호를 강화한다. 본 연구에서는 노드 샘플러가 MHA와 선형층으로 확률분포를 생성하고 RL 유사 손실로 확률을 갱신한다. 핵심노드 보존으로 인해 재구성 불가능한 핵심 정보를 마스킹해 학습 효율이 향상된다.
노드 샘플러(Node Sampler): — 그래프 인코더 출력으로부터 각 노드의 마스킹 확률을 계산하는 모듈로, Multi-Head Attention, 선형층, softmax로 확률 벡터를 만들고 샘플링율 rho에 따라 N_key를 결정한다. 샘플러는 L_NS 손실로 비차별적 노드를 보조로 남기고 예측이 어려운 노드를 핵심으로 할당하도록 학습된다. 이 모듈은 그래프 구조의 중복성 부족을 고려해 마스킹 신호의 효율을 높이는 핵심 요소다.
그래프 기반 RAG(GraphRAG): — 검색 증강 생성(RAG) 방식의 확장으로 그래프 구조의 외부 지식을 검색해 LLM 입력으로 제공하는 프레임워크다. 그래프의 노드와 엣지를 부분 서브그래프로 취득해 LLM의 추론에 필요한 구조적 관계를 보존하고자 한다. 본 논문은 GraphRAG에서 비매개변수형 리트리버와 고정 LLM 조합의 실용성에 주목해 그래프 임베딩 정렬 문제를 해결했다.

Adaptive-masking(AGE)을 이용한 GraphRAG용 그래프 임베딩

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

언어 모델과 지식 그래프를 결합하는 파라미터 효율적 학습 및 통합 그래프 방법론.

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

언어 모델과 지식 그래프를 결합하는 파라미터 효율적 학습 및 통합 그래프 방법론.