Repository: Lightning-AI/torchmetrics

Issue:
Converting from half precision to `bfloat16` in Pearson correlation coefficient cause numerical errors
## 🐛 Bug

The lack of half-precision op support on CPUs has lead to the decision to convert to `bfloat16` in the calculation of Pearson correlation coefficient. (#1813)

However this leads to broadcasting errors when the values are (meaningfully) small.


### To Reproduce

```py
>>> import torch
>>> a = torch.rand((32,), dtype=torch.half)
>>> b = torch.rand((32,), dtype=torch.half)
>>> from torchmetrics.functional import pearson_corrcoef
>>> pearson_corrcoef(a, b)
tensor(0.2240)
>>> pearson_corrcoef(a * 1e-1, b * 1e-1)
tensor(0.2251)
>>> pearson_corrcoef(a * 1e-2, b * 1e-2)
tensor(0.2213)
>>> pearson_corrcoef(a * 1e-3, b * 1e-3)
tensor(0.)
>>> pearson_corrcoef(a * 1e-4, b * 1e-4)
tensor(nan)
```

Which generally makes sense since there are only 5 exponential bits in `float16`, and some are lost in the operations. However this is not obvious during debugging since 0.001 is not *that* small...

### Expected behavior

A warning when the dynamic range is problematic, such as [what SciPy does](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.pearsonr.html#scipy-stats-pearsonr).

Or there may be some way to preserve more precision?

### Environment

- TorchMetrics version (pypi v1.0.1):
- Python & PyTorch Version (3.10, 2.0.1):
- Any other relevant information such as OS (tested on Linux_x64 and macOS_arm64):



Discussion:
Hi! thanks for your contribution!, great first issue!
Hi @ziw-liu, thanks for raising this issue :]
I looked into it and it actually seems from the next version of Pytorch (v2.1.0) all operations involved in calculating the Pearson corr coef will have native support for `half` precision. So at that point we could remove the `bfloat16` conversion.

That said, it is not really the underlying problem here. For values in the 1e-3 and 1e-4 range you are also simply running into problems with what can be done in `half` precision. In particular, it is the final computation that is a problem
https://github.com/Lightning-AI/torchmetrics/blob/74e8c0679c4d20433a07757f6b9c3c43ade2b52d/src/torchmetrics/functional/regression/pearson.py#L101
Essentially, the variance of the input is so small that the multiplication of the variance is running out of precision:
![billede](https://github.com/Lightning-AI/torchmetrics/assets/24896311/70bb600c-2774-411a-b428-0e95e09a45cd)
The calculation is only working in `half` precision when the individual variances are larger than 1e-3 to 1e-4.



@SkafteNicki Thanks for looking into this!

The best way to handle small variance input seems to be a `UserWarning` then?

Resolving patch:
diff --git a/src/torchmetrics/functional/regression/pearson.py b/src/torchmetrics/functional/regression/pearson.py
--- a/src/torchmetrics/functional/regression/pearson.py
+++ b/src/torchmetrics/functional/regression/pearson.py
@@ -11,12 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import math
 from typing import Tuple
 
 import torch
 from torch import Tensor
 
 from torchmetrics.functional.regression.utils import _check_data_shape_to_num_outputs
+from torchmetrics.utilities import rank_zero_warn
 from torchmetrics.utilities.checks import _check_same_shape
 
 
@@ -100,6 +102,15 @@ def _pearson_corrcoef_compute(
         var_x = var_x.bfloat16()
         var_y = var_y.bfloat16()
 
+    bound = math.sqrt(torch.finfo(var_x.dtype).eps)
+    if (var_x < bound).any() or (var_y < bound).any():
+        rank_zero_warn(
+            "The variance of predictions or target is close to zero. This can cause instability in Pearson correlation"
+            "coefficient, leading to wrong results. Consider re-scaling the input if possible or computing using a"
+            f"larger dtype (currently using {var_x.dtype}).",
+            UserWarning,
+        )
+
     corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()
     return torch.clamp(corrcoef, -1.0, 1.0)
 


Test patch:
diff --git a/tests/unittests/regression/test_pearson.py b/tests/unittests/regression/test_pearson.py
--- a/tests/unittests/regression/test_pearson.py
+++ b/tests/unittests/regression/test_pearson.py
@@ -140,3 +140,12 @@ def test_final_aggregation_function(shapes):
     output = _final_aggregation(input_fn(), input_fn(), input_fn(), input_fn(), input_fn(), torch.randint(10, shapes))
     assert all(isinstance(out, torch.Tensor) for out in output)
     assert all(out.ndim == input_fn().ndim - 1 for out in output)
+
+
+@pytest.mark.parametrize(("dtype", "scale"), [(torch.float16, 1e-4), (torch.float32, 1e-8), (torch.float64, 1e-16)])
+def test_pearsons_warning_on_small_input(dtype, scale):
+    """Check that a user warning is raised for small input."""
+    preds = scale * torch.randn(100, dtype=dtype)
+    target = scale * torch.randn(100, dtype=dtype)
+    with pytest.warns(UserWarning, match="The variance of predictions or target is close to zero.*"):
+        pearson_corrcoef(preds, target)