加速UMAP：预计算KNN后度量不符问题排查与优化建议

问题背景 #

特征数达10M但观测值仅数百个，测试后发现cosine和correlation度量比euclidean效果更好，但UMAP运行时多数时间仅用单核，仅末尾几秒使用多核，推测是度量计算环节未并行化。尝试通过预计算KNN利用NumPy多核能力，但修改后的代码无报错，输出结果却基于默认的euclidean度量，而非预期的cosine。

原始代码 #

import umap
from sklearn.preprocessing import StandardScaler

metric = 'cosine' # alternatively 'correlation'
scaler = StandardScaler()
ip_std = scaler.fit_transform(ip_mat)

# Start UMAP
reducer = umap.UMAP(n_components=n_components, n_neighbors=n_neighbors, metric=metric)
umap_embed = reducer.fit_transform(ip_std)

修改后代码 #

import umap
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import pairwise_distances
from scipy.spatial.distance import cosine
from umap.umap_ import nearest_neighbors

# Start precomputed_knn

scaler = StandardScaler()
eiip_std = scaler.fit_transform(eiip_mat)

dist_cosine = 1 - pairwise_distances(eiip_std, metric="cosine")
precomputed_knn = nearest_neighbors(dist_cosine, metric="cosine", \
                                      metric_kwds=None, angular=False, \
                                      n_neighbors=n_neighbors, random_state=42)
# Start UMAP
reducer = umap.UMAP(n_components=n_components, precomputed_knn=precomputed_knn)
umap_embed = reducer.fit_transform(eiip_std)

return umap_embed

代码错误分析 #

1. 距离/相似度矩阵混淆：UMAP的nearest_neighbors需要输入距离矩阵，但你计算的dist_cosine = 1 - pairwise_distances(..., metric="cosine")是相似度矩阵，直接传入会导致KNN的近邻逻辑完全错误。
2. 未指定UMAP全局度量：即便传入预计算KNN，UMAP后续的布局优化环节默认仍使用euclidean度量，必须显式指定metric参数匹配预计算逻辑。
3. nearest_neighbors的metric参数错误：传入预计算矩阵时，必须设置metric="precomputed"，否则函数会忽略你传入的矩阵，重新计算cosine距离覆盖结果。

修正后的代码 #

import umap
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import pairwise_distances
from umap.umap_ import nearest_neighbors

scaler = StandardScaler()
eiip_std = scaler.fit_transform(eiip_mat)

# 直接计算cosine距离矩阵，无需转成相似度
dist_cosine = pairwise_distances(eiip_std, metric="cosine")
# 传入预计算矩阵时，metric必须设为"precomputed"
precomputed_knn = nearest_neighbors(
    dist_cosine,
    metric="precomputed",
    n_neighbors=n_neighbors,
    random_state=42
)

# 显式指定metric为"cosine"，确保全流程匹配预计算逻辑
reducer = umap.UMAP(
    n_components=n_components,
    precomputed_knn=precomputed_knn,
    metric="cosine"
)
umap_embed = reducer.fit_transform(eiip_std)

return umap_embed

优化建议 #

• 激活多核计算：sklearn.metrics.pairwise_distances默认依赖OpenBLAS/MKL的多核优化，可通过设置环境变量export OMP_NUM_THREADS=8（根据CPU核心数调整）最大化利用多核资源。
• correlation度量适配：如果使用correlation度量，直接用pairwise_distances(eiip_std, metric="correlation")即可，该函数返回的1 - 皮尔逊相关系数刚好符合UMAP的距离输入要求。
• 减少冗余内存占用：预计算KNN后，最新版UMAP允许fit_transform传入None替代原始数据，可尝试reducer.fit_transform(None)降低内存负载。
• 验证KNN正确性：通过precomputed_knn[0]查看近邻索引、precomputed_knn[1]查看距离值，手动核对样本近邻是否符合cosine距离的预期结果。

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