大规模对称矩阵G计算的时间与内存优化方案问询

我现在需要计算一个元素为标量的矩阵G，它的元素计算规则是：
$$G_{ij} = \sum_{k=1}^n \sum_{l=1}^n M_k M_l \cdot k(R_k - R_l, \gamma(Z_i - Z_j))$$
其中$k$为Matern核函数。

我要处理的场景是n>10000，d>30，目前写的代码开销极大，运行耗时非常久。想请教大家，在不使用GPU的前提下，怎么才能最快完成这个计算，同时尽量降低内存占用？

当前实现代码 #

import numpy as np
from sklearn.gaussian_process.kernels import Matern
from tqdm import tqdm
from joblib import Parallel, delayed

# Pre-flattened computation to minimize data transfer overhead
def precompute_differences(R, Z):
    n, d        = R.shape
    R_diff_flat = (R[:, None, :] - R[None, :, :]).reshape(n * n, d)
    Z_diff      = Z[:, None, :] - Z[None, :, :]
    return R_diff_flat, Z_diff

def compute_G_row(i, R_diff_flat, Z_diff, W, gamma_val, kernel, n, d):
    """
    Compute the i-th row for j >= i and store them in a temporary array.
    """
    row_values = np.zeros(n)
    for j in range(i, n):
        Z_ij   = gamma_val * Z_diff[i, j].reshape(1, d)
        K_flat = kernel(R_diff_flat, Z_ij)
        K_ij   = K_flat.reshape(n, n)
        row_values[j] = np.sum(W * K_ij)
    return i, row_values

def compute_G(M, gamma, R, Z, nu=1.5, length_scale=1.0, use_parallel=True):
    """
    Compute the G matrix with fewer kernel evaluations by exploiting symmetry:
    G[i,j] = G[j,i]. We only compute for j >= i, then mirror the result.
    """
    R = np.asarray(R)
    Z = np.asarray(Z)
    M = np.asarray(M).reshape(-1, 1)  # ensure (n,1)
    n, d = R.shape

    # Precompute data
    R_diff_flat, Z_diff = precompute_differences(R, Z)
    W = M @ M.T  # Weight matrix

    G = np.zeros((n, n))

    kernel = Matern(length_scale=length_scale, nu=nu)

    if use_parallel and n > 1:
        # Parallel computation
        results = Parallel(n_jobs=-1)(
            delayed(compute_G_row)(i, R_diff_flat, Z_diff, W, gamma, kernel, n, d)
            for i in tqdm(range(n), desc="Computing G matrix")
        )
    else:
        # Single-threaded computation
        results = []
        for i in tqdm(range(n), desc="Computing G matrix"):
            row_values = np.zeros(n)
            for j in range(i, n):
                Z_ij   = gamma * Z_diff[i, j].reshape(1, d)
                K_flat = kernel(R_diff_flat, Z_ij)
                K_ij   = K_flat.reshape(n, n)
                row_values[j] = np.sum(W * K_ij)
            results.append((i, row_values))

    # Sort and fill final G by symmetry
    results.sort(key=lambda x: x[0])
    for i, row_vals in results:
        for j in range(i, n):
            G[i, j] = row_vals[j]
            G[j, i] = row_vals[j]  # mirror for symmetry

    # Delete auxiliary variables to save memory
    del R_diff_flat, Z_diff, W, kernel, results

    # Optional checks
    is_symmetric = np.allclose(G, G.T, atol=1e-8)
    eigenvalues = np.linalg.eigvalsh(G)
    is_semi_positive_definite = np.all(eigenvalues >= -1e-8)
    print(f"G is semi-positive definite: {is_semi_positive_definite}")
    print(f"G is symmetric: {is_symmetric}")

    # Delete all local auxiliary variables except G to save memory
    local_vars = list(locals().keys())
    for var_name in local_vars:
        if var_name not in ["G"]:
            del locals()[var_name]

    return G

测试示例 #

# Example usage:
if __name__ == "__main__":
    __spec__ = None
    n = 20
    d = 10
    gamma = 0.9
    R = np.random.rand(n, d)
    Z = np.random.rand(n, d)
    M = np.random.rand(n, 1)

    G = compute_G(M, gamma, R, Z, nu=1.5, length_scale=1.0, use_parallel=True)
    print("G computed with shape:", G.shape)

备注：内容来源于stack exchange，提问作者domath