仅有效计算对称外积矩阵的唯一元素答案

【问题标题】：Efficiently compute only the unique elements of a symmetric outer product matrix仅有效计算对称外积矩阵的唯一元素
【发布时间】：2023-04-09 10:17:01
【问题描述】：

我有一大组 N d 维向量（驻留在矩阵中），我通过取自外积（即每个向量 i 乘以自身）来提升它们。对于每个向量，这会产生一个对称矩阵，其中 (d+1) 选择 2 个唯一条目。对于整个数据，这是一个 N x d x d 张量。我只想计算唯一的 (d+1) 从每个张量切片的下对角线中选择 2 个条目并将它们存储在一个向量中。我想在 Python 中以尽可能少的内存占用和尽可能快的速度做到这一点——包括使用 C 绑定。

如果您使用标准 numpy 方法执行此操作，它会分配每个矩阵的整体。这大约是实际所需内存复杂度的两倍。

对于这里的规模感，考虑 N = 20k 和 d = 20k 的情况。然后 N * d^2 * ~8bytes per element = (2*10^4)^3 * 8 bytes = 64 TB。

如果我们只计算编码唯一条目的向量，我们有 (20001 choose 2) * 20k * 8 = 200010000 * 20000 * 8 bytes = 32 TB。

有没有一种快速的方法来做到这一点，而无需求助于缓慢的方法（例如在 python 中编写我自己的外部产品）？

编辑：我会注意到Create array of outer products in numpy中提出了类似的问题

我已经知道如何使用 einsum 进行计算（如上述问题）。但是，如果没有额外的（d 选择 2）计算和分配，就没有得到任何答案

编辑 2：这个线程How to exploit symmetry in outer product in Numpy (or other Python solutions)? 提出了一个相关的问题，但没有解决内存复杂性。最佳答案仍然会为每个外部产品分配一个 d x d 数组。

这个线程Numpy Performance - Outer Product of a vector with its transpose 还解决了自外积的计算考虑，但没有达到内存高效的解决方案。

编辑 3：如果想分配整个数组然后提取元素，np.tril_indices 或scipy.spatial.distance.squareform 可以解决问题。

【问题讨论】：

scipy.spatial.distance 可以使用condensed 和squareform 距离数组。 squareform 是二维冗余版本。 condensed 将上下三角形值映射到一维数组。
您是否建议考虑外部产品以便将其传递给scipy.spatial.distance.pdist？这会产生预期的结果，但我还不清楚如何进行分解
您是否希望将唯一的元素作为一维数组/向量或在 N x d x d 张量中获取，其中上层 diag 元素与下层元素重复？

标签： python numpy matrix-multiplication

【解决方案1】：

不确定您希望输出的确切方式，但始终可以选择使用 Numba：

import numpy as np
import numba as nb

# Computes unique pairwise products
@nb.njit(parallel=True)
def self_outer_unique(a):
    n, d = a.shape
    out = np.empty((n, (d * d + d) // 2), dtype=a.dtype)
    for i in nb.prange(n):
        for j1 in range(d):
            for j2 in range(j1, d):
                idx = j1 * (2 * d - j1 + 1) // 2 + j2 - j1
                out[i, idx] = a[i, j1] * a[i, j2]
    return out

这将为您提供一个数组，其中包含每行上的所有唯一产品（即完整输出的扁平上三角形）。

import numpy as np

a = np.arange(12).reshape(4, 3)
print(a)
# [[ 0  1  2]
#  [ 3  4  5]
#  [ 6  7  8]
#  [ 9 10 11]]
print(self_outer_unique(a))
# [[  0   0   0   1   2   4]
#  [  9  12  15  16  20  25]
#  [ 36  42  48  49  56  64]
#  [ 81  90  99 100 110 121]]

在性能方面，它比使用 NumPy 计算完整的外积要快，尽管从中重新创建完整的数组需要更长的时间。

import numpy as np

def np_self_outer(a):
    return a[:, :, np.newaxis] * a[:, np.newaxis, :]

def my_self_outer(a):
    b = self_outer_unique(a)
    n, d = a.shape
    b_full = np.zeros((n, d, d), dtype=a.dtype)
    idx0 = np.arange(n)[:, np.newaxis]
    idx1, idx2 = np.triu_indices(d)
    b_full[idx0, idx1, idx2] = b
    b_full += np.triu(b_full, 1).transpose(0, 2, 1)
    return b_full

n, d = 1000, 100
a = np.arange(n * d).reshape(n, d)
print(np.all(np_self_outer(a) == my_self_outer(a)))
# True

%timeit np_self_outer(a)
# 24.6 ms ± 248 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit self_outer_unique(a)
# 6.32 ms ± 69.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit my_self_outer(a)
# 124 ms ± 770 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

【讨论】：