衡量和改进语义搜索中的检索质量
| 时间: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 hub 的 Qdrant/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
标准模式 vs 精确搜索
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=16 和 ef_construct=100。让我们尝试将它们增加到 m=32 和 ef_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。