优化文本嵌入,大幅提升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 einops

import torch
from sentence_transformers import SentenceTransformer

model = 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 = 64

wikipedia_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-cpu
import faiss
index = 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.9633528  0.729192   0.63353264 0.62068397 0.512541   0.43155164]]

我们最相似的文本在数据库中的索引是1,相似性得分为0.96(最高是1.0)。

# results with d=64
print(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=768
print(indices)
print(distances)
>> [[1 2 4 3 0 5]]
>> [[0.92466116 0.645744   0.54405797 0.54004824 0.39331824 0.37972206]]

6 MRL & 量化

如果我们想要进一步压缩我们的嵌入,可以使用MRL和二进制向量量化。二进制量化将嵌入向量中所有大于零的数字转换为一,其余的转换为零。

使用二进制量化,一个维度为 d 的嵌入向量只需要 d / 8 字节的内存,这比32位浮点数的 d * 4 字节减少了32倍。然而,这种减少是以性能为代价的。

7 结语

在训练过程中,嵌入模型采用了套娃损失函数,以优化多个嵌入维度。

通过套娃表示学习,LLM用户可以在减少文本嵌入大小和接受轻微性能损失之间进行权衡。

较小的嵌入向量占用的内存更少,计算量也更小,长期来看有助于节省成本。同时,它们的计算速度也更快,因此具有更高的检索速度,这对于像RAG这样的应用程序来说尤其重要。

最后的最后

感谢你们的阅读和喜欢,我收藏了很多技术干货,可以共享给喜欢我文章的朋友们,如果你肯花时间沉下心去学习,它们一定能帮到你。

因为这个行业不同于其他行业,知识体系实在是过于庞大,知识更新也非常快。作为一个普通人,无法全部学完,所以我们在提升技术的时候,首先需要明确一个目标,然后制定好完整的计划,同时找到好的学习方法,这样才能更快的提升自己。

这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费

大模型知识脑图

为了成为更好的 AI大模型 开发者,这里为大家提供了总的路线图。它的用处就在于,你可以按照上面的知识点去找对应的学习资源,保证自己学得较为全面。
在这里插入图片描述

经典书籍阅读

阅读AI大模型经典书籍可以帮助读者提高技术水平,开拓视野,掌握核心技术,提高解决问题的能力,同时也可以借鉴他人的经验。对于想要深入学习AI大模型开发的读者来说,阅读经典书籍是非常有必要的。

在这里插入图片描述

实战案例

光学理论是没用的,要学会跟着一起敲,要动手实操,才能将自己的所学运用到实际当中去,这时候可以搞点实战案例来学习。

在这里插入图片描述

面试资料

我们学习AI大模型必然是想找到高薪的工作,下面这些面试题都是总结当前最新、最热、最高频的面试题,并且每道题都有详细的答案,面试前刷完这套面试题资料,小小offer,不在话下

在这里插入图片描述

640套AI大模型报告合集

这套包含640份报告的合集,涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师,还是对AI大模型感兴趣的爱好者,这套报告合集都将为您提供宝贵的信息和启示。

在这里插入图片描述

这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值