如何为核k-means算法中的RBF核函数实现矩阵运算加速

你实现RBF核与核距离用到的公式如下：

现有代码存在的问题 #

• 语法错误：kernel函数中res = np.exp(-*up/gamma)存在多余的*符号，无法正常运行
• 变量未定义：get_gamma函数中的length变量没有声明，按你的计算逻辑应该是X、Y样本数量的乘积，缺失会直接报错
• 执行效率极低：双层for循环的时间复杂度为O(nmd)（n为X样本数、m为Y样本数、d为特征维度），处理784维的图像类数据时，样本量稍大就会出现严重的性能瓶颈
• 功能限制：当前kernel仅支持单个样本对的核值计算，无法批量输出核矩阵，也不适配核k-means需要批量计算样本与聚类中心距离的场景
• 数值稳定性隐患：没有对gamma做最小值截断，当样本距离普遍为0时会出现除零错误

矩阵运算优化方案 #

核心利用numpy的广播机制与矩阵运算替代循环，使用欧氏距离平方的矩阵化计算公式：
||X-Y||² = 行方向sum(X²) + 列方向sum(Y²) - 2 * X @ Y.T
可以一次性计算所有样本对的距离平方，时间复杂度降低到numpy底层优化的矩阵运算级别，784维数据下性能可以提升数百倍。

优化后完整代码 #

import numpy as np

def get_gamma(X, Y):
    # X shape: (n_samples, n_features), Y shape: (m_samples, n_features)
    n = X.shape[0]
    m = Y.shape[0]
    # 批量计算所有样本对的距离平方和
    sum_x2 = np.sum(np.square(X), axis=1).reshape(-1, 1)  # shape (n, 1)
    sum_y2 = np.sum(np.square(Y), axis=1).reshape(1, -1)  # shape (1, m)
    dist_sq = sum_x2 + sum_y2 - 2 * X @ Y.T
    # 替代原来的双层循环求和，gamma取所有距离平方的平均值
    gamma = np.mean(dist_sq)
    # 加极小值避免除零
    return max(gamma, 1e-10)

def kernel(X, Y, gamma):
    # 批量输出核矩阵，shape (n_samples, m_samples)
    sum_x2 = np.sum(np.square(X), axis=1).reshape(-1, 1)
    sum_y2 = np.sum(np.square(Y), axis=1).reshape(1, -1)
    dist_sq = sum_x2 + sum_y2 - 2 * X @ Y.T
    return np.exp(-dist_sq / gamma)

def kernel_distance(X, Y):
    gamma = get_gamma(X, Y)
    a = kernel(X, X, gamma).diagonal().reshape(-1, 1)  # 取K(x,x)的值
    b = kernel(Y, Y, gamma).diagonal().reshape(1, -1)  # 取K(y,y)的值
    c = kernel(X, Y, gamma)
    # 批量输出所有两两样本的核距离矩阵
    return np.sqrt(a + b - 2 * c)

优化说明 #

• 所有循环逻辑都替换为矩阵运算，完全去掉了Python层面的for循环，适配高维大数据量场景
• 支持批量输入，kernel函数直接输出(n, m)维度的核矩阵，kernel_distance直接输出所有样本对的距离矩阵，完全适配核k-means的批量计算需求
• 新增了gamma的极小值截断，避免除零错误
• 如果仅需要计算单个样本对的距离，也可以直接传入形状为(1, 784)的样本数组，不需要修改代码逻辑

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