衡量和改进语义搜索中的检索质量

时间:30分钟难度:中等

语义搜索管道的质量取决于它们使用的嵌入(embeddings)。如果你的模型无法正确表示输入数据,那么相似的对象在向量空间中可能会彼此相距甚远。毫无疑问,在这种情况下,搜索结果会很差。然而,过程中还有一个组件也可能降低搜索结果的质量,那就是近似最近邻(ANN)算法本身。

在本教程中,我们将展示如何衡量语义检索的质量,以及如何调整 Qdrant 中使用的 ANN 算法 HNSW 的参数以获得最佳结果。

嵌入质量

嵌入的质量是一个单独教程的主题。简而言之,它通常通过基准进行衡量和比较,例如 Massive Text Embedding Benchmark (MTEB)。评估过程本身非常直接,基于人工构建的地面真实数据集。我们有一组查询和一组我们期望为每个查询接收的文档。在评估过程中,我们获取一个查询,在向量空间中找到最相似的文档,并将它们与地面真实进行比较。在这种设置中,查找最相似的文档被实现为完全的 kNN 搜索,没有任何近似。因此,我们可以衡量嵌入本身的质量,而不受 ANN 算法的影响。

检索质量

嵌入质量确实是语义搜索质量中最重要的因素。然而,像 Qdrant 这样的向量搜索引擎并不执行纯粹的 kNN 搜索。相反,它们使用近似最近邻(ANN)算法,这些算法比精确搜索快得多,但可能返回次优结果。我们还可以衡量这种近似的检索质量,这也对整体搜索质量有所贡献。

质量指标

有多种方法可以量化语义搜索的质量。其中一些,例如 Precision@k,基于前 k 个搜索结果中相关文档的数量。另一些,例如 Mean Reciprocal Rank (MRR),考虑了搜索结果中第一个相关文档的位置。DCG 和 NDCG 指标则基于文档的相关性分数。

如果我们将搜索管道视为一个整体,我们可以全部使用这些指标。嵌入质量评估也是如此。然而,对于 ANN 算法本身,任何基于相关性分数或排名的指标都不适用。向量搜索中的排名依赖于查询和文档在向量空间中的距离,但由于近似,距离不会改变,因为函数仍然相同。

因此,通过前 k 个搜索结果中相关文档的数量(例如 precision@k)来衡量 ANN 算法的质量才有意义。它计算为前 k 个搜索结果中相关文档的数量除以 k。在仅测试 ANN 算法的情况下,我们可以使用精确的 kNN 搜索作为地面真实,其中 k 是固定的。这将衡量 ANN 算法近似精确搜索的程度

衡量搜索结果的质量

让我们建立一个 Qdrant 中 ANN 算法的质量评估。我们首先以标准方式调用搜索端点以获取近似搜索结果。然后,我们将调用精确搜索端点以获取精确匹配,最后在精度方面比较这两个结果。

在开始之前,让我们创建一个集合,填充一些数据,然后开始我们的评估。我们将使用与从 Hugging Face hub 加载数据集教程中相同的数据集,即来自 Hugging Face hubQdrant/arxiv-titles-instructorxl-embeddings。我们将以流模式下载它,因为我们只使用其中的一部分。

from datasets import load_dataset

dataset = load_dataset(
    "Qdrant/arxiv-titles-instructorxl-embeddings", split="train", streaming=True
)

我们需要一些数据进行索引,另一些用于测试目的。让我们获取前 50000 个项目用于训练,接下来的 1000 个项目用于测试。

dataset_iterator = iter(dataset)
train_dataset = [next(dataset_iterator) for _ in range(60000)]
test_dataset = [next(dataset_iterator) for _ in range(1000)]

现在,让我们创建一个集合并索引训练数据。此集合将以默认配置创建。请注意,它可能与你的集合设置不同,始终重要的是测试与你稍后在生产中使用的完全相同的配置。

from qdrant_client import QdrantClient, models

client = QdrantClient("https://:6333")
client.create_collection(
    collection_name="arxiv-titles-instructorxl-embeddings",
    vectors_config=models.VectorParams(
        size=768,  # Size of the embeddings generated by InstructorXL model
        distance=models.Distance.COSINE,
    ),
)

我们现在准备索引训练数据。上传记录将触发索引过程,该过程将构建 HNSW 图。索引过程可能需要一些时间,具体取决于数据集的大小,但在从 upsert 端点接收响应后,你的数据将立即可用于搜索。只要索引尚未完成,HNSW 尚未构建,Qdrant 将执行精确搜索。我们必须等到索引完成后才能确保执行近似搜索。

client.upload_points(  # upload_points is available as of qdrant-client v1.7.1
    collection_name="arxiv-titles-instructorxl-embeddings",
    points=[
        models.PointStruct(
            id=item["id"],
            vector=item["vector"],
            payload=item,
        )
        for item in train_dataset
    ]
)

while True:
    collection_info = client.get_collection(collection_name="arxiv-titles-instructorxl-embeddings")
    if collection_info.status == models.CollectionStatus.GREEN:
        # Collection status is green, which means the indexing is finished
        break

Qdrant 具有内置的精确搜索模式,可用于衡量搜索结果的质量。在此模式下,Qdrant 为每个查询执行完整的 kNN 搜索,没有任何近似。它不适用于高负载的生产环境,但非常适合评估 ANN 算法及其参数。可以通过在搜索请求中将 exact 参数设置为 True 来触发它。我们将简单地使用测试数据集中的所有示例作为查询,并比较近似搜索结果与精确搜索结果。让我们创建一个带有参数 k 的辅助函数,这样我们就可以计算不同 k 值的 precision@k

def avg_precision_at_k(k: int):
    precisions = []
    for item in test_dataset:
        ann_result = client.query_points(
            collection_name="arxiv-titles-instructorxl-embeddings",
            query=item["vector"],
            limit=k,
        ).points
    
        knn_result = client.query_points(
            collection_name="arxiv-titles-instructorxl-embeddings",
            query=item["vector"],
            limit=k,
            search_params=models.SearchParams(
                exact=True,  # Turns on the exact search mode
            ),
        ).points

        # We can calculate the precision@k by comparing the ids of the search results
        ann_ids = set(item.id for item in ann_result)
        knn_ids = set(item.id for item in knn_result)
        precision = len(ann_ids.intersection(knn_ids)) / k
        precisions.append(precision)
    
    return sum(precisions) / len(precisions)

计算 precision@5 就像调用具有相应参数的函数一样简单

print(f"avg(precision@5) = {avg_precision_at_k(k=5)}")

响应

avg(precision@5) = 0.9935999999999995

正如我们所看到的,近似搜索与精确搜索的精度相当高。然而,在某些场景下,我们需要更高的精度并可以接受更高的延迟。HNSW 具有很高的可调性,我们可以通过更改其参数来提高精度。

调整 HNSW 参数

HNSW 是一个分层图,其中每个节点都有一组指向其他节点的链接。每个节点的边数称为 m 参数。它的值越大,搜索精度越高,但所需的空间也越大。ef_construct 参数是在索引构建期间要考虑的邻居数量。同样,值越大,精度越高,但索引时间越长。这些参数的默认值为 m=16ef_construct=100。让我们尝试将它们增加到 m=32ef_construct=200,看看它如何影响精度。当然,我们需要等到索引完成后才能执行搜索。

client.update_collection(
    collection_name="arxiv-titles-instructorxl-embeddings",
    hnsw_config=models.HnswConfigDiff(
        m=32,  # Increase the number of edges per node from the default 16 to 32
        ef_construct=200,  # Increase the number of neighbours from the default 100 to 200
    )
)

while True:
    collection_info = client.get_collection(collection_name="arxiv-titles-instructorxl-embeddings")
    if collection_info.status == models.CollectionStatus.GREEN:
        # Collection status is green, which means the indexing is finished
        break

相同的函数可用于计算平均 precision@5

print(f"avg(precision@5) = {avg_precision_at_k(k=5)}")

响应

avg(precision@5) = 0.9969999999999998

精度显然有所提高,我们也知道如何控制它。然而,精度与搜索延迟和内存要求之间存在权衡。在某些特定情况下,我们可能希望尽可能提高精度,现在我们知道如何做到这一点。

总结

评估检索质量是评估语义搜索性能的关键方面。为了获得最佳搜索结果质量,衡量检索质量至关重要。Qdrant 提供内置的精确搜索模式,可用于衡量 ANN 算法本身的质量,甚至可以作为 CI/CD 管道的一部分以自动化方式进行。

再次强调,嵌入的质量是最重要的因素。HNSW 在精度方面做得相当不错,并且在需要时可以参数化和调整。还有其他一些 ANN 算法可用,例如 IVF*,但它们通常在质量和性能方面不如 HNSW

此页面有用吗?

感谢您的反馈!🙏

很抱歉听到这个消息。😔 你可以在 GitHub 上编辑此页面,或者创建一个 GitHub 问题。