# Weighted PCA: Memory requirements when multiplying two large matrices

All–

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:

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()

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 pca_lowrank implementation.

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.

Best,

–Davis

Welcome to the forums!

If we change that line, should still yield the same result for what you’re doing:

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

We can check this with:

print(t.isclose(orig, new).all())

Which returns True.

Keep in mind that diagonalizing a vector and matmuling that on another matrix is effectively the same as row wise multiplication.

2 Likes

Thanks so much @J_Johnson! 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!!