Pencarian Corong dengan Penyematan Matryoshka
Ketika membangun sistem pencarian vektor yang efisien, salah satu tantangan utama adalah mengelola biaya penyimpanan sambil mempertahankan latensi dan pemanggilan yang dapat diterima. Model penyematan modern menghasilkan vektor dengan ratusan atau ribuan dimensi, menciptakan penyimpanan yang signifikan dan overhead komputasi untuk vektor mentah dan indeks.
Secara tradisional, kebutuhan penyimpanan dikurangi dengan menerapkan metode kuantisasi atau pengurangan dimensi sebelum membangun indeks. Sebagai contoh, kita dapat menghemat penyimpanan dengan menurunkan presisi menggunakan Product Quantization (PQ) atau jumlah dimensi menggunakan Principal Component Analysis (PCA). Metode-metode ini menganalisis seluruh kumpulan vektor untuk menemukan vektor yang lebih ringkas yang mempertahankan hubungan semantik antar vektor.
Meskipun efektif, pendekatan standar ini mengurangi presisi atau dimensi hanya sekali dan pada skala tunggal. Namun, bagaimana jika kita dapat mempertahankan beberapa lapisan detail secara bersamaan, seperti piramida representasi yang semakin presisi?
Masuklah ke dalam penyematan Matryoshka. Dinamai berdasarkan boneka sarang Rusia (lihat ilustrasi), konstruksi cerdas ini menyematkan beberapa skala representasi dalam satu vektor. Tidak seperti metode pasca-pemrosesan tradisional, embeddings Matryoshka mempelajari struktur multi-skala ini selama proses pelatihan awal. Hasilnya luar biasa: tidak hanya penyematan penuh yang menangkap semantik input, tetapi setiap awalan subset bersarang (paruh pertama, kuartal pertama, dll.) memberikan representasi yang koheren, meskipun tidak terlalu detail.
Dalam buku catatan ini, kami membahas cara menggunakan embedding Matryoshka dengan Milvus untuk pencarian semantik. Kami mengilustrasikan sebuah algoritma yang disebut "pencarian corong" yang memungkinkan kita untuk melakukan pencarian kemiripan pada sebagian kecil dari dimensi penyematan kita tanpa penurunan yang drastis dalam mengingat.
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
Memuat Model Penyematan Matryoshka
Alih-alih menggunakan model embedding standar seperti sentence-transformers/all-MiniLM-L12-v2kami menggunakan model dari Nomic yang dilatih secara khusus untuk menghasilkan embedding Matryoshka.
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>
Memuat Dataset, Menyematkan Item, dan Membangun Basis Data Vektor
Kode berikut ini adalah modifikasi dari halaman dokumentasi "Pencarian Film dengan Pengubah Kalimat dan Milvus". Pertama, kita memuat dataset dari HuggingFace. Dataset ini berisi sekitar 35 ribu entri, masing-masing sesuai dengan film yang memiliki artikel Wikipedia. Kita akan menggunakan bidang Title dan PlotSummary dalam contoh ini.
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
})
Selanjutnya, kita sambungkan ke basis data Milvus Lite, tentukan skema datanya, dan buat koleksi dengan skema ini. Kita akan menyimpan embedding yang belum dinormalisasi dan seperenam pertama dari embedding di bidang yang terpisah. Alasannya adalah karena kita membutuhkan 1/6 pertama dari embedding Matryoshka untuk melakukan pencarian kemiripan, dan 5/6 embedding yang tersisa untuk menentukan peringkat dan meningkatkan hasil pencarian.
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 saat ini tidak mendukung pencarian melalui subset dari embedding, jadi kami memecah embedding menjadi dua bagian: kepala mewakili subset awal vektor untuk diindeks dan dicari, dan ekor adalah sisanya. Model dilatih untuk pencarian kemiripan jarak kosinus, jadi kami menormalkan embedding kepala. Namun, untuk menghitung kemiripan untuk subset yang lebih besar di kemudian hari, kita perlu menyimpan norma dari penyisipan kepala, sehingga kita dapat menormalkannya kembali sebelum menggabungkannya ke ekor.
Untuk melakukan pencarian melalui 1/6 pertama dari penyematan, kita perlu membuat indeks pencarian vektor di bidang head_embedding. Nantinya, kita akan membandingkan hasil "pencarian corong" dengan pencarian vektor biasa, dan juga membuat indeks pencarian di atas penyematan penuh.
Yang penting, kita menggunakan metrik jarak COSINE dan bukan IP, karena jika tidak, kita harus melacak norma-norma penyematan, yang akan menyulitkan implementasi (hal ini akan lebih masuk akal setelah algoritme pencarian corong dijelaskan).
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)
Terakhir, kami mengkodekan ringkasan plot untuk semua 35 ribu film dan memasukkan penyematan yang sesuai ke dalam basis data.
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]
Melakukan Pencarian Corong
Sekarang mari kita mengimplementasikan "pencarian corong" dengan menggunakan 1/6 pertama dari dimensi penyematan Matryoshka. Saya memiliki tiga film yang akan diambil dan telah membuat ringkasan plot saya sendiri untuk menanyakan database. Kami menyematkan kueri, lalu melakukan pencarian vektor pada bidang head_embedding, mengambil 128 kandidat hasil.
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"],
)
Pada titik ini, kita telah melakukan pencarian pada ruang vektor yang jauh lebih kecil, dan oleh karena itu kemungkinan besar telah menurunkan latensi dan mengurangi kebutuhan penyimpanan untuk indeks dibandingkan dengan pencarian pada ruang penuh. Mari kita periksa 5 kecocokan teratas untuk setiap kueri:
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
Seperti yang bisa kita lihat, recall mengalami penurunan akibat pemotongan sematan selama pencarian. Pencarian corong memperbaiki hal ini dengan trik yang cerdas: kita dapat menggunakan sisa dimensi penyematan untuk memberi peringkat ulang dan memangkas daftar kandidat kita untuk memulihkan kinerja pencarian tanpa menjalankan pencarian vektor tambahan yang mahal.
Untuk memudahkan penjelasan algoritme pencarian corong, kami mengubah hit pencarian Milvus untuk setiap kueri menjadi bingkai data 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]
Sekarang, untuk melakukan pencarian corong, kami melakukan iterasi pada himpunan bagian yang semakin besar dari sematan. Pada setiap iterasi, kami mengubah peringkat kandidat berdasarkan kesamaan baru dan memangkas sebagian kecil dari peringkat terendah.
Untuk memperjelas hal ini, dari langkah sebelumnya kami telah mengambil 128 kandidat dengan menggunakan 1/6 dari dimensi embedding dan kueri. Langkah pertama dalam melakukan pencarian corong adalah menghitung ulang kemiripan antara kueri dan kandidat menggunakan 1/3 dimensi pertama. 64 kandidat terbawah dipangkas. Kemudian kami mengulangi proses ini dengan 2/3 dimensi pertama, dan kemudian semua dimensi, secara berurutan memangkas menjadi 32 dan 16 kandidat.
# 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
Kami telah mampu memulihkan daya ingat tanpa melakukan pencarian vektor tambahan! Secara kualitatif, hasil ini tampaknya memiliki daya ingat yang lebih tinggi untuk "Raiders of the Lost Ark" dan "The Shining" daripada pencarian vektor standar dalam tutorial, "Pencarian Film menggunakan Milvus dan Sentence Transformers", yang menggunakan model penyematan yang berbeda. Namun, tidak dapat menemukan "Ferris Bueller's Day Off", yang akan kita bahas nanti di buku catatan ini. (Lihat makalah Pembelajaran Representasi Matryoshka untuk eksperimen dan pembandingan yang lebih kuantitatif).
Membandingkan Pencarian Corong dengan Pencarian Biasa
Mari kita bandingkan hasil pencarian corong kita dengan pencarian vektor standar pada dataset yang sama dengan model penyematan yang sama. Kami melakukan pencarian pada penyematan penuh.
# 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
Dengan pengecualian hasil untuk "Seorang remaja berpura-pura sakit agar tidak masuk sekolah...", hasil di bawah pencarian corong hampir sama dengan pencarian lengkap, meskipun pencarian corong dilakukan pada ruang pencarian 128 dimensi vs 768 dimensi untuk pencarian biasa.
Menyelidiki Kegagalan Pemanggilan Kembali Pencarian Corong untuk Hari Libur Ferris Bueller
Mengapa pencarian corong tidak berhasil mendapatkan Ferris Bueller's Day Off? Mari kita periksa apakah film tersebut ada dalam daftar kandidat asli atau salah disaring.
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
Kami melihat bahwa masalahnya adalah daftar kandidat awal tidak cukup besar, atau lebih tepatnya, hasil yang diinginkan tidak cukup mirip dengan kueri pada tingkat perincian tertinggi. Mengubahnya dari 128 ke 256 menghasilkan pengambilan yang berhasil. Kita harus membuat sebuah aturan praktis untuk mengatur jumlah kandidat pada kumpulan yang ditahan untuk mengevaluasi secara empiris pertukaran antara recall dan latensi.
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
Apakah urutannya penting? Penyematan awalan vs akhiran.
Model ini dilatih untuk melakukan pencocokan awalan yang lebih kecil secara rekursif dengan baik dari sematan. Apakah urutan dimensi yang kita gunakan penting? Sebagai contoh, bisakah kita juga mengambil himpunan bagian dari sematan yang merupakan sufiks? Dalam percobaan ini, kami membalik urutan dimensi dalam embedding Matryoshka dan melakukan pencarian corong.
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
Daya ingat jauh lebih buruk daripada pencarian corong atau pencarian biasa seperti yang diharapkan (model penyematan dilatih dengan pembelajaran kontras pada awalan dari dimensi penyematan, bukan akhiran).
Ringkasan
Berikut ini adalah perbandingan hasil pencarian kami di seluruh metode:
