IVFFlat（Inverted File with Flat Storage）索引算法

IVFFlat 索引算法介绍概述IVFFlatInverted File with Flat Storage是IVF算法的一个变种它在IVF的基础上保持了原始向量的精确存储。与IVFADC使用量化压缩不同IVFFlat在每个聚类中完整存储原始向量因此在查询时能够提供与FLAT算法相同的精确距离计算。这使得IVFFlat在需要高精度的场景下特别有价值。基本原理IVFFlat结合了IVF的空间划分优势和FLAT的精确计算能力IVF聚类使用K-means等算法将向量空间划分为多个聚类精确存储在每个聚类中完整存储原始向量不进行量化压缩局部搜索查询时只搜索最相关的几个聚类但在这些聚类内进行精确距离计算这种策略通过限制搜索范围来提升查询性能同时保持距离计算的精确性。算法架构数据结构IVFFlat Index Structure: ├── Cluster Centers (k × d) │ ├── Center 1: [c11, c12, ..., c1d] │ ├── Center 2: [c21, c22, ..., c2d] │ └── ... │ └── Center k: [ck1, ck2, ..., ckd] └── Inverted Lists (k lists) ├── List 1: [(id1, v1_full), (id2, v2_full), ...] ├── List 2: [(id3, v3_full), (id4, v4_full), ...] └── ... └── List k: [(idn, vn_full), ...]查询策略聚类选择计算查询向量与所有聚类中心的距离候选聚类选择距离最近的nprobe个聚类精确搜索在选定的聚类内使用FLAT算法进行精确搜索结果排序对所有候选进行精确距离计算并排序数学表示给定向量集合V{v1,v2,...,vn}V \{v_1, v_2, ..., v_n\}V{v1​,v2​,...,vn​}聚类中心集合C{c1,c2,...,ck}C \{c_1, c_2, ..., c_k\}C{c1​,c2​,...,ck​}。聚类分配cluster(vi)arg⁡min⁡cj∈C∥vi−cj∥ \text{cluster}(v_i) \arg\min_{c_j \in C} \|v_i - c_j\|cluster(vi​)argcj​∈Cmin​∥vi​−cj​∥查询聚类选择selected_clusters(q){cj∈C∣∥q−cj∥ 是最小的nprobe个值之一} \text{selected\_clusters}(q) \{c_j \in C | \|q - c_j\| \text{ 是最小的nprobe个值之一}\}selected_clusters(q){cj​∈C∣∥q−cj​∥是最小的nprobe个值之一}精确距离计算dexact(q,vi)∥q−vi∥,其中 vi∈selected_clusters(q) d_{exact}(q, v_i) \|q - v_i\|, \text{其中 } v_i \in \text{selected\_clusters}(q)dexact​(q,vi​)∥q−vi​∥,其中vi​∈selected_clusters(q)时间复杂度分析构建时间O(n⋅d⋅k⋅I)O(n \cdot d \cdot k \cdot I)O(n⋅d⋅k⋅I)其中k是聚类数I是K-means迭代次数查询时间O(nprobe⋅nk⋅d)O(\text{nprobe} \cdot \frac{n}{k} \cdot d)O(nprobe⋅kn​⋅d)平均情况下空间复杂度O(n⋅dk⋅d)O(n \cdot d k \cdot d)O(n⋅dk⋅d)存储原始向量和聚类中心参数调优关键参数n_clusters聚类数量影响搜索精度和查询速度的平衡典型值[4, 1024]通常选择n\sqrt{n}n​的数量级聚类越多每个聚类越小搜索精度越高nprobe查询时搜索的聚类数量直接影响查询精度和速度典型值[1, 20]通常设置为n_clusters\sqrt{\text{n\_clusters}}n_clusters​nprobe越大精度越高速度越慢metric距离度量欧几里得距离适合稠密向量余弦距离适合需要方向相似性的场景算法实现importnumpyasnpfromsklearn.clusterimportKMeansfromtypingimportList,Tuple,DictimportheapqclassIVFFlatIndex:def__init__(self,n_clusters:int100,nprobe:int10,metric:streuclidean):self.n_clustersn_clusters self.nprobenprobe self.metricmetric self.cluster_centersNoneself.inverted_lists{}# cluster_id - list of (id, vector)self.vector_to_cluster{}# vector_id - cluster_iddeffit(self,vectors:np.ndarray):构建IVFFlat索引# 使用K-means聚类kmeansKMeans(n_clustersself.n_clusters,max_iter20,tol1e-4,random_state42)labelskmeans.fit_predict(vectors)self.cluster_centerskmeans.cluster_centers_ self.inverted_lists{}# 构建倒排列表fori,vectorinenumerate(vectors):cluster_idlabels[i]ifcluster_idnotinself.inverted_lists:self.inverted_lists[cluster_id][]self.inverted_lists[cluster_id].append((i,vector))self.vector_to_cluster[i]cluster_iddef_distance(self,v1:np.ndarray,v2:np.ndarray)-float:计算距离ifself.metriceuclidean:returnnp.linalg.norm(v1-v2)elifself.metriccosine:return1-np.dot(v1,v2)/(np.linalg.norm(v1)*np.linalg.norm(v2))else:raiseValueError(fUnsupported metric:{self.metric})defsearch(self,query:np.ndarray,k:int10)-Tuple[List[int],List[float]]:搜索最近邻ifnotself.inverted_lists:return[],[]# 计算与所有聚类中心的距离cluster_distances[]forcluster_id,centerinenumerate(self.cluster_centers):distanceself._distance(query,center)cluster_distances.append((distance,cluster_id))# 选择最近的nprobe个聚类cluster_distances.sort()selected_clusters[cluster_idfor_,cluster_idincluster_distances[:self.nprobe]]# 在选定的聚类中收集候选candidates[]forcluster_idinselected_clusters:ifcluster_idinself.inverted_lists:forvector_id,vectorinself.inverted_lists[cluster_id]:distanceself._distance(query,vector)candidates.append((distance,vector_id))# 使用堆来高效获取前k个结果iflen(candidates)k:# 使用nlargest获取最大的k个距离最小的k个top_kheapq.nsmallest(k,candidates,keylambdax:x[0])else:top_ksorted(candidates,keylambdax:x[0])[:k]# 提取结果indices[idxfor_,idxintop_k]distances[distfordist,_intop_k]returnindices,distancesdefget_cluster_info(self)-Dict:获取聚类信息info{n_clusters:self.n_clusters,cluster_sizes:{},total_vectors:0}forcluster_id,vectorsinself.inverted_lists.items():info[cluster_sizes][cluster_id]len(vectors)info[total_vectors]len(vectors)returninfo查询优化策略1. 聚类距离预计算classOptimizedIVFFlatIndex(IVFFlatIndex):def__init__(self,n_clusters:int100,nprobe:int10,metric:streuclidean):super().__init__(n_clusters,nprobe,metric)self.cluster_distance_cache{}defsearch(self,query:np.ndarray,k:int10)-Tuple[List[int],List[float]]:优化的搜索方法# 预计算聚类距离cluster_distances[]forcluster_id,centerinenumerate(self.cluster_centers):# 使用缓存cache_key(cluster_id,tuple(query))ifcache_keyinself.cluster_distance_cache:distanceself.cluster_distance_cache[cache_key]else:distanceself._distance(query,center)self.cluster_distance_cache[cache_key]distance cluster_distances.append((distance,cluster_id))# 选择最近的聚类cluster_distances.sort()selected_clusters[cluster_idfor_,cluster_idincluster_distances[:self.nprobe]]# 并行搜索选定的聚类candidates[]forcluster_idinselected_clusters:ifcluster_idinself.inverted_lists:cluster_candidatesself._search_cluster(query,cluster_id)candidates.extend(cluster_candidates)# 使用堆优化结果选择returnself._select_top_k(candidates,k)def_search_cluster(self,query:np.ndarray,cluster_id:int)-List[Tuple[float,int]]:搜索单个聚类candidates[]forvector_id,vectorinself.inverted_lists[cluster_id]:distanceself._distance(query,vector)candidates.append((distance,vector_id))returncandidatesdef_select_top_k(self,candidates:List[Tuple[float,int]],k:int)-Tuple[List[int],List[float]]:选择前k个结果iflen(candidates)k:candidates.sort()indices[idxfor_,idxincandidates]distances[distfordist,_incandidates]returnindices,distances# 使用堆选择前k个top_kheapq.nsmallest(k,candidates,keylambdax:x[0])indices[idxfor_,idxintop_k]distances[distfordist,_intop_k]returnindices,distances2. 动态nprobe调整classAdaptiveIVFFlatIndex(IVFFlatIndex):defsearch(self,query:np.ndarray,k:int10,min_accuracy:float0.95)-Tuple[List[int],List[float]]:自适应nprobe的搜索current_nprobe1best_indices,best_distances[],[]whilecurrent_nprobeself.n_clusters:# 使用当前nprobe进行搜索indices,distancessuper().search(query,k,current_nprobe)# 评估结果质量ifself._evaluate_result_quality(query,indices,distances)min_accuracy:returnindices,distances current_nprobe*2# 如果所有聚类都无法满足精度要求返回最佳结果returnbest_indices,best_distancesdef_evaluate_result_quality(self,query:np.ndarray,indices:List[int],distances:List[float])-float:评估结果质量# 这里可以实现更复杂的质量评估逻辑# 简单实现检查距离是否单调递增iflen(distances)2:return1.0foriinrange(1,len(distances)):ifdistances[i]distances[i-1]:return0.0return1.0优缺点分析优点精确距离计算提供与FLAT算法相同的精确距离查询性能提升通过聚类限制搜索范围内存使用可控相比图结构内存使用更稳定实现简单算法直观易于理解和实现参数调优相对简单主要参数影响相对直观缺点聚类边界问题向量可能位于聚类边界影响搜索精度内存开销大需要完整存储所有原始向量构建时间长K-means聚类需要较长时间动态更新困难插入新向量可能需要重新聚类适用场景高精度要求需要精确距离计算的场景中等规模数据10万到1000万向量的数据集内存充足环境能够承受完整存储原始向量的内存开销查询性能要求高需要在保证精度的同时获得较好的查询性能与其他算法的比较算法查询时间准确性内存使用距离计算精度FLATO(n⋅d)O(n \cdot d)O(n⋅d)100%O(n⋅d)O(n \cdot d)O(n⋅d)精确HNSWO(log⁡n⋅M⋅d)O(\log n \cdot M \cdot d)O(logn⋅M⋅d)95-99%O(n⋅log⁡n⋅M)O(n \cdot \log n \cdot M)O(n⋅logn⋅M)近似IVFO(nprobe⋅nk⋅d)O(\text{nprobe} \cdot \frac{n}{k} \cdot d)O(nprobe⋅kn​⋅d)90-95%O(n⋅dk⋅d)O(n \cdot d k \cdot d)O(n⋅dk⋅d)近似(ADC)IVFFlatO(nprobe⋅nk⋅d)O(\text{nprobe} \cdot \frac{n}{k} \cdot d)O(nprobe⋅kn​⋅d)95-98%O(n⋅dk⋅d)O(n \cdot d k \cdot d)O(n⋅dk⋅d)精确性能优化批量查询支持批量查询优化并行聚类搜索多线程并行搜索不同聚类缓存优化缓存热点聚类和查询结果增量更新支持增量式聚类更新总结IVFFlat算法通过结合IVF的空间划分优势和FLAT的精确计算能力在向量相似性搜索中实现了性能与精度的良好平衡。它特别适合那些需要精确距离计算但又希望获得较好查询性能的场景。尽管在内存使用和构建时间方面存在一些挑战但其在高精度要求下的稳定表现使其成为许多向量数据库系统中的重要索引选择。通过合理的参数调优和查询优化策略IVFFlat能够在各种应用场景中提供优秀的搜索性能。