Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

기술 요약: KGERMAR

문제 정의
장문맥 언어 모델링(Long-context language modeling)은 단순히 컨텍스트 창을 확장하는 것을 넘어, 수천 개의 토큰에 걸쳐 엔티티의 상태와 관계를 일관되게 이해하는 데 근본적인 어려움을 겪고 있습니다. 표준적인 접근 방식들은 두 가지 주요 한계점을 가집니다. 첫째, 확장된 어텐션 모델(예: LongLoRA, YARN)은 긴 컨텍스트의 중간 부분에 있는 정보가 체계적으로 과소 주목받는 "lost-in-the-middle" 저하 현상을 겪습니다. 둘째, 메모리 증강 검색 시스템(예: MemLong, ERMAR)은 오직 의미적 유사성(semantic similarity)에만 의존합니다. 이는 관련 용어가 포함된 구절을 찾는 데는 효과적이지만, 단순히 관련 용어를 우연히 언급한 구절과 동일한 엔티티 간의 명시적인 인과 관계나 구조적 관계를 설명하는 구절을 구분하지 못합니다. 이러한 한계는 장치 오류 상태의 진행 과정을 추적하는 것과 같이, 쿼리와 어휘적 중복은 없더라도 깊은 구조적 연결을 공유하는 상황에서의 구조적 추론이 필요한 시나리오에서 매우 치명적입니다.

방법론
본 논문은 추론 과정에서 동적이고 컨텍스트 특화적인 지식 그래프를 구축하여 그래프 구조 임베딩과 텍extual semantics를 통합하는 프레임워크인 KGERMAR(Knowledge Graph Enhanced Memory-Augmented Retrieval)을 제안합니다. 이 시스템은 네 가지 통합된 구성 요소로 작동합니다:

1. 컨텍스트 지식 그래프 구축:

   • 추출: 시스템은 CoNLL-2003으로 미세 조정된 BERT-large 모델을 사용하여 실시간 개체명 인식(NER)을 수행하고, TACRED로 미세 조정된 BERT-base 모델을 사용하여 관계 추출(Relation Extraction)을 수행합니다.
   • 그래프 구축: 슬라이딩 윈도우 내에서 엔티티와 관계(예: CAUSEDBY, PARTOF, USEDFOR)를 추출합니다. 품질 보장을 위해 문자열 및 임베딩 유사도를 통한 엔티티 통합(entity consolidation), 분류기 확신도와 빈도를 결합한 관계 신뢰도 점수 산정, 그리고 그래프 필터링 기법을 채택합니다.
   • 동적 특성: 정적 방식(예: ConceptNet)과 달리, 이 그래프는 입력 텍스트 자체로부터 구축되어, 사전 구축된 지식 베이스에는 존재하지 않는 도메인 특화 및 세션 특화 엔티티(예: 특정 프린터 모델 및 고유 에러 코드)를 포착합니다.

2. 그래프 강화 임베딩 생성:

   • 인코딩: 추출된 엔티티는 다중 홉 메시지 전달(multi-hop message passing)을 통해 개별 속성과 구조적 역할을 모두 포착하도록 관계형 그래프 합성곱 네트워크(R-GCN)를 사용하여 인코딩됩니다.
   • 융합: 이러한 그래프 구조 임베딩은 교차 모달 어텐션(cross-modal attention)을 통해 기본 언어 모델에서 유도된 텍스트 엔티티 임베딩과 융합됩니다. 이를 통해 표면적 의미론과 더 깊은 관계적 패턴을 모두 포착하는 하이브리드 표현을 생성합니다.

3. 다중 구성 요소 메모리 아키텍처:

   • KGERMAR은 세 가지 특화된 메모리 뱅크를 유지합니다:
     • 컨텍스트 메모리(Contextual Memory): ERMAR을 따라 언어 모델 레이어의 토큰 레벨 키-값 쌍을 저장합니다.
     • 의미 메모리(Semantic Memory): 중복되지 않는 청크들의 밀집 텍스트 임베딩을 저장합니다(BGE-M3 사용).
     • 구조 메모리(Structural Memory): 융합된 텍스트-그래프 엔티티 임베딩을 저장합니다.
   • 모든 뱅크는 점진적으로 업데이트되며, 용량에 도달하면 LRU 제거(eviction) 방식으로 관리됩니다.

4. 하이브리드 검색 메커니즘:

   • 주어진 쿼리에 대해 시스템은 세 가지 뱅크 모두에서 병렬로 Top-K 검색을 수행합니다.
   • 검색된 표현은 학습된 융합 가중치($\lambda_1, \lambda_2, \lambda_3$)를 사용하여 결합되어 최종 점수를 생성합니다.
   • 검색된 표현은 지정된 상위 레이어에서 검색 인과 어텐션(retrieval causal attention)을 통해 언어 모델에 주입됩니다.

학습 목표
모델은 다음 네 가지 목적 함수를 동시에 최적화하기 위해 멀티태스크 학습을 채택합니다:

• 언어 모델링 손실 ($L_{LM}$): 검색된 컨텍스트를 사용한 다음 토큰 예측.
• 지식 그래프 임베딩 손실 ($L_{KG}$): R-GCN을 학습시키기 위한 이진 교차 엔트로피 기반의 링크 예측.
• 검색 랭킹 손실 ($L_{retrieval}$): 관련 있는 컨텍스트가 무관한 컨텍스트보다 높은 점수를 받도록 하는 마진 기반 랭킹.
• 교차 모달 정렬 손실 ($L_{alignment}$): 텍스트와 구조적 표현 사이의 일관성을 보장.

주요 기여

1. 동적 지식 그래프 구축: 입력 텍스트에서 도메인 특화 엔티티와 관계를 실시간으로 추출하는 파이프라인을 통해, 정적 지식 베이스에 없는 전문적인 콘텐츠에 적응할 수 있게 합니다.
2. 하이브리드 다중 구성 요소 메모리: 컨텍스트, 의미, 구조적 표현을 통합하는 아키텍처입니다. 특히 구조 메모리 뱅크는 순수 의미적 유사성보다 더 보완적이고 종종 더 변별력 있는 엔티티 관계를 포착합니다.
3. 경험적 검증: 구조적 지식을 통합하는 것이 장문맥 모델링에서 상당한 개선을 가져온다는 것을 입증하였으며, 강력한 베이스라인 모델들과 비교하여 최대 8.5% 낮은 퍼플렉시티(perplexity)와 2~2.5배 더 나은 메모리 효율성을 달s성했습니다.

실험 결과
프레임워크는 SlimPajama(미세 조정을 위한 84.7K 예시)와 세 가지 벤치마크 데이터셋인 WikiText-103, PG-19, Proof-pile에서 평가되었습니다.

• 퍼플렉시티(Perplexity): KGERMAR은 1K에서 32K 토큰에 이르는 컨텍스트 길이 전반에서 모든 메모리 증강 3B 모델 중 가장 낮은 퍼플렉시티를 기록했습니다. WikiText-103에서 16K 토큰 기준 7.14의 퍼플렉시티를 기록하며 ERMAR(7.80) 및 MemLong을 크게 앞질렀습니다.
• 인-컨텍스트 학습(ICL): 다섯 가지 NLU 태스크(SST-2, MR, Subj, SST-5, MPQA)에서 KGERMAR은 4-shot(78.6%) 및 20-shot(81.2%) 설정 모두에서 ERMAR과 MemLong을 능가하는 가장 높은 평균 정확도를 달성했습니다. 특히 엔티티 수준의 변별력이 요구되는 태스크(예: Subj, MPQA)에서 향상이 두드러졌습니다.
• 메모리 효율성: KGERMAR은 2~2.5배 더 낮은 메모리 사용량을 유지하며 우수한 메모리 효율성을 보여주었습니다. 32K 토큰에서 ERMAR보다 36%, MemLong보다 41% 적은 피크 메모리를 사용했습니다. 토큰당 정규화된 메모리는 컨텍스트 길이가 길어질수록 개선되었으며, 이는 표준 어텐션 메커니즘의 이차 함수적 메모리 증가를 완화하는 능력을 입증합니다.
• 스케일링 동작: KGERMAR은 실시간 그래프 구축으로 인해 토큰당 레이턴시(latency)가 약간 높지만, 시퀀스 길이에 따라 거의 선형적인 레이턴시 증가를 보여, ERMAR보다 장문맥 배포 시 예측 가능성이 더 높습니다.

의의
본 논문은 KGERMAR이 고정된 정적 지식 베이스에 대한 의존에서 도메인 적응형 지식 표현으로 전환함으로써 메모리 증강 언어 모델링을 발전시킨다고 주장합니다. 입력 컨텍스트로부터 지식 그래프를 동적으로 구축함으로써, 시스템은 의미적 유사성만으로는 포착할 수 없는 엔티티의 상태와 인과적 진행 과정을 추론할 수 있습니다. 이 접근 방식은 높은 모델링 품질과 현저히 줄어든 메모리 오버헤드 사이의 균형을 제공하며, 구조적 추론이 매우 중요한 기술 지원, 과학 문헌, 법률 문서 분석 등의 애플리케이션에 대한 실질적인 솔루션을 제공합니다.