3.4 向量索引与检索优化（下） - 第一部分

本节导读：深入探讨向量索引的高级优化技术和检索性能调优方法，掌握企业级索引系统的构建与维护。

学习目标

• 掌握高级索引算法的原理和实现
• 理解向量检索的性能优化策略
• 学会索引质量评估和调优方法
• 能够构建高性能的索引系统
• 了解索引技术的最新发展趋势

高级索引算法

1. 乘积量化（Product Quantization, PQ）

基本原理

乘积量化通过将高维向量分解为多个低维子向量的量化，大幅减少内存占用和计算复杂度。

核心思想：

• 将高维向量空间分解为多个低维子空间
• 每个子空间独立进行量化编码
• 通过码本压缩存储，显著降低内存需求
• 查询时通过距离表快速计算近似距离

实现代码

import numpy as np
from sklearn.cluster import KMeans

class ProductQuantizationIndex:
    """乘积量化索引实现"""
    
    def __init__(self, m: int = 8, k: int = 256):
        self.m = m
        self.k = k
        self.subspace_codebooks = []
        self.index_vectors = None
        self.codes = None
        self.vector_dim = None
    
    def fit(self, vectors: np.ndarray):
        """构建PQ索引"""
        self.index_vectors = vectors
        n_samples, n_dimensions = vectors.shape
        self.vector_dim = n_dimensions
        
        # 确保维度可以被m整除
        if n_dimensions % self.m != 0:
            raise ValueError(f"维度 {n_dimensions} 不能被 {self.m} 整除")
        
        # 计算每个子空间的维度
        subspace_dim = n_dimensions // self.m
        
        # 为每个子空间学习码本
        for i in range(self.m):
            # 提取子空间
            start_idx = i * subspace_dim
            end_idx = (i + 1) * subspace_dim
            subspace_vectors = vectors[:, start_idx:end_idx]
            
            # 使用K-means聚类
            kmeans = KMeans(n_clusters=self.k, random_state=42)
            kmeans.fit(subspace_vectors)
            
            # 存储码本
            self.subspace_codebooks.append(kmeans.cluster_centers_)
            
            # 编码向量
            if self.codes is None:
                self.codes = np.zeros((n_samples, self.m), dtype=np.uint8)
            
            # 为每个样本分配最近的聚类中心
            for j in range(n_samples):
                distances = np.linalg.norm(subspace_vectors[j] - kmeans.cluster_centers_, axis=1)
                closest_cluster = np.argmin(distances)
                self.codes[j, i] = closest_cluster
    
    def search(self, query_vector: np.ndarray, top_k: int = 5):
        """PQ搜索"""
        n_samples = self.index_vectors.shape[0]
        n_dimensions = self.vector_dim
        
        # 计算距离表
        distance_table = np.zeros((self.m, self.k))
        
        for i in range(self.m):
            start_idx = i * (n_dimensions // self.m)
            end_idx = (i + 1) * (n_dimensions // self.m)
            
            query_subspace = query_vector[start_idx:end_idx]
            
            # 计算查询向量与每个码本的距离
            for j in range(self.k):
                codebook = self.subspace_codebooks[i][j]
                distance = np.linalg.norm(query_subspace - codebook)
                distance_table[i, j] = distance
        
        # 计算每个样本的近似距离
        approximate_distances = np.zeros(n_samples)
        
        for i in range(n_samples):
            total_distance = 0
            for j in range(self.m):
                cluster_id = self.codes[i, j]
                total_distance += distance_table[j, cluster_id]
            approximate_distances[i] = total_distance
        
        # 获取top-k结果
        top_indices = np.argsort(approximate_distances)[:top_k]
        
        return list(zip(top_indices, approximate_distances[top_indices]))
    
    def get_memory_usage(self):
        """获取内存使用情况"""
        # 码本内存
        codebook_memory = sum(book.shape[0] * book.shape[1] * 4 for book in self.subspace_codebooks)
        
        # 编码内存
        codes_memory = self.codes.shape[0] * self.codes.shape[1]
        
        total_memory = codebook_memory + codes_memory
        return {
            'codebooks_mb': codebook_memory / (1024 * 1024),
            'codes_mb': codes_memory / (1024 * 1024),
            'total_mb': total_memory / (1024 * 1024)
        }
​

性能特点

优势：

• 内存占用大幅减少：压缩比可达10-100倍
• 查询速度快：距离预计算，避免重复计算
• 适合超大规模数据：内存敏感型应用的首选

局限性：

• 精度有一定损失：量化过程会引入近似误差
• 参数调优复杂：需要平衡子空间数量和聚类数量

2. 层次可聚类树（HNSW with Clustering）

基本原理

结合HNSW和聚类技术，构建层次化的索引结构，提升查询性能。

核心思想：

• 全局聚类：将数据空间划分为多个聚类区域
• 局部索引：在每个聚类区域内构建HNSW索引
• 全局路由：使用全局索引快速定位到相关聚类
• 分层搜索：先全局后局部的两阶段搜索策略

实现架构

import numpy as np
from sklearn.cluster import KMeans

class HierarchicalClusterHNSW:
    """层次聚类HNSW索引"""
    
    def __init__(self, M: int = 16, ef: int = 32, n_clusters: int = 10):
        self.M = M  # HNSW连接数
        self.ef = ef  # HNSW搜索宽度
        self.n_clusters = n_clusters
        self.kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        self.cluster_indexes = {}
        self.global_index = None
        self.cluster_centers = None
        self.cluster_labels = None
    
    def fit(self, vectors: np.ndarray):
        """构建层次索引"""
        # 第一层：全局聚类
        self.kmeans.fit(vectors)
        self.cluster_labels = self.kmeans.labels_
        self.cluster_centers = self.kmeans.cluster_centers_
        
        # 第二层：构建局部HNSW索引
        for cluster_id in range(self.n_clusters):
            cluster_vectors = vectors[self.cluster_labels == cluster_id]
            
            if len(cluster_vectors) > 0:
                # 构建簇内HNSW索引
                cluster_index = SimpleHNSWIndex(M=self.M, ef=self.ef)
                cluster_index.fit(cluster_vectors)
                self.cluster_indexes[cluster_id] = cluster_index
        
        # 构建全局索引（使用簇中心）
        self.global_index = SimpleHNSWIndex(M=self.M, ef=self.ef)
        self.global_index.fit(self.cluster_centers)
    
    def search(self, query_vector: np.ndarray, top_k: int = 5):
        """层次搜索"""
        # 1. 在全局索引中找到最近的簇
        global_results = self.global_index.search(query_vector, min(top_k * 2, self.n_clusters))
        
        # 2. 在相关聚类中搜索
        candidates = []
        for cluster_idx, distance in global_results:
            cluster_id = cluster_idx
            
            if cluster_id in self.cluster_indexes:
                cluster_index = self.cluster_indexes[cluster_id]
                local_results = cluster_index.search(query_vector, top_k)
                
                # 将局部索引转换为全局索引，加上距离权重
                for local_idx, local_distance in local_results:
                    # 找到原始索引ID
                    global_vector_idx = np.where(self.cluster_labels == cluster_id)[0][local_idx]
                    weighted_distance = local_distance + distance
                    candidates.append((global_vector_idx, weighted_distance))
        
        # 3. 排序并返回top-k
        candidates.sort(key=lambda x: x[1])
        return candidates[:top_k]
​

特点：

• 分层搜索，速度快：全局索引大幅减少搜索范围
• 准确性高：保留HNSW的精确搜索特性
• 内存效率好：避免了单个大规模HNSW的高内存消耗
• 扩展性强：易于添加新的聚类和索引

检索质量优化

1. 多阶段检索策略

策略概述

多阶段检索通过分层的检索策略，在性能和精度之间取得平衡：

阶段1：快速检索

• 使用索引进行快速粗粒度检索
• 返回较大候选集（如100个结果）
• 计算成本低，速度快

阶段2：精细过滤

• 对候选集应用过滤条件
• 排除不相关结果
• 提高结果相关性

阶段3：重排序

• 使用更精确的算法对结果重排序
• 优化最终结果质量
• 计算成本相对较高

实现框架

class MultiStageRetriever:
    """多阶段检索器"""
    
    def __init__(self):
        self.indexes = {}
        self.filters = {}
        self.rankers = {}
    
    def add_stage(self, stage_name: str, index, 
                  filter_func=None, 
                  ranker_func=None):
        """添加检索阶段"""
        self.indexes[stage_name] = index
        self.filters[stage_name] = filter_func or (lambda x: True)
        self.rankers[stage_name] = ranker_func or (lambda x: x[1])
    
    def search(self, query: dict, stages: list):
        """多阶段检索"""
        results = {}
        
        for stage_name in stages:
            if stage_name not in self.indexes:
                continue
            
            # 阶段1：快速检索
            index = self.indexes[stage_name]
            
            # 根据阶段调整检索数量
            if 'coarse' in stage_name:
                top_k = 100
            elif 'medium' in stage_name:
                top_k = 50
            else:
                top_k = 20
            
            candidate_results = index.search(query['vector'], top_k)
            
            # 阶段2：过滤
            filter_func = self.filters[stage_name]
            filtered_results = []
            
            for doc_id, score in candidate_results:
                if filter_func({
                    'doc_id': doc_id, 
                    'score': score, 
                    'query': query
                }):
                    filtered_results.append((doc_id, score))
            
            # 阶段3：重排序
            if len(filtered_results) > 0:
                ranker_func = self.rankers[stage_name]
                ranked_results = sorted(filtered_results, key=ranker_func, reverse=True)
                
                results[stage_name] = ranked_results[:10]
        
        return results
    
    def fusion_results(self, results: dict, strategy: str = 'weighted'):
        """结果融合"""
        if strategy == 'weighted':
            return self._weighted_fusion(results)
        elif strategy == 'rank_based':
            return self._rank_based_fusion(results)
        else:
            raise ValueError(f"不支持的融合策略: {strategy}")
    
    def _weighted_fusion(self, results: dict):
        """加权融合"""
        doc_scores = {}
        stage_names = list(results.keys())
        stage_weights = {name: 1.0/len(stage_names) for name in stage_names}
        
        # 收集所有文档ID
        all_docs = set()
        for stage_results in results.values():
            for doc_id, _ in stage_results:
                all_docs.add(doc_id)
        
        # 计算加权分数
        for doc_id in all_docs:
            total_score = 0
            valid_stages = 0
            
            for stage_name in stage_names:
                for result_doc_id, score in results[stage_name]:
                    if result_doc_id == doc_id:
                        total_score += score * stage_weights[stage_name]
                        valid_stages += 1
                        break
            
            if valid_stages > 0:
                doc_scores[doc_id] = total_score / valid_stages
        
        return sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)
​

2. 查询优化技术

查询扩展

查询扩展通过增加相关词汇或语义变体来提高检索覆盖率：

class QueryExpansion:
    """查询扩展器"""
    
    def __init__(self, embedding_model, expansion_method: str = 'similarity'):
        self.embedding_model = embedding_model
        self.expansion_method = expansion_method
        self.similarity_threshold = 0.7
        self.expansion_limit = 5
        self.context_vectors = None
    
    def set_context(self, context_vectors: np.ndarray):
        """设置上下文向量"""
        self.context_vectors = context_vectors
    
    def expand_query(self, original_query: str, **kwargs):
        """扩展查询"""
        # 获取查询向量
        query_vector = self.embedding_model.encode([original_query])[0]
        
        if self.expansion_method == 'similarity':
            return self._similarity_expansion(query_vector, original_query)
        elif self.expansion_method == 'semantic':
            return self._semantic_expansion(query_vector, original_query)
        else:
            return [original_query]
    
    def _similarity_expansion(self, query_vector: np.ndarray, original_query: str):
        """基于相似度的查询扩展"""
        if self.context_vectors is None:
            return [original_query]
        
        # 计算相似度
        from sklearn.metrics.pairwise import cosine_similarity
        similarities = cosine_similarity([query_vector], self.context_vectors)[0]
        
        # 获取最相似的文档
        top_indices = np.argsort(similarities)[-self.expansion_limit:]
        top_similarities = similarities[top_indices]
        
        # 生成扩展查询
        expanded_queries = [original_query]
        
        for idx, similarity in zip(top_indices, top_similarities):
            if similarity > self.similarity_threshold:
                expanded_queries.append(f"{original_query} related_content")
        
        return expanded_queries
​

索引质量评估

1. 评估指标体系

检索质量评估

class IndexEvaluator:
    """索引质量评估器"""
    
    def __init__(self):
        self.metrics = {}
    
    def evaluate_quality(self, test_queries: list, index, ground_truth: dict):
        """评估索引质量"""
        results = {
            'recall': [],
            'precision': [],
            'map': [],
            'ndcg': [],
            'latency': []
        }
        
        for query in test_queries:
            # 执行检索
            start_time = time.time()
            retrieval_results = index.search(query['vector'], top_k=10)
            end_time = time.time()
            
            # 计算指标
            recall = self._calculate_recall(retrieval_results, ground_truth[query['id']])
            precision = self._calculate_precision(retrieval_results, ground_truth[query['id']])
            ap = self._calculate_average_precision(retrieval_results, ground_truth[query['id']])
            ndcg = self._calculate_ndcg(retrieval_results, ground_truth[query['id']])
            latency = end_time - start_time
            
            # 记录结果
            results['recall'].append(recall)
            results['precision'].append(precision)
            results['map'].append(ap)
            results['ndcg'].append(ndcg)
            results['latency'].append(latency)
        
        # 计算平均指标
        evaluation_results = {
            'mean_recall': np.mean(results['recall']),
            'mean_precision': np.mean(results['precision']),
            'mean_average_precision': np.mean(results['map']),
            'mean_ndcg': np.mean(results['ndcg']),
            'mean_latency': np.mean(results['latency']),
            'throughput': 1.0 / np.mean(results['latency'])  # QPS
        }
        
        return evaluation_results
    
    def _calculate_recall(self, retrieval_results: list, ground_truth: list, k: int = 10):
        """计算召回率"""
        retrieved_set = set([idx for idx, _ in retrieval_results[:k]])
        truth_set = set(ground_truth)
        
        return len(retrieved_set & truth_set) / len(truth_set)
    
    def _calculate_precision(self, retrieval_results: list, ground_truth: list, k: int = 10):
        """计算准确率"""
        retrieved_set = set([idx for idx, _ in retrieval_results[:k]])
        truth_set = set(ground_truth)
        
        if len(retrieved_set) == 0:
            return 0.0
        
        return len(retrieved_set & truth_set) / len(retrieved_set)
    
    def _calculate_average_precision(self, retrieval_results: list, ground_truth: list):
        """计算平均精度"""
        retrieved_set = set([idx for idx, _ in retrieval_results])
        truth_set = set(ground_truth)
        
        precision_sum = 0
        relevant_count = 0
        
        for i, (idx, _) in enumerate(retrieval_results):
            if idx in truth_set:
                relevant_count += 1
                precision_sum += relevant_count / (i + 1)
        
        if len(truth_set) == 0:
            return 0.0
        
        return precision_sum / len(truth_set)
    
    def _calculate_ndcg(self, retrieval_results: list, ground_truth: list, k: int = 10):
        """计算NDCG"""
        def dcg_at_k(relevance_scores: list, k: int):
            dcg = 0
            for i in range(min(k, len(relevance_scores))):
                dcg += relevance_scores[i] / (np.log2(i + 2))
            return dcg
        
        # 生成相关性得分
        relevance_scores = []
        for idx, _ in retrieval_results[:k]:
            relevance_scores.append(1.0 if idx in ground_truth else 0.0)
        
        # 计算DCG
        dcg = dcg_at_k(relevance_scores, k)
        
        # 计算理想DCG
        ideal_relevance_scores = [1.0] * min(k, len(ground_truth))
        idcg = dcg_at_k(ideal_relevance_scores, k)
        
        return dcg / idcg if idcg > 0 else 0.0
​

最佳实践总结

1. 索引选择策略

def select_optimal_index_strategy(data_characteristics: dict, query_patterns: dict):
    """选择最优索引策略"""
    
    strategy_recommendations = {
        'small_dataset': {
            'data_size': '< 10K',
            'recommended_index': 'HNSW',
            'parameters': {'M': 16, 'ef': 32},
            'reasoning': '数据量小，使用HNSW可以获得最佳精度'
        },
        'medium_dataset': {
            'data_size': '10K - 100K',
            'recommended_index': 'HNSW with Clustering',
            'parameters': {'M': 24, 'ef': 48, 'n_clusters': 10},
            'reasoning': '中等数据量，分层索引提供较好的性能平衡'
        },
        'large_dataset': {
            'data_size': '100K - 1M',
            'recommended_index': 'Product Quantization',
            'parameters': {'m': 8, 'k': 256},
            'reasoning': '大数据量，PQ可以大幅减少内存使用'
        },
        'very_large_dataset': {
            'data_size': '> 1M',
            'recommended_index': 'Distributed HNSW',
            'parameters': {'shards': 16, 'M': 16, 'ef': 32},
            'reasoning': '超大数据量，分布式架构提供水平扩展能力'
        }
    }
    
    # 基于数据规模选择策略
    data_size = data_characteristics.get('size', 0)
    
    if data_size < 10000:
        strategy = strategy_recommendations['small_dataset']
    elif data_size < 100000:
        strategy = strategy_recommendations['medium_dataset']
    elif data_size < 1000000:
        strategy = strategy_recommendations['large_dataset']
    else:
        strategy = strategy_recommendations['very_large_dataset']
    
    # 基于查询模式调整参数
    if query_patterns.get('real_time_requirement', False):
        strategy['parameters']['ef'] = min(strategy['parameters']['ef'] * 2, 128)
    
    if query_patterns.get('accuracy_critical', False):
        strategy['parameters']['M'] = min(strategy['parameters']['M'] * 2, 64)
    
    return strategy
​

2. 性能调优原则

• 数据驱动：基于实际数据特征和查询模式选择策略
• 渐进式优化：从简单到复杂逐步优化，避免激进改动
• 监控先行：建立完善的监控机制，及时发现性能问题
• 持续改进：定期评估和优化索引性能
• 成本意识：在性能和成本之间找到平衡点

总结

向量索引与检索优化是RAG系统性能提升的核心技术。通过本章的学习，你应该掌握：

1. 高级索引算法：乘积量化、层次聚类HNSW等高效索引技术
2. 检索质量优化：多阶段检索、查询扩展、结果融合等优化策略
3. 索引质量评估：完整的评估体系和基准测试框架
4. 企业级索引系统：分布式架构、监控维护、自动优化等企业级特性

在实际应用中，建议采用系统化的方法：从基础索引开始，根据业务需求逐步升级到更复杂的索引技术，并通过持续的监控和优化确保系统性能。记住，没有"最佳"的索引策略，只有最适合当前业务需求的索引选择。

高质量的索引将为RAG系统提供高效、可靠的检索能力，是整个系统成功的关键保障。

延伸阅读

• 官方文档：FAISS技术文档、Milvus索引指南
• 相关章节：本教程 3.5 节数据库性能调优

关键词：RAG知识库实战, 向量索引, 检索优化, 乘积量化, HNSW, 性能评估, 分布式索引
难度：高级
预计阅读：25分钟