마트료시카 임베딩을 사용한 퍼널 검색
효율적인 벡터 검색 시스템을 구축할 때, 한 가지 핵심 과제는 허용 가능한 지연 시간과 리콜을 유지하면서 스토리지 비용을 관리하는 것입니다. 최신 임베딩 모델은 수백 또는 수천 개의 차원을 가진 벡터를 출력하기 때문에 원시 벡터와 인덱스에 상당한 스토리지와 계산 오버헤드를 발생시킵니다.
전통적으로, 인덱스 구축 직전에 양자화 또는 차원 축소 방법을 적용하여 저장소 요구 사항을 줄입니다. 예를 들어, 제품 정량화(PQ)를 사용해 정밀도를 낮추거나 주성분 분석(PCA)을 사용해 차원 수를 줄임으로써 저장 공간을 절약할 수 있습니다. 이러한 방법은 전체 벡터 집합을 분석하여 벡터 간의 의미 관계를 유지하는 더 간결한 벡터 집합을 찾습니다.
이러한 표준 접근 방식은 효과적이기는 하지만 정밀도나 차원을 단 한 번만 감소시킵니다. 하지만 점점 더 정밀하게 표현되는 피라미드처럼 여러 계층의 디테일을 동시에 유지할 수 있다면 어떨까요?
마트료시카 임베딩을 소개합니다. 러시아 중첩 인형(그림 참조)에서 이름을 따온 이 기발한 구조는 단일 벡터 내에 여러 스케일의 표현을 임베딩합니다. 기존의 사후 처리 방법과 달리, 마트료시카 임베딩은 초기 학습 과정에서 이 다중 스케일 구조를 학습합니다. 그 결과는 놀랍습니다. 전체 임베딩이 입력 의미를 포착할 뿐만 아니라 중첩된 각 하위 집합 접두사(전반부, 1/4 등)가 덜 상세하더라도 일관성 있는 표현을 제공합니다.
이 노트에서는 의미 검색을 위해 Milvus와 함께 Matryoshka 임베딩을 사용하는 방법을 살펴봅니다. "깔때기 검색"이라는 알고리즘을 통해 임베딩 차원의 작은 하위 집합에 대한 유사성 검색을 수행하면서 회상률을 급격히 떨어뜨리지 않고도 검색을 수행할 수 있는 방법을 설명합니다.
import functools
from datasets import load_dataset
import numpy as np
import pandas as pd
import pymilvus
from pymilvus import MilvusClient
from pymilvus import FieldSchema, CollectionSchema, DataType
from sentence_transformers import SentenceTransformer
import torch
import torch.nn.functional as F
from tqdm import tqdm
마트료시카 임베딩 모델 로드
다음과 같은 표준 임베딩 모델을 사용하는 대신 sentence-transformers/all-MiniLM-L12-v2와 같은 표준 임베딩 모델을 사용하는 대신, Matryoshka 임베딩을 생성하도록 특별히 훈련된 Nomic의 모델을 사용합니다.
model = SentenceTransformer(
# Remove 'device='mps' if running on non-Mac device
"nomic-ai/nomic-embed-text-v1.5",
trust_remote_code=True,
device="mps",
)
<All keys matched successfully>
데이터 세트 로드, 항목 임베딩 및 벡터 데이터베이스 구축하기
다음 코드는 문서 페이지 "문장 트랜스포머와 밀버스를 사용한 영화 검색"의 코드를 수정한 것입니다. 먼저 HuggingFace에서 데이터 세트를 로드합니다. 여기에는 약 35,000개의 항목이 포함되어 있으며, 각 항목은 Wikipedia 문서가 있는 영화에 해당합니다. 이 예제에서는 Title 및 PlotSummary 필드를 사용하겠습니다.
ds = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries", split="train")
print(ds)
Dataset({
features: ['Release Year', 'Title', 'Origin/Ethnicity', 'Director', 'Cast', 'Genre', 'Wiki Page', 'Plot', 'PlotSummary'],
num_rows: 34886
})
다음으로 Milvus Lite 데이터베이스에 연결하여 데이터 스키마를 지정하고 이 스키마로 컬렉션을 생성합니다. 정규화되지 않은 임베딩과 임베딩의 첫 번째 여섯 번째를 모두 별도의 필드에 저장합니다. 그 이유는 유사도 검색을 수행하기 위해 Matryoshka 임베딩의 첫 1/6이 필요하고, 검색 결과의 순위를 재조정하고 개선하기 위해 나머지 5/6의 임베딩이 필요하기 때문입니다.
embedding_dim = 768
search_dim = 128
collection_name = "movie_embeddings"
client = MilvusClient(uri="./wiki-movie-plots-matryoshka.db")
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
# First sixth of unnormalized embedding vector
FieldSchema(name="head_embedding", dtype=DataType.FLOAT_VECTOR, dim=search_dim),
# Entire unnormalized embedding vector
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
]
schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
client.create_collection(collection_name=collection_name, schema=schema)
Milvus는 현재 임베딩의 하위 집합에 대한 검색을 지원하지 않기 때문에 임베딩을 두 부분으로 나누어 헤드는 색인 및 검색할 벡터의 초기 하위 집합을 나타내고 테일은 나머지 부분입니다. 이 모델은 코사인 거리 유사성 검색을 위해 훈련되었으므로 헤드 임베딩을 정규화합니다. 그러나 나중에 더 큰 하위 집합에 대한 유사도를 계산하려면 헤드 임베딩의 규범을 저장해야 하므로 꼬리에 결합하기 전에 정규화를 해제할 수 있습니다.
임베딩의 첫 1/6을 통해 검색을 수행하려면 head_embedding 필드에 벡터 검색 인덱스를 생성해야 합니다. 나중에 '퍼널 검색'의 결과를 일반 벡터 검색과 비교할 것이므로 전체 임베딩에 대해서도 검색 인덱스를 구축할 것입니다.
중요한 것은 IP 거리 메트릭이 아닌 COSINE 을 사용하는 것인데, 그렇지 않으면 임베딩 규범을 추적해야 하므로 구현이 복잡해지기 때문입니다(퍼널 검색 알고리즘을 설명하면 더 이해가 쉬워질 것입니다).
index_params = client.prepare_index_params()
index_params.add_index(
field_name="head_embedding", index_type="FLAT", metric_type="COSINE"
)
index_params.add_index(field_name="embedding", index_type="FLAT", metric_type="COSINE")
client.create_index(collection_name, index_params)
마지막으로 35,000개의 모든 영화에 대한 줄거리 요약을 인코딩하고 해당 임베딩을 데이터베이스에 입력합니다.
for batch in tqdm(ds.batch(batch_size=512)):
# This particular model requires us to prefix 'search_document:' to stored entities
plot_summary = ["search_document: " + x.strip() for x in batch["PlotSummary"]]
# Output of embedding model is unnormalized
embeddings = model.encode(plot_summary, convert_to_tensor=True)
head_embeddings = embeddings[:, :search_dim]
data = [
{
"title": title,
"head_embedding": head.cpu().numpy(),
"embedding": embedding.cpu().numpy(),
}
for title, head, embedding in zip(batch["Title"], head_embeddings, embeddings)
]
res = client.insert(collection_name=collection_name, data=data)
100%|██████████| 69/69 [05:57<00:00, 5.18s/it]
퍼널 검색 수행하기
이제 마트료시카 임베딩 차원의 첫 1/6을 사용해 '퍼널 검색'을 구현해 보겠습니다. 검색할 영화 세 편을 염두에 두고 있으며 데이터베이스 쿼리를 위한 줄거리 요약을 직접 작성했습니다. 쿼리를 임베드한 다음 head_embedding 필드에서 벡터 검색을 수행하여 128개의 결과 후보를 검색합니다.
queries = [
"An archaeologist searches for ancient artifacts while fighting Nazis.",
"A teenager fakes illness to get off school and have adventures with two friends.",
"A young couple with a kid look after a hotel during winter and the husband goes insane.",
]
# Search the database based on input text
def embed_search(data):
embeds = model.encode(data)
return [x for x in embeds]
# This particular model requires us to prefix 'search_query:' to queries
instruct_queries = ["search_query: " + q.strip() for q in queries]
search_data = embed_search(instruct_queries)
# Normalize head embeddings
head_search = [x[:search_dim] for x in search_data]
# Perform standard vector search on first sixth of embedding dimensions
res = client.search(
collection_name=collection_name,
data=head_search,
anns_field="head_embedding",
limit=128,
output_fields=["title", "head_embedding", "embedding"],
)
이 시점에서 훨씬 더 작은 벡터 공간에서 검색을 수행했기 때문에 전체 공간에서 검색할 때보다 지연 시간이 줄어들고 인덱스에 대한 저장 공간 요구 사항이 줄어들었을 가능성이 높습니다. 각 쿼리에 대해 상위 5개 일치 항목을 살펴보겠습니다:
for query, hits in zip(queries, res):
rows = [x["entity"] for x in hits][:5]
print("Query:", query)
print("Results:")
for row in rows:
print(row["title"].strip())
print()
Query: An archaeologist searches for ancient artifacts while fighting Nazis.
Results:
"Pimpernel" Smith
Black Hunters
The Passage
Counterblast
Dominion: Prequel to the Exorcist
Query: A teenager fakes illness to get off school and have adventures with two friends.
Results:
How to Deal
Shorts
Blackbird
Valentine
Unfriended
Query: A young couple with a kid look after a hotel during winter and the husband goes insane.
Results:
Ghostkeeper
Our Vines Have Tender Grapes
The Ref
Impact
The House in Marsh Road
보시다시피, 검색 중에 임베딩을 잘라낸 결과 리콜이 저하되었습니다. 퍼널 검색은 이 문제를 영리한 방법으로 해결합니다. 나머지 임베딩 차원을 사용하여 후보 목록의 순위를 재조정하고 잘라내어 비용이 많이 드는 벡터 검색을 추가로 실행하지 않고도 검색 성능을 회복할 수 있습니다.
퍼널 검색 알고리즘을 쉽게 설명하기 위해 각 쿼리에 대한 Milvus 검색 히트를 Pandas 데이터 프레임으로 변환합니다.
def hits_to_dataframe(hits: pymilvus.client.abstract.Hits) -> pd.DataFrame:
"""
Convert a Milvus search result to a Pandas dataframe. This function is specific to our data schema.
"""
rows = [x["entity"] for x in hits]
rows_dict = [
{"title": x["title"], "embedding": torch.tensor(x["embedding"])} for x in rows
]
return pd.DataFrame.from_records(rows_dict)
dfs = [hits_to_dataframe(hits) for hits in res]
이제 퍼널 검색을 수행하기 위해 임베딩의 점점 더 큰 하위 집합을 반복합니다. 반복할 때마다 새로운 유사성에 따라 후보의 순위를 재조정하고 가장 낮은 순위를 가진 일부 후보를 제거합니다.
이를 구체적으로 설명하기 위해 이전 단계에서는 임베딩 및 쿼리 차원의 1/6을 사용하여 128개의 후보를 검색했습니다. 퍼널 검색을 수행하는 첫 번째 단계는 첫 번째 1/3의 차원을 사용하여 쿼리와 후보 간의 유사성을 다시 계산하는 것입니다. 하위 64개의 후보가 제거됩니다. 그런 다음 처음 2/3의 차원과 모든 차원에 대해 이 프로세스를 반복하여 32개와 16개의 후보를 연속적으로 정리합니다.
# An optimized implementation would vectorize the calculation of similarity scores across rows (using a matrix)
def calculate_score(row, query_emb=None, dims=768):
emb = F.normalize(row["embedding"][:dims], dim=-1)
return (emb @ query_emb).item()
# You could also add a top-K parameter as a termination condition
def funnel_search(
df: pd.DataFrame, query_emb, scales=[256, 512, 768], prune_ratio=0.5
) -> pd.DataFrame:
# Loop over increasing prefixes of the embeddings
for dims in scales:
# Query vector must be normalized for each new dimensionality
emb = torch.tensor(query_emb[:dims] / np.linalg.norm(query_emb[:dims]))
# Score
scores = df.apply(
functools.partial(calculate_score, query_emb=emb, dims=dims), axis=1
)
df["scores"] = scores
# Re-rank
df = df.sort_values(by="scores", ascending=False)
# Prune (in our case, remove half of candidates at each step)
df = df.head(int(prune_ratio * len(df)))
return df
dfs_results = [
{"query": query, "results": funnel_search(df, query_emb)}
for query, df, query_emb in zip(queries, dfs, search_data)
]
for d in dfs_results:
print(d["query"], "\n", d["results"][:5]["title"], "\n")
An archaeologist searches for ancient artifacts while fighting Nazis.
0 "Pimpernel" Smith
1 Black Hunters
29 Raiders of the Lost Ark
34 The Master Key
51 My Gun Is Quick
Name: title, dtype: object
A teenager fakes illness to get off school and have adventures with two friends.
21 How I Live Now
32 On the Edge of Innocence
77 Bratz: The Movie
4 Unfriended
108 Simon Says
Name: title, dtype: object
A young couple with a kid look after a hotel during winter and the husband goes insane.
9 The Shining
0 Ghostkeeper
11 Fast and Loose
7 Killing Ground
12 Home Alone
Name: title, dtype: object
추가 벡터 검색을 수행하지 않고도 리콜을 복원할 수 있었습니다! 질적으로 이러한 결과는 다른 임베딩 모델을 사용하는 튜토리얼 '밀버스와 문장 트랜스포머를 사용한 영화 검색'의 표준 벡터 검색보다 '로스트 아크의 침입자'와 '샤이닝'에 대한 리콜이 더 높은 것으로 보입니다. 하지만 "페리스 뷸러의 하루"는 찾을 수 없었는데, 이 부분은 나중에 노트북에서 다시 다룰 예정입니다. (더 많은 정량적 실험과 벤치마킹은 Matryoshka 표현 학습 논문을 참조하세요.)
퍼널 검색과 일반 검색 비교
동일한 임베딩 모델을 사용하는 동일한 데이터 세트에서 퍼널 검색의 결과를 표준 벡터 검색과 비교해 보겠습니다. 전체 임베딩에 대해 검색을 수행합니다.
# Search on entire embeddings
res = client.search(
collection_name=collection_name,
data=search_data,
anns_field="embedding",
limit=5,
output_fields=["title", "embedding"],
)
for query, hits in zip(queries, res):
rows = [x["entity"] for x in hits]
print("Query:", query)
print("Results:")
for row in rows:
print(row["title"].strip())
print()
Query: An archaeologist searches for ancient artifacts while fighting Nazis.
Results:
"Pimpernel" Smith
Black Hunters
Raiders of the Lost Ark
The Master Key
My Gun Is Quick
Query: A teenager fakes illness to get off school and have adventures with two friends.
Results:
A Walk to Remember
Ferris Bueller's Day Off
How I Live Now
On the Edge of Innocence
Bratz: The Movie
Query: A young couple with a kid look after a hotel during winter and the husband goes insane.
Results:
The Shining
Ghostkeeper
Fast and Loose
Killing Ground
Home Alone
"한 십대가 학교를 빠지기 위해 병을 속이고..."에 대한 결과를 제외하면, 퍼널 검색의 경우 128차원의 검색 공간에서 퍼널 검색을 수행한 반면 일반 검색의 경우 768차원에서 검색을 수행했음에도 불구하고 전체 검색과 거의 동일한 결과를 얻을 수 있습니다.
페리스 뷸러의 휴무일에 대한 퍼널 검색 리콜 실패 조사하기
왜 퍼널 검색이 페리스 뷸러의 휴일 검색에 성공하지 못했을까요? 원래 후보 목록에 있었던 것인지 아니면 실수로 필터링이 잘못되었는지 살펴보겠습니다.
queries2 = [
"A teenager fakes illness to get off school and have adventures with two friends."
]
# Search the database based on input text
def embed_search(data):
embeds = model.encode(data)
return [x for x in embeds]
instruct_queries = ["search_query: " + q.strip() for q in queries2]
search_data2 = embed_search(instruct_queries)
head_search2 = [x[:search_dim] for x in search_data2]
# Perform standard vector search on subset of embeddings
res = client.search(
collection_name=collection_name,
data=head_search2,
anns_field="head_embedding",
limit=256,
output_fields=["title", "head_embedding", "embedding"],
)
for query, hits in zip(queries, res):
rows = [x["entity"] for x in hits]
print("Query:", queries2[0])
for idx, row in enumerate(rows):
if row["title"].strip() == "Ferris Bueller's Day Off":
print(f"Row {idx}: Ferris Bueller's Day Off")
Query: A teenager fakes illness to get off school and have adventures with two friends.
Row 228: Ferris Bueller's Day Off
초기 후보 목록이 충분히 크지 않았거나, 오히려 원하는 검색어가 가장 높은 수준의 세부 수준에서 쿼리와 충분히 유사하지 않다는 것이 문제임을 알 수 있습니다. 128 에서 256 으로 변경하면 검색에 성공합니다. 리콜과 지연 시간 간의 균형을 경험적으로 평가하기 위해 홀드아웃 세트의 후보자 수를 설정하는 경험 법칙을 만들어야 합니다.
dfs = [hits_to_dataframe(hits) for hits in res]
dfs_results = [
{"query": query, "results": funnel_search(df, query_emb)}
for query, df, query_emb in zip(queries2, dfs, search_data2)
]
for d in dfs_results:
print(d["query"], "\n", d["results"][:7]["title"].to_string(index=False), "\n")
A teenager fakes illness to get off school and have adventures with two friends.
A Walk to Remember
Ferris Bueller's Day Off
How I Live Now
On the Edge of Innocence
Bratz: The Movie
Unfriended
Simon Says
순서가 중요한가요? 접두사 임베딩 대 접미사 임베딩.
이 모델은 재귀적으로 더 작은 접두사 임베딩을 잘 매칭하도록 학습되었습니다. 우리가 사용하는 차원의 순서가 중요할까요? 예를 들어, 접미사인 임베딩의 하위 집합도 사용할 수 있을까요? 이 실험에서는 마트료시카 임베딩의 차원 순서를 반대로 하여 깔때기 검색을 수행합니다.
client = MilvusClient(uri="./wikiplots-matryoshka-flipped.db")
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
FieldSchema(name="head_embedding", dtype=DataType.FLOAT_VECTOR, dim=search_dim),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
]
schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
client.create_collection(collection_name=collection_name, schema=schema)
index_params = client.prepare_index_params()
index_params.add_index(
field_name="head_embedding", index_type="FLAT", metric_type="COSINE"
)
client.create_index(collection_name, index_params)
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
- Avoid using `tokenizers` before the fork if possible
- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
for batch in tqdm(ds.batch(batch_size=512)):
plot_summary = ["search_document: " + x.strip() for x in batch["PlotSummary"]]
# Encode and flip embeddings
embeddings = model.encode(plot_summary, convert_to_tensor=True)
embeddings = torch.flip(embeddings, dims=[-1])
head_embeddings = embeddings[:, :search_dim]
data = [
{
"title": title,
"head_embedding": head.cpu().numpy(),
"embedding": embedding.cpu().numpy(),
}
for title, head, embedding in zip(batch["Title"], head_embeddings, embeddings)
]
res = client.insert(collection_name=collection_name, data=data)
100%|██████████| 69/69 [05:50<00:00, 5.08s/it]
# Normalize head embeddings
flip_search_data = [
torch.flip(torch.tensor(x), dims=[-1]).cpu().numpy() for x in search_data
]
flip_head_search = [x[:search_dim] for x in flip_search_data]
# Perform standard vector search on subset of embeddings
res = client.search(
collection_name=collection_name,
data=flip_head_search,
anns_field="head_embedding",
limit=128,
output_fields=["title", "head_embedding", "embedding"],
)
dfs = [hits_to_dataframe(hits) for hits in res]
dfs_results = [
{"query": query, "results": funnel_search(df, query_emb)}
for query, df, query_emb in zip(queries, dfs, flip_search_data)
]
for d in dfs_results:
print(
d["query"],
"\n",
d["results"][:7]["title"].to_string(index=False, header=False),
"\n",
)
An archaeologist searches for ancient artifacts while fighting Nazis.
"Pimpernel" Smith
Black Hunters
Raiders of the Lost Ark
The Master Key
My Gun Is Quick
The Passage
The Mole People
A teenager fakes illness to get off school and have adventures with two friends.
A Walk to Remember
How I Live Now
Unfriended
Cirque du Freak: The Vampire's Assistant
Last Summer
Contest
Day One
A young couple with a kid look after a hotel during winter and the husband goes insane.
Ghostkeeper
Killing Ground
Leopard in the Snow
Stone
Afterglow
Unfaithful
Always a Bride
예상대로 퍼널 검색이나 일반 검색보다 리콜률이 훨씬 떨어집니다(임베딩 모델은 접미사가 아닌 임베딩 차원의 접두사에 대한 대조 학습을 통해 학습되었습니다).
요약
다음은 여러 방법의 검색 결과를 비교한 것입니다:
