Differentiable 1-D, 2-D covariance (numpy.cov) clone

[LemonPi]

🚀 Feature

A differentiable way to calculate covariance for a tensor of random variables similar to numpy.cov.

Motivation

Sometimes the covariance of the data can be used as a norm (as opposed to the implicit identity matrix in the standard x'y norm we have x'Sy where S is the covariance matrix, for example between x and y). It would be good to have a differentiable way of calculating these to allow backpropagation.

Pitch

Here is my partial clone of numpy.cov that's tested with pytest against what numpy.cov finds.

def cov(x, rowvar=False, bias=False, ddof=None, aweights=None):
    """Estimates covariance matrix like numpy.cov"""
    # ensure at least 2D
    if x.dim() == 1:
        x = x.view(-1, 1)

    # treat each column as a data point, each row as a variable
    if rowvar and x.shape[0] != 1:
        x = x.t()

    if ddof is None:
        if bias == 0:
            ddof = 1
        else:
            ddof = 0

    w = aweights
    if w is not None:
        if not torch.is_tensor(w):
            w = torch.tensor(w, dtype=torch.float)
        w_sum = torch.sum(w)
        avg = torch.sum(x * (w/w_sum)[:,None], 0)
    else:
        avg = torch.mean(x, 0)

    # Determine the normalization
    if w is None:
        fact = x.shape[0] - ddof
    elif ddof == 0:
        fact = w_sum
    elif aweights is None:
        fact = w_sum - ddof
    else:
        fact = w_sum - ddof * torch.sum(w * w) / w_sum

    xm = x.sub(avg.expand_as(x))

    if w is None:
        X_T = xm.t()
    else:
        X_T = torch.mm(torch.diag(w), xm).t()

    c = torch.mm(X_T, xm)
    c = c / fact

    return c.squeeze()

Testing

def assert_same_cov(A, w=None):
    c1 = np.cov(A, rowvar=False, aweights=w)
    c2 = cov(torch.tensor(A, dtype=torch.float), aweights=w)
    assert np.linalg.norm(c2.numpy() - c1) < 1e-6


def test_cov():
    a = [1, 2, 3, 4]
    assert_same_cov(a)
    A = [[1, 2], [3, 4]]
    assert_same_cov(A)

    assert_same_cov(a, w=[1, 1, 1, 1])
    assert_same_cov(a, w=[2, 0.5, 3, 1])

    assert_same_cov(A, w=[1, 1])
    assert_same_cov(A, w=[2, 0.5])

Alternatives

I don't think there's an existing method to compute the covariance in pytorch - are there alternatives? And is this an OK format for computation (instead of a forward and backward method)?

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @mruberry @rgommers @heitorschueroff

[zuoxingdong]

Any update on this?

[umanwizard]

Sorry, I do not work on PyTorch anymore. @soumith can you remove the triaged label so this gets put back in the oncall queue?

[zuoxingdong]

I guess we might need a C++ implementation if we want to port this function to the core? @soumith

[ProGamerGov]

It'd be insanely useful to have a PyTorch version of numpy.cov!

[turmeric-blend]

any updates on having numpy.cov version in PyTorch?

[mauricett]

I suggest to calculate a covariance matrix like:

def cov(X):
    D = X.shape[-1]
    mean = torch.mean(X, dim=-1).unsqueeze(-1)
    X = X - mean
    return 1/(D-1) * X @ X.transpose(-1, -2)

Works with batches of matrices, i.e. an input with shape (batchsize, m, n) will give you covariance matrices with shape (batchsize, m, m). And you can calculate the correlation matrix similarly:

def corr(X, eps=1e-08):
    D = X.shape[-1]
    std = torch.std(X, dim=-1).unsqueeze(-1)
    mean = torch.mean(X, dim=-1).unsqueeze(-1)
    X = (X - mean) / (std + eps)
    return 1/(D-1) * X @ X.transpose(-1, -2)

[ProGamerGov]

I have an RGB covariance matrix calculation function for image datasets here: https://github.com/ProGamerGov/dream-creator/blob/master/data_tools/calc_cm.py#L61-L67

Encase anyone has the same use case for covariance matrices.

[mauricett]

I have an RGB covariance matrix calculation function for image datasets here: https://github.com/ProGamerGov/dream-creator/blob/master/data_tools/calc_cm.py#L61-L67

Encase anyone has the same use case for covariance matrices.

You can do the same with my code, by the way, by moving the reshape outside of the covariance function.

img = img.reshape(-1, 3).T
cov(img)

I'll think about covering special cases like "1D covariance" and then maybe make a pull request.

[xmodar]

I think @mauricett's implementation is the way to go but let's add support for complex tensors as well.

def cov(tensor, rowvar=True, bias=False):
    """Estimate a covariance matrix (np.cov)"""
    tensor = tensor if rowvar else tensor.transpose(-1, -2)
    tensor = tensor - tensor.mean(dim=-1, keepdim=True)
    factor = 1 / (tensor.shape[-1] - int(not bool(bias)))
    return factor * tensor @ tensor.transpose(-1, -2).conj()

As for corrcoef(), I prefer to follow numpy's implementation since torch.std_mean() doesn't support complex tensors.

def corrcoef(tensor, rowvar=True):
    """Get Pearson product-moment correlation coefficients (np.corrcoef)"""
    covariance = cov(tensor, rowvar=rowvar)
    variance = covariance.diagonal(0, -1, -2)
    if variance.is_complex():
        variance = variance.real
    stddev = variance.sqrt()
    correlation = covariance
    correlation /= stddev.unsqueeze(-1)
    correlation /= stddev.unsqueeze(-2)
    if correlation.is_complex():
        correlation.real.clip_(-1, 1)
        correlation.imag.clip_(-1, 1)
    else:
        correlation.clip_(-1, 1)
    return correlation

If you want to add support for aweights and fweights to cov(), see this full implementation.

[umanwizard]

@gchanan , @cpuhrsch , I fear this has fallen through the cracks since it was unassigned to me. Can you please put it in the proper issue triage queue (I don't have permissions to change tags as I am no longer in the organization).