这个矩阵计算可以在没有中间 3D 矩阵的情况下实现或近似吗？

Question

给定一个 NxN 矩阵 W，我正在寻找由以下公式给出的 NxN 矩阵 C：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 内存不足，或者在 k 上使用 for 循环然后很慢。我不知道如何绕过这些。如果没有像这样的大型中间矩阵，我该如何实现这个计算或它的近似值？

非常感谢任何语言的建议、伪代码或任何 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)