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

user_15885288 0
Viewed 103

Given an NxN matrix W, I'm looking to calculate an NxN matrix C given by the equation in this link: https://i.stack.imgur.com/dY7rY.png, or in LaTeX

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

I have tried to implement this in PyTorch but I've either encountered memory problems by constructing an intermediate NxNxN 3D matrix which, for large N, causes my GPU to run out of memory, or used a for-loop over k which is then very slow. I can't work out how I can get round these. How might I implement this calculation, or an approximation of it, without a large intermediate matrix like this?

Suggestions, pseudocode in any language or an implementation in any of Python/Numpy/PyTorch would be much appreciated.

pytorch linear-algebra matrix python c++
2 Answers

The formula can be simplified as

C_ij = max_k ( W_ik M_kj)

where

M = W * W - N * W

with N the size of the matrix W and W * W the usual product.

Then, in the formula above, for every i, j there is an independent maximum to be computed. Without knowing further properties of W, it is in general not possible to further simplify the problem. So, after computing the matrix M, you can do a loop over i and j, and compute the maximum.

user_87699850

The first solution using Numba (You can do the same using Cython or plain C) would be to formulate the problem using simple loops.

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 provided a simplification which scales far better for larger array sizes. This leads to the following, where I also optimized away a small temporary array.

@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

This can be optimized further by partial loop unrolling and optimizing the array access. Some compilers may do this automatically.

@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

Timings

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)
user_40457740
Related