你的RAG检索太“笨”？用K-Means聚类来“调教”

今天我们要聊的话题，可能会让不少RAG开发者产生共鸣：为什么辛苦搭建的检索系统，返回的结果总是差强人意，要么内容高度重复，要么无法精准命中用户的意图？其实，问题的症结往往不在于模型本身，而在于检索策略。

如果你也觉得自己的RAG系统有点“傻气”，那么今天这篇文章可能会帮你找到一个全新的突破口——用K-Means聚类“调教”它。

我们会从K-Means算法的底层原理讲起，再结合目前表现优秀的BGE-M3嵌入模型，一步步教你如何将文本向量化、聚类分析，最后落地到RAG系统的两个实战场景：一是通过聚类提升检索结果的多样性，二是实现智能化的查询路由，让系统知道该去哪个“知识库”里翻答案。

文章里所有代码都已经调通，可以直接跑。配套的代码仓库也已开放，方便大家动手实践。

话不多说，我们开始吧。

前言

在自然语言处理（NLP）领域，检索增强生成（RAG）技术近年来迅速崛起，成为解决知识密集型任务的重要工具。通过将外部知识检索与生成模型结合，RAG能在回答复杂问题时提供更准确、更有上下文相关性的内容。然而，很多开发者在实际使用中都会遇到一个共同的痛点：检索结果过于单一，不够“聪明”，导致生成内容缺乏多样性，甚至没能覆盖用户真正想要的答案。

问题出在哪里？传统检索方法大多基于简单的相似度排序——余弦相似度或欧几里得距离。这种方式的优点是简单高效，但它有一个天然的短板：倾向于返回语义高度相似的文档。在某些场景下，这就造成了“信息冗余”，完全没能捕捉数据内部的潜在语义结构。

而K-Means聚类算法，作为一种经典的无监督学习方法，刚好可以帮你“调教”这个过程。通过对文本数据进行分组和语义挖掘，它能显著提升检索的多样性和智能性。

接下来，我们一步步深入：从K-Means的基本原理，到如何用它做文本向量聚类，最后是RAG系统中的两个实战案例。不管你是NLP开发者还是RAG爱好者，这篇文章都会提供实用的技术洞见和可运行的代码。

一、K-Means聚类的基本原理

有监督学习 vs 无监督学习

在深入算法之前，先来理清一个基础概念：有监督学习和无监督学习的区别。

有监督学习：算法使用带有标签（正确答案）的数据训练，学习输入和输出之间的映射关系，目标是预测新数据的标签或值。典型算法包括分类（如决策树、支持向量机）和回归（如线性回归），应用场景包括垃圾邮件检测、图像识别、股票价格预测。

无监督学习：算法处理的数据没有标签，模型需要通过识别数据中的内在结构和模式来学习，目标是发现隐藏的规律或分组。典型算法包括聚类（如K-Means、DBSCAN）和降维（如PCA），应用包括客户细分、异常检测、主题发现。

K-Means属于无监督学习，目标是把数据集划分为K个簇，使得每个数据点与最近的簇中心距离最小。

K-Means算法步骤详解

数学上，K-Means的优化目标可以表示为最小化所有数据点到其所属簇中心的距离平方和（即簇内平方和，WCSS）：

其中，K是预设的簇数量，C_j是第j个簇中的数据点集合，x是数据点，μ_j是第j个簇的中心向量。

算法通过以下步骤逐步逼近局部最优解。

1️⃣ 初始化：选择初始簇中心

初始中心的选择对最终聚类结果和收敛速度影响很大。常用的有两种方式：

• 随机选择法：从数据集中随机挑选K个点作为初始簇中心。
• K-Means++ 初始化：通过更“聪明”的选择策略，让初始中心分布更均匀，有效提高算法找到全局最优解的可能性。K-Means++的步骤是：随机选第一个中心；然后计算每个数据点到最近已选中心的距离，距离越远的点被选中的概率越大；重复直到选出K个中心，K-Means++也是scikit-learn的默认初始化方式。

2️⃣ 分配：将数据点分配到最近的簇

每个数据点被分配到距离它最近的簇中心所在的簇。距离通常使用欧几里得距离平方计算。

实现代码中，利用广播机制高效计算所有数据点到所有中心的距离平方矩阵，然后为每个点分配最近的簇编号。

3️⃣ 更新：重新计算簇中心

新的簇中心是该簇内所有数据点的均值，即求每个簇的“质心”。

4️⃣ 迭代：重复直到收敛

不断重复分配和更新步骤，直到满足以下任一条件：簇中心变化小于某个阈值、数据点簇分配不再变化、达到最大迭代次数。

5️⃣ 时间复杂度

单次迭代复杂度为O(n·K·d)，其中n是数据点数量，K是簇数量，d是数据维度；总复杂度为O(t·n·K·d)，t是迭代次数，通常t较小。

K-Means算法参数详解

在scikit-learn的KMeans类中，关键参数包括：n_clusters（簇数量）、init（初始化方法，默认'k-means++'）、n_init（运行次数，默认10，取WCSS最小的结果）、max_iter（最大迭代次数，默认300）、tol（收敛阈值，默认1e-4）、random_state（随机种子，确保可重复性）等。

K-Means的数学目标

K-Means的核心是优化目标函数WCSS，算法通过交替执行分配步骤（固定中心，优化分配）和更新步骤（固定分配，优化中心），逐步减小WCSS，直到收敛。多次运行并选择最优结果可以缓解局部最优问题。

二、文本向量生成与聚类实践

在开始实践之前，先安装好必要的依赖库：

pip install numpy scikit-learn pandas matplotlib FlagEmbedding

2.1 使用BGE-M3模型生成文本向量

BGE-M3是一个强大的嵌入模型，能够为文本生成密集向量。假设我们有两个句子："What is BGE M3?" 和 "Defination of BM25"，在本地部署后，可以用以下代码提取向量：

from FlagEmbedding import BGEM3FlagModel
import numpy as np

model_path = r"C:UserskDesktopBaiduSyncdiskbaidu_sync_documentshf_modelsbge-m3"
model = BGEM3FlagModel(model_path, use_fp16=True)

sentences = ["What is BGE M3?", "Defination of BM25"]
output = model.encode(
    sentences=sentences,
    batch_size=12,
    max_length=8192,
    return_dense=True,
    return_sparse=False,
    return_colbert_vecs=False
)
dense_vecs = output['dense_vecs']
print(dense_vecs.shape)  # 输出: (2, 1024)

输出显示dense_vecs的形状为(2, 1024)，表示每个句子被编码为一个1024维的向量。这些向量捕捉了句子的语义信息，为聚类提供了高质量的输入。

2.2 将K-Means应用于句子向量

有了向量表示，就可以直接输入K-Means算法了。假设希望将两个句子分为两类：

from sklearn.cluster import KMeans
kmeans = KMeans(n_clusters=2, random_state=42)
clusters = kmeans.fit_predict(dense_vecs)
print(f"聚类结果: {clusters}")

运行后，结果可能为[0, 1]，表示第一个句子属于簇0，第二个句子属于簇1。簇中心可以通过kmeans.cluster_centers_获取，形状为(2, 1024)。这些簇中心可以被视为每个簇的代表性向量，帮助我们理解每个簇的语义主题。

2.3 余弦相似度与K-Means聚类

scikit-learn的K-Means默认使用欧几里得距离，但在处理文本向量时，余弦相似度往往更合适。通过向量归一化，可以让标准K-Means等效于使用余弦相似度：

from sklearn.preprocessing import normalize
normalized_vecs = normalize(dense_vecs, norm='l2')
cos_kmeans = KMeans(n_clusters=2, random_state=42)
cos_clusters = cos_kmeans.fit_predict(normalized_vecs)
print(f"基于余弦相似度的聚类结果: {cos_clusters}")

原理在于：当向量归一化后，它们都位于单位超球面上，此时欧几里得距离的平方与余弦相似度之间存在数学关系：||x - y||² = 2(1 - cos(x,y))。因此，在归一化向量上最小化欧几里得距离，等同于最大化余弦相似度。这在文本聚类中往往更符合需求，因为余弦相似度更能捕捉文本的主题信息，而不受向量绝对大小的影响。

2.4 完整案例：新闻标题聚类

下面通过一个完整的案例，展示如何将BGE-M3和K-Means结合用于实际应用。我们对20条不同主题的新闻标题进行聚类分析，看看算法能否自动识别出潜在的主题分组。

from FlagEmbedding import BGEM3FlagModel
import numpy as np
from sklearn.cluster import KMeans
from sklearn.preprocessing import normalize
from sklearn.metrics import silhouette_score
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

# 加载模型
model_path = r"C:UserskDesktopBaiduSyncdiskbaidu_sync_documentshf_modelsbge-m3"
model = BGEM3FlagModel(model_path, use_fp16=True)

# 样本新闻标题集合（涵盖科技、体育、整治、娱乐和健康五个主题）
news_titles = [
    # 科技主题
    "苹果发布最新iPhone 15系列，搭载A17芯片",
    "谷歌推出新一代人工智能助手，支持自然语言理解",
    "特斯拉自动驾驶技术获突破，事故率降低30%",
    "微软宣布收购AI初创公司，强化云服务能力",
    # 体育主题
    "梅西在巴黎首秀进球，球迷欢呼雀跃",
    "东京奥运会闭幕，中国队金牌榜位列第二",
    "NBA季后赛：湖人击败热火，夺得总冠军",
    "国足世预赛不敌日本队，出线形势严峻",
    # 整治主题
    "中美元首通话，就双边关系交换意见",
    "欧盟通过新气候法案，承诺2050年实现碳中和",
    "联合国大会召开，各国领导人讨论全球治理",
    "英国宣布脱欧后新贸易政策，加强与亚洲合作",
    # 娱乐主题
    "新电影《沙丘》全球热映，票房突破4亿美元",
    "流行歌手泰勒·斯威夫特发布新专辑，粉丝热情高涨",
    "网飞热门剧集《鱿鱼游戏》创收视纪录",
    "奥斯卡颁奖典礼举行，《无依之地》获最佳影片",
    # 健康主题
    "新研究发现常规锻炼可降低阿尔茨海默病风险",
    "全球新冠疫苗接种突破30亿剂，发展中国家覆盖率仍低",
    "医学专家建议减少超加工食品摄入，降低慢性病风险",
    "心理健康问题在年轻人中上升，专家呼吁加强关注"
]

# 生成文本向量
news_vectors = model.encode[news_titles]('dense_vecs')
# 对向量进行归一化，准备基于余弦相似度的聚类
normalized_vectors = normalize(news_vectors, norm='l2')
# 使用归一化向量进行K-means聚类
kmeans = KMeans(n_clusters=5, init='k-means++', n_init=10, random_state=42)
clusters = kmeans.fit_predict(normalized_vectors)

# 创建结果DataFrame
results_df = pd.DataFrame({'title': news_titles, 'cluster': clusters})
# 打印每个簇的新闻标题
for cluster_id in range(5):
    print(f"\n=== 簇 {cluster_id} ===")
    cluster_titles = results_df[results_df['cluster'] == cluster_id]['title']
    for title in cluster_titles:
        print(f"- {title}")
# 计算聚类评估指标
silhouette_a vg = silhouette_score(normalized_vectors, clusters)
print(f"\n聚类轮廓系数: {silhouette_a vg:.4f}")

输出结果展示了聚类情况：算法成功地将新闻标题按主题分组，识别出了科技、体育、整治、娱乐和健康五个不同主题。虽然轮廓系数不算高，但考虑到新闻标题的简洁性和主题之间的潜在关联，实际聚类效果还是不错的。

结果分析与应用

这个例子展示了如何将文本数据转换为向量并应用K-Means聚类。尽管没有提供任何标签信息，算法成功地将新闻标题按主题分组。这种聚类方法在实际中有多种用途：

• 内容推荐系统：根据用户阅读的内容所属簇，推荐同一簇内的其他文章。
• 自动归档与组织：为大量文档自动生成主题分类，方便管理。
• 热点话题发现：通过分析哪些簇包含更多近期文章，识别热门话题。
• 舆情分析：分析不同主题簇的情感倾向和数量变化。
• 文本摘要：从每个簇中选择最具代表性的文本作为该簇的摘要。
• 构建知识图谱：将每个簇视为一个概念，簇内文本作为该概念的实例。

三、RAG实战：K-Means的妙用

3.1 多样性增强的检索方法

多样性增强检索的核心思路是：通过K-Means聚类，将候选文档分为几个语义簇，然后从每个簇中选出最具代表性的文档。这样可以避免传统相似度排序导致的信息冗余，确保RAG系统获得不同角度的信息。

from FlagEmbedding import BGEM3FlagModel
import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity

# 加载模型
model_path = r"C:UserskDesktopBaiduSyncdiskbaidu_sync_documentshf_modelsbge-m3"
model = BGEM3FlagModel(model_path, use_fp16=True)

def diversity_enhanced_retrieval(query, doc_vectors, doc_texts, top_k=5, diversity_clusters=3):
    # 编码查询并计算相似度
    query_vec = model.encode[[query]]['dense_vecs'](0)
    similarities = cosine_similarity[[query_vec], doc_vectors](0)
    # 选择相似度最高的top_n个文档作为候选集
    top_n = min(top_k * 3, len(doc_vectors))
    candidate_indices = np.argsort[similarities][-top_n:](::-1)
    candidate_vectors = doc_vectors[candidate_indices]
    # 执行K-means聚类
    kmeans = KMeans(n_clusters=min(diversity_clusters, len(candidate_vectors)),
                    random_state=42, n_init=10)
    clusters = kmeans.fit_predict(candidate_vectors)
    # 从每个簇中选最佳文档
    selected_indices = []
    cluster_dict = {}
    for idx, cluster_id in enumerate(clusters):
        cluster_dict.setdefault(cluster_id, []).append((candidate_indices[idx], similarities[candidate_indices[idx]]))
    for cluster_id in range(min(diversity_clusters, len(cluster_dict))):
        if cluster_dict.get(cluster_id):
            best_doc = max[cluster_dict[cluster_id], key=lambda x: x[1]](0)
            selected_indices.append(best_doc)
    # 补充不足的文档
    remaining = [i for i in candidate_indices if i not in selected_indices]
    if len(selected_indices) < top_k and remaining:
        remaining_similarities = [similarities[i] for i in remaining]
        extra_indices = [remaining[i] for i in np.argsort[remaining_similarities](-top_k + len(selected_indices):)]
        selected_indices.extend(extra_indices)
    return selected_indices[:top_k]

测试结果表明，这种多样性检索方法能确保RAG系统获取不同角度的信息，避免回答过于单一，大大提高了生成内容的全面性。

3.2 聚类感知的查询路由系统

在包含多种专业领域文档的大型RAG系统中，基于聚类的路由可以智能选择最相关的知识库。其核心思想是：对每个知识库中的文档进行聚类，提取代表性语义簇；当用户提出查询时，将查询向量与各知识库的簇中心进行比较，选择最相关的知识库，从而提高检索效率和准确性。

测试结果显示，查询“感冒的症状有哪些？”被正确路由到医学知识库，而“AI开发最好的编程语言是什么？”则被路由到技术知识库。这种方法尤其适用于拥有多个垂直领域知识库的场景，比如一个包含医学、法律和技术文档的问答系统。

当然，这里只是提供一些启发。其实K-Means在RAG场景下还有很多其他应用，比如在多轮对话中利用聚类方式载入历史对话，效果往往不错。大家可以自由探索，相信会有不少收获。

总结与展望

技术全景图

• 无监督学习：K-Means聚类算法提供了一种无需标签即可对数据进行分组的有效方法。
• 文本嵌入模型：BGE-M3模型能生成高质量的文本向量，为聚类等下游任务提供强大的特征表示。
• 检索增强生成（RAG）：通过优化检索过程，聚类技术可以显著提升RAG系统的性能和生成内容的多样性。

学习汇总

• 掌握了K-Means聚类算法的基本原理和实现步骤。
• 学会了使用BGE-M3模型生成文本的密集向量表示。
• 了解了如何在NLP任务中应用K-Means进行文本聚类。
• 探索了如何利用聚类技术增强RAG系统的检索效果，包括提升检索结果的多样性和实现智能查询路由。

动手挑战

• 尝试修改新闻标题聚类案例中的n_clusters参数，观察聚类结果的变化，并分析不同K值对结果的影响。
• 基于提供的多样性增强检索代码，尝试使用不同的diversity_clusters值，分析检索结果的多样性变化。
• 扩展聚类感知的查询路由系统，添加更多的知识库，并设计更复杂的查询场景进行测试。