Spearman's Correlation

I need to use a rank-based correlation (Spearman’s Correlation) to compute my loss.
Could anyone please help me with the implementation of Spearman’s Correlation between two vectors in Pytorch?

In case someone is trying to look for a solution:

def cov(m):
    # m = m.type(torch.double)  # uncomment this line if desired
    fact = 1.0 / (m.shape[-1] - 1)  # 1 / N
    m -= torch.mean(m, dim=(1, 2), keepdim=True)
    mt = torch.transpose(m, 1, 2)  # if complex: mt = m.t().conj()
    return fact * m.matmul(mt).squeeze()

def compute_rank_correlation(x, y):
    x, y = rankmin(x), rankmin(y)
    return corrcoef(x, y)

def corrcoef(x, y):
    batch_size = x.shape[0]
    x = torch.stack((x, y), 1)
    # calculate covariance matrix of rows
    c = cov(x)
    # normalize covariance matrix
    d = torch.diagonal(c, dim1=1, dim2=2)
    stddev = torch.pow(d, 0.5)
    stddev = stddev.repeat(1, 2).view(batch_size, 2, 2)
    c = c.div(stddev)
    c = c.div(torch.transpose(stddev, 1, 2))
    return c[:, 1, 0]

There is small difference between the above implementation and scipy’s spearmanr implementation:

  1. This operate on batches, that is, it expects input to be of dimension (N, x) where N is the batch size for both parameters x and y.
  2. scipy uses rank average while this uses rank min to compute the ranking before corrcoef. There is 10^-3 disagreement and although I suspect that this is the cause, this could possibly be.

Please feel free to share your improvements on this as it could be useful for other people and me

1 Like

Hi haltaha!

Presumably you will use your “loss” to train your network with
backpropagation by calling loss.backward().

Because the ranks of your data values are discrete integers, the
Spearman’s-rank-correlation loss you propose calculating will not
be (usefully) differentiable, so you will not be able to backpropagate
nor train.

You don’t show us the code for rankmin(), but presumably buried
in there somewhere is a non-differentiable call that returns the indices
for the sorted values of x (sometimes referred to as argsort()).


K. Frank

You are correct, thank you.


Here is simplified version for computation of Spearman’s rank correlation coefficient. It’s non differentiable, but much faster than scipy.stats.spearmanr version.

def _get_ranks(x: torch.Tensor) -> torch.Tensor:
    tmp = x.argsort()
    ranks = torch.zeros_like(tmp)
    ranks[tmp] = torch.arange(len(x))
    return ranks

def spearman_correlation(x: torch.Tensor, y: torch.Tensor):
    """Compute correlation between 2 1-D vectors
        x: Shape (N, )
        y: Shape (N, )
    x_rank = _get_ranks(x)
    y_rank = _get_ranks(y)
    n = x.size(0)
    upper = 6 * torch.sum((x_rank - y_rank).pow(2))
    down = n * (n ** 2 - 1.0)
    return 1.0 - (upper / down)

Time comparison

x = torch.rand(1000)
y = torch.rand(1000)
%timeit spearman_correlation(x, y)
206 µs ± 15.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit spearmanr(x_n, y_n)[0]
592 µs ± 2.22 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Hi, I found a differential version at post:

Copying from above post:
" A recent paper, Fast Differentiable Sorting and Ranking , introduced a novel method for differentiable sorting and ranking, with the added bonus of O(nlogn) complexity (I would encourage reading the paper to learn more). We can leverage their open sourced code google-research/fast-soft-sort in order to implement a differentiable version of the Spearman"