긴 문맥 처리를 위한 쿼리 집중형 및 메모리 인식 재순위화 모델 (QRRanker)

기존 임베딩 모델의 정보 압축 한계와 LLM 기반 재순위화 모델의 불안정한 점수 생성 문제를 동시에 해결한다. 텍스트 생성 과정 없이 어텐션 가중치만 활용하여 연산 효율을 높였으며, 요약 정보를 활용한 메모리 기법으로 아주 긴 소설이나 대화 기록에서도 정확한 정보를 찾아낼 수 있다.

Transformer의 Self-Attention은 시퀀스 내 토큰 간의 관계를 파악하는 핵심 기제다. 하지만 기존 RAG 시스템은 문서를 고정된 크기의 벡터로 압축하는 임베딩 방식에 의존하며, 이는 복잡한 쿼리-문서 상호작용을 모두 담지 못하는 '기하학적 병목' 현상을 일으킨다. 또한 기존 LLM 재순위화 모델은 점수를 텍스트로 생성하게 시키는데, 이는 출력이 불안정하고 정밀한 점수 차이를 표현하기 어렵다는 한계가 있다.

QRRanker는 LLM이 학습 과정에서 이미 특정 어텐션 헤드를 통해 질문과 관련된 정보를 찾는 '검색 능력'을 내재화하고 있다는 점에 주목한다. 모델에게 답을 쓰라고 시키는 대신, 질문 토큰이 각 문서 토큰에 보내는 어텐션 에너지의 크기를 직접 추출하여 관련성 지표로 사용한다. 이는 모델의 내부 판단 근거를 직접 들여다보는 것과 같아 더 정확하고 정밀한 점수 산출이 가능하다.

특히 여러 문서를 한 번에 입력하는 리스트웨이즈 방식을 채택하여 문서 간의 상대적 중요도를 동시에 비교한다. 여기에 문서들의 요약본을 미리 보여주는 '메모리' 기법을 더해, 아주 긴 소설이나 대화 기록에서도 길을 잃지 않고 전역적인 맥락과 세부 정보를 동시에 고려하여 최적의 문서를 찾아낸다.

QR-head 선정 및 점수 계산 메커니즘을 구축한다. 먼저 NarrativeQA 데이터셋에서 샘플을 추출하여 질문 토큰과 정답 문서 토큰 간의 어텐션 합이 가장 높은 상위 16개 헤드를 QR-head로 지정한다. 질문 토큰과 문서 토큰 간의 어텐션 행렬을 입력으로 하여, 질문 토큰들이 문서 토큰들에 보낸 가중치 값을 모두 합산하고 질문 길이로 나누는 연산을 수행해 해당 헤드에서의 문서별 관련성 점수를 얻는다. 이 값이 높을수록 질문에 답하기 위해 해당 문서가 중요함을 의미한다.

리스트웨이즈 학습을 위해 질문과 50개의 후보 문서를 하나의 프롬프트로 구성한다. 계산된 원시 어텐션 점수 S를 입력으로 하여 (S - 최소값) / (최대값 - 최소값) 연산을 수행해 0에서 1 사이의 정규화된 점수를 얻고, 이를 통해 서로 다른 헤드나 샘플 간의 점수 편차를 줄여 학습 안정성을 확보한다. 이후 정답 문서와 오답 문서의 점수 분포를 입력으로 하여 정답 문서의 지수 함수 값을 전체 문서의 지수 함수 합으로 나누는 연산을 수행해 정답일 확률 값을 얻고, 이 확률의 로그 값을 최대화하도록 모델을 최적화한다.

긴 문맥 처리를 위해 메모리 강화 기법을 적용한다. 각 문서 블록의 요약 텍스트를 입력으로 하여 후보 문서 리스트 앞에 결합하는 연산을 수행해 확장된 입력 시퀀스를 얻고, 이를 통해 모델이 각 문서를 평가하기 전 전체적인 맥락 정보를 미리 참조할 수 있게 한다. 또한 효율성을 위해 36개 레이어 중 24번 레이어 이후를 제거하는 절단 기법을 적용하여 추론 속도를 높인다.

LoCoMo 벤치마크에서 기존 SOTA 모델들을 제치고 최고 성능을 기록했다. GPT-4o-mini를 생성기로 사용했을 때 Overall F1 57.03점을 달성하며, 복잡한 그래프 기반 메모리 모델들보다 우수한 성능을 보였다. 특히 Wikipedia QA 및 긴 서사 QA 데이터셋에서도 Qwen-Reranker-4B 및 GroupRank-32B를 능가하는 결과를 보였다.

4B 규모의 작은 모델임에도 불구하고 32B 모델보다 높은 Recall@k를 기록하여 모델 크기 대비 효율성을 입증했다. Wikipedia 데이터셋인 Musique와 HotpotQA에서 QRRanker는 각각 Recall@3 70.19, 95.05를 기록하며 기존 방식들을 큰 폭으로 앞질렀다.

효율성 분석 결과, 상위 레이어를 제거한 QRRanker(middle) 버전은 Qwen3-Reranker-4B 대비 지연 시간을 약 50% 단축하면서도 성능 하락이 거의 없었다. TFLOPs 기준 연산량 또한 기존 재순위화 모델 대비 약 30% 이상 감소하여 실무 적용 가능성을 높였다.

QRRanker는 Qwen3-4B-Instruct를 백본으로 사용하며, 36개 레이어 중 중간 레이어에 위치한 16개의 QR-head를 활용하도록 설계되었다. 핵심 점수 계산 수식은 질문 토큰들이 특정 문서 토큰들에 할당한 어텐션 가중치의 평균값으로 정의된다. 이는 모델이 텍스트를 생성하기 전 내부적으로 수행하는 정보 취합 과정을 직접적인 랭킹 신호로 변환한 것이다.

학습 시에는 Group Contrastive Loss를 사용하여 샘플 내의 여러 정답 문서를 동시에 최적화한다. 이는 단일 정답만 고려하던 기존의 포인트웨이즈 방식보다 긴 문맥 내 다중 증거 탐색에 훨씬 유리한 구조를 제공한다. 또한 레이어 절단 기법을 통해 24번 레이어 이후를 제거해도 성능이 유지됨을 확인했는데, 이는 검색 및 관련성 판단에 필요한 핵심 정보가 모델의 중간 단계에서 이미 충분히 처리됨을 시사한다.

메모리 인식 구조를 위해 블록 기반 요약(Block-based Summary)과 이벤트 중심 요약(Event-centric Summary) 두 가지 전략을 사용한다. 긴 소설의 경우 시간 흐름에 따른 블록 요약을, 대화의 경우 주요 사건 중심의 요약을 프롬프트에 주입하여 모델이 세부 문서를 평가할 때 전역적인 가이드라인으로 삼게 한다.

현재 평가는 4B 규모의 단일 모델군에 집중되어 있어, 더 큰 규모나 다른 아키텍처를 가진 모델로의 일반화 가능성에 대한 추가 검증이 필요하다. 또한 학습 과정에서 사람이 직접 라벨링한 데이터가 아닌 모델이 생성한 'Silver' 데이터에 의존하고 있어 잠재적인 라벨 노이즈 문제가 존재할 수 있다.