Fix RIN inverse for location-scale likelihood parameters

CHANGELOG.md

@@ -14,6 +14,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 **Fixed**
 
 - Fixed compatibility with pandas 3.1: fitting and predicting models on series with a daily frequency failed with a `TypeError`, and `MIDAS` accepted a series whose frequency was already the target low frequency. [#3289](https://github.com/unit8co/darts/pull/3289) by [Nicolò Angileri](https://github.com/nicoloangileri).
+- 🔴 Fixed the scale parameter of `GaussianLikelihood`, `LaplaceLikelihood`, `CauchyLikelihood` and `GumbelLikelihood` with `use_reversible_instance_norm`: the inverse normalization was applied before the softplus, so the predicted scale followed the level of the series instead of its spread. Models trained this way in earlier versions should be retrained. Combining `use_reversible_instance_norm` with any other likelihood except `QuantileRegression` now raises a `ValueError` at model creation. [#3306](https://github.com/unit8co/darts/pull/3306) by [Cassian433](https://github.com/Cassian433).
 
 **Dependencies**
 

darts/models/components/layer_norm_variants.py

@@ -135,3 +135,18 @@ def inverse(self, x: torch.Tensor) -> torch.Tensor:
         x = x * self.stdev.view(self.stdev.shape + (1,))
         x = x + self.mean.view(self.mean.shape + (1,))
         return x
+
+    def inverse_scale(self, x: torch.Tensor) -> torch.Tensor:
+        """Undo normalization of a softplus scale parameter in raw output space.
+
+        ``x`` has shape ``(batch_size, output_chunk_length, n_targets, 1)``.
+        Scales use the absolute inverse affine weight and are not shifted.
+        """
+        scale = nn.functional.softplus(x) * self.stdev.unsqueeze(-1)
+        if self.affine:
+            scale = scale / (self.affine_weight + self.eps * self.eps).abs().view(
+                self.affine_weight.shape + (1,)
+            )
+        scale = scale.clamp_min(torch.finfo(scale.dtype).tiny)
+        # Stable inverse softplus, including for very small and very large scales.
+        return scale + torch.log(-torch.expm1(-scale))

darts/models/forecasting/pl_forecasting_module.py

@@ -59,7 +59,24 @@ def forward_wrapper(self, x_in: PLModuleInput, *args, **kwargs):
         x_in = x_in.replace(past_target=self.rin(x_in.past_target.clone()))
         out = forward(self, x_in, *args, **kwargs)
         # inverse transform target output back to original scale
-        return out.replace(prediction=self.rin.inverse(out.prediction))
+        rin_parameters = (
+            None if self.likelihood is None else self.likelihood._rin_parameters
+        )
+        if rin_parameters is None or "scale" not in rin_parameters:
+            prediction = self.rin.inverse(out.prediction)
+        else:
+            prediction = torch.cat(
+                [
+                    (
+                        self.rin.inverse_scale(out.prediction[..., i : i + 1])
+                        if parameter == "scale"
+                        else self.rin.inverse(out.prediction[..., i : i + 1])
+                    )
+                    for i, parameter in enumerate(rin_parameters)
+                ],
+                dim=-1,
+            )
+        return out.replace(prediction=prediction)
 
     return forward_wrapper
 

darts/models/forecasting/torch_forecasting_model.py

@@ -384,6 +384,21 @@ def _validate_model_params(cls, **kwargs):
                 ),
             )
 
+        likelihood = kwargs.get("likelihood")
+        if (
+            kwargs.get("use_reversible_instance_norm")
+            and likelihood is not None
+            and likelihood._rin_parameters is None
+        ):
+            raise_log(
+                ValueError(
+                    "`use_reversible_instance_norm` only supports `QuantileRegression`, "
+                    "`GaussianLikelihood`, `LaplaceLikelihood`, `CauchyLikelihood` and "
+                    "`GumbelLikelihood` as `likelihood`, "
+                    f"received `{type(likelihood).__name__}`."
+                ),
+            )
+
     @classmethod
     def _extract_torch_model_params(cls, **kwargs):
         """extract params from model creation to set up TorchForecastingModels"""

darts/tests/models/components/test_layer_norm_variants.py

@@ -16,6 +16,10 @@
     RINorm,
     RMSNorm,
 )
+from darts.utils.likelihood_models.torch import (
+    MIN_CAUCHY_GAMMA_SAMPLING,
+    CauchyLikelihood,
+)
 
 
 class TestLayerNormVariants:
@@ -45,3 +49,43 @@ def test_rin(self):
         rin = RINorm(input_dim=3, affine=True)
         with pytest.raises(RuntimeError):
             x_norm = rin(x)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+    @pytest.mark.parametrize("affine", [False, True])
+    def test_rin_inverse_scale_extremes(self, dtype, affine):
+        rin = RINorm(input_dim=1, affine=affine).to(dtype=dtype)
+        if affine:
+            with torch.no_grad():
+                rin.affine_weight.fill_(-2.5)
+                rin.affine_bias.fill_(1000.0)
+        rin(torch.tensor([[[10.0], [30.0]]], dtype=dtype))
+        raw = (
+            torch.tensor([-1000.0, -100.0, -2.0, 0.0, 25.0, 1000.0], dtype=dtype)
+            .reshape(1, 6, 1, 1)
+            .requires_grad_()
+        )
+        output = rin.inverse_scale(raw)
+        factor = rin.stdev.unsqueeze(-1)
+        if affine:
+            factor = factor / (rin.affine_weight + rin.eps**2).abs().unsqueeze(-1)
+        expected = (torch.nn.functional.softplus(raw) * factor).clamp_min(
+            torch.finfo(dtype).tiny
+        )
+        assert torch.isfinite(output).all()
+        torch.testing.assert_close(
+            torch.nn.functional.softplus(output), expected, atol=0, rtol=1e-5
+        )
+        # Cauchy's existing gamma floor still applies after the raw-space inverse.
+        _, gamma = CauchyLikelihood()._params_from_output(
+            torch.cat([torch.zeros_like(output), output], dim=-1)
+        )
+        torch.testing.assert_close(
+            gamma,
+            expected.squeeze(-1).clamp_min(MIN_CAUCHY_GAMMA_SAMPLING),
+            atol=0,
+            rtol=1e-5,
+        )
+        output.sum().backward()
+        assert torch.isfinite(raw.grad).all()
+        if affine:
+            assert torch.isfinite(rin.affine_weight.grad).all()

darts/tests/models/forecasting/test_torch_forecasting_model.py

@@ -13,6 +13,7 @@
 import math
 import os
 from dataclasses import fields
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import patch
 
@@ -58,6 +59,7 @@
 )
 from darts.models.components.layer_norm_variants import RINorm
 from darts.models.forecasting.global_baseline_models import _GlobalNaiveModel
+from darts.models.forecasting.pl_forecasting_module import io_processor
 from darts.tests.conftest import tfm_kwargs, tfm_kwargs_dev
 from darts.tests.parametrize_helpers import param_product
 from darts.utils.data.torch_datasets._data_module import TorchDataModule
@@ -70,6 +72,8 @@
     TorchTrainingDataset,
 )
 from darts.utils.data.torch_datasets.utils import (
+    PLModuleInput,
+    PLModuleOutput,
     TorchInferenceSample,
     TorchTrainingSample,
     _batch_collate_fn_predict,
@@ -78,7 +82,9 @@
 from darts.utils.likelihood_models.torch import (
     CauchyLikelihood,
     GaussianLikelihood,
+    GumbelLikelihood,
     LaplaceLikelihood,
+    PoissonLikelihood,
     QuantileRegression,
     TorchLikelihood,
 )
@@ -2317,6 +2323,124 @@ def test_rin(self, model_config):
         assert isinstance(model_rin_mv.model.rin, RINorm)
         assert model_rin_mv.model.rin.input_dim == self.multivariate_series.n_components
 
+    @pytest.mark.parametrize(
+        "likelihood",
+        [
+            GaussianLikelihood(),
+            LaplaceLikelihood(),
+            CauchyLikelihood(),
+            GumbelLikelihood(),
+            QuantileRegression([0.1, 0.5, 0.9]),
+        ],
+    )
+    def test_rin_likelihood_equivariance(self, likelihood):
+        rng = np.random.default_rng(42)
+        series = TimeSeries.from_values(rng.normal(size=(48, 2)))
+        model = DLinearModel(
+            input_chunk_length=12,
+            output_chunk_length=3,
+            likelihood=likelihood,
+            use_reversible_instance_norm=True,
+            n_epochs=2,
+            random_state=42,
+            **tfm_kwargs,
+        )
+        model.fit(series)
+        reference = model.predict(3, predict_likelihood_parameters=True).values()
+        for factor, shift in [(1.0, 1000.0), (100.0, 0.0)]:
+            prediction = model.predict(
+                3,
+                series=series * factor + shift,
+                predict_likelihood_parameters=True,
+            ).values()
+            expected = reference * factor
+            if isinstance(likelihood, QuantileRegression):
+                expected += shift
+            else:
+                expected[:, ::2] += shift
+            np.testing.assert_allclose(prediction, expected, rtol=1e-3, atol=1e-4)
+
+    @pytest.mark.parametrize(
+        "likelihood",
+        [
+            None,
+            PoissonLikelihood(),
+            GaussianLikelihood(),
+            LaplaceLikelihood(),
+            CauchyLikelihood(),
+            GumbelLikelihood(),
+            QuantileRegression([0.1, 0.5, 0.9]),
+        ],
+    )
+    @pytest.mark.parametrize("affine_weight", [None, 2.5, -0.5])
+    def test_rin_likelihood_raw_output(self, likelihood, affine_weight):
+        # Exercise the shared forward wrapper without training, with two batches/targets.
+        rin = RINorm(input_dim=2, affine=affine_weight is not None).double()
+        if rin.affine:
+            with torch.no_grad():
+                rin.affine_weight.copy_(torch.tensor([affine_weight, 1.5]))
+                rin.affine_bias.copy_(torch.tensor([3.0, -4.0]))
+        past_target = torch.arange(24, dtype=torch.float64).reshape(2, 6, 2)
+        n_params = 1 if likelihood is None else likelihood.num_parameters
+        raw = (
+            torch.linspace(-3, 3, 2 * 3 * 2 * n_params, dtype=torch.float64)
+            .reshape(2, 3, 2, n_params)
+            .requires_grad_()
+        )
+        module = SimpleNamespace(
+            use_reversible_instance_norm=True, rin=rin, likelihood=likelihood
+        )
+
+        @io_processor
+        def forward(self, x_in):
+            return PLModuleOutput(prediction=raw)
+
+        output = forward(module, PLModuleInput(past_target=past_target)).prediction
+        if likelihood is None or isinstance(
+            likelihood, QuantileRegression | PoissonLikelihood
+        ):
+            torch.testing.assert_close(output, rin.inverse(raw))
+        else:
+            location, scale = likelihood._params_from_output(output)
+            expected_location = rin.inverse(raw[..., :1]).squeeze(-1)
+            factor = rin.stdev
+            if rin.affine:
+                factor = factor / (rin.affine_weight + rin.eps**2).abs()
+            expected_scale = torch.nn.functional.softplus(raw[..., 1]) * factor
+            torch.testing.assert_close(location, expected_location)
+            torch.testing.assert_close(scale, expected_scale)
+        output.sum().backward()
+        assert torch.isfinite(raw.grad).all()
+        if rin.affine:
+            assert torch.isfinite(rin.affine_weight.grad).all()
+            assert torch.isfinite(rin.affine_bias.grad).all()
+
+    @pytest.mark.parametrize("rin", [True, {"affine": False}])
+    def test_rin_unsupported_likelihood(self, rin):
+        with pytest.raises(
+            ValueError,
+            match=(
+                "`use_reversible_instance_norm` only supports `QuantileRegression`, "
+                "`GaussianLikelihood`, `LaplaceLikelihood`, `CauchyLikelihood` and "
+                r"`GumbelLikelihood` as `likelihood`, received `PoissonLikelihood`\."
+            ),
+        ):
+            DLinearModel(
+                input_chunk_length=12,
+                output_chunk_length=3,
+                likelihood=PoissonLikelihood(),
+                use_reversible_instance_norm=rin,
+                **tfm_kwargs,
+            )
+        # The likelihood remains available when RIN is disabled.
+        DLinearModel(
+            input_chunk_length=12,
+            output_chunk_length=3,
+            likelihood=PoissonLikelihood(),
+            use_reversible_instance_norm=False,
+            **tfm_kwargs,
+        )
+
     @pytest.mark.parametrize("use_mc_dropout", [False, True])
     def test_mc_dropout_active(self, use_mc_dropout):
         """Test that model activates dropout ."""

darts/utils/likelihood_models/torch.py

@@ -97,6 +97,10 @@ def _check_in_open_0_1_intvl(param, param_name=""):
 
 
 class TorchLikelihood(Likelihood, ABC):
+    # RIN inverse for each raw output channel; scale channels use softplus.
+    # None indicates that the likelihood does not support RIN.
+    _rin_parameters: tuple[str, ...] | None = None
+
     def __init__(
         self,
         likelihood_type: LikelihoodType,
@@ -232,6 +236,8 @@ def _get_equality_attrs(likelihood, ignore_attrs):
 
 
 class GaussianLikelihood(TorchLikelihood):
+    _rin_parameters = ("loc", "scale")
+
     def __init__(
         self, prior_mu=None, prior_sigma=None, prior_strength=1.0, beta_nll=0.0
     ):
@@ -562,6 +568,8 @@ def _params_from_output(
 
 
 class CauchyLikelihood(TorchLikelihood):
+    _rin_parameters = ("loc", "scale")
+
     def __init__(self, prior_xzero=None, prior_gamma=None, prior_strength=1.0):
         """
         Cauchy Distribution.
@@ -876,6 +884,8 @@ def _params_from_output(
 
 
 class GumbelLikelihood(TorchLikelihood):
+    _rin_parameters = ("loc", "scale")
+
     def __init__(self, prior_mu=None, prior_beta=None, prior_strength=1.0):
         """
         Gumbel distribution.
@@ -978,6 +988,8 @@ def _params_from_output(
 
 
 class LaplaceLikelihood(TorchLikelihood):
+    _rin_parameters = ("loc", "scale")
+
     def __init__(self, prior_mu=None, prior_b=None, prior_strength=1.0):
         """
         Laplace distribution.
@@ -1180,6 +1192,10 @@ def __init__(self, quantiles: list[float] | None = None):
         # ignore additional attributes
         self.ignore_attrs_equality += ["_median_idx", "first", "quantiles_tensor"]
 
+    @property
+    def _rin_parameters(self) -> tuple[str, ...]:
+        return ("loc",) * self.num_parameters
+
     def sample(self, model_output: torch.Tensor) -> torch.Tensor:
         """
         Sample uniformly between [0, 1] (for each batch example) and return the linear interpolation between the fitted