优化文本嵌入 大幅优化RAG检索速度

1 简介

文本嵌入技术能够将文字消息转换成高维向量示意的数字，提供了一种了解和解决文本数据的新模式，协助咱们更好地理解和解决文本数据。

这些向量，也就是数字数组，能够捕捉文本的深层特色，进而允许多种运行。比如了解语义、启动文本分类、聚类、消息检索，甚至优化搜查结果排序等。

传统上，嵌入向量的维度是固定的，理论取2的幂次方，大小介于64到4096之间。

如今，有了套娃嵌入技术，咱们可以依据不同的运行需求，灵敏调整嵌入向量的维度。这样做的好处是显而易见的：不只能够缩小存储需求，降落老本，还能大幅优化检索效率。

2 文本嵌入

从输入字符串到句子嵌入

咱们先定义一个词汇表，这个表把一切或者输入的字符，包括字母、不凡符号、短词和子词，都映射到整数值。比如：

{"a": 1,"b": 2,"c": 3,..."z": 26,"the": 27," ": 28}

经过标志化解决后，咱们可以将令牌（token）列表输入到编码器模型中。这个模型经过少量数据的训练，能够将每个令牌转换为高维数值向量嵌入。

例如，OpenAI的text-embedding-3-large模型的嵌入向量输入维度为3072。

假构想要取得单个句子嵌入，咱们须要从多个令牌嵌入中提敞开息。经常出现的做法是，对一切令牌嵌入求平均值。

3 套娃嵌入（Matryoshka Representation Learning）

套娃嵌入（Matryoshka Representation Learning）是一种先进的文本示意技术，由华盛顿大学、谷歌钻研院和哈佛大学的学者们在2022年宣布的论文《Matryoshka Representation Learning》中初次提出。

套娃嵌入技术能够在繁多的嵌入向量中嵌入多个档次的消息。

打个比如，它不是只训练一个繁多维度为1024的嵌入向量，而是同时优化一组不同大小的维度，如1024、512、256、128、64等。

这样的设计让嵌入向量像套娃一样，外层蕴含着较为概括的消息，而内层则逐渐蕴含更粗疏的消息。这种结构让咱们能够在简直不影响功能的状况下，依据实践需求来调整嵌入向量的长度，从而更好地顺应各种不同的运行环境。

4 套娃嵌入的关键性

假定咱们要在向量数据库中存储一少量文本嵌入向量。每个嵌入有 d 个维度。每个维度都是一个32位的浮点数。这样算上去，存储空间就须要n * d * 4 个字节。

假设咱们想要计算这些向量的相似性，如点积或余弦相似性（只是归一化的点积），维度 d 越高，须要做的数学计算量就越多。

点积公式

有了MRL技术，假设咱们更看重节俭内存和提高解决速度，从而缩小老本，那咱们或者只取前64个维度来用。假设咱们谋求最佳的功能，那就用上一切的维度。当然，也可以选用一个折中的维度数。

总的来说，MRL技术让LLM用户能够在嵌入向量的存储老本和功能之间找到一个平衡点。

5 Nomic AI的MRL运行

Nomic的套娃文本嵌入模型nomic-embed-text-v1.5​是经常使用 matryoshka_dims = [768,512,256,128,64] 训练的。该模型在Hugging Face上地下可用。

这个编码器模型还允许多种前缀，比如[search_query, search_document, classification, clustering]，这象征着它能针对搜查查问、搜查文档、文本分类和聚类等特定义务，提供更为精准的嵌入结果。

以下是nomic-embed-text-v1.5在大规模文本嵌入基准（MTEB）上的体现：

让咱们经常使用PyTorch和Sentence Transformers库在Python中成功该模型：

!pip install torch sentence_transformers einopsimport torchfrom sentence_transformers import SentenceTransformermodel = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5",device=device,trust_remote_code=True,prompts={"search_query": "search_query: ","search_document": "search_document: ","classification": "classification: ","clustering": "clustering: ",},)def embed_sentences(model: SentenceTransformer,sentences: list[str],prompt_name: str,matryoshka_dim: int,device: str,):assert matryoshka_dim <= 768, "maximum dimension for nomic-embed-text-v1.5 is 768"embeddings = model.encode(sentences, prompt_name=prompt_name, device=device, convert_to_tensor=True)embeddings = torch.nn.functional.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],))embeddings = embeddings[:, :matryoshka_dim]embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)return embeddings.cpu()

经常使用 matryoshka_dim 参数，可以将原本768维的嵌入向量启动截断，而后归一化新的嵌入向量。

如今，可以设置咱们希冀的维度，对维基百科上的一些文本内容以及关系疑问启动编码，以供检索增强生成（RAG）的运行场景经常使用:

matryoshka_dim = 64wikipedia_texts = ["The dog (Canis familiaris or Canis lupus familiaris) is a domesticated descendant of the wolf.","Albert Einstein was born in Ulm in the Kingdom of Württemberg in the German Empire, on 14 March 1879.","Einstein excelled at physics and mathematics from an early age, and soon acquired the mathematical expertise normally only found in a child several years his senior.","Werner Karl Heisenberg was a German theoretical physicist, one of the main pioneers of the theory of quantum mechanics, and a principal scientist in the Nazi nuclear weapons program during World War II.","Steven Paul Jobs (February 24, 1955 - October 5, 2011) was an American businessman, inventor, and investor best known for co-founding the technology giant Apple Inc.","The cat (Felis catus), commonly referred to as the domestic cat or house cat, is the only domesticated species in the family Felidae.",]question = ["Where was Albert Einstein born?"]question_embedding = embed_sentences(model,sentences=question,prompt_name="search_query",matryoshka_dim=matryoshka_dim,device=device,)document_embeddings = embed_sentences(model,sentences=wikipedia_texts,prompt_name="search_document",matryoshka_dim=matryoshka_dim,device=device,)

print(f"document_embeddings.shape: {document_embeddings.shape}")print(f"question_embedding.shape:{question_embedding.shape}")>> document_embeddings.shape: torch.Size([6, 64])>> question_embedding.shape:torch.Size([1, 64])

咱们可以用散点图可视化套娃文本嵌入的前两个维度。不过，须要留意的是，这个嵌入模型并没有专门针对二维展现启动优化。

散点图展现了维基百科文本和关系疑问的套娃嵌入结果

接上去，将咱们的文档嵌入存储在向量数据库中。这里经常使用的是Faiss。Faiss是Meta Research的开源库，用于高效相似性搜查和密集向量的聚类。

!pip install faiss-cpuimport faissindex = faiss.IndexFlatIP(matryoshka_dim)index.add(document_embeddings)

经过“准确搜查内积”的方法，咱们构建了一个名为IndexFlatIP的向量数据库，它经常使用的是点积相似性度量。由于咱们经常使用的嵌入向量曾经过归一化解决，所以点积和余弦相似性在这种状况下是等价的。

index 如今是一个蕴含六个文本嵌入的向量数据库：

print(index.ntotal)>> 6

搜查与咱们的疑问最相似的嵌入，并检索前k个结果：

distances, indices = index.search(question_embedding, k=6)print(indices)print(distances)>> [[1 2 3 4 0 5]]>> [[0.96335280.7291920.63353264 0.62068397 0.5125410.43155164]]

咱们最相似的文本在数据库中的索引是1，相似性得分为0.96（最高是1.0）。

# results with d=64print(question)print(wikipedia_texts[1])>> ['Where was Albert Einstein born?']>> 'Albert Einstein was born in Ulm in the Kingdom of Württemberg in the German Empire, on 14 March 1879.'

这里也用matryoshka_dim=768从新运转了代码，获取了相似的结果。但是，更高的维度须要更多的内存和更多的计算。

# results with d=768print(indices)print(distances)>> [[1 2 4 3 0 5]]>> [[0.92466116 0.6457440.54405797 0.54824 0.39331824 0.37972206]]

6 MRL & 量化

假设咱们想要进一步紧缩咱们的嵌入，可以经常使用MRL和二进制向量量化。二进制量化将嵌入向量中一切大于零的数字转换为一，其他的转换为零。

从完整的嵌入向量到小巧的二进制版本

经常使用二进制量化，一个维度为 d 的嵌入向量只有要 d / 8​ 字节的内存，这比32位浮点数的 d * 4 字节缩小了32倍。但是，这种缩小是以功能为代价的。

7 结语

在训练环节中，嵌入模型驳回了套娃损失函数，以优化多个嵌入维度。

经过套娃示意学习，LLM用户可以在缩小文本嵌入大小和接受细微功能损失之间启动掂量。

较小的嵌入向量占用的内存更少，计算量也更小，常年来看有助于节俭老本。同时，它们的计算速度也更快，因此具备更高的检索速度，这关于像RAG这样的运行程序来说尤其关键。

本文转载自​​，作者：