如何为Sklearn KMeans预设聚类标签与嵌入样本的对应关系？

为什么扩展数据集后标签会变化 #

• KMeans的聚类标签本身是无意义的类别标识符，不具备固定语义。默认情况下，算法每次运行会随机初始化质心（即便用k-means++也存在随机性），导致最终生成的标签编号无法固定。
• 扩展数据集后，数据分布变化，质心的收敛位置也会改变，进一步加剧标签的不确定性。

如何预先定义特定嵌入向量对应的聚类标签 #

直接让KMeans在训练时指定某类对应固定标签做不到，但可以通过两种方式实现类似需求：

1. 基于已知样本的标签映射 #

如果你有已知类别代表样本的嵌入向量，可以先训练KMeans，再将聚类结果与已知样本匹配，重新映射标签：

from sklearn.cluster import KMeans
import numpy as np

# 假设target_embedding是你想要对应标签3的类别的代表向量
target_embedding = np.array([[...]])  # 替换为你的目标类嵌入向量
all_embeddings = np.array([...])  # 所有数据集的嵌入向量

# 设置random_state固定初始质心，减少训练随机性
kmeans = KMeans(n_clusters=5, random_state=42)
pred_labels = kmeans.fit_predict(all_embeddings)

# 找到与target_embedding最近的质心索引
closest_centroid_idx = np.argmin(np.linalg.norm(kmeans.cluster_centers_ - target_embedding, axis=1))

# 构建标签映射：将目标质心对应的标签改为3，其余标签分配未被占用的编号
label_mapping = {closest_centroid_idx: 3}
remaining_labels = [i for i in range(kmeans.n_clusters) if i != 3]
current_idx = 0
for idx in range(kmeans.n_clusters):
    if idx not in label_mapping:
        label_mapping[idx] = remaining_labels[current_idx]
        current_idx += 1

# 应用映射得到固定标签
fixed_labels = [label_mapping[label] for label in pred_labels]

• 核心：random_state能保证相同输入下训练结果稳定，再通过匹配已知样本的质心完成标签映射。

2. 初始化质心强制绑定目标类别 #

如果你能确定目标类的质心位置，可以直接在KMeans初始化时指定该质心，确保对应类别被分配到目标标签（标签顺序与初始化质心列表的索引对应）：

# 假设target_centroid是你想要对应标签3的类别的质心
target_centroid = np.array([[...]])
# 生成其他类的质心（可通过k-means++初始化其他样本得到）
other_embeddings = all_embeddings[...].copy()  # 排除目标类样本的数据集
other_centroids = KMeans(n_clusters=4, random_state=42).fit(other_embeddings).cluster_centers_

# 将目标质心放在索引3的位置，组合成初始化质心列表
initial_centroids = np.vstack([other_centroids[:3], target_centroid, other_centroids[3:]])

# 用自定义质心初始化KMeans，n_init必须设为1避免重新初始化
kmeans = KMeans(n_clusters=5, init=initial_centroids, n_init=1, random_state=42)
fixed_labels = kmeans.fit_predict(all_embeddings)

• 注意：这种方式要求你能准确获取目标类的质心，适合有明确类别代表的场景。

额外建议 #

• 扩数据后务必设置random_state，保证重复训练的结果一致，这是维持标签稳定性的基础。
• 如果需要长期复用聚类结果，建议用joblib.dump保存训练好的KMeans模型，后续直接加载模型预测，避免重复训练带来的标签变化。

内容的提问来源于stack exchange，提问作者Shashank Shandilya