BM25 뒤에 리랭커를 붙이면 얼마나 좋아질까? AutoRAG로 직접 재 봤어요

공개 RAG 검색 데이터셋 AutoRAG(질문 114개, 문단 720개)에서 BM25 상위 10개를 cheon-reranker-0.6b-v1로 다시 줄 세웠더니 nDCG@10이 0.857에서 0.916으로, MRR@10이 0.822에서 0.898로 올랐어요. 데이터 받기부터 지표 계산까지 파이썬 코드를 전부 붙였고, 후보를 20개로 늘렸더니 오히려 떨어진 이유도 같이 적었어요.

RAG를 한 번이라도 붙여 봤다면 이런 장면이 낯설지 않을 거예요. 질문에 딱 맞는 문단이 검색 결과에 있긴 있는데, 1등이 아니라 4등이나 8등에 가 있어요. LLM에 상위 세 개만 넘기면 정답이 잘려 나가고, 스무 개를 다 넘기면 컨텍스트가 지저분해지죠.

리랭커는 이 순서를 바로잡는 모델이에요. 1차 검색이 후보를 넉넉하게 가져오면, 리랭커가 질문과 후보를 같이 읽고 줄을 다시 세워요. 말로만 하면 얼마나 좋아지는지 감이 잘 안 와서, 누구나 받을 수 있는 공개 데이터셋으로 직접 재 봤어요. 돌린 코드는 전부 아래에 붙여 뒀어요.

결과부터

AutoRAG 검색 데이터셋의 질문 114개에 BM25를 돌리고, 상위 10개를 cheon-reranker-0.6b-v1로 다시 정렬했어요.

방식 nDCG@10 MRR@10 정답이 1위인 질문
BM25 0.857 0.822 85개 (74.6%)
BM25 → 리랭커 (상위 10개) 0.916 0.898 96개 (84.2%)
BM25와 BM25 → 리랭커AutoRAG 검색 데이터셋, 질문 114개. 1위 적중률은 정답 문단이 1위에 온 질문의 비율이에요.
BM25BM25 → 리랭커 (상위 10개)0.00.20.40.60.81.0 BM25 · nDCG@10: 0.8570.857BM25 · MRR@10: 0.8220.822BM25 · 1위 적중률: 0.7460.746BM25 → 리랭커 (상위 10개) · nDCG@10: 0.9160.916BM25 → 리랭커 (상위 10개) · MRR@10: 0.8980.898BM25 → 리랭커 (상위 10개) · 1위 적중률: 0.8420.842 nDCG@10MRR@101위 적중률

질문별로 보면 정답 순위가 올라간 게 19개, 내려간 게 2개, 그대로인 게 93개였어요. 내려간 두 건은 뒤에서 따로 열어 볼게요. 생각보다 재밌는 이유가 있었거든요.

데이터: AutoRAG 검색 셋

AutoRAG 논문 팀이 공개한 RAG 평가 데이터예요. MTEB에도 AutoRAGRetrieval이라는 이름으로 들어가 있어서 Hugging Face에서 바로 받을 수 있어요(MIT 라이선스). 금융·공공·법률·커머스 PDF를 페이지 단위로 자른 문단 720개와 사람이 만든 질문 114개로 되어 있고, 질문마다 정답 문단이 딱 하나씩 붙어 있어요. 문단 길이는 평균 825자 정도예요.

분야 문단 질문
법률 282 37
공공 152 29
커머스 176 26
금융 110 22

받는 코드는 이게 전부예요.

from datasets import load_dataset

name = "mteb/AutoRAGRetrieval"
corpus = load_dataset(name, "corpus", split="test")
queries = load_dataset(name, "queries", split="test")
qrels = load_dataset(name, "qrels", split="test")

doc_ids = corpus["_id"]
doc_text = dict(zip(corpus["_id"], corpus["text"]))
answer = dict(zip(qrels["query-id"], qrels["corpus-id"]))  # 질문마다 정답 문단 하나

print(len(corpus), len(queries))  # 720 114

1단계: BM25로 후보 뽑기

한국어는 조사가 단어에 붙어 있어서, 띄어쓰기로만 자르면 '예비인가를'과 '예비인가'가 서로 다른 단어가 돼요. 그래서 형태소 분석기 Kiwi로 먼저 자르고 BM25를 돌렸어요. BM25 파라미터는 rank_bm25 기본값(k1=1.5, b=0.75)을 그대로 썼어요.

import numpy as np
from kiwipiepy import Kiwi
from rank_bm25 import BM25Okapi

kiwi = Kiwi()

def tokenize(text):
    return [token.form for token in kiwi.tokenize(text)]

bm25 = BM25Okapi([tokenize(text) for text in corpus["text"]])

def bm25_search(query, k=100):
    scores = bm25.get_scores(tokenize(query))
    order = np.argsort(-scores, kind="stable")[:k]
    return [doc_ids[i] for i in order]

runs = {q["_id"]: bm25_search(q["text"]) for q in queries}

솔직히 BM25 점수를 처음 보고 조금 놀랐어요. nDCG@10이 0.857이면 꽤 높은 편이거든요. 질문이 문서에 나온 표현을 그대로 가져다 쓴 경우가 많아서, 단어만 맞춰도 정답이 잘 걸리는 데이터예요. 정답이 상위 10개 안에 들어온 질문이 110개(96.5%), 20개 안으로 넓히면 112개(98.2%)였어요.

리랭커는 후보에 없는 정답을 끌어올 수 없으니 이 숫자가 사실상 천장이에요. 그러니까 이번 실험의 진짜 질문은 "BM25만으로도 이미 잘 되는 데이터에서 리랭커가 할 일이 남아 있나"였던 셈이에요.

2단계: 리랭커로 다시 줄 세우기

cheon-reranker는 흔히 보는 cross-encoder와 조금 달라요. 질문-문단 쌍마다 따로 점수를 매기지 않고, 질문과 후보 전부를 한 시퀀스에 넣어서 한 번에 읽어요. 그 위에서 후보 벡터끼리 서로 참고하는 작은 트랜스포머가 한 번 더 돌고 나서 점수를 내요. 여기에는 위치 인코딩이 없어서 후보 순서를 섞어도 점수는 똑같아요.

그래서 입력을 만들 때 각 후보가 시퀀스의 몇 번째 토큰부터 몇 번째까지인지(doc_spans)를 같이 넘겨야 해요. 모델 카드의 예제를 함수 하나로 묶으면 이렇게 돼요.

import torch
from transformers import AutoModel, AutoTokenizer

repo = "cheonai/cheon-reranker-0.6b-v1"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModel.from_pretrained(repo, trust_remote_code=True, torch_dtype=torch.float32)
model = model.eval().to("cuda")

MAX_TOKENS = model.config.max_sequence_length             # 12,288
MAX_CANDIDATES = model.config.max_candidates_per_context  # 50

@torch.no_grad()
def rerank(query, passages):
    assert len(passages) <= MAX_CANDIDATES
    ids = tok(f"query: {query.strip()}\n", add_special_tokens=False)["input_ids"]
    # 전부 한 시퀀스에 들어가도록 후보 하나가 쓸 수 있는 토큰을 똑같이 나눠요
    budget = (MAX_TOKENS - len(ids)) // len(passages)
    spans = []
    for text in passages:
        piece = tok(f"candidate: {text.strip()}", add_special_tokens=False)["input_ids"][:budget]
        spans.append((len(ids), len(ids) + len(piece)))
        ids += piece
    input_ids = torch.tensor([ids], device=model.device)
    scores = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids), doc_spans=spans)
    return scores.float().cpu().tolist()

budget 줄을 눈여겨봐 주세요. 이 모델은 질문과 후보 전부가 12,288토큰 안에 들어가야 해서, 후보 하나가 쓸 수 있는 토큰을 똑같이 나누고 넘치는 뒷부분은 잘라요. 이 한 줄이 뒤에서 결과를 꽤 크게 바꿔요.

BM25 결과를 받아 상위 k개만 다시 정렬하는 부분은 이게 다예요. 점수는 확률이 아니라 상대적인 순서 신호라서 정렬에만 써요.

def rerank_run(runs, k=10):
    out = {}
    for q in queries:
        candidates = runs[q["_id"]][:k]
        scores = rerank(q["text"], [doc_text[d] for d in candidates])
        order = sorted(range(len(candidates)), key=lambda i: -scores[i])
        out[q["_id"]] = [candidates[i] for i in order] + runs[q["_id"]][k:]
    return out

reranked = rerank_run(runs, k=10)

측정은 저희 GPU 서버에서 위와 같은 입력 형식, fp32로 돌렸어요. 질문 하나에 모델 호출이 한 번이고, 이 데이터에서 후보 10개를 넣으면 시퀀스 길이가 중간값 기준 6,636토큰쯤 돼요.

지표: nDCG@10과 MRR@10

정답 문단이 하나뿐이라 식이 아주 단순해져요. 질문 qq의 정답이 rqr_q등에 있다고 하면 이렇게 돼요.

nDCG@10=1∣Q∣∑q∈Q1[rq≤10]log⁡2(rq+1)\mathrm{nDCG@10} = \frac{1}{|Q|} \sum_{q \in Q} \frac{\mathbb{1}[r_q \le 10]}{\log_2(r_q + 1)}
MRR@10=1∣Q∣∑q∈Q1[rq≤10]rq\mathrm{MRR@10} = \frac{1}{|Q|} \sum_{q \in Q} \frac{\mathbb{1}[r_q \le 10]}{r_q}

원래 nDCG는 관련도가 여러 단계인 경우까지 다루는 지표예요. DCG@k=∑i=1k2reli−1log⁡2(i+1)\mathrm{DCG@k} = \sum_{i=1}^{k} \frac{2^{rel_i} - 1}{\log_2(i + 1)}를 가장 이상적인 순서일 때의 DCG(IDCG)로 나눈 값인데, 정답이 하나고 관련도가 1이면 IDCG가 1이 돼서 위 식으로 줄어들어요.

두 지표가 순위를 얼마나 깎는지 나란히 놓으면 성격 차이가 보여요.

정답 순위 nDCG@10 MRR@10
1위 1.000 1.000
2위 0.631 0.500
3위 0.500 0.333
5위 0.387 0.200
10위 0.289 0.100
11위 이하, 또는 없음 0 0

MRR은 2위만 돼도 반토막이 나요. LLM에 문단을 하나만 넘기는 구성이라면 MRR 쪽이 체감에 더 가까워요. 계산 코드는 이렇습니다.

import math

def rank_of(ranked, doc_id):
    return ranked.index(doc_id) + 1 if doc_id in ranked else None

def report(run, k=10):
    ranks = [rank_of(run[q], answer[q]) for q in run]
    found = [r for r in ranks if r is not None and r <= k]
    scores = {
        f"nDCG@{k}": sum(1 / math.log2(r + 1) for r in found) / len(ranks),
        f"MRR@{k}": sum(1 / r for r in found) / len(ranks),
        "Hit@1": sum(r == 1 for r in ranks) / len(ranks),
    }
    return {name: round(value, 3) for name, value in scores.items()}

print(report(runs))      # {'nDCG@10': 0.857, 'MRR@10': 0.822, 'Hit@1': 0.746}
print(report(reranked))  # {'nDCG@10': 0.916, 'MRR@10': 0.898, 'Hit@1': 0.842}

분야별로 보면

분야별 nDCG@10괄호 안은 질문 수예요.
BM25BM25 → 리랭커 (상위 10개)0.00.20.40.60.81.0 BM25 · 법률 (37): 0.8540.854BM25 · 공공 (29): 0.8060.806BM25 · 커머스 (26): 0.8730.873BM25 · 금융 (22): 0.9100.910BM25 → 리랭커 (상위 10개) · 법률 (37): 0.9410.941BM25 → 리랭커 (상위 10개) · 공공 (29): 0.8840.884BM25 → 리랭커 (상위 10개) · 커머스 (26): 0.9280.928BM25 → 리랭커 (상위 10개) · 금융 (22): 0.8990.899 법률 (37)공공 (29)커머스 (26)금융 (22)

가장 크게 오른 건 법률이에요(0.854 → 0.941). 판결문이나 법령 문단은 같은 조문 번호와 용어가 여러 문단에 반복돼서, 단어가 많이 겹친다고 정답이 되지는 않아요. 이 모델이 원래 법률·행정 문서 검색을 염두에 두고 만든 모델이기도 하고요. 공공(0.806 → 0.884)과 커머스(0.873 → 0.928)도 꽤 올랐어요.

금융만 살짝 내려갔어요(0.910 → 0.899). 질문이 22개밖에 없는데, 그중 두 건이 1위에서 2위로 밀린 영향이에요.

올라간 질문, 내려간 질문

올라간 쪽부터 하나 볼게요.

라이브커머스 이용률이 가장 높은 서비스는 무엇이며 그 이유는 무엇인가요?

BM25의 1위는 한 이커머스 솔루션 소개 자료였어요. "자사몰에서 라이브 커머스를 운영하는 것에는 어떤 장점이 있을까요?" 같은 문장이 들어 있어서 질문과 단어가 많이 겹쳤거든요. 정작 정답인 산업 보고서 페이지("라이브커머스에서는 네이버쇼핑라이브의 이용률이 압도적으로 높습니다. 포털에서 바로 접근할 수 있고, 구매 편의성이 높아…")는 8위에 있었어요. 리랭커는 이걸 1위로 올렸어요. 단어가 겹치는 문단이 아니라 질문에 실제로 답하는 문단을 고른 거죠.

2024년에 새로 만드는 'K-스카우터' 제도를 묻는 질문도 BM25에서는 정답이 10위였는데, 리랭커를 거치고 1위가 됐어요.

이제 내려간 두 건이에요. 둘 다 금융 분야고, 둘 다 1위에서 2위로 한 칸 내려갔어요.

은행법에 의거 예비인가를 신청할 수 있는지와, 그 경우 금융위원회가 검토했어야 하는 요건은 무엇인가요?

정답으로 지정된 문단은 은행법 제8조의 인가 요건 같은 관련 조문을 모아 둔 참고 페이지예요. 리랭커가 1위로 올린 건 같은 문서의 두 쪽 앞 페이지로, "은행업 본인가를 받으려는 자는 예비인가를 신청할 수 있으며…"라고 질문의 앞부분에 정면으로 답하는 문단이었어요. 질문은 두 가지를 한꺼번에 묻는데 정답은 하나만 지정돼 있으니, 어느 쪽을 1위로 둬도 반쯤은 맞는 셈이에요.

현행 계약형 퇴직연금제도와 기금형 퇴직연금제도의 주요 차이점은 무엇인가요?

이 건은 더 재밌어요. 정답은 한 포럼 책자에 실린 축사의 첫 페이지인데, 핵심 문장인 "현행 계약형 제도는 DC제도에 가입된 노동자의 수급권 보호에 한계가 있기 때문입니다"가 페이지 경계에서 둘로 쪼개져 있어요. 리랭커가 1위로 고른 다음 페이지에는 그 문장의 뒷부분과 함께 기금형 제도가 왜 필요한지에 대한 설명이 이어져요. 점수로는 틀린 거지만, 사람이 봐도 쉽게 틀렸다고 하기 어려운 선택이었어요.

정답을 하나만 지정한 데이터는 이렇게 "같이 맞는 문단"을 틀린 것으로 쳐요. 문서를 페이지 단위로 자른 데이터라면 더 그렇고요.

후보를 20개로 늘리면 더 좋아질까?

BM25 상위 20개로 넓히면 정답이 두 개 더 들어와요(110개 → 112개). 천장이 올라가니 점수도 오를 거라고 기대했는데, 결과는 반대였어요.

설정 후보 하나당 최대 토큰 잘린 문단 nDCG@10 MRR@10
BM25만 - - 0.857 0.822
상위 10개 약 1,220 7% 0.916 0.898
상위 20개 약 610 64% 0.888 0.862
상위 10개, 610토큰으로 자름 약 610 62% 0.890 0.865

이유는 토큰 예산이에요. 이 데이터의 문단은 후보 형식으로 넣으면 평균 585토큰, 긴 건 1,683토큰이에요. 후보가 10개면 하나당 약 1,220토큰까지 쓸 수 있어서 잘리는 문단이 7%뿐이지만, 20개면 하나당 약 610토큰이라 64%가 뒷부분을 잃어요.

범인이 후보 수인지 자르기인지 가려 보려고, 후보는 10개 그대로 두고 문단만 20개일 때처럼 610토큰으로 잘라서 한 번 더 돌려 봤어요(위 rerank()에서 budget을 후보 20개 기준으로 계산하면 돼요). 0.890이 나왔어요. 후보를 20개로 늘렸을 때(0.888)와 거의 같아요. 점수가 떨어진 건 거의 전부 문단을 잘라서였던 거죠.

적어도 이 모델과 이 데이터에서는 후보를 많이 넣는 것보다 후보를 온전히 넣는 게 더 중요했어요. 문단이 긴 데이터라면 후보 수를 무작정 늘리기 전에, 문단이 잘리지 않는 선이 어디까지인지부터 보는 게 좋겠어요.

해 보고 나서

BM25가 이미 강한 데이터였는데도, 리랭커 한 단으로 정답이 1위에 오는 질문이 85개에서 96개로 늘었어요. LLM에 상위 한두 개만 넘기는 RAG라면 이 11개가 곧 답변이 달라지는 질문이에요.

남은 네 건은 정답이 BM25 상위 10개 밖(17위, 19위, 25위, 96위)에 있어서 리랭커가 손댈 수 없었어요. 여기서 더 올리려면 리랭커가 아니라 1차 검색 쪽을 손봐야 해요.

그리고 질문 114개는 작은 셋이에요. 금융 22개처럼 쪼개면 한두 문제로 점수가 크게 흔들려요. 분야별 숫자는 방향 정도로만 봐 주세요.

직접 돌려 보려면

필요한 패키지는 이 정도예요.

pip install datasets kiwipiepy rank_bm25 transformers torch

위 코드를 순서대로 이어 붙이면 같은 실험을 해 볼 수 있어요. GPU 종류나 라이브러리 버전에 따라 소수점 셋째 자리가 조금 달라질 수는 있어요. 모델은 CC-BY-NC-4.0 라이선스라 연구와 평가에는 자유롭게 쓰실 수 있고, 서비스에 넣으려면 문의해 주세요. 설치 없이 감만 보고 싶다면 리랭커 플레이그라운드에 질문과 문단을 붙여 넣어 보셔도 돼요.