Repository: Lightning-AI/torchmetrics

Issue:
Self-Implemented Metrics Behaved Weirdly in DDP Mode
## 🐛 Bug

Self-implemented nDCG@5 is larger than nDCG@1 in DDP mode. What I was trying to do here is to record the sum of sample-wise scores **_s_** and the number of samples **_n_**. At the end of batch/epoch, I can obtain the mean of nDCG@k by dividing **_s_** by **_n_**.

My implementation looks very similar to the example on [IMPLEMENTING A METRIC](https://torchmetrics.readthedocs.io/en/stable/pages/implement.html). However, It failed in distributed settings.

`Epoch 8:  60%|██████    | 1541/2555 [04:58<03:16,  5.17it/s, loss=0.00354, v_num=213259, train_ndcg@1=0.656, train_ndcg@3=0.675, train_ndcg@5=0.698,...`

### To Reproduce

<details>
  <summary>Self-implemented Metric</summary>

```python
from torchmetrics.functional.retrieval import retrieval_normalized_dcg
from torchmetrics import Metric
import torch
from torch import Tensor

class NDCG(Metric):
    full_state_update = False

    def __init__(self, top_k: int):
        super().__init__()
        self.top_k = top_k
        self.add_state("ndcg", default=torch.tensor(0, dtype=torch.float), dist_reduce_fx="sum")
        self.add_state("n", default=torch.tensor(0), dist_reduce_fx="sum")

    def update(self, preds: Tensor, target: Tensor):
        # preds and target are 2-D tensors.
        assert preds.shape == target.shape

        self.ndcg += torch.stack([retrieval_normalized_dcg(p, t, k=self.top_k) for p, t in zip(preds, target)]).sum()
        self.n += len(preds)

    def compute(self):
        return self.ndcg / self.n
```

</details>

### Expected behavior

The documentation of self-implemented metric looked too simple and probably missed some details of distributed metrics.

I noticed that you used list and in `compute()` you call mean() to get the average score of a given batch. Can you please explain the difference between these two implementations?

### Environment
- TorchMetrics version (and how you installed TM, e.g. conda, pip, build from source):
0.10.3 (conda-forge)
- Python version:
3.10.12
PyTorch version (e.g., 1.0):
1.13.1
- Any other relevant information such as OS (e.g., Linux):
pytorch-lightning: 1.9.4
Linux: 5.15.0-78-generic

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi @donglihe-hub, thanks for reporting this issue.
I investigated it a bit, but to me it seems that maybe your intuition behind nDCG@k is not correct. It is not impossible for nDCG@5 to be higher than nDCG@1, which is probably why you are seeing the results you are seeing.

Here is a example, which I computed both using torchmetrics and sklearn (just to confirm the behavior):
```python
import torch
from torchmetrics.functional.retrieval import retrieval_normalized_dcg
from sklearn.metrics import ndcg_score

target = torch.tensor([[1, 0, 0, 1, 0, 0, 1, 1, 1, 0]])
preds = torch.tensor([[0.54, 0.89, 0.50, 0.27, 0.33, 0.35, 0, 0.75, 0.53, 0.14]])

print([retrieval_normalized_dcg(preds, target, top_k=k) for k in [1,3,5]])
#prints [tensor(0.), tensor(0.5307), tensor(0.5296)]

print([ndcg_score(target, preds, k=k) for k in [1,3,5]])
#prints [0.0, 0.5307212739772434, 0.529634717214042]
```
that said I will be updating the documentation regarding implementing a metric in the near future.
Update:
I suddenly realize that I saw 64 because `dist_sync_on_step` is set to False. I will continue testing the logic of the metric.
\----------------------------------------------------------------------------------------------------------------------
@SkafteNicki Thank you for your reply. I think I might have some clues why this happened.

In the [IMPLEMENTING A METRIC](https://torchmetrics.readthedocs.io/en/stable/pages/implement.html) it said calling compute() does the following internally:
```
1. Syncs metric states between processes.
2. Reduce gathered metric states.
3. Calls the user defined compute() method on the gathered metric states.
...
```

I printed the number of samples in compute() and found that during training process the number is equal to the batch size on **a single GPU**, say 64 on each step. However, what I was expected to see is 64 * 4 = 256 as I used 4 GPUs if I understand the documentation correctly. While in validation, the metric states are correctly synced, I have 4000 samples for validation and the number is exactly the same.

**Results:**

Training:
```
len(self.ndcg): 64
```

Validation:
```
ndcg size: torch.Size([4000])
```

To make it clear, I implemented nDCG in two ways. One stores the results in a single list (I saw torchmetrics wrote Precision@K in this way before 0.11) and the other in two variables, i.e., sum of sample-wise nDCG and the number of samples. Here I printed the length of the list.

**Logs of DDP:**
```
...
[rank: 3] Global seed set to 1337
Initializing distributed: GLOBAL_RANK: 3, MEMBER: 4/4
...
2023-08-26 13:40:26,393 INFO:
distributed_backend=nccl
All distributed processes registered. Starting with 4 processes
...
```

Here are the related codes:
**Lightning Codes**
```
...
def training_step(self, batch: Tensor, batch_idx: int):
        x, y = batch
        logits = self.network(x)
        loss = self.loss_fn(logits, y)
        # train_metrics is an instance of MetricCollection
        metrics = self.train_metrics(torch.sigmoid(logits).detach(), y.long())
        self.log("train_loss", loss, sync_dist=True)
        self.log_dict(metrics, sync_dist=True, prog_bar=True)
...
def validation_step(self, batch: Tensor, batch_idx: int):
    x, y = batch
    logits = self.network(x)
    loss = self.loss_fn(logits, y)
    # valid_metrics is also an instance of MetricCollection
    self.valid_metrics.update(torch.sigmoid(logits).detach(), y.long())
    self.log("val_loss", loss, sync_dist=True)
    return loss

def on_validation_epoch_end(self):
    self.log_dict(self.valid_metrics.compute(), sync_dist=True, prog_bar=True)
    self.valid_metrics.reset()
...
```

**Metric Codes:**
```
class NDCG(Metric):
    def __init__(self, top_k: int):
        super().__init__()
        self.add_state("ndcg", default=[], dist_reduce_fx="cat")
...
    def compute(self):
        """Performs stacking on ndcg if neccesary"""
        if isinstance(self.ndcg, list):
            print(f"len(self.ndcg): {len(self.ndcg)}")
        else:
            print(f"ndcg size: {self.ndcg.shape}")
        ndcg = torch.stack(self.ndcg) if isinstance(self.ndcg, list) else self.ndcg
        return ndcg.mean()
```
> ```python
> #prints [0.0, 0.5307212739772434, 0.529634717214042]
> ```

Hi, Nicki. I tested my metric on Colab and obtained the same results.

```python
!pip install -q torchmetrics
import torchmetrics
from torchmetrics import Metric
from torchmetrics.functional.retrieval import retrieval_normalized_dcg

class NDCG(Metric):
    is_differentiable = False
    higher_is_better = True
    full_state_update = False

    def __init__(self, top_k: int):
        super().__init__()
        self.top_k = top_k
        # the range of ndcg is from 0 to 1
        self.add_state("ndcg", default=torch.tensor(0, dtype=torch.float), dist_reduce_fx="sum")
        self.add_state("n", default=torch.tensor(0), dist_reduce_fx="sum")

    def update(self, preds, target):
        assert preds.shape == target.shape
        self.ndcg += torch.stack([retrieval_normalized_dcg(p, t, top_k=self.top_k) for p, t in zip(preds, target)]).sum()
        self.n += len(preds)

    def compute(self):
        return self.ndcg / self.n

target = torch.tensor([[1, 0, 0, 1, 0, 0, 1, 1, 1, 0]])
preds = torch.tensor([[0.54, 0.89, 0.50, 0.27, 0.33, 0.35, 0, 0.75, 0.53, 0.14]])

m_1 = NDCG(1)
m_3 = NDCG(3)
m_5 = NDCG(5)
m_1.update(preds, target)
m_3.update(preds, target)
m_5.update(preds, target)
m_1.compute(), m_3.compute(), m_5.compute()
# (tensor(0.), tensor(0.5307), tensor(0.5296))
```

I realize there is one edge case, though it is unlikely to happen, where ndcg@5 is larger than ndcg@3 or even ndcg@1.
In the following case when all predicted scores are 0.1, ndcg@5 is larger than ndcg@3. Note that scikit-learn's ndcg averages ties, while torchmetrics's ndcg seems to break ties randomly.

```python
import torch
from torchmetrics.functional.retrieval import retrieval_normalized_dcg
from sklearn.metrics import ndcg_score

# the same target in your example
target = torch.tensor([[1, 0, 0, 1, 0, 0, 1, 1, 1, 0]])
preds = torch.tensor([[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]])

print([retrieval_normalized_dcg(preds, target, top_k=k) for k in [1,3,5]])
#prints [tensor(1.), tensor(0.4693), tensor(0.4852)]

print([ndcg_score(target, preds, k=k) for k in [1,3,5]])
#prints [0.5, 0.5, 0.5]

# Another example given in scikit-learn ndcg
target = torch.tensor([[10, 0, 0, 1, 5]])
preds = torch.tensor([[0.1, 0, 0, 0, 0.1]])

print([retrieval_normalized_dcg(preds, target, top_k=k) for k in [1,3,5]])
# prints [tensor(1.), tensor(0.9634), tensor(0.9917)]

print([ndcg_score(target, preds, k=k) for k in [1,3,5]])
# prints [0.75, 0.9080160192504324, 0.9279733094794905]
```
@donglihe-hub thanks for the example :]
You are completly right, there is a issue with our implementation that it does not deal with ties in the same way as sklearn does. That we need to fix.
@SkafteNicki I forgot to say thank you for correcting my understanding of nDCG in my last reply. Thank you for your help :)

As you have solved my issue, I will close it now. Thank you again for solving my problem! Hope you have a nice day!

Resolving patch:
diff --git a/src/torchmetrics/functional/retrieval/ndcg.py b/src/torchmetrics/functional/retrieval/ndcg.py
--- a/src/torchmetrics/functional/retrieval/ndcg.py
+++ b/src/torchmetrics/functional/retrieval/ndcg.py
@@ -19,10 +19,53 @@
 from torchmetrics.utilities.checks import _check_retrieval_functional_inputs
 
 
-def _dcg(target: Tensor) -> Tensor:
-    """Compute Discounted Cumulative Gain for input tensor."""
-    denom = torch.log2(torch.arange(target.shape[-1], device=target.device) + 2.0)
-    return (target / denom).sum(dim=-1)
+def _tie_average_dcg(target: Tensor, preds: Tensor, discount_cumsum: Tensor) -> Tensor:
+    """Translated version of sklearns `_tie_average_dcg` function.
+
+    Args:
+        target: ground truth about each document relevance.
+        preds: estimated probabilities of each document to be relevant.
+        discount_cumsum: cumulative sum of the discount.
+
+    Returns:
+        The cumulative gain of the tied elements.
+
+    """
+    _, inv, counts = torch.unique(-preds, return_inverse=True, return_counts=True)
+    ranked = torch.zeros_like(counts, dtype=torch.float32)
+    ranked.scatter_add_(0, inv, target.to(dtype=ranked.dtype))
+    ranked = ranked / counts
+    groups = counts.cumsum(dim=0) - 1
+    discount_sums = torch.zeros_like(counts, dtype=torch.float32)
+    discount_sums[0] = discount_cumsum[groups[0]]
+    discount_sums[1:] = discount_cumsum[groups].diff()
+    return (ranked * discount_sums).sum()
+
+
+def _dcg_sample_scores(target: Tensor, preds: Tensor, top_k: int, ignore_ties: bool) -> Tensor:
+    """Translated version of sklearns `_dcg_sample_scores` function.
+
+    Args:
+        target: ground truth about each document relevance.
+        preds: estimated probabilities of each document to be relevant.
+        top_k: consider only the top k elements
+        ignore_ties: If True, ties are ignored. If False, ties are averaged.
+
+    Returns:
+        The cumulative gain
+
+    """
+    discount = 1.0 / (torch.log2(torch.arange(target.shape[-1], device=target.device) + 2.0))
+    discount[top_k:] = 0.0
+
+    if ignore_ties:
+        ranking = preds.argsort(descending=True)
+        ranked = target[ranking]
+        cumulative_gain = (discount * ranked).sum()
+    else:
+        discount_cumsum = discount.cumsum(dim=-1)
+        cumulative_gain = _tie_average_dcg(target, preds, discount_cumsum)
+    return cumulative_gain
 
 
 def retrieval_normalized_dcg(preds: Tensor, target: Tensor, top_k: Optional[int] = None) -> Tensor:
@@ -59,15 +102,12 @@ def retrieval_normalized_dcg(preds: Tensor, target: Tensor, top_k: Optional[int]
     if not (isinstance(top_k, int) and top_k > 0):
         raise ValueError("`top_k` has to be a positive integer or None")
 
-    sorted_target = target[torch.argsort(preds, dim=-1, descending=True)][:top_k]
-    ideal_target = torch.sort(target, descending=True)[0][:top_k]
-
-    ideal_dcg = _dcg(ideal_target)
-    target_dcg = _dcg(sorted_target)
+    gain = _dcg_sample_scores(target, preds, top_k, ignore_ties=False)
+    normalized_gain = _dcg_sample_scores(target, target, top_k, ignore_ties=True)
 
     # filter undefined scores
-    all_irrelevant = ideal_dcg == 0
-    target_dcg[all_irrelevant] = 0
-    target_dcg[~all_irrelevant] /= ideal_dcg[~all_irrelevant]
+    all_irrelevant = normalized_gain == 0
+    gain[all_irrelevant] = 0
+    gain[~all_irrelevant] /= normalized_gain[~all_irrelevant]
 
-    return target_dcg.mean()
+    return gain.mean()


Test patch:
diff --git a/tests/unittests/retrieval/test_ndcg.py b/tests/unittests/retrieval/test_ndcg.py
--- a/tests/unittests/retrieval/test_ndcg.py
+++ b/tests/unittests/retrieval/test_ndcg.py
@@ -15,6 +15,7 @@
 
 import numpy as np
 import pytest
+import torch
 from sklearn.metrics import ndcg_score
 from torch import Tensor
 from torchmetrics.functional.retrieval.ndcg import retrieval_normalized_dcg
@@ -185,3 +186,15 @@ def test_arguments_functional_metric(self, preds: Tensor, target: Tensor, messag
             exception_type=ValueError,
             kwargs_update=metric_args,
         )
+
+
+def test_corner_case_with_tied_scores():
+    """See issue: https://github.com/Lightning-AI/torchmetrics/issues/2022."""
+    target = torch.tensor([[10, 0, 0, 1, 5]])
+    preds = torch.tensor([[0.1, 0, 0, 0, 0.1]])
+
+    for k in [1, 3, 5]:
+        assert torch.allclose(
+            retrieval_normalized_dcg(preds, target, top_k=k),
+            torch.tensor([ndcg_score(target, preds, k=k)], dtype=torch.float32),
+        )