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

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

그래프 데이터를 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 입력에 맞추는 데 초점을 두어 전통적 그래프 학습 과제인 노드 분류나 링크 예측의 직접적 성능 향상에 대한 이론적 보장은 제공하지 않는다.