4.1 基础召回算法 — AI搜索技术内幕 向量检索算法

本节导读：掌握AI搜索中最基础的召回算法原理，包括精确检索、近似检索和分层检索等核心技术，为构建高效召回系统奠定基础。

学习目标

• 掌握精确检索（Exact Search）的数学原理和实现方法
• 理解近似检索（Approximate Nearest Neighbor, ANN）的核心思想和常用算法
• 学习分层检索（Hierarchical Retrieval）的架构设计和优化策略
• 熟悉不同召回算法的性能特点和适用场景
• 能够根据实际业务需求选择合适的召回算法

核心概念

精确检索 vs 近似检索

在AI搜索系统中，召回算法主要分为精确检索和近似检索两大类：

精确检索（Exact Search）：

• 在高维向量空间中找到与查询向量最相似的k个向量
• 保证结果的准确性，但计算复杂度高
• 时间复杂度：O(d×n)，其中d是向量维度，n是向量总数
• 适用于小规模数据集或对精度要求极高的场景

近似检索（ANN）：

• 以牺牲一定精度为代价，大幅提升检索速度
• 通过各种索引结构和算法设计，找到"足够好"的近似结果
• 时间复杂度：O(log n)或O(1)，与向量总数无关
• 适用于大规模数据集和实时检索场景

召回算法的评估指标

召回算法的性能通常通过以下指标进行评估：

• 召回率（Recall）：检索到的相关文档数 / 总相关文档数
• 精确率（Precision）：检索到的相关文档数 / 检索到的总文档数
• 延迟（Latency）：单次查询的平均响应时间
• 吞吐量（Throughput）：单位时间内能够处理的查询数量
• 内存使用（Memory Usage）：算法运行所需的内存空间

环境准备 / 前置知识

必需的Python库

pip install numpy faiss-cpu matplotlib scikit-learn pandas
​

核心依赖说明

• numpy：数值计算基础库
• faiss-cpu：Facebook的相似性搜索库（CPU版本）
• scikit-learn：机器学习工具包，包含传统检索算法
• matplotlib：数据可视化库
• pandas：数据处理和分析库

理论基础要求

• 线性代数：向量运算、矩阵运算、内积计算
• 数据结构：树结构、图结构、哈希表
• 算法分析：时间复杂度、空间复杂度分析
• 概率统计：概率分布、期望值、方差计算

分步实战

步骤 1：精确检索算法实现

精确检索是AI搜索的基础，让我们从最基本的暴力检索开始：

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import time

class ExactSearch:
    """精确检索类"""
    
    def __init__(self, vectors):
        """
        初始化精确检索器
        
        Args:
            vectors: 向量矩阵，形状为(n_samples, n_features)
        """
        self.vectors = vectors
        self.n_samples = vectors.shape[0]
        self.n_features = vectors.shape[1]
        print(f"初始化精确检索器，向量数量: {self.n_samples}, 维度: {self.n_features}")
    
    def cosine_search(self, query_vector, k=10):
        """
        使用余弦相似度进行检索
        
        Args:
            query_vector: 查询向量，形状为(1, n_features)
            k: 返回最相似的k个向量
            
        Returns:
            indices: 最相似向量的索引
            similarities: 对应的相似度分数
        """
        # 计算余弦相似度
        similarities = cosine_similarity(query_vector, self.vectors)[0]
        
        # 获取最相似的k个索引
        top_k_indices = np.argsort(similarities)[::-1][:k]
        top_k_similarities = similarities[top_k_indices]
        
        return top_k_indices, top_k_similarities
    
    def euclidean_search(self, query_vector, k=10):
        """
        使用欧氏距离进行检索
        
        Args:
            query_vector: 查询向量，形状为(1, n_features)
            k: 返回最相似的k个向量
            
        Returns:
            indices: 最相似向量的索引
            distances: 对应的距离分数（越小越好）
        """
        # 计算欧氏距离
        distances = np.sqrt(np.sum((self.vectors - query_vector)**2, axis=1))
        
        # 距离越小越相似，所以取前k个最小的距离
        top_k_indices = np.argsort(distances)[:k]
        top_k_distances = distances[top_k_indices]
        
        return top_k_indices, top_k_distances

# 生成测试数据
def generate_test_data(n_samples=10000, n_features=128):
    """生成测试向量数据"""
    np.random.seed(42)
    vectors = np.random.randn(n_samples, n_features)
    # 归一化向量
    vectors = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
    return vectors

# 测试精确检索
print("=== 精确检索测试 ===")
vectors = generate_test_data(5000, 128)
searcher = ExactSearch(vectors)

# 创建查询向量
query_vector = np.random.randn(1, 128)
query_vector = query_vector / np.linalg.norm(query_vector)

# 测试余弦相似度搜索
start_time = time.time()
indices_cosine, similarities_cosine = searcher.cosine_search(query_vector, k=10)
cosine_time = time.time() - start_time
print(f"余弦相似度搜索耗时: {cosine_time:.4f}秒")
print(f"前5个相似向量索引: {indices_cosine[:5]}")
print(f"前5个相似度分数: {similarities_cosine[:5]:.4f}")

# 测试欧氏距离搜索
start_time = time.time()
indices_euclidean, distances_euclidean = searcher.euclidean_search(query_vector, k=10)
euclidean_time = time.time() - start_time
print(f"欧氏距离搜索耗时: {euclidean_time:.4f}秒")
print(f"前5个相似向量索引: {indices_euclidean[:5]}")
print(f"前5个距离分数: {distances_euclidean[:5]:.4f}")
​

步骤 2：近似检索算法实现

精确检索在大规模数据集上效率低下，让我们实现几种常用的近似检索算法：

import faiss
from sklearn.neighbors import BallTree

class ApproximateSearch:
    """近似检索类"""
    
    def __init__(self, vectors):
        """
        初始化近似检索器
        
        Args:
            vectors: 向量矩阵，形状为(n_samples, n_features)
        """
        self.vectors = vectors.astype('float32')
        self.n_samples = vectors.shape[0]
        self.n_features = vectors.shape[1]
        
        # 初始化各种索引
        self._build_indices()
    
    def _build_indices(self):
        """构建各种检索索引"""
        print("构建近似检索索引...")
        
        # FAISS 索引
        # 使用IVF倒排文件索引
        nlist = min(100, int(np.sqrt(self.n_samples)))  # 聚类中心数量
        quantizer = faiss.IndexFlatL2(self.n_features)  # 量化器
        self.ivf_index = faiss.IndexIVFFlat(quantizer, self.n_features, nlist)
        self.ivf_index.train(self.vectors)
        self.ivf_index.add(self.vectors)
        
        # 使用HNSW索引
        self.hnsw_index = faiss.IndexHNSWFlat(self.n_features, 32)  # M参数为32
        self.hnsw_index.add(self.vectors)
        
        # 使用PQ乘积量化
        nbits = 8  # 每个向量用8个int8表示
        pq_index = faiss.IndexPQ(self.n_features, nbits, 8)  # 8个子空间
        pq_index.train(self.vectors)
        pq_index.add(self.vectors)
        
        # Scikit-learn的Ball Tree
        self.ball_tree = BallTree(self.vectors)
        
        print("索引构建完成")
    
    def search_ivf(self, query_vector, k=10, nprobe=10):
        """
        使用IVF索引进行检索
        
        Args:
            query_vector: 查询向量
            k: 返回数量
            nprobe: 搜索的聚类中心数量
            
        Returns:
            indices: 检索到的向量索引
            distances: 对应的距离
        """
        query_vector = query_vector.astype('float32')
        self.ivf_index.nprobe = nprobe  # 设置搜索的聚类中心数量
        distances, indices = self.ivf_index.search(query_vector, k)
        return indices[0], distances[0]
    
    def search_hnsw(self, query_vector, k=10):
        """
        使用HNSW索引进行检索
        
        Args:
            query_vector: 查询向量
            k: 返回数量
            
        Returns:
            indices: 检索到的向量索引
            distances: 对应的距离
        """
        query_vector = query_vector.astype('float32')
        distances, indices = self.hnsw_index.search(query_vector, k)
        return indices[0], distances[0]
    
    def search_ball_tree(self, query_vector, k=10):
        """
        使用Ball Tree进行检索
        
        Args:
            query_vector: 查询向量
            k: 返回数量
            
        Returns:
            indices: 检索到的向量索引
            distances: 对应的距离
        """
        distances, indices = self.ball_tree.query(query_vector, k=k)
        return indices[0], distances[0]

# 测试近似检索
print("\n=== 近似检索测试 ===")
approx_searcher = ApproximateSearch(vectors)

# 创建查询向量
query_vector = np.random.randn(1, 128).astype('float32')
query_vector = query_vector / np.linalg.norm(query_vector)

# 测试不同检索算法
algorithms = [
    ("IVF", approx_searcher.search_ivf),
    ("HNSW", approx_searcher.search_hnsw),
    ("Ball Tree", approx_searcher.search_ball_tree),
]

for name, search_func in algorithms:
    start_time = time.time()
    if name == "IVF":
        indices, distances = search_func(query_vector, k=10, nprobe=10)
    elif name == "HNSW":
        indices, distances = search_func(query_vector, k=10)
    else:
        indices, distances = search_func(query_vector, k=10)
    
    elapsed_time = time.time() - start_time
    print(f"{name}搜索耗时: {elapsed_time:.4f}秒")
    print(f"前5个相似向量索引: {indices[:5]}")
    print(f"前5个距离分数: {distances[:5]:.4f}")
​

步骤 3：分层检索架构实现

分层检索是处理超大规模数据集的有效策略，让我们实现一个完整的分层检索系统：

class HierarchicalSearch:
    """分层检索系统"""
    
    def __init__(self, vectors, layer_ratios=[0.1, 0.3, 0.6]):
        """
        初始化分层检索系统
        
        Args:
            vectors: 向量矩阵
            layer_ratios: 各层的数据比例
        """
        self.vectors = vectors
        self.layer_ratios = layer_ratios
        self.layers = []
        self._build_layers()
    
    def _build_layers(self):
        """构建分层结构"""
        total_samples = len(self.vectors)
        
        for i, ratio in enumerate(self.layer_ratios):
            layer_size = int(total_samples * ratio)
            layer_vectors = self.vectors[:layer_size]
            
            if i == 0:
                # 第一层：使用精确检索
                layer = ExactSearch(layer_vectors)
            else:
                # 其他层：使用近似检索
                layer = ApproximateSearch(layer_vectors)
            
            self.layers.append({
                'size': layer_size,
                'vectors': layer_vectors,
                'searcher': layer
            })
            
            print(f"第{i+1}层: {layer_size}个向量，使用{'精确' if i == 0 else '近似'}检索")
    
    def search(self, query_vector, k=10, layer_weights=None):
        """
        分层检索
        
        Args:
            query_vector: 查询向量
            k: 返回数量
            layer_weights: 各层权重
            
        Returns:
            combined_indices: 合并后的索引
            combined_scores: 合并后的分数
        """
        if layer_weights is None:
            # 默认权重：底层权重更高
            layer_weights = [0.7, 0.2, 0.1]
        
        results = []
        
        # 在每一层进行检索
        for i, layer in enumerate(self.layers):
            layer_searcher = layer['searcher']
            
            if isinstance(layer_searcher, ExactSearch):
                indices, scores = layer_searcher.cosine_search(query_vector, k=k)
            else:
                indices, scores = layer_searcher.search_hnsw(query_vector, k=k)
            
            # 根据层数调整分数
            adjusted_scores = scores * layer_weights[i]
            
            # 添加层ID和位置
            for idx, score in zip(indices, adjusted_scores):
                results.append({
                    'global_idx': idx,
                    'layer_score': score,
                    'layer_id': i
                })
        
        # 按分数排序并返回前k个结果
        results.sort(key=lambda x: x['layer_score'], reverse=True)
        top_k_results = results[:k]
        
        combined_indices = [r['global_idx'] for r in top_k_results]
        combined_scores = [r['layer_score'] for r in top_k_results]
        
        return combined_indices, combined_scores

# 测试分层检索
print("\n=== 分层检索测试 ===")
hierarchical_searcher = HierarchicalSearch(vectors)

# 创建查询向量
query_vector = np.random.randn(1, 128).astype('float32')
query_vector = query_vector / np.linalg.norm(query_vector)

# 测试分层检索
start_time = time.time()
indices, scores = hierarchical_searcher.search(query_vector, k=10)
hierarchical_time = time.time() - start_time

print(f"分层检索耗时: {hierarchical_time:.4f}秒")
print(f"前5个相似向量索引: {indices[:5]}")
print(f"前5个调整后的分数: {scores[:5]:.4f}")
​