Ricerca a imbuto con Matryoshka Embeddings
Quando si costruiscono sistemi di ricerca vettoriale efficienti, una sfida fondamentale è la gestione dei costi di memorizzazione, mantenendo una latenza e un richiamo accettabili. I moderni modelli di embedding producono vettori con centinaia o migliaia di dimensioni, creando un significativo overhead di archiviazione e calcolo per il vettore grezzo e l'indice.
Tradizionalmente, i requisiti di memorizzazione vengono ridotti applicando un metodo di quantizzazione o di riduzione della dimensionalità appena prima di costruire l'indice. Ad esempio, è possibile risparmiare memoria riducendo la precisione con la quantizzazione del prodotto (PQ) o il numero di dimensioni con l'analisi delle componenti principali (PCA). Questi metodi analizzano l'intero insieme di vettori per trovarne uno più compatto che mantenga le relazioni semantiche tra i vettori.
Pur essendo efficaci, questi approcci standard riducono la precisione o la dimensionalità solo una volta e su un'unica scala. Ma cosa succederebbe se potessimo mantenere più livelli di dettaglio contemporaneamente, come una piramide di rappresentazioni sempre più precise?
Ecco le Matryoshka embeddings. Questi costrutti intelligenti, che prendono il nome dalle matrioske russe (vedi illustrazione), incorporano più scale di rappresentazione all'interno di un singolo vettore. A differenza dei metodi tradizionali di post-elaborazione, le Matryoshka embeddings apprendono questa struttura multi-scala durante il processo di addestramento iniziale. Il risultato è notevole: non solo l'embedding completo cattura la semantica dell'input, ma ogni prefisso annidato del sottoinsieme (prima metà, primo quarto, ecc.) fornisce una rappresentazione coerente, anche se meno dettagliata.
In questo quaderno esaminiamo come utilizzare gli embedding Matryoshka con Milvus per la ricerca semantica. Illustriamo un algoritmo chiamato "funnel search" che ci permette di eseguire ricerche di similarità su un piccolo sottoinsieme di dimensioni dell'embedding senza un drastico calo del richiamo.
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
Modello di incorporamento a matrioska
Invece di utilizzare un modello di incorporamento standard, come ad esempio sentence-transformers/all-MiniLM-L12-v2utilizziamo un modello di Nomic addestrato appositamente per produrre 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>
Caricare il set di dati, incorporare gli elementi e costruire il database di vettori
Il codice seguente è una modifica di quello della pagina di documentazione "Movie Search with Sentence Transformers and Milvus". Per prima cosa, carichiamo il dataset da HuggingFace. Contiene circa 35.000 voci, ognuna delle quali corrisponde a un film con un articolo di Wikipedia. In questo esempio utilizzeremo i campi Title e 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
})
Successivamente, ci colleghiamo a un database Milvus Lite, specifichiamo lo schema dei dati e creiamo una collezione con questo schema. Memorizzeremo sia l'incorporazione non normalizzata sia il primo sesto dell'incorporazione in campi separati. Il motivo è che abbiamo bisogno del primo 1/6 dell'embedding Matryoshka per eseguire una ricerca di somiglianza, e dei restanti 5/6 dell'embedding per eseguire un reranking e migliorare i risultati della ricerca.
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 non supporta attualmente la ricerca su sottoinsiemi di incorporazioni, quindi dividiamo le incorporazioni in due parti: la testa rappresenta il sottoinsieme iniziale del vettore da indicizzare e cercare, mentre la coda è il resto. Il modello è stato addestrato per la ricerca di somiglianze a distanza coseno, quindi normalizziamo gli embeddings della testa. Tuttavia, per calcolare in seguito le somiglianze per sottoinsiemi più ampi, è necessario memorizzare la norma dell'embedding della testa, in modo da poterla normalizzare prima di unirla alla coda.
Per eseguire la ricerca attraverso il primo 1/6 dell'incorporazione, dovremo creare un indice di ricerca vettoriale sul campo head_embedding. In seguito, confronteremo i risultati della "ricerca a imbuto" con una normale ricerca vettoriale, e quindi creeremo un indice di ricerca anche sull'intero embedding.
È importante notare che utilizziamo la metrica di distanza COSINE anziché IP, perché altrimenti dovremmo tenere traccia delle norme di incorporamento, il che complicherebbe l'implementazione (questo avrà più senso una volta descritto l'algoritmo di ricerca a imbuto).
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)
Infine, codifichiamo i riassunti delle trame di tutti i 35k film e inseriamo le corrispondenti incorporazioni nel database.
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]
Esecuzione della ricerca a imbuto
Ora implementiamo una "ricerca a imbuto" utilizzando il primo 1/6 delle dimensioni dell'embedding Matryoshka. Ho in mente tre film da recuperare e ho prodotto il mio riassunto della trama per interrogare il database. Incorporiamo le query, quindi eseguiamo una ricerca vettoriale sul campo head_embedding, recuperando 128 candidati risultati.
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"],
)
A questo punto, abbiamo eseguito una ricerca su uno spazio vettoriale molto più piccolo e quindi è probabile che la latenza e i requisiti di memorizzazione dell'indice siano ridotti rispetto alla ricerca sull'intero spazio. Esaminiamo le prime 5 corrispondenze per ciascuna query:
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
Come possiamo vedere, il richiamo ha sofferto come conseguenza del troncamento delle incorporazioni durante la ricerca. La ricerca a imbuto risolve questo problema con un trucco intelligente: possiamo usare il resto delle dimensioni dell'embedding per riordinare e sfoltire il nostro elenco di candidati e recuperare le prestazioni di recupero senza eseguire ulteriori costose ricerche vettoriali.
Per facilitare l'esposizione dell'algoritmo di ricerca funnel, convertiamo i risultati della ricerca Milvus per ogni query in un dataframe 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]
Ora, per eseguire la ricerca a imbuto, iteriamo sui sottoinsiemi sempre più grandi delle incorporazioni. A ogni iterazione, i candidati vengono classificati in base alle nuove somiglianze e viene eliminata una frazione di quelli classificati più in basso.
Per fare un esempio concreto, dalla fase precedente abbiamo recuperato 128 candidati utilizzando 1/6 delle dimensioni dell'embedding e della query. Il primo passo per eseguire la ricerca a imbuto consiste nel ricalcolare le somiglianze tra le query e i candidati utilizzando il primo 1/3 delle dimensioni. Gli ultimi 64 candidati vengono eliminati. Quindi si ripete il processo con i primi 2/3 delle dimensioni e poi con tutte le dimensioni, riducendo successivamente a 32 e 16 i candidati.
# 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
Siamo riusciti a ripristinare il richiamo senza eseguire ulteriori ricerche vettoriali! Qualitativamente, questi risultati sembrano avere un richiamo più elevato per "I predatori dell'arca perduta" e "Shining" rispetto alla ricerca vettoriale standard del tutorial "Movie Search using Milvus and Sentence Transformers", che utilizza un modello di incorporamento diverso. Tuttavia, non riesce a trovare "Ferris Bueller's Day Off", su cui torneremo più avanti nel quaderno. (Si veda il documento Matryoshka Representation Learning per ulteriori esperimenti quantitativi e benchmark).
Confronto tra la ricerca a imbuto e la ricerca regolare
Confrontiamo i risultati della nostra ricerca a imbuto con una ricerca vettoriale standard sullo stesso set di dati con lo stesso modello di incorporazione. Eseguiamo una ricerca sulle incorporazioni complete.
# 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
Con l'eccezione dei risultati per "Un adolescente finge una malattia per non andare a scuola...", i risultati della ricerca a imbuto sono quasi identici a quelli della ricerca completa, anche se la ricerca a imbuto è stata eseguita su uno spazio di ricerca di 128 dimensioni contro le 768 della ricerca normale.
Analisi del fallimento della ricerca a imbuto per Ferris Bueller's Day Off
Perché la ricerca a imbuto non è riuscita a recuperare Ferris Bueller's Day Off? Verifichiamo se era presente nell'elenco originale dei candidati o se è stato erroneamente filtrato.
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
Vediamo che il problema era che l'elenco iniziale dei candidati non era abbastanza ampio, o meglio, l'hit desiderato non era abbastanza simile alla query al massimo livello di granularità. Cambiando da 128 a 256 il risultato è stato recuperato con successo. Dovremmo stabilire una regola empirica per impostare il numero di candidati su un set di attesa per valutare empiricamente il compromesso tra richiamo e latenza.
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
L'ordine è importante? Incorporazione di prefissi e suffissi.
Il modello è stato addestrato per ottenere buoni risultati con i prefissi ricorsivamente più piccoli degli embeddings. L'ordine delle dimensioni che utilizziamo è importante? Per esempio, potremmo prendere anche sottoinsiemi di embeddings che sono suffissi? In questo esperimento, invertiamo l'ordine delle dimensioni nelle Matryoshka embeddings ed eseguiamo una ricerca a imbuto.
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
Il richiamo è molto più scarso rispetto alla ricerca a imbuto o alla ricerca normale, come ci si aspettava (il modello di incorporazione è stato addestrato tramite apprendimento contrastivo sui prefissi delle dimensioni di incorporazione, non sui suffissi).
Sintesi
Ecco un confronto dei risultati della ricerca tra i vari metodi:
