Introducing Contextual Retrieval
TMTAI 모델이 특정 맥락에서 유용하려면 배경 지식에 접근할 수 있어야 합니다. 예를 들어, 고객 지원 챗봇은 사용되는 특정 비즈니스에 대한 지식이 필요하고, 법률 분석 봇은 방대한 과거 사례에 대한 지식이 필요합니다.
개발자들은 일반적으로 Retrieval-Augmented Generation(RAG)이라는 방법을 사용해 AI 모델의 지식을 확장합니다. RAG는 지식 베이스에서 관련 정보를 검색해 사용자의 프롬프트에 추가함으로써 모델의 응답을 크게 향상시킵니다. 문제는 기존 RAG 솔루션이 정보를 인코딩할 때 맥락을 제거해, 시스템이 지식 베이스에서 관련 정보를 검색하지 못하는 경우가 많다는 점입니다.
이 글에서는 RAG의 검색 단계를 획기적으로 개선하는 방법을 소개합니다. 이 방법은 “Contextual Retrieval(맥락 기반 검색)”이라 하며, Contextual Embeddings와 Contextual BM25라는 두 가지 하위 기법을 사용합니다. 이 방법은 검색 실패를 49% 줄이고, 재정렬(reranking)과 결합하면 67%까지 줄일 수 있습니다. 이는 검색 정확도의 상당한 향상으로, 다운스트림 작업의 성능 향상으로 직결됩니다.
Claude와 함께 Contextual Retrieval 솔루션을 쉽게 배포할 수 있습니다.
단순히 더 긴 프롬프트를 사용하는 것에 대하여
가장 간단한 해결책이 최선일 때도 있습니다. 지식 베이스가 200,000 토큰(약 500페이지 분량)보다 작다면, 전체 지식 베이스를 모델에 주는 프롬프트에 그냥 포함시키면 되며, RAG나 유사한 방법이 필요 없습니다.
몇 주 전, Claude용 프롬프트 캐싱을 출시해 이 접근법을 훨씬 빠르고 비용 효율적으로 만들었습니다. 개발자는 이제 자주 사용하는 프롬프트를 API 호출 간에 캐시할 수 있어, 지연 시간을 2배 이상 줄이고 비용을 최대 90%까지 절감할 수 있습니다.(prompt caching cookbook)
하지만 지식 베이스가 커지면 더 확장성 있는 솔루션이 필요합니다. 이때 Contextual Retrieval이 등장합니다.
RAG 기초: 더 큰 지식 베이스로 확장하기
컨텍스트 윈도우에 들어가지 않는 대형 지식 베이스에는 RAG가 일반적인 해결책입니다. RAG는 다음과 같은 전처리 과정을 거칩니다:
- 지식 베이스(“코퍼스”)를 수백 토큰 이하의 작은 텍스트 청크로 분할합니다.
- 임베딩 모델을 사용해 이 청크들을 의미를 담은 벡터 임베딩으로 변환합니다.
- 이 임베딩을 벡터 데이터베이스에 저장해 의미적 유사성으로 검색할 수 있게 합니다.
실행 시, 사용자가 모델에 쿼리를 입력하면 벡터 데이터베이스가 쿼리와 의미적으로 가장 유사한 청크를 찾아냅니다. 그런 다음, 가장 관련성 높은 청크를 생성 모델에 전달할 프롬프트에 추가합니다.
임베딩 모델은 의미적 관계를 포착하는 데 뛰어나지만, 중요한 정확 일치를 놓칠 수 있습니다. 다행히 이런 상황에 도움이 되는 오래된 기법이 있습니다. BM25(Best Matching 25)는 어휘적 일치를 사용해 정확한 단어나 구를 찾는 랭킹 함수입니다. 고유 식별자나 기술 용어가 포함된 쿼리에 특히 효과적입니다.
BM25는 TF-IDF(단어 빈도-역문서 빈도) 개념을 기반으로 합니다. TF-IDF는 단어가 컬렉션 내 문서에서 얼마나 중요한지 측정합니다. BM25는 여기에 문서 길이와 용어 빈도에 포화 함수를 적용해, 일반 단어가 결과를 지배하지 않도록 개선합니다.
BM25가 의미 임베딩이 실패하는 상황에서 어떻게 성공하는지 예를 들어보겠습니다. 사용자가 기술 지원 데이터베이스에서 “Error code TS-999”를 쿼리하면, 임베딩 모델은 일반적인 오류 코드에 대한 내용을 찾을 수 있지만, 정확히 “TS-999”에 대한 일치는 놓칠 수 있습니다. BM25는 이 특정 텍스트 문자열을 찾아 관련 문서를 식별합니다.
RAG 솔루션은 임베딩과 BM25 기법을 결합해 다음과 같은 단계로 더 정확하게 관련 청크를 검색할 수 있습니다:
- 지식 베이스(“코퍼스”)를 수백 토큰 이하의 작은 텍스트 청크로 분할합니다.
- 이 청크들에 대해 TF-IDF 인코딩과 의미 임베딩을 생성합니다.
- BM25를 사용해 정확 일치 기반 상위 청크를 찾습니다.
- 임베딩을 사용해 의미적 유사성 기반 상위 청크를 찾습니다.
- (3)과 (4)의 결과를 랭크 퓨전 기법으로 결합하고 중복을 제거합니다.
- 상위 K개의 청크를 프롬프트에 추가해 응답을 생성합니다.
BM25와 임베딩 모델을 모두 활용하면, 기존 RAG 시스템은 정확한 용어 일치와 넓은 의미 이해를 균형 있게 제공해 더 포괄적이고 정확한 결과를 얻을 수 있습니다.
임베딩과 BM25를 모두 사용하는 표준 RAG 시스템. TF-IDF(단어 빈도-역문서 빈도)는 단어의 중요도를 측정하며 BM25의 기반이 됩니다.
이 접근법은 한 번에 프롬프트에 담을 수 없는 방대한 지식 베이스로 비용 효율적으로 확장할 수 있게 해줍니다. 하지만 이런 기존 RAG 시스템에는 중요한 한계가 있습니다. 바로 맥락을 자주 파괴한다는 점입니다.
기존 RAG의 맥락 문제
기존 RAG에서는 효율적인 검색을 위해 문서를 작은 청크로 분할하는 것이 일반적입니다. 이 방식은 많은 애플리케이션에서 잘 작동하지만, 개별 청크에 충분한 맥락이 없을 때 문제가 발생할 수 있습니다.
예를 들어, 금융 정보(예: 미국 SEC 공시) 모음이 지식 베이스에 있고, 다음과 같은 질문을 받았다고 가정해봅시다: “ACME Corp의 2023년 2분기 매출 성장률은 얼마였나요?”
관련 청크에는 _“회사의 매출은 전 분기 대비 3% 증가했습니다.”_라는 내용이 있을 수 있습니다. 하지만 이 청크만으로는 어떤 회사에 대한 것인지, 어떤 기간에 대한 것인지 명확하지 않아, 올바른 정보를 검색하거나 효과적으로 사용할 수 없습니다.
Contextual Retrieval 소개
Contextual Retrieval은 각 청크에 청크별 설명 맥락을 앞에 붙여 임베딩(“Contextual Embeddings”)과 BM25 인덱스(“Contextual BM25”)를 생성함으로써 이 문제를 해결합니다.
SEC 공시 모음 예시로 돌아가 보겠습니다. 청킹이 다음과 같이 변환됩니다:
original_chunk = "회사의 매출은 전 분기 대비 3% 증가했습니다."
contextualized_chunk = "이 청크는 ACME corp의 2023년 2분기 실적에 대한 SEC 공시에서 가져온 것입니다; 전 분기 매출은 3억 1,400만 달러였습니다. 회사의 매출은 전 분기 대비 3% 증가했습니다."맥락을 활용해 검색을 개선하는 다른 접근법도 과거에 제안된 바 있습니다. 예를 들어, adding generic document summaries to chunks(https://aclanthology.org/W02-0405.pdf)(실험 결과 효과 미미), hypothetical document embedding, HyDE, summary-based indexing(평가 결과 성능 낮음) 등이 있습니다. 이 방법들은 본문에서 제안하는 방식과 다릅니다.
Contextual Retrieval 구현
수천, 수백만 개의 청킹을 수작업으로 처리하는 것은 현실적으로 불가능합니다. Contextual Retrieval을 구현하기 위해 Claude를 사용합니다. Claude 3 Haiku에 다음과 같은 프롬프트를 작성해 각 청크에 대해 간결하고 청크별 맥락을 생성하도록 했습니다:
<document>
{{WHOLE_DOCUMENT}}
</document>
여기 전체 문서 내에서 위치를 설명하고 싶은 청크가 있습니다
<chunk>
{{CHUNK_CONTENT}}
</chunk>
검색을 개선하기 위해 이 청크를 전체 문서 내에서 어떻게 위치시킬 수 있을지 간결한 맥락을 제공해주세요. 맥락만 간단히 답변해주세요.생성된 맥락 텍스트(보통 50~100토큰)는 임베딩과 BM25 인덱스를 만들기 전에 청크 앞에 붙입니다.
실제 전처리 흐름은 다음과 같습니다:
Contextual Retrieval은 검색 정확도를 높이는 전처리 기법입니다.
Contextual Retrieval을 사용해보고 싶다면 바로 시작할 수 있습니다.
Contextual Retrieval의 비용을 줄이는 Prompt Caching 활용
Contextual Retrieval은 Claude의 특별한 프롬프트 캐싱 기능 덕분에 저렴하게 구현할 수 있습니다. 프롬프트 캐싱을 사용하면 각 청크마다 참조 문서를 매번 전달할 필요가 없습니다. 한 번만 문서를 캐시에 로드한 뒤, 이후에는 캐시된 내용을 참조하면 됩니다. 800토큰 청크, 8,000토큰 문서, 50토큰 맥락 지시문, 청크당 100토큰 맥락을 가정하면, 맥락화된 청크를 생성하는 일회성 비용은 백만 문서 토큰당 $1.02입니다.
방법론
다양한 지식 도메인(코드베이스, 소설, ArXiv 논문, 과학 논문), 임베딩 모델, 검색 전략, 평가 지표에 걸쳐 실험을 진행했습니다. 각 도메인별 질문과 답변 예시는 부록 II에서 확인할 수 있습니다.
아래 그래프는 상위 임베딩 구성(Gemini Text 004)과 상위 20개 청크 검색 시 모든 지식 도메인에서의 평균 성능을 보여줍니다. 평가 지표로는 1- recall@20(상위 20개 청크 내에서 검색되지 못한 관련 문서의 비율)을 사용했습니다. 전체 결과는 부록에서 확인할 수 있습니다. 맥락화는 모든 임베딩-소스 조합에서 성능을 향상시켰습니다.
성능 향상
실험 결과:
- Contextual Embeddings는 상위 20개 청크 검색 실패율을 35% 감소시켰습니다 (5.7% → 3.7%).
- Contextual Embeddings와 Contextual BM25를 결합하면 상위 20개 청크 검색 실패율이 49% 감소했습니다 (5.7% → 2.9%).
Contextual Embedding과 Contextual BM25를 결합하면 상위 20개 청크 검색 실패율이 49% 감소합니다.
구현 시 고려사항
Contextual Retrieval을 구현할 때 고려해야 할 점은 다음과 같습니다:
- 청크 경계: 문서를 어떻게 분할할지 고민하세요. 청크 크기, 경계, 중첩 여부가 검색 성능에 영향을 미칠 수 있습니다.
- 임베딩 모델: Contextual Retrieval은 테스트한 모든 임베딩 모델에서 성능을 향상시켰지만, 일부 모델이 더 큰 이득을 볼 수 있습니다. Gemini와 Voyage 임베딩이 특히 효과적이었습니다.
- 맞춤형 맥락화 프롬프트: 제공한 일반 프롬프트도 잘 작동하지만, 도메인이나 사용 사례에 맞게 프롬프트를 맞추면 더 나은 결과를 얻을 수 있습니다(예: 지식 베이스 내 다른 문서에만 정의된 용어집 포함).
- 청크 수: 컨텍스트 윈도우에 더 많은 청크를 넣으면 관련 정보를 포함할 확률이 높아집니다. 하지만 정보가 너무 많으면 모델에 방해가 될 수 있으니 한계가 있습니다. 5, 10, 20개 청크를 실험한 결과 20개가 가장 성능이 좋았으나, 사용 사례에 따라 실험해보길 권장합니다.
항상 평가를 실행하세요: 응답 생성 시 맥락화된 청크를 전달하고, 무엇이 맥락이고 무엇이 청크인지 구분하면 성능이 향상될 수 있습니다.
Reranking으로 성능 추가 향상
마지막 단계로, Contextual Retrieval을 또 다른 기법과 결합해 성능을 더욱 높일 수 있습니다. 기존 RAG에서는 AI 시스템이 지식 베이스에서 잠재적으로 관련 있는 정보 청크를 찾습니다. 대형 지식 베이스에서는 이 초기 검색이 수백 개의 다양한 관련성·중요도를 가진 청크를 반환하는 경우가 많습니다.
Reranking(재정렬)은 가장 관련성 높은 청크만 모델에 전달하도록 필터링하는 데 널리 쓰이는 기법입니다. Reranking은 더 나은 응답을 제공하고, 모델이 처리하는 정보가 줄어 비용과 지연 시간도 감소시킵니다. 주요 단계는 다음과 같습니다:
- 초기 검색으로 잠재적으로 관련성 높은 상위 청크(예: 150개)를 가져옵니다.
- 상위 N개 청크와 사용자 쿼리를 reranking 모델에 전달합니다.
- reranking 모델이 각 청크의 쿼리 관련성과 중요도를 점수화해 상위 K개(예: 20개)를 선택합니다.
- 상위 K개 청크를 컨텍스트로 모델에 전달해 최종 결과를 생성합니다.
Contextual Retrieval과 Reranking을 결합해 검색 정확도를 극대화합니다.
성능 향상
여러 reranking 모델이 시중에 나와 있습니다. 우리는 Cohere reranker로 테스트를 진행했습니다. Voyage도 reranker를 제공하지만, 시간상 테스트하지 못했습니다. 다양한 도메인에서 실험한 결과, reranking 단계를 추가하면 검색이 더욱 최적화됨을 확인했습니다.
구체적으로, Reranked Contextual Embedding과 Contextual BM25는 상위 20개 청크 검색 실패율을 67% 감소시켰습니다(5.7% → 1.9%).
Reranked Contextual Embedding과 Contextual BM25는 상위 20개 청크 검색 실패율을 67% 감소시킵니다.
비용 및 지연 시간 고려사항
Reranking의 한 가지 중요한 고려사항은, 많은 청크를 reranking할 때 비용과 지연 시간에 미치는 영향입니다. reranking은 실행 시 추가 단계가 들어가므로, 모든 청크를 병렬로 점수화하더라도 약간의 지연이 추가됩니다. 더 나은 성능을 위해 더 많은 청크를 reranking할지, 더 낮은 지연과 비용을 위해 적게 할지 트레이드오프가 있습니다. 사용 사례에 맞는 최적의 균형을 찾기 위해 다양한 설정을 실험해보길 권장합니다.
결론
위에서 설명한 모든 기법(임베딩 모델, BM25 사용, Contextual Retrieval 사용, reranker 사용, 상위 K개 결과 수)을 다양한 데이터셋에 걸쳐 조합해 대규모 테스트를 진행했습니다. 요약은 다음과 같습니다:
- 임베딩+BM25가 임베딩 단독보다 낫다.
- Voyage와 Gemini가 테스트한 임베딩 중 가장 우수하다.
- 상위 20개 청크를 모델에 전달하는 것이 10개, 5개보다 효과적이다.
- 청크에 맥락을 추가하면 검색 정확도가 크게 향상된다.
- reranking이 없을 때보다 있는 것이 낫다.
- 이 모든 이점은 누적된다: 성능을 극대화하려면, Voyage나 Gemini의 contextual embedding, contextual BM25, reranking 단계를 결합하고, 20개 청크를 프롬프트에 추가하면 된다.
지식 베이스를 다루는 모든 개발자에게 이 방법을 실험해 새로운 수준의 성능을 경험해보길 권장합니다.
Appendix I
아래는 데이터셋, 임베딩 제공자, 임베딩 외 BM25 사용 여부, Contextual Retrieval 사용 여부, reranking 사용 여부에 따른 Retrievals @ 20 결과 분포입니다.
Retrievals @ 10 및 @ 5의 분포와 각 데이터셋별 질문·답변 예시는 부록 II에서 확인할 수 있습니다.
1 minus recall @ 20 결과(데이터셋 및 임베딩 제공자별)