💻 computer science

InfiniRetri: Attention-Guided Retrieval for Enhancing LLMs' Long-Context Processing

InfiniRetri는 외부 임베딩 모델이나 비용이 많이 드는 재학습 없이도, 사전 학습된 LLM의 내재된 어텐션 메커니즘을 활용하여 임의의 길이의 입력으로부터 관련 컨텍스트를 반복적으로 추출함으로써 효율적이고 매우 정확한 롱 컨텍스트 검색 및 추론을 가능하게 하는 훈련 불필요(training-free) 방식입니다.

원저자: XiaoJu YE, Fang Wang, FangFang Wang, YuanRui Ren

게시일 2026-10-01

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: XiaoJu YE, Fang Wang, FangFang Wang, YuanRui Ren

원본 논문은 CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

대규모 언어 모델로 알려진 현대의 인공지능 시스템은 인간의 언어를 이해하고 생성하는 데 있어 놀라운 능력을 갖추게 되었습니다. 이들은 이야기를 쓰고, 수학 문제를 풀며, 복잡한 질문에 답할 수 있습니다. 그러나 이 시스템들은 중대한 물리적 한계에 직면해 있습니다. 바로 한 번에 '작업 기억'에 담을 수 있는 정보의 양이 제한되어 있다는 점입니다. 사용자가 모델에게 소설 한 권이나 법적 기록 모음과 같이 매우 긴 문서를 읽도록 요청할 때, 시스템은 종번 어려움을 겪습니다. 텍스트가 길어짐에 따라 모델은 문서 앞부분의 세부 사항을 잊어버리거나 핵심 논점을 놓치곤 합니다. 이는 단어들을 서로 연결하는 모델의 수학적 과정이 텍스트가 커질수록 믿기 힘들 정도로 비용이 많이 들고 느려지기 때문에 발생합니다. 이를 해결하기 위해 연구자들은 두 가지 주요 접근 방식을 시도해 왔습니다. 하나는 비용이 많이 들고 시간이 오래 걸리는 재학습을 통해 모델이 더 긴 텍스트를 처리하는 법을 강제로 배우게 하는 것이고, 다른 하나는 모델이 텍수를 읽기 전에 관련 부분을 검색하는 보조 시스템을 사용하는 것인데, 이 보조 시스템은 종종 깊은 의미를 놓치는 단순한 단어 매칭에 의존한다는 단점이 있습니다.

중국 원화 대학(WenHua College)의 연구진은 재학습도 필요 없고 외부 검색 도구도 필요 없는 다른 해결책을 제안했습니다. 그들은 모델이 이미 긴 텍스트 내에서 적절한 정보를 찾아내는 능력을 이미 갖추고 있으며, 단지 그 능력을 사용하도록 유도될 필요가 있을 뿐이라는 것을 발견했습니다. 모델의 내부 어텐션(attention) 메커니즘이 어떻게 작동하는지 관찰한 결과, 연구진은 모델이 방대한 문서 속에 깊이 파묻혀 있더라도 질문에 대한 답을 포함하고 있는 특정 문장들에 자연스럽게 주의를 집중한다는 사실을 발견했습니다. 이러한 패턴은 모델의 처리 네트워크 중 마지막 레이어에서 가장 뚜렷하게 나타납니다. 연구진은 이러한 자연스러운 행동을 활용하는 '인피니트리트리(InfiniRetri)'라는 방법을 구축했습니다. 텍스트를 읽을 수 없는 작은 조각들로 압축하거나 별도의 검색 엔진에 의존하는 대신, 이 방법은 긴 문서를 청크(chunk, 덩어리) 단위로 반복해서 읽으며, 모델이 자신의 내부적 초점을 바탕으로 가장 관련성이 높은 문장을 식별하도록 요청하고, 오직 그 완전한 문장들만을 작고 효율적인 메모리 캐시에 저장합니다.

이 과정은 마치 두꺼운 책을 페이지별로 넘기며 읽는 세심한 독자와 같습니다. 연구진은 긴 텍스트를 관리 가능한 섹션으로 나눕니다. 각 섹션에 대해, 그들은 텍스트와 사용자의 질문을 함께 모델에 입력합니다. 그러면 모델은 자신의 내부 메커니즘이 해당 섹션의 서로 다른 부분들에 얼마나 많은 주의를 기울이는지 살펴봅니다. 연구진은 이 어텐션의 지도를 추출하며, 특히 질문의 단어들이 텍스트의 단어들과 어떻게 연결되는지를 조사했습니다. 그들은 모델의 어텐션 점수가 정확한 신호로서 작동하여, 어떤 문장이 답을 담고 있는지를 정확히 강조해 준다는 것을 발견했습니다. 시스템은 그다음 상위 점수를 받은 문장들을 선택하여 가벼운 캐시에 저장합니다. 이 순환 과정은 문서의 모든 청크에 대해 반복됩니다. 문서 전체가 처리될 때쯤이면, 캐시에는 질문에 답하는 데 필요한 가장 중요한 문장들만 남게 되어 나머지 노이즈를 효과적으로 걸러내게 됩니다. 마지막으로, 모델은 이 작고 정제된 문장 모음을 읽고 최종 답변을 생성합니다.

이 접근 방식의 결과는 놀랍습니다. 모델이 거대한 문서 속에서 숨겨진 사실을 찾아야 하는 테스트에서, 이 방법을 사용한 매우 작은 규모의 모델은 텍스트에 백만 단어가 넘게 포함되어 있음에도 불구하고 완벽한 정확도를 달ach했습니다. 이는 텍스트가 32,000 단어를 초과할 때 종종 실패했던 기존 방식들에 비해 상당한 개선입니다. 또한 이 방법은 복잡한 훈련이나 외부 데이터베이스에 의존하는 다른 고급 기술들을 능가하며, 여러 문서를 가로질러 정보를 찾아야 하는 질문을 해결하는 데 매우 효과적임이 입증되었습니다. 놀랍게도, 이러한 높은 수준의 성능은 일반 소비자용 하드웨어에서도 실행 가능한 크기인 2기가바이트의 컴퓨터 메모리만을 사용하여 달성되었습니다. 연구진은 이 방법이 명확한 질문이 있을 때 가장 잘 작동한다고 언급했는데, 이는 시스템이 주의를 집중하기 위해 질문에 의존하기 때문입니다. 전체 내용을 요약하는 것(즉, 특정 사실을 찾는 것이 아니라 모든 것을 압축하는 것이 목표인 작업)에는 개선 정도가 덜 나타났지만, 사실상 무제 limit한 길이의 텍스트로부터 정밀한 정보를 추출하는 능력은 값비싼 훈련 없이도 가능하다는 점에서 중요한 진전을 의미합니다. 이 연구는 모델 고유의 내부 신호를 신뢰함으로써, 우리가 어떤 크기의 텍스트에 대해서도 읽고 추론할 수 있는 능력을 끌어낼 수 있음을 시사합니다.

기술 요약: InfiniRetri

문제 정의

대규모 언어 모델(LLM)은 입력 길이가 길어짐에 따라 성능이 지속적으로 저하되는 문제를 겪으며, 이는 직접적인 검색(retrieval)과 긴 문서에 대한 복잡한 추론 모두에 영향을 미칩니다. 기존의 해결책들은 다음과 같은 상당한 한계에 직면해 있습니다:

컨텍스트 창 확장(Context Window Extension): 계속 학습(continued training)을 통해 컨텍스트 창을 확장하는 것은 어텐션 메커니즘의 이차 복잡도(quadratic complexity)로 인해 비용이 매우 많이 듭니다.
프롬프트 압축(Prompt Compression): 현재의 방법들(예: SelectiveContext, LLMLingua)은 종종 토큰 단위로 작동하여 의미적 단위(semantic units)를 버리고 읽을 수 없는 텍스트를 생성하며, 이는 별도의 학습 없이 성능을 저하시킵니다.
KV 캐시 압축(KV Cache Compression): H2O나 SnapKV와 같은 방식은 부분적인 Key-Value 상태만을 유지하며, 모든 상태를 학습 없이 압축하는 데 실패하고 세밀한 어텐션 패턴을 간과하는 경우가 많습니다.
검색 증강 생성(RAG): 표준 RAG는 표면적인 어휘적 관련성만을 포착하는 외부 임베딩 모델에 의존하며, 이로 인해 의미적 격차를 메우거나 여러 검색된 구절 사이의 연관성을 구축하는 데 어려움을 겪습니다.

방법론: InfiniRetri

저자들은 사전 학습된 LLM의 내재된 어텐션 메커니즘을 활용하여, 임의로 긴 입력으로부터 관련 정보를 검색하고 캐싱하는 학습 불필요(training-free) 방식인 InfiniRetri를 제안합니다. 이 접근 방식은 어텐션 정렬 검색(Attention-Aligned Retrieval, AAR) 패턴의 발견을 기반으로 합니다.

핵심 통찰: 어텐션 정렬 검색 (AAR)

계층별 시각화 및 분석을 통해 저자들은 다음을 관찰했습니다:

검색 능력: 사전 학습된 LLM(특히 마지막 트랜스포머 레이어)의 질문-대-컨텍스트(question-to-context) 어텐션 부행렬(sub-matrix)은 쿼리와 관련된 콘텐츠를 식별하기 위한 정밀하고 해석 가능한 신호를 제공합니다. 모델은 어떤 컨텍스트 구간이 관련이 있는지 본질적으로 "알고" 있습니다.
레이어 깊이: 이러한 검색 패턴은 깊이가 깊어질수록 강화되며, 관련 있는 컨텍스트와 관련 없는 컨텍스트 사이의 어텐션이 가장 변별력이 높은 마지막 트랜스포머 레이어에서 정점에 도달합니다.

워크플로(Workflow)

InfiniRetri는 다음 5단계로 구성된 반복 루프를 통해 긴 컨텍스트를 처리합니다:

청킹(Chunking): 긴 컨텍스트를 의미적 무결성을 보존하기 위해 문장 경계에서 세그먼트로 나눕니다.
병합(Merging): 이전에 검색된 문장들로 구성된 가벼운 외부 캐시를 현재 청크 앞에 붙입니다.
LLM 추론(LLM Inference): 병합된 컨텍스트와 사용자 질문을 LLM에 입력합니다.
어텐션 내 검색(Retrieval in Attention):
- 마지막 레이어의 어텐션 행렬을 추출합니다.
- (어텐션 싱크를 제외한) 질문-대-컨텍스트 부행렬을 슬라이싱합니다.
- 1-D sum-pooling을 사용하여 어텐션 점수를 구절(phrase) 단위로 집계합니다.
- 집계된 점수를 바탕으로 상위 K개의 가장 관련 있는 위치를 식별합니다.
캐시 업데이트(Cache Update): 개별 토큰을 저장하는 대신, 상위 K개 토큰을 포함하는 완전한 문장을 원래 컨텍스트에서 추출하여 캐시에 추가합니다.

모든 청크를 처리한 후, 누적된 캐시와 원래의 질문을 사용하여 최종 답변을 생성합니다.

주요 기여

AAR 패턴 식별: 본 논문은 사전 학습된 LLM의 마지막 레이어에 있는 질문-대-컨텍스트 어텐션 부행렬이 정확하고 신뢰할 수 있으며 해석 가능한 검색 신호 역할을 한다는 것을 밝혀냈으며, 이를 통해 외부 임베딩 모델의 필요성을 제거했습니다.
InfiniRetri 프레임워크: 긴 컨텍스트를 가벼운 캐시로 반복적으로 증류하는 학습 불필요 방식입니다. 이는 프롬프트 압축의 목표와 일치하면서도, 더 나은 해석 가능성과 의미적 일관성을 위해 문장 단위로 작동합니다.
포괄적인 평가: 이 방법은 프롬프트 압축, KV 캐시 압축, RAG라는 세 가지 계열의 베이스라인과 비교 평가되어 여러 벤치마크에서 우수한 성능을 입증했습니다.

실험 결과

InfiniRetri는 LongBench, LongBenchV2, $\infty$ Bench, 그리고 Needle-In-a-Haystack (NIH) 테스트를 통해 평가되었습니다.

LongBench 성능:
- 프롬프트 압축: 이전의 학습 불필요 SOTA인 DAC보다 평균 점수에서 11.4% 앞섰으며(42.05 vs. 37.76), 압축률은 현저히 낮았습니다(0.26 vs. 0.5).
- KV 캐시 압축: 오프라인 캘리브레이션이 필요한 강력한 기존 방식인 CompilerKV를 Qwen2-7B 기준으로 전체에서 15.9% 능가했습니다. 특히, 다중 문서 QA 작업에서 238.7%의 개선(13.40 $\to$ 45.38)을 달성했습니다.
- RAG: 선도적인 임베딩 기반 검색기(BGE-M3, Stella-v5)와 고급 RAG 시스템(GraphRAG, RQ-RAG)을 모두 능가했습니다. 오직 풀 스테이지 학습의 이점을 가진 MemoRAG에만 뒤처졌습니다.
초장기 컨텍스트 능력:
- NIH 테스트: 0.5B 파라미터 모델로 100만 토큰 이상의 환경에서 100% 정확도를 달고, 단 2GB의 GPU 메모리만을 요구했습니다.
- $\infty$ Bench: 1.5B 모델이 초장기 검색 작업(214K 토큰)에서 100% 정확도를 달성하여, KV 검색 방법을 사용하는 더 큰 모델들을 능가했습니다.
- LongBenchV2: InfiniRetri를 사용한 7B 모델이 긴 텍스트 작업에서 72B 모델과 대등한 성능에 도달했습니다.
복잡도: 이 방법은 컨텍스트를 관련 문장의 작은 캐시로 가지치기함으로써, 추론 복잡도를 이차 복잡도 $O(N^2)$ 에서 약 선형 복잡도 $O(mN)$ ( $m$ 은 청크 크기)으로 줄입니다.

의의 및 주장

본 논문은 InfiniRetri가 LLM의 내재적 능력에 기반한 정보 검색의 새로운 패러다임을 구축한다고 주장합니다. LLM 자체의 어텐션 모듈을 활용함으로써, 외부 검색 구성 요소나 학습 없이 쿼리와 문서 사이의 의미적 격차를 해소합니다.

저자들은 다음을 강조합니다:

해석 가능성: 토큰 단위 압축과 달리, InfiniRetri는 완전한 문장을 유지하므로 검색 과정이 투명하고 인간이 읽을 수 있습니다.
효율성: 정확도를 유지하거나 향상시키면서도 메모리 오버헤드(예: 1M 토큰에 2GB)와 계산 비용을 획기적으로 줄입니다.
일반성: 이 방법은 모든 어텐션 기반 LLM에 적용 가능하며 이론적으로 무한한 컨텍스트 길력을 지원합니다.

인정된 한계점:
논문은 InfiniRetri가 미리 정의된 질문을 필요로 한다는 점(질문 생성이 없는 개방형 시나리오에서의 제한)과, 현재의 단일 패스 검색 메커니즘을 넘어 여러 윈도우에 걸친 정보를 집계해야 하는 요약 작업에서 이득이 제한적이라는 점을 언급합니다. 또한, 성능은 하이퍼파라미터 튜닝, 특히 타겟 답변 길이에 대한 구절 크기에 민다는 점을 명시하고 있습니다.