Weighted PCA: Memory requirements when multiplying two large matrices

DVigneault: <p>All–</p>
<p>Say I have a grayscale (non-binary) image with some object, and I’d like to determine its major axis. I could binarize and perform PCA on the foreground components only, but I’m wondering if I can instead perform a weighted PCA based on the pixel values, similar to what is described here [1]. Here’s the implementation I’ve come up with based on that post:</p>
<pre><code class="lang-auto">import pylab as plt
import torch as t
import numpy as np

def weighted_pca(input_mat, weights):
    weights_sum = weights.sum()
    weights_doubled = weights.repeat((2, 1))
    mean = t.sum(input_mat*weights_doubled, axis=1) / weights_sum
    matrix_centered = input_mat - mean.unsqueeze(1).float()
    print(weights.shape, matrix_centered.shape)
    # This line attempts to allocate a massive amount of memory (256 GB when s=512) and fails.
    inner = t.matmul(t.diag(weights), matrix_centered.transpose(1, 0))
    # Using sparse matrices also fails.
    # inner = t.sparse.mm(t.diag(weights).to_sparse_csr(), matrix_centered.transpose(1, 0))
    weighted_cov = t.matmul(matrix_centered, inner) / weights_sum
    u,e,v = t.linalg.svd(weighted_cov)
    return u, e, v, mean


if __name__ == "__main__":
    s = 512
    # Dummy binary data -- real data would vary between 0 and 1.
    img = t.diag(t.ones((s)))
    points = np.stack(np.meshgrid(*(np.arange(s), np.arange(s)))).reshape(2, -1)

    u, e, v, mean = weighted_pca(t.Tensor(points).cuda(), img.flatten().cuda())

    plt.imshow(img, origin="lower")
    plt.quiver(*mean, u[0,0], u[0,1])
    plt.show()
</code></pre>
<p>This works great for a small matrix size (say 256 x 256) but requires allocation of a very large amount of memory during calculation of the covariance matrix for larger sizes (say 512 x 512).  I tried converting the diagonal matrix to a sparse matrix, but this didn’t help.  I’m wondering if anyone can see a memory efficient way of performing this calculation, or maybe if there’s another way I could formulate this problem to take advantage of the built in <code>pca_lowrank</code> implementation.</p>
<p>Happy to provide more details if necessary – thanks so much in advance to everyone who has built this incredible software and maintains this incredibly helpful forum.</p>
<p>Best,</p>
<p>–Davis</p>
<p>[1] <a href="https://stats.stackexchange.com/questions/113485/weighted-principal-components-analysis" class="inline-onebox" rel="noopener nofollow ugc">pca - Weighted principal components analysis - Cross Validated</a></p>

J_Johnson: <p>Welcome to the forums!</p>
<p>If we change that line, should still yield the same result for what you’re doing:</p>
<pre><code class="lang-auto">v = 4
dummy_weight = t.rand(v)
dummy_matrix_centered = t.rand((2,v))

#old function
orig = t.matmul(t.diag(dummy_weight), dummy_matrix_centered.T)

#new function
new = (dummy_weight*dummy_matrix_centered).T
</code></pre>
<p>We can check this with:</p>
<pre><code class="lang-auto">print(t.isclose(orig, new).all())
</code></pre>
<p>Which returns True.</p>
<p>Keep in mind that diagonalizing a vector and matmuling that on another matrix is effectively the same as row wise multiplication.</p>

DVigneault: <p>Thanks so much <a class="mention" href="/u/j_johnson">@J_Johnson</a>!  That does the trick – I must have been staring at this too long, seems so obvious now that you point that out.  Really appreciate your help!!  <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>