ColBERT가 개발자가 RAG의 한계를 극복하도록 돕는 방법

검색 증강 생성(RAG)은 이제 생성 인공 지능(AI) 애플리케이션의 표준 부분 입니다. 벡터 데이터베이스에서 검색된 관련 컨텍스트로 애플리케이션 프롬프트를 보완하면 정확도가 크게 향상되고 환각이 줄어들 수 있습니다. 이는 벡터 검색 결과의 관련성 증가가 RAG 애플리케이션의 품질과 직접적인 상관관계가 있음을 의미합니다.

LLM(대형 언어 모델)이 컨텍스트 창을 늘려도 RAG가 여전히 인기를 끌고 관련성이 점점 높아지는 데는 두 가지 이유가 있습니다.

1. LLM 응답 시간과 가격은 모두 컨텍스트 길이에 따라 선형적으로 증가합니다.

2. LLM은 여전히 방대한 맥락에서 검색 과 추론 에 어려움을 겪고 있습니다.

   
   

그러나 RAG는 마술 지팡이가 아닙니다. 특히, 가장 일반적인 디자인인 DPR(Dense Passage Retrieval)은 쿼리와 구절을 모두 단일 임베딩 벡터로 표현하고 직접적인 코사인 유사성을 사용하여 관련성을 평가합니다. 이는 DPR이 모든 관련 검색어를 인식하기 위해 광범위한 교육을 받은 임베딩 모델에 크게 의존한다는 것을 의미합니다.

불행하게도 기성 모델은 훈련 데이터에 일반적으로 포함되지 않는 이름을 포함한 특이한 용어로 인해 어려움을 겪습니다. DPR은 또한 청킹 전략에 과민한 경향이 있는데, 이로 인해 관련 없는 정보가 많이 주변에 있으면 관련 구절을 놓칠 수 있습니다. 이 모든 것은 애플리케이션 개발자에게 "처음에 제대로 해야 하는" 부담을 안겨줍니다. 왜냐하면 일반적으로 실수로 인해 인덱스를 처음부터 다시 작성해야 하기 때문입니다.

ColBERT를 통해 DPR의 과제 해결

ColBERT는 DPR의 문제를 실질적으로 해결하는 BERT 언어 모델을 사용하여 구절 관련성을 평가하는 새로운 방법입니다. 첫 번째 ColBERT 논문 의 다이어그램은 이것이 왜 그렇게 흥미로운지 보여줍니다.

이는 ColBERT의 성능을 MS-MARCO 데이터 세트에 대한 다른 최첨단 솔루션과 비교합니다. (MS-MARCO는 Microsoft가 가장 관련성이 높은 부분을 직접 채점한 Bing 쿼리 집합입니다. 이는 더 나은 검색 벤치마크 중 하나입니다.) 낮고 오른쪽이 더 좋습니다.

간단히 말해서, ColBERT는 대기 시간이 약간 증가하는 대신 훨씬 더 복잡한 솔루션 분야를 쉽게 능가합니다.

이를 테스트하기 위해 데모를 만들고 ada002 DPR과 ColBERT를 모두 사용하여 1,000개 이상의 Wikipedia 기사를 색인화했습니다. 나는 ColBERT가 특이한 검색어에 대해 훨씬 더 나은 결과를 제공한다는 것을 발견했습니다.

다음 스크린샷은 DPR이 Abraham Lincoln의 동료인 William H. Herndon의 특이한 이름을 인식하지 못하는 반면 ColBERT는 Springfield 기사에서 참조를 찾는 것을 보여줍니다. 또한 ColBERT의 2위 결과는 다른 William에 대한 것이지만 DPR의 결과는 관련이 없습니다.

ColBERT는 종종 밀집된 기계 학습 전문 용어로 설명되지만 실제로는 매우 간단합니다. 단 몇 줄의 Python 및 CQL( Cassandra Query Language )을 사용하여 DataStax Astra DB 에서 ColBERT 검색 및 채점을 구현하는 방법을 보여 드리겠습니다.

큰 아이디어

구절을 단일 "임베딩" 벡터로 바꾸는 기존의 단일 벡터 기반 DPR 대신 ColBERT는 구절의 각 토큰에 대해 상황에 따라 영향을 받는 벡터를 생성합니다. ColBERT는 마찬가지로 쿼리의 각 토큰에 대한 벡터를 생성합니다.

(토큰화는 LLM에서 처리하기 전에 입력을 단어의 일부로 나누는 것을 의미합니다. OpenAI 팀의 창립 멤버인 Andrej Karpathy가 방금 이것이 어떻게 작동하는지에 대한 뛰어난 비디오를 공개했습니다 .)

그런 다음 각 문서의 점수는 문서 임베딩에 임베딩된 각 쿼리의 최대 유사성의 합계입니다.

 def maxsim(qv, document_embeddings): return max(qv @ dv for dv in document_embeddings) def score(query_embeddings, document_embeddings): return sum(maxsim(qv, document_embeddings) for qv in query_embeddings)

(@는 내적에 대한 PyTorch 연산자이며 벡터 유사성의 가장 일반적인 척도 입니다.)

그게 다입니다. Python 네 줄로 ColBERT 채점을 구현할 수 있습니다! 이제 당신은 ColBERT에 대해 X(이전의 Twitter)에 게시하는 사람들의 99%보다 더 잘 이해하고 있습니다.

ColBERT 논문의 나머지 부분은 다음을 다룹니다.

1. 주어진 데이터 세트에 대해 최상의 임베딩을 생성하기 위해 BERT 모델을 어떻게 미세 조정합니까?
2. 여기에 표시된 (상대적으로 비용이 많이 드는) 점수를 계산하는 문서 세트를 어떻게 제한합니까?

첫 번째 질문은 선택 사항이며 이 글의 범위를 벗어납니다. 사전 훈련된 ColBERT 체크포인트를 사용하겠습니다. 그러나 두 번째는 DataStax Astra DB와 같은 벡터 데이터베이스를 사용하는 것이 간단합니다.

Astra DB의 ColBERT

RAGatouille 이라는 인기 있는 ColBERT용 Python 올인원 라이브러리가 있습니다. 그러나 이는 정적 데이터세트를 가정합니다. RAG 애플리케이션의 강력한 기능 중 하나는 실시간으로 동적으로 변화하는 데이터 에 응답하는 것입니다. 따라서 대신 Astra의 벡터 인덱스를 사용하여 각 하위 벡터에 가장 적합한 후보로 점수를 매기는 데 필요한 문서 세트의 범위를 좁힐 것입니다.

ColBERT를 RAG 애플리케이션에 추가하는 경우 수집과 검색이라는 두 단계가 있습니다.

음식물 섭취

각 문서 청크에는 여러 개의 임베딩이 연결되어 있으므로 두 개의 테이블이 필요합니다.

 CREATE TABLE chunks ( title text, part int, body text, PRIMARY KEY (title, part) ); CREATE TABLE colbert_embeddings ( title text, part int, embedding_id int, bert_embedding vector<float, 128>, PRIMARY KEY (title, part, embedding_id) ); CREATE INDEX colbert_ann ON colbert_embeddings(bert_embedding) WITH OPTIONS = { 'similarity_function': 'DOT_PRODUCT' };

ColBERT 라이브러리( pip install colbert-ai )를 설치하고 사전 학습된 BERT 체크포인트를 다운로드한 후 다음 테이블에 문서를 로드할 수 있습니다.

 from colbert.infra.config import ColBERTConfig from colbert.modeling.checkpoint import Checkpoint from colbert.indexing.collection_encoder import CollectionEncoder from cassandra.concurrent import execute_concurrent_with_args from db import DB def encode_and_save(title, passages): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encoder = CollectionEncoder(cf, cp) # encode_passages returns a flat list of embeddings and a list of how many correspond to each passage embeddings_flat, counts = encoder.encode_passages(passages) # split up embeddings_flat into a nested list start_indices = [0] + list(itertools.accumulate(counts[:-1])) embeddings_by_part = [embeddings_flat[start:start+count] for start, count in zip(start_indices, counts)] # insert into the database for part, embeddings in enumerate(embeddings_by_part): execute_concurrent_with_args(db.session, db.insert_colbert_stmt, [(title, part, i, e) for i, e in enumerate(embeddings)])

(저는 DB 로직을 전용 모듈에 캡슐화하는 것을 좋아합니다. 전체 소스는 GitHub 저장소 에서 액세스할 수 있습니다.)

검색

그런 다음 검색은 다음과 같습니다.

 def retrieve_colbert(query): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encode = lambda q: cp.queryFromText([q])[0] query_encodings = encode(query) # find the most relevant documents for each query embedding. using a set # handles duplicates so we don't retrieve the same one more than once docparts = set() for qv in query_encodings: rows = db.session.execute(db.query_colbert_ann_stmt, [list(qv)]) docparts.update((row.title, row.part) for row in rows) # retrieve these relevant documents and score each one scores = {} for title, part in docparts: rows = db.session.execute(db.query_colbert_parts_stmt, [title, part]) embeddings_for_part = [tensor(row.bert_embedding) for row in rows] scores[(title, part)] = score(query_encodings, embeddings_for_part) # return the source chunk for the top 5 return sorted(scores, key=scores.get, reverse=True)[:5]

가장 관련성이 높은 문서 부분( db.query_colbert_ann_stmt )에 대해 실행되는 쿼리는 다음과 같습니다.

 SELECT title, part FROM colbert_embeddings ORDER BY bert_embedding ANN OF ? LIMIT 5

기본 그 이상: RAGStack

이 기사와 링크된 저장소는 ColBERT의 작동 방식을 간략하게 소개합니다. 오늘 자신의 데이터로 이를 구현하고 즉각적인 결과를 확인할 수 있습니다. AI의 모든 것과 마찬가지로 모범 사례는 매일 바뀌고 있으며 새로운 기술이 끊임없이 등장하고 있습니다.

최신 기술을 더 쉽게 따라갈 수 있도록 DataStax는 LangChain 및 LlamaIndex를 활용하는 프로덕션 지원 RAG 라이브러리인 RAGStack 에 이 기능과 기타 개선 사항을 적용하고 있습니다. 우리의 목표는 개발자에게 새로운 기능으로의 발전을 제어할 수 있는 RAG 애플리케이션용 일관된 라이브러리를 제공하는 것입니다. 기술과 라이브러리의 수많은 변경 사항을 따라갈 필요 없이 단일 스트림이 있으므로 애플리케이션 구축에 집중할 수 있습니다. 지금 바로 RAGStack을 사용하여 LangChain 및 LlamaIndex에 대한 모범 사례를 통합할 수 있습니다. ColBERT와 같은 발전은 향후 릴리스에서 RAGstack에 적용될 예정입니다.

작성자: 조나단 엘리스(Jonathan Ellis), DataStax