4.2 向量检索算法 — LightRAG 向量模块

本节导读：深入掌握向量检索的核心算法和工程实现，构建高效精准的语义检索系统

学习目标

• 理解向量检索的基本原理和重要性
• 掌握常用的相似度计算方法
• 学习高效的向量检索算法实现
• 了解向量检索的性能优化策略

核心概念

向量检索的定义

向量检索（Vector Retrieval）是在高维向量空间中根据相似性快速找到最接近目标向量的技术。在LightRAG中，向量检索模块负责从大规模向量库中快速、准确地找到与查询向量最相似的候选文档。

向量检索的重要性

查询向量 → 向量检索 → 候选排序 → 最终结果
    ↓          ↓         ↓        ↓
语义理解    快速匹配   相关度计算   结果呈现
​

向量检索的主要挑战

1. 高维空间：向量维度通常在数百到数千维
2. 大规模数据：向量库可能包含数百万到数十亿向量
3. 实时性要求：需要在毫秒级响应时间内完成检索
4. 准确性要求：确保检索结果的语义相关性
5. 资源约束：有限的计算和内存资源

环境准备 / 前置知识

技术栈要求

• Python 3.8+：主要编程语言
• NumPy：高性能数值计算
• Faiss：高效的向量检索库
• SciPy：科学计算库
• scikit-learn：机器学习库

依赖安装

pip install numpy faiss-cpu scikit-learn scipy h5py
​

分步实战

步骤 1：基础相似度计算实现

import numpy as np
from typing import List, Tuple, Dict
from scipy.spatial.distance import cosine
from sklearn.metrics.pairwise import cosine_similarity

class SimilarityCalculator:
    """相似度计算器"""
    
    def __init__(self):
        self.similarity_methods = {
            'cosine': self._cosine_similarity,
            'euclidean': self._euclidean_similarity,
            'manhattan': self._manhattan_similarity
        }
    
    def compute_similarity(self, vec1: np.ndarray, vec2: np.ndarray, method: str = 'cosine') -> float:
        """计算两个向量的相似度"""
        if method not in self.similarity_methods:
            raise ValueError(f"不支持的方法: {method}")
        
        return self.similarity_methods[method](vec1, vec2)
    
    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """余弦相似度"""
        vec1_norm = vec1 / (np.linalg.norm(vec1) + 1e-8)
        vec2_norm = vec2 / (np.linalg.norm(vec2) + 1e-8)
        return np.dot(vec1_norm, vec2_norm)
    
    def _euclidean_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """欧氏距离相似度"""
        distance = np.linalg.norm(vec1 - vec2)
        max_distance = np.sqrt(2 * len(vec1))
        return 1 - (distance / max_distance)

# 使用示例
similarity_calc = SimilarityCalculator()
vec1 = np.array([0.5, 0.2, 0.8, 0.3])
vec2 = np.array([0.4, 0.3, 0.7, 0.2])

similarity = similarity_calc.compute_similarity(vec1, vec2, 'cosine')
print(f"余弦相似度: {similarity:.4f}")
​

步骤 2：基础向量检索实现

import heapq
from typing import List, Tuple, Any

class BasicVectorRetriever:
    """基础向量检索器"""
    
    def __init__(self, vectors: np.ndarray, method: str = 'cosine'):
        self.vectors = vectors
        self.similarity_calc = SimilarityCalculator()
        self.method = method
        self.size = vectors.shape[0]
    
    def search(self, query: np.ndarray, top_k: int = 10) -> List[Tuple[int, float]]:
        """搜索最相似的向量"""
        if top_k > self.size:
            top_k = self.size
        
        heap = []
        
        for idx, vector in enumerate(self.vectors):
            similarity = self.similarity_calc.compute_similarity(query, vector, self.method)
            
            if len(heap) < top_k:
                heapq.heappush(heap, (similarity, idx))
            else:
                if similarity > heap[0][0]:
                    heapq.heappop(heap)
                    heapq.heappush(heap, (similarity, idx))
        
        results = sorted(heap, key=lambda x: x[0], reverse=True)
        return [(idx, similarity) for similarity, idx in results]
    
    def get_vector_info(self) -> Dict[str, Any]:
        """获取向量库信息"""
        return {
            'size': self.size,
            'memory_usage_mb': self.vectors.nbytes / (1024 * 1024)
        }

# 使用示例
if __name__ == "__main__":
    np.random.seed(42)
    vector_size = 1000
    vector_dim = 128
    vector_db = np.random.randn(vector_size, vector_dim)
    
    retriever = BasicVectorRetriever(vector_db, method='cosine')
    query_vector = np.random.randn(vector_dim)
    top_k = 5
    
    results = retriever.search(query_vector, top_k)
    print(f"搜索结果 (top-{top_k}):")
    for idx, similarity in results:
        print(f"  索引 {idx}: 相似度 {similarity:.4f}")
​

步骤 3：基于Faiss的高效向量检索

import faiss
import time

class FaissVectorRetriever:
    """基于Faiss的高效向量检索器"""
    
    def __init__(self, vectors: np.ndarray, method: str = 'cosine', nlist: int = 100):
        self.vectors = vectors
        self.dimension = vectors.shape[1]
        self.method = method
        self.nlist = nlist
        
        # 构建Faiss索引
        self.index = self._build_index()
        self.index.train(vectors)
        self.index.add(vectors)
        
        # 缓存机制
        self.query_cache = {}
        self.cache_size = 1000
    
    def _build_index(self) -> faiss.Index:
        """构建Faiss索引"""
        if self.method == 'cosine':
            # 余弦相似度：使用内积索引，归一化向量
            quantizer = faiss.IndexFlatIP(self.dimension)
            index = faiss.IndexIVFFlat(quantizer, self.dimension, self.nlist, faiss.METRIC_INNER_PRODUCT)
            
            # 归一化向量
            normalized_vectors = self.vectors / (np.linalg.norm(self.vectors, axis=1, keepdims=True) + 1e-8)
            self.vectors = normalized_vectors
            
        elif self.method == 'l2':
            # 欧氏距离
            quantizer = faiss.IndexFlatL2(self.dimension)
            index = faiss.IndexIVFFlat(quantizer, self.dimension, self.nlist, faiss.METRIC_L2)
        else:
            raise ValueError(f"不支持的方法: {self.method}")
        
        return index
    
    def search(self, query: np.ndarray, top_k: int = 10) -> List[Tuple[int, float]]:
        """使用Faiss搜索"""
        # 检查缓存
        cache_key = tuple(query.tolist())
        if cache_key in self.query_cache:
            return self.query_cache[cache_key]
        
        # 归一化查询向量
        if self.method == 'cosine':
            query = query / (np.linalg.norm(query) + 1e-8)
        
        # 执行搜索
        distances, indices = self.index.search(query.reshape(1, -1), top_k)
        
        # 转换为相似度分数
        if self.method == 'cosine':
            scores = distances[0]
        else:
            max_distance = np.sqrt(2 * self.dimension)
            scores = 1 - (distances[0] / max_distance)
        
        results = []
        for idx, score in zip(indices[0], scores):
            if idx >= 0:
                results.append((int(idx), float(score)))
        
        # 更新缓存
        if len(self.query_cache) < self.cache_size:
            self.query_cache[cache_key] = results
        
        return results
    
    def get_index_info(self) -> Dict:
        """获取索引信息"""
        return {
            'index_type': type(self.index).__name__,
            'dimension': self.dimension,
            'nlist': self.nlist,
            'ntotal': self.index.ntotal,
            'cache_size': len(self.query_cache)
        }

# 使用示例
if __name__ == "__main__":
    # 创建大规模向量库
    np.random.seed(42)
    vector_size = 10000
    vector_dim = 128
    vector_db = np.random.randn(vector_size, vector_dim)
    
    # 初始化Faiss检索器
    print("初始化Faiss检索器...")
    faiss_retriever = FaissVectorRetriever(vector_db, method='cosine', nlist=100)
    
    # 测试查询
    query_vector = np.random.randn(vector_dim)
    top_k = 5
    
    # 单次搜索
    results = faiss_retriever.search(query_vector, top_k)
    print(f"Faiss搜索结果 (top-{top_k}):")
    for idx, similarity in results:
        print(f"  索引 {idx}: 相似度 {similarity:.4f}")
    
    # 索引信息
    print("\n索引信息:")
    for key, value in faiss_retriever.get_index_info().items():
        print(f"  {key}: {value}")
​

常见问题 FAQ

Q1：如何选择合适的相似度计算方法？

A：根据使用场景选择：

• 余弦相似度：最常用，适合高维数据和语义搜索
• 欧氏距离：适合需要考虑向量幅度的情况
• 曼哈顿距离：适合高维稀疏数据

Q2：如何处理大规模向量库？

A：使用Faiss处理大规模向量库：

• 分片处理：将向量库分成多个分片，分别构建索引
• 降维技术：使用PCA等降维方法减少维度
• 缓存机制：缓存常用查询结果

最佳实践

1. 索引选择：根据数据规模选择合适的索引结构
2. 缓存机制：合理使用缓存提升重复查询性能
3. 批量处理：使用批处理提高查询效率
4. 内存管理：监控内存使用，避免内存泄漏

本节小结

本节介绍了LightRAG向量检索模块的核心算法和工程实现。通过基础相似度计算、向量检索实现、Faiss高效检索等步骤，我们掌握了构建高效精准语义检索系统的关键技术。

关键要点：

1. 向量检索是语义检索的核心技术
2. Faiss可以大幅提升检索性能
3. 缓存和批量处理能显著优化性能
4. 需要根据具体场景选择合适的参数

下一节将探讨向量模块与图模块的融合策略算法，完成LightRAG双层检索架构的完整实现。

关键词：向量检索，相似度计算，Faiss，LightRAG，语义搜索
难度：进阶
预计阅读：40分钟