Pesquisa em funil com embeddings Matryoshka
Ao criar sistemas de pesquisa vetorial eficientes, um dos principais desafios é gerir os custos de armazenamento, mantendo uma latência e uma recuperação aceitáveis. Os modelos modernos de incorporação produzem vectores com centenas ou milhares de dimensões, criando um armazenamento significativo e uma sobrecarga computacional para o vetor bruto e o índice.
Tradicionalmente, os requisitos de armazenamento são reduzidos através da aplicação de um método de quantização ou de redução da dimensionalidade imediatamente antes da construção do índice. Por exemplo, podemos poupar armazenamento reduzindo a precisão utilizando a Quantização de Produtos (PQ) ou o número de dimensões utilizando a Análise de Componentes Principais (PCA). Estes métodos analisam todo o conjunto de vectores para encontrar um conjunto mais compacto que mantenha as relações semânticas entre os vectores.
Embora eficazes, estas abordagens padrão reduzem a precisão ou a dimensionalidade apenas uma vez e numa única escala. Mas e se pudéssemos manter várias camadas de pormenor em simultâneo, como uma pirâmide de representações cada vez mais precisas?
Eis os embeddings Matryoshka. Com o nome das bonecas russas (ver ilustração), estas construções inteligentes incorporam várias escalas de representação num único vetor. Ao contrário dos métodos tradicionais de pós-processamento, os encaixes Matryoshka aprendem esta estrutura multi-escala durante o processo de formação inicial. O resultado é notável: não só a incorporação completa capta a semântica da entrada, como cada prefixo de subconjunto aninhado (primeira metade, primeiro quarto, etc.) fornece uma representação coerente, embora menos pormenorizada.
Neste caderno, examinamos a forma de utilizar as incrustações Matryoshka com o Milvus para a pesquisa semântica. Ilustramos um algoritmo chamado "pesquisa em funil" que nos permite efetuar pesquisas por semelhança num pequeno subconjunto das nossas dimensões de incorporação sem uma queda drástica na recuperação.
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
Carregar o modelo de incorporação Matryoshka
Em vez de utilizar um modelo de incorporação padrão, como sentence-transformers/all-MiniLM-L12-v2utilizamos um modelo da Nomic treinado especialmente para produzir 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>
Carregar o conjunto de dados, incorporar itens e construir uma base de dados de vectores
O código seguinte é uma modificação do código da página de documentação "Movie Search with Sentence Transformers and Milvus". Primeiro, carregamos o conjunto de dados do HuggingFace. Ele contém cerca de 35 mil entradas, cada uma correspondendo a um filme com um artigo na Wikipédia. Neste exemplo, usaremos os campos 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
})
Em seguida, ligamo-nos a uma base de dados Milvus Lite, especificamos o esquema de dados e criamos uma coleção com este esquema. Iremos armazenar a incorporação não normalizada e o primeiro sexto da incorporação em campos separados. A razão para isto é que precisamos do primeiro 1/6 da incorporação Matryoshka para efetuar uma pesquisa de semelhança e os restantes 5/6 das incorporações para reordenar e melhorar os resultados da pesquisa.
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)
Atualmente, o Milvus não suporta a pesquisa em subconjuntos de sequências, pelo que dividimos as sequências em duas partes: a cabeça representa o subconjunto inicial do vetor a indexar e pesquisar, e a cauda é o restante. O modelo é treinado para a pesquisa de semelhança de distância de cosseno, por isso normalizamos os embeddings da cabeça. No entanto, para calcular semelhanças para subconjuntos maiores mais tarde, precisamos armazenar a norma da incorporação da cabeça, para que possamos desnormalizá-la antes de juntar à cauda.
Para efetuar a pesquisa através do primeiro 1/6 da incorporação, teremos de criar um índice de pesquisa vetorial sobre o campo head_embedding. Mais tarde, compararemos os resultados da "pesquisa em funil" com uma pesquisa vetorial normal e, por isso, criaremos também um índice de pesquisa sobre a incorporação completa.
É importante notar que utilizamos a métrica de distância COSINE em vez da métrica de distância IP, porque, de outro modo, teríamos de registar as normas de incorporação, o que complicaria a implementação (isto fará mais sentido quando o algoritmo de pesquisa em funil for descrito).
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)
Por fim, codificamos os resumos dos enredos de todos os 35 mil filmes e introduzimos os correspondentes embeddings na base de dados.
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]
Efetuar a pesquisa em funil
Vamos agora implementar uma "pesquisa em funil" utilizando o primeiro 1/6 das dimensões do embedding Matryoshka. Tenho três filmes em mente para recuperar e produzi o meu próprio resumo do enredo para consultar a base de dados. Incorporamos as consultas e, em seguida, efectuamos uma pesquisa vetorial no campo head_embedding, obtendo 128 candidatos a resultados.
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"],
)
Nesta altura, fizemos uma pesquisa num espaço vetorial muito mais pequeno e, por isso, é provável que tenhamos reduzido a latência e os requisitos de armazenamento do índice em relação à pesquisa no espaço total. Vamos examinar as 5 principais correspondências para cada consulta:
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
Como podemos ver, a recordação foi afetada como consequência do truncamento dos embeddings durante a pesquisa. A pesquisa em funil resolve isto com um truque inteligente: podemos usar o resto das dimensões de incorporação para classificar e podar a nossa lista de candidatos para recuperar o desempenho da recuperação sem executar quaisquer pesquisas vectoriais adicionais dispendiosas.
Para facilitar a exposição do algoritmo de pesquisa de funil, convertemos os resultados da pesquisa Milvus para cada consulta num quadro de dados 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]
Agora, para efetuar a pesquisa de funil, iteramos sobre os subconjuntos cada vez maiores dos embeddings. Em cada iteração, classificamos novamente os candidatos de acordo com as novas semelhanças e eliminamos uma fração dos candidatos com a classificação mais baixa.
Para concretizar isto, a partir do passo anterior recuperámos 128 candidatos utilizando 1/6 das dimensões de incorporação e de consulta. O primeiro passo para efetuar a pesquisa em funil consiste em recalcular as semelhanças entre as consultas e os candidatos utilizando o primeiro 1/3 das dimensões. Os últimos 64 candidatos são eliminados. Em seguida, repetimos este processo com os primeiros 2/3 das dimensões e, depois, com todas as dimensões, reduzindo sucessivamente para 32 e 16 candidatos.
# 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
Conseguimos recuperar a memória sem efetuar quaisquer pesquisas vectoriais adicionais! Em termos qualitativos, estes resultados parecem ter uma recuperação mais elevada para "Raiders of the Lost Ark" e "The Shining" do que a pesquisa vetorial padrão no tutorial, "Movie Search using Milvus and Sentence Transformers", que utiliza um modelo de incorporação diferente. No entanto, não consegue encontrar "Ferris Bueller's Day Off", ao qual voltaremos mais tarde no caderno. (Ver o artigo Aprendizagem de representações Matryoshka para mais experiências quantitativas e avaliação comparativa).
Comparação da pesquisa em funil com a pesquisa normal
Vamos comparar os resultados da nossa pesquisa em funil com uma pesquisa vetorial normal no mesmo conjunto de dados com o mesmo modelo de incorporação. Realizamos uma pesquisa nos embeddings completos.
# 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
Com exceção dos resultados para "Um adolescente finge estar doente para não ir à escola...", os resultados da pesquisa em funil são quase idênticos aos da pesquisa completa, apesar de a pesquisa em funil ter sido efectuada num espaço de pesquisa de 128 dimensões contra 768 dimensões para a pesquisa regular.
Investigando a falha de recordação da pesquisa de funil para Ferris Bueller's Day Off
Por que é que a pesquisa de funil não conseguiu recuperar o filme Ferris Bueller's Day Off? Vamos examinar se ele estava ou não na lista de candidatos original ou se foi filtrado por engano.
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
Vemos que o problema é que a lista inicial de candidatos não era suficientemente grande, ou melhor, o resultado desejado não é suficientemente semelhante à consulta no nível mais alto de granularidade. Alterá-la de 128 para 256 resulta numa recuperação bem sucedida. Devemos criar uma regra geral para definir o número de candidatos num conjunto retido para avaliar empiricamente o compromisso entre a recuperação e a latência.
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
A ordem é importante? Embeddings de prefixo vs sufixo.
O modelo foi treinado para ter um bom desempenho na correspondência de prefixos recursivamente mais pequenos dos embeddings. Será que a ordem das dimensões que utilizamos é importante? Por exemplo, será que também poderíamos pegar em subconjuntos dos embeddings que são sufixos? Nesta experiência, invertemos a ordem das dimensões nos embeddings Matryoshka e efectuamos uma pesquisa de funil.
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
A recuperação é muito mais fraca do que a pesquisa em funil ou a pesquisa regular, como esperado (o modelo de incorporação foi treinado por aprendizagem contrastiva em prefixos das dimensões de incorporação, e não em sufixos).
Resumo
Aqui está uma comparação dos nossos resultados de pesquisa entre métodos:
