spearman_corrcoef

[AhmedHussKhalifa]

Hi,
I am trying to calculate the spearman coefficient during training on the GPU, and it is slowing down. Is there any advice to speed it up?

torchmetrics.functional import spearman_corrcoef

[mhorlacher]

You could try my custom implementation for PCC and change it to SCC.

import torch
import torchmetrics

def pearson_corrcoef(x, y, dim=-1):
    x = x - torch.unsqueeze(torch.mean(x, dim), dim)
    y = y - torch.unsqueeze(torch.mean(y, dim), dim)
    return torch.sum(x * y, dim) / torch.sqrt(torch.sum(x ** 2, dim) * torch.sum(y ** 2, dim))

class PearsonCorrCoeff(torchmetrics.MeanMetric):
    def __init__(self, dim=-1, reduction=torch.mean, *args, **kwargs):
        super(PearsonCorrCoeff,  self).__init__(*args, **kwargs)

        self.dim = dim
        self.reduction = reduction

    def update(self, y: torch.Tensor, y_pred: torch.Tensor):
        assert y.shape == y_pred.shape

        pcc = pearson_corrcoef(y, y_pred, dim=self.dim)
        pcc = torch.nan_to_num(pcc, 0.0) # replace nan's with 0 (this might underestimate the pcc)

        reduced_pcc = self.reduction(pcc)

        # update (i.e. take mean)
        super().update(reduced_pcc)

[Borda]

The slowness is expected — Spearman correlation requires ranking (sorting), and sorting on GPU is inherently slower than vectorized arithmetic (which is why Pearson is fast).

Options to speed it up:

1. Compute on CPU (counterintuitive but often faster for ranking):

from torchmetrics.functional.regression import spearman_corrcoef

score = spearman_corrcoef(preds.cpu(), target.cpu())

2. Compute less frequently — epoch-only:

class MyModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.spearman = SpearmanCorrCoef()

    def validation_step(self, batch, batch_idx):
        self.spearman.update(preds, target)  # just update, don't compute

    def on_validation_epoch_end(self):
        self.log("spearman", self.spearman.compute())
        self.spearman.reset()

3. Subsample for large N:

if preds.shape[0] > 10000:
    idx = torch.randperm(preds.shape[0])[:10000]
    score = spearman_corrcoef(preds[idx], target[idx])

4. Approximate with differentiable sorting (if used as a loss):

# pip install torchsort
import torchsort

def fast_spearman(preds, target):
    pred_ranks = torchsort.soft_rank(preds.unsqueeze(0), regularization_strength=1e-2).squeeze(0)
    target_ranks = torchsort.soft_rank(target.unsqueeze(0), regularization_strength=1e-2).squeeze(0)
    pred_ranks = pred_ranks - pred_ranks.mean()
    target_ranks = target_ranks - target_ranks.mean()
    return (pred_ranks * target_ranks).sum() / (pred_ranks.norm() * target_ranks.norm())

Docs: SpearmanCorrCoef