MCC Multilabel with average macro

🚀 Feature

The current behaviour of the multilabel MCC is to calculate the overall MCC by adding the TP, TN, FP & FN of each label (C matrices of shape (2,2)). This does, however, lead to an unintended behaviour in which the metric is computed as a micro average and never as macro.

Motivation

There are disbalanced multilabel classification tasks which might rely on the MCC as their primary metric. By not offering the option to select the average on this metric, this often leads to a not representative metric regarding the problem.

Pitch

Add the possibility to compute macro average on MCC.

Alternatives

I've done this implementation for my code, although it is a bit messy, but to give an idea of what we are looking for

class MacroMCCFromBinary(Metric):
    """Macro-averaged MCC: mean of per-label binary MCCs (mirrors STL heads)."""

    def __init__(self, num_labels: int):
        super().__init__()
        self.num_labels = num_labels
        self._binary_mccs = torch.nn.ModuleList(
            [MatthewsCorrCoef(task="binary") for _ in range(num_labels)]
        )

    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
        for i, mcc in enumerate(self._binary_mccs):
            mcc.update(preds[:, i], target[:, i])

    def compute(self) -> torch.Tensor:
        return torch.stack([mcc.compute() for mcc in self._binary_mccs]).mean()

    def reset(self) -> None:
        for mcc in self._binary_mccs:
            mcc.reset()

Additional context

I think part of the fix is located on https://github.com/Lightning-AI/torchmetrics/blob/master/src/torchmetrics/functional/classification/matthews_corrcoef.py

There must be some kind of fix.

Start by reading src/torchmetrics/functional/classification/matthews_corrcoef.py and trace how multilabel MCC currently aggregates label confusion matrices. Check the surrounding classification metric APIs for the expected averaging pattern, then verify that MCC supports selecting macro averaging while preserving the current behavior when no option is selected.