Pruning RAG context down to what the answer actually needs

TL;DR Highlight

kapa.ai가 소형 LLM을 활용해 RAG 파이프라인의 retrieved chunk 68%를 걸러내면서도 recall 96%를 유지하고 쿼리 비용을 1/3 줄인 방법을 공개했다.

Who Should Read

RAG 기반 AI 어시스턴트나 검색 파이프라인을 운영 중인데 LLM 호출 비용이 부담스럽고 context 최적화를 고민 중인 백엔드/ML 엔지니어.

Core Mechanics

kapa.ai는 기술 문서, API 레퍼런스, 포럼, 지원 스레드 등 대규모 지식베이스에 답변하는 AI 어시스턴트를 만드는데, retriever가 찾아온 chunk들이 쿼리 비용의 약 2/3를 차지한다. 답변 텍스트, 대화 이력, 시스템 프롬프트를 합친 것보다 retrieved chunk 비용이 더 크다.
기존 파이프라인은 retriever → reranker → generator(비싼 LLM) 순서였는데, reranker와 generator 사이에 '소형 LLM pruner'를 한 단계 추가했다. 이 pruner가 질문과 모든 chunk를 한꺼번에 보고 불필요한 chunk를 버린다.
결과적으로 chunk의 68%를 제거하면서 recall은 96%를 유지했고, pruner 자체 비용을 감안해도 쿼리당 총 비용이 약 1/3 감소했다. chunk 하나를 줄일 때마다 쿼리 비용이 약 4% 절감된다.
rerank 점수에 threshold를 설정해 자르는 방식(예: 0.7 이상만 유지)은 두 가지 이유로 실패했다. 첫째, rerank 점수는 순서를 매기는 것이지 절대적 척도가 아니라서 쿼리 간 비교가 안 된다(Cohere도 이를 공식 인정). 둘째, pointwise cross-encoder(각 chunk를 개별적으로 점수 매기는 방식)는 chunk 간 관계를 보지 못한다.
pointwise scoring의 핵심 한계는 '청크 하나만 봐서는 관련성을 판단할 수 없다'는 것이다. 예를 들어 audit log 관련 질문에서 두 번째 chunk가 'audit log'를 직접 언급하지 않아 낮은 점수를 받지만, 첫 번째 chunk와 함께 봐야 비로소 답의 절반이 된다. 개별 점수가 아니라 chunk 집합 전체가 질문에 답할 수 있는지를 봐야 한다.
anchor document 방식(Sinhababu et al.)도 시도했다. Essential부터 Unrelated까지 수준별 합성 chunk를 ranking에 심어서 reranker의 절대적 기준을 만들려 했지만, 간접적으로만 관련된 chunk가 명백히 무관한 chunk보다 낮은 점수를 받는 문제는 해결하지 못했다. anchor를 낮게 설정해야 이런 chunk를 살릴 수 있어서 pruning 효과가 거의 없었다.
최종 해결책은 LLM에게 5단계 척도(ESSENTIAL / HIGH / MEDIUM / LOW / UNRELATED)로 모든 chunk를 한 번에 채점하게 하는 'listwise grading'이다. 질문과 모든 chunk를 동시에 보기 때문에 chunk 간 관계, 멀티파트 질문 분산, 간접 관련성을 모두 고려할 수 있다.
에이전트 환경에서 이 최적화가 특히 중요한데, 에이전트는 tool call마다 결과가 context에 쌓여서 빠르게 context가 커진다. pruning으로 context를 줄이면 에이전트가 추가 작업을 위한 공간을 확보할 수 있다.

Evidence

소형 LLM에 pruning을 맡기는 것에 대한 신뢰 문제가 제기됐다. 어렵게 검색해 낸 정보를 소형 LLM이 걸러내는 과정에서 지적 병목이 생길 수 있다는 우려로, 소형 모델이 복잡한 기술 문서의 미묘한 관련성을 제대로 판단할 수 있는지에 대한 의문이다.
희귀하지만 결정적인 chunk가 pruning에서 탈락할 위험성에 대한 날카로운 지적이 있었다. 예를 들어 약물 A와 B의 상호작용을 묻는 질문에서 경고문이 B의 금기 항목에 A를 언급하지 않고 쓰여 있다면, semantic pruning이 그 chunk를 먼저 버릴 수 있다. 반론으로 구조 기반 pruning(개별 문장 점수가 아닌 섹션/레코드 단위로 타입별 제거)이 더 안전할 수 있다는 대안이 제시됐다.
reranker를 완전히 pruner로 대체하는 실험을 했는지에 대한 질문이 있었다. 원문에서는 reranker와 generator 사이에 pruner를 추가하는 방식으로 설명했는데, 두 단계를 통합하는 가능성에 대한 궁금증이다.
관련 없거나 느슨하게 연결된 context가 오히려 최종 답변 품질을 떨어뜨린다는 실용적 관점의 댓글이 있었다. 이는 pruning의 정당성을 지지하는 현장 경험으로, 많은 context가 항상 좋은 것은 아님을 시사한다.
context가 에너지와 처리 시간에 얼마나 민감한지를 언급하며, 이 접근법이 비용 절감뿐 아니라 에너지 효율 측면에서도 의미 있다는 의견이 나왔다. 또한 LLM이 'knob'(파라미터를 의미하는 단어)를 사용한 것을 보고 이 글이 AI로 작성된 것 아닌지 의심하는 재미있는 댓글도 있었다.

How to Apply

RAG 파이프라인에서 retrieved chunk 비용이 전체 LLM 비용의 절반 이상을 차지하는 경우, reranker와 generator 사이에 소형 LLM pruner 단계를 추가해 chunk를 5단계로 채점하고 하위 단계를 버리면 전체 비용을 최대 1/3까지 줄일 수 있다.
여러 tool call을 연속으로 수행하는 에이전트를 구현하는 경우, 각 tool call 결과를 context에 그대로 쌓지 말고 listwise pruner를 통해 현재 질문에 필요한 chunk만 남기면 context window를 효율적으로 사용하고 이후 tool call을 위한 공간을 확보할 수 있다.
reranker 점수 threshold로 context를 자르는 방식을 이미 쓰고 있다면, 이 방법은 멀티파트 질문이나 간접 관련 chunk에서 실패한다. 질문과 모든 chunk를 동시에 보는 listwise LLM 채점 방식으로 교체하면 chunk 간 보완 관계를 고려한 더 정확한 pruning이 가능하다.
의료, 법률처럼 희귀하지만 결정적인 예외 정보가 중요한 도메인에서 RAG를 구축하는 경우, 점수 기반 개별 chunk pruning 대신 섹션/레코드 단위 구조 기반 pruning을 적용하거나 recall 테스트 셋을 '희귀 chunk가 정답인 케이스'로 구성해 pruner의 실제 성능을 검증해야 한다.

Terminology

listwise grading검색된 chunk들을 하나씩 개별로 평가하는 게 아니라 질문과 모든 chunk를 한꺼번에 LLM에 넣고 '이 질문을 답하는 데 각 chunk가 얼마나 필요한가'를 동시에 채점하는 방식.

pointwise cross-encoder질문-chunk 쌍을 하나씩 독립적으로 점수 매기는 reranker 방식. 다른 chunk들과의 관계를 전혀 고려하지 않아서 두 chunk가 함께 있어야 의미 있는 경우를 놓친다.

recall정답을 내는 데 필요한 chunk를 얼마나 빠짐없이 포함했는지 비율. recall 96%는 필요한 chunk 100개 중 96개를 유지했다는 뜻.

rerankerretriever가 넓게 건져온 수백 개의 후보 chunk를 질문과의 관련성 순으로 다시 정렬해주는 모델. 보통 cross-encoder 방식을 사용한다.

RAGRetrieval-Augmented Generation의 약자. LLM이 답변을 생성할 때 외부 지식베이스에서 관련 문서를 검색해 context로 제공하는 방식.

anchor documentreranker 점수를 절대적 척도로 만들기 위해 '이 chunk는 Essential이다', '이 chunk는 Unrelated다'처럼 관련성 수준이 미리 정해진 합성 chunk를 ranking에 심는 기법.