BERTopic处理80万波兰语推文时UMAP余弦度量崩溃问题

临时稳定方案 #

直接改用euclidean度量替代余弦度量，该方案已验证可稳定处理80万+数据量，集成到BERTopic的代码如下：

from umap import UMAP
from sentence_transformers import SentenceTransformer
from bertopic import BERTopic

# 配置使用欧氏距离的UMAP模型
umap_model = UMAP(n_neighbors=15, n_components=5, metric='euclidean')
embedding_model = SentenceTransformer("sdadas/st-polish-paraphrase-from-distilroberta")
topic_model = BERTopic(embedding_model=embedding_model, umap_model=umap_model)
topics, probs = topic_model.fit_transform(docs)

余弦度量的修复尝试方向 #

若必须使用余弦度量，可尝试以下方法：

• 降低邻居数量：将n_neighbors从15下调至5-10，减少余弦距离计算的复杂度，降低内存访问异常概率
• 增量式UMAP拟合：将嵌入向量分批传入UMAP进行拟合（需UMAP版本≥0.5.0），示例代码：

import numpy as np
from umap import UMAP

# 拆分嵌入为批次
def split_batch(arr, batch_size):
    for i in range(0, arr.shape[0], batch_size):
        yield arr[i:i+batch_size]

umap_model = UMAP(n_neighbors=15, n_components=5, metric='cosine')
# 分批拟合
for batch in split_batch(embeddings, batch_size=100000):
    umap_model.fit(batch)
# 转换全量嵌入
reduced_embeddings = umap_model.transform(embeddings)

• 源码编译UMAP并启用OpenMP：卸载现有UMAP后从源码编译安装，开启OpenMP优化可修复部分内存访问错误：

  pip uninstall -y umap-learn
  pip install git+https://github.com/lmcinnes/umap.git --no-binary umap-learn
  

  注意：Linux环境需提前安装libgomp1依赖，Windows需确保编译工具链支持OpenMP。

问题说明 #

该段错误源于UMAP在大样本量下计算余弦距离时的底层内存访问异常，和可用内存大小无直接关联（大内存服务器同样触发），非确定性表现是因为UMAP的随机初始化特性。

内容的提问来源于stack exchange，提问作者Julek Sienkiewicz