Recherche en entonnoir avec Matryoshka Embeddings
Lors de la construction de systèmes de recherche vectorielle efficaces, l'un des principaux défis consiste à gérer les coûts de stockage tout en maintenant une latence et un rappel acceptables. Les modèles d'intégration modernes produisent des vecteurs comportant des centaines ou des milliers de dimensions, ce qui entraîne une surcharge de stockage et de calcul importante pour le vecteur brut et l'index.
Traditionnellement, les besoins en stockage sont réduits par l'application d'une méthode de quantification ou de réduction de la dimensionnalité juste avant la construction de l'index. Par exemple, nous pouvons économiser de l'espace de stockage en réduisant la précision à l'aide de la quantification par produit (PQ) ou le nombre de dimensions à l'aide de l'analyse en composantes principales (ACP). Ces méthodes analysent l'ensemble des vecteurs pour en trouver un plus compact qui préserve les relations sémantiques entre les vecteurs.
Bien qu'efficaces, ces approches standard ne réduisent la précision ou la dimensionnalité qu'une seule fois et à une seule échelle. Mais que se passerait-il si nous pouvions maintenir simultanément plusieurs couches de détails, comme une pyramide de représentations de plus en plus précises ?
C'est là qu'interviennent les encastrements Matryoshka. Nommées d'après les poupées russes gigognes (voir illustration), ces constructions astucieuses intègrent plusieurs échelles de représentation dans un seul vecteur. Contrairement aux méthodes traditionnelles de post-traitement, les Matryoshka embeddings apprennent cette structure multi-échelle au cours du processus d'apprentissage initial. Le résultat est remarquable : non seulement l'intégration complète capture la sémantique de l'entrée, mais chaque préfixe de sous-ensemble imbriqué (première moitié, premier quart, etc.) fournit une représentation cohérente, bien que moins détaillée.
Dans ce carnet, nous examinons comment utiliser les encastrements Matryoshka avec Milvus pour la recherche sémantique. Nous illustrons un algorithme appelé "funnel search" qui nous permet d'effectuer une recherche de similarité sur un petit sous-ensemble de nos dimensions d'encastrement sans baisse drastique du rappel.
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
Modèle d'intégration de la charge Matryoshka
Au lieu d'utiliser un modèle d'intégration standard tel que sentence-transformers/all-MiniLM-L12-v2nous utilisons un modèle de Nomic entraîné spécialement pour produire des embeddings 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>
Chargement de l'ensemble de données, intégration des éléments et construction de la base de données vectorielle
Le code suivant est une modification de celui de la page de documentation "Movie Search with Sentence Transformers and Milvus". Tout d'abord, nous chargeons le jeu de données de HuggingFace. Il contient environ 35k entrées, chacune correspondant à un film ayant un article Wikipédia. Nous utiliserons les champs Title et PlotSummary dans cet exemple.
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
})
Ensuite, nous nous connectons à une base de données Milvus Lite, nous spécifions le schéma de données et nous créons une collection avec ce schéma. Nous stockerons à la fois l'intégration non normalisée et le premier sixième de l'intégration dans des champs distincts. La raison en est que nous avons besoin du premier 1/6e de l'intégration de Matryoshka pour effectuer une recherche de similarité, et des 5/6e restants de l'intégration pour le reclassement et l'amélioration des résultats de la recherche.
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 ne prend pas actuellement en charge la recherche sur des sous-ensembles d'enchâssements, c'est pourquoi nous divisons les enchâssements en deux parties : la tête représente le sous-ensemble initial du vecteur à indexer et à rechercher, et la queue est le reste. Le modèle est formé pour la recherche de similarité par distance cosinusoïdale, nous normalisons donc les intégrations de tête. Cependant, afin de calculer ultérieurement les similarités pour des sous-ensembles plus importants, nous devons stocker la norme de l'intégration de la tête, afin de pouvoir la dé-normaliser avant de la joindre à la queue.
Pour effectuer une recherche via le premier 1/6e de l'intégration, nous devrons créer un index de recherche vectorielle sur le champ head_embedding. Plus tard, nous comparerons les résultats de la "recherche en entonnoir" à ceux d'une recherche vectorielle classique, et construirons donc également un index de recherche sur l'ensemble de l'intégration.
Il est important de noter que nous utilisons la métrique de distance COSINE plutôt que IP, car sinon nous devrions tenir compte des normes d'intégration, ce qui compliquerait la mise en œuvre (cela aura plus de sens une fois que l'algorithme de recherche en entonnoir aura été décrit).
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)
Enfin, nous codons les résumés de l'intrigue pour l'ensemble des 35 000 films et entrons les encastrements correspondants dans la base de données.
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]
Recherche en entonnoir
Mettons maintenant en œuvre une "recherche en entonnoir" en utilisant le premier 1/6e des dimensions de l'encastrement de Matryoshka. J'ai en tête trois films à récupérer et j'ai produit mon propre résumé de l'intrigue pour interroger la base de données. Nous intégrons les requêtes, puis nous effectuons une recherche vectorielle sur le champ head_embedding, ce qui nous permet d'obtenir 128 résultats.
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"],
)
À ce stade, nous avons effectué une recherche sur un espace vectoriel beaucoup plus petit et nous sommes donc susceptibles d'avoir réduit la latence et les besoins de stockage de l'index par rapport à une recherche sur l'espace complet. Examinons les 5 premiers résultats pour chaque requête :
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
Comme nous pouvons le constater, le rappel a souffert de la troncature des encastrements lors de la recherche. La recherche en entonnoir résout ce problème grâce à une astuce astucieuse : nous pouvons utiliser le reste des dimensions d'intégration pour reclasser et élaguer notre liste de candidats afin de récupérer les performances d'extraction sans effectuer de recherches vectorielles supplémentaires coûteuses.
Pour faciliter la présentation de l'algorithme de recherche en entonnoir, nous convertissons les résultats de la recherche Milvus pour chaque requête en un cadre de données 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]
Pour effectuer une recherche en entonnoir, nous itérons sur les sous-ensembles de plus en plus grands des encastrements. À chaque itération, nous classons les candidats en fonction des nouvelles similarités et nous éliminons une partie des candidats les moins bien classés.
Concrètement, à l'étape précédente, nous avons récupéré 128 candidats en utilisant 1/6 des dimensions de l'intégration et de la requête. La première étape de la recherche en entonnoir consiste à recalculer les similitudes entre les requêtes et les candidats en utilisant le premier tiers des dimensions. Les 64 candidats les plus bas sont éliminés. Nous répétons ensuite ce processus avec les deux premiers tiers des dimensions, puis avec toutes les dimensions, en réduisant successivement le nombre de candidats à 32 et à 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
Nous avons pu rétablir le rappel sans effectuer de recherches vectorielles supplémentaires ! Qualitativement, ces résultats semblent avoir un meilleur rappel pour "Raiders of the Lost Ark" et "The Shining" que la recherche vectorielle standard du tutoriel, "Movie Search using Milvus and Sentence Transformers", qui utilise un modèle d'intégration différent. Cependant, elle ne parvient pas à trouver "Ferris Bueller's Day Off", sur lequel nous reviendrons plus loin dans le carnet. (Voir l'article Matryoshka Representation Learning pour plus d'expériences quantitatives et d'analyses comparatives).
Comparaison entre la recherche en entonnoir et la recherche normale
Comparons les résultats de notre recherche en entonnoir à une recherche vectorielle standard sur le même ensemble de données avec le même modèle d'intégration. Nous effectuons une recherche sur les encastrements complets.
# 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
À l'exception des résultats pour "Un adolescent fait semblant d'être malade pour ne pas aller à l'école...", les résultats de la recherche en entonnoir sont presque identiques à ceux de la recherche complète, même si la recherche en entonnoir a été effectuée sur un espace de recherche de 128 dimensions contre 768 dimensions pour la recherche normale.
Enquête sur l'échec du rappel de la recherche en entonnoir pour Ferris Bueller's Day Off
Pourquoi la recherche en entonnoir n'a-t-elle pas permis de retrouver Ferris Bueller's Day Off ? Examinons s'il figurait ou non dans la liste originale des candidats ou s'il a été filtré par erreur.
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
Nous constatons que le problème réside dans le fait que la liste initiale des candidats n'était pas assez grande, ou plutôt que l'occurrence souhaitée n'était pas assez similaire à la requête au niveau de granularité le plus élevé. En passant de 128 à 256, la recherche est couronnée de succès. Nous devrions établir une règle empirique pour fixer le nombre de candidats sur un ensemble retenu afin d'évaluer empiriquement le compromis entre le rappel et la latence.
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'ordre a-t-il de l'importance ? Intégration des préfixes et des suffixes.
Le modèle a été entraîné à bien faire correspondre les préfixes récursivement plus petits des enchâssements. L'ordre des dimensions que nous utilisons a-t-il une importance ? Par exemple, pourrions-nous également prendre des sous-ensembles d'enchâssements qui sont des suffixes ? Dans cette expérience, nous inversons l'ordre des dimensions dans les encastrements de Matryoshka et effectuons une recherche en entonnoir.
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
Le rappel est beaucoup plus faible que la recherche en entonnoir ou la recherche normale, comme prévu (le modèle d'intégration a été formé par apprentissage contrastif sur les préfixes des dimensions d'intégration, et non sur les suffixes).
Résumé
Voici une comparaison des résultats de nos recherches entre les différentes méthodes :
