Attention 레이어, 긴 문서 검색의 숨은 비밀

초록

본 논문은 사전 학습된 대형 언어 모델(LLM)의 어텐션 스코어를 활용해, 별도 학습 없이 긴 문서에서 관련 구간을 효율적으로 찾아내는 Retrieval 모델인 AttentionRetriever를 제안한다. 어텐션을 이용한 문장‑레벨 스코어링에 엔터티 그래프 기반 범위 결정 방식을 결합해, 컨텍스트·인과·쿼리 의존성을 모두 만족한다. 100k 토큰 규모의 신규 데이터셋을 포함한 10개 벤치마크에서 기존 sparse·dense 모델을 크게 앞서며, 3B 파라미터 규모 모델만으로도 dense 검색 수준의 효율성을 달성한다.

상세 분석

AttentionRetriever는 기존 RAG 파이프라인이 직면한 “긴 문서 검색” 문제를 두 가지 핵심 관점에서 재해석한다. 첫째, 트랜스포머의 어텐션은 본질적으로 쿼리‑키의 내적을 통해 토큰 간 유사도를 계산하므로, 별도 인코더 없이도 “검색 스코어”를 제공한다는 점을 이용한다. 논문은 LLaMA‑3.2‑3B, Qwen‑2.5‑3B, Mistral‑7B 등 세 모델에 대해 층별 어텐션 정확도를 정량화했으며, 특히 후반부 층이 “인과적 의존성”을 구축하고, 초기 층이 “독립적 서브쿼리”를 포착한다는 계층적 패턴을 발견했다. 이는 어텐션이 단순한 의미 유사도뿐 아니라 문맥 흐름을 단계적으로 축적한다는 기존 가설을 실험적으로 입증한다.

둘째, 어텐션 스코어만으로는 “배경 정보”와 같이 직접적인 쿼리와의 유사도가 낮은 구간을 놓치기 쉽다. 이를 보완하기 위해 논문은 엔터티 기반 그래프를 도입한다. 문서 내 각 엔터티를 노드로, 해당 엔터티가 등장하는 문장을 연결하는 간단한 bipartite 구조를 만든 뒤, 쿼리와 연관된 엔터티를 랭킹한다. 높은 랭크를 받은 엔터티가 포함된 모든 문장을 최종 후보 집합에 추가함으로써, 쿼리‑의존성(배경 정보) 문제를 해결한다. 이 접근법은 복잡한 관계 추출 없이도 엔터티 수준의 연결성을 활용해 검색 범위를 자동으로 확장한다는 점에서 효율적이다.

효율성 측면에서는 “Cascading KV Cache”와 같은 어텐션 근사 기법을 그대로 적용해 100k 토큰 문서에서도 메모리·시간 비용을 크게 낮춘다. 실험 결과, 3B 규모 모델만 사용했음에도 기존 7B‑10B dense 모델 대비 2‑3배 빠른 추론 속도를 보였으며, 평균 정밀도·재현율에서 10~15%p 상승을 기록했다. 특히 새로 구축한 100k‑토큰 평균 길이의 LongDoc 데이터셋에서 가장 높은 MAP@10을 달성했으며, 다중 문서 베이스의 기존 벤치마크에서도 경쟁력 있는 성능을 유지했다.

한계점으로는 (1) 어텐션 스코어가 레이어마다 변동성이 커서 최적 레이어 선택이 모델‑특화적이라는 점, (2) 엔터티 그래프 구축 시 엔터티 인식 정확도에 크게 의존한다는 점, (3) 현재는 쿼리‑문서 쌍에 대한 단일 스코어링만 지원해 복합 질의(예: 다중 조건 결합)에는 추가 설계가 필요하다는 점을 언급한다. 향후 연구에서는 레이어‑별 어텐션 가중치를 학습적으로 조정하거나, 엔터티‑관계 추출을 결합해 보다 정교한 범위 결정 메커니즘을 탐색할 여지가 있다.