如果没有中间 3D 矩阵，是否可以实现或近似计算此矩阵？

Question

给定一个 NxN 矩阵 W，我想计算一个 NxN 矩阵 C，该矩阵由 link 中的方程式给出：https://i.stack.imgur.com/dY7rY.png，或在 LaTeX

中

$$C_{ij} = \max_k \bigg\{ \sum_l \bigg( W_{ik}W_{kl}W_{lj} - W_{ik}W_{kj} \bigg)  \bigg\}.$$

我曾尝试在 PyTorch 中实现这一点，但我要么通过构建中间 NxNxN 3D 矩阵遇到内存问题，对于大 N，导致我的 GPU 运行内存不足，或者使用了for-loop over k 然后非常慢。我不知道如何解决这些问题。如果没有像这样的大型中间矩阵，我如何实现此计算或它的近似值？

任何语言的建议、伪代码或任何 Python/Numpy/PyTorch 的实现都将不胜感激。

Answer 1

公式可以简化为

C_ij = max_k ( W_ik M_kj)

哪里

M = W * W - N * W

与 N 矩阵的大小 W 和 W * W 通常的乘积。

然后，在上面的公式中，对于每个i、j，都有一个独立的最大值要计算。在不知道 W 的更多属性的情况下，通常不可能进一步简化问题。因此，在计算完矩阵 M 之后，您可以对 i 和 j 进行循环，然后计算最大值。

Answer 2

第一个使用 Numba 的解决方案（您可以使用 Cython 或纯 C 来做同样的事情）是使用简单的循环来表述问题。

import numpy as np
import numba as nb

@nb.njit(fastmath=True,parallel=True)
def calc_1(W):
    C=np.empty_like(W)
    N=W.shape[0]

    for i in nb.prange(N):
        TMP=np.empty(N,dtype=W.dtype)
        for j in range(N):
            for k in range(N):
                acc=0
                for l in range(N):
                    acc+=W[i,k]*W[k,l]*W[l,j]-W[i,k]*W[k,j]
                TMP[k]=acc
            C[i,j]=np.max(TMP)
    return C

Francesco 提供了一种简化方法，可以更好地适应更大的数组大小。这导致了以下内容，我还优化了一个小的临时数组。

@nb.njit(fastmath=True,parallel=True)
def calc_2(W):
    C=np.empty_like(W)
    N=W.shape[0]
    M = np.dot(W,W) - N * W

    for i in nb.prange(N):
        for j in range(N):
            val=W[i,0]*M[0,j]
            for k in range(1,N):
                TMP=W[i,k]*M[k,j]
                if TMP>val:
                    val=TMP
            C[i,j]=val
    return C

这可以通过部分循环展开和优化数组访问来进一步优化。一些编译器可能会自动执行此操作。

@nb.njit(fastmath=True,parallel=True)
def calc_3(W):
    C=np.empty_like(W)
    N=W.shape[0]
    W=np.ascontiguousarray(W)
    M = np.dot(W.T,W.T) - W.shape[0] * W.T

    for i in nb.prange(N//4):
        for j in range(N):
            val_1=W[i*4+0,0]*M[j,0]
            val_2=W[i*4+1,0]*M[j,0]
            val_3=W[i*4+2,0]*M[j,0]
            val_4=W[i*4+3,0]*M[j,0]
            for k in range(1,N):
                TMP_1=W[i*4+0,k]*M[j,k]
                TMP_2=W[i*4+1,k]*M[j,k]
                TMP_3=W[i*4+2,k]*M[j,k]
                TMP_4=W[i*4+3,k]*M[j,k]
                if TMP_1>val_1:
                    val_1=TMP_1
                if TMP_2>val_2:
                    val_2=TMP_2
                if TMP_3>val_3:
                    val_3=TMP_3
                if TMP_4>val_4:
                    val_4=TMP_4

            C[i*4+0,j]=val_1
            C[i*4+1,j]=val_2
            C[i*4+2,j]=val_3
            C[i*4+3,j]=val_4

    #Remainder
    for i in range(N//4*4,N):
        for j in range(N):
            val=W[i,0]*M[j,0]
            for k in range(1,N):
                TMP=W[i,k]*M[j,k]
                if TMP>val:
                    val=TMP
            C[i,j]=val
    return C

时间

W=np.random.rand(100,100)
%timeit calc_1(W)
#16.8 ms ± 131 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit calc_2(W)
#449 µs ± 25.7 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit calc_3(W)
#259 µs ± 47.4 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

W=np.random.rand(2000,2000)
#Temporary array would be 64GB in this case
%timeit calc_2(W)
#5.37 s ± 174 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit calc_3(W)
#596 ms ± 30.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如果没有中间 3D 矩阵，是否可以实现或近似计算此矩阵？

Can this matrix calculation be implemented or approximated without an intermediate 3D matrix?

c++

python

matrix

linear-algebra

pytorch