From 46bb8bf41afda8da9032a15ab24943ec29841ef0 Mon Sep 17 00:00:00 2001 From: Cassian433 <262026629+Cassian433@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:07:19 +0530 Subject: [PATCH] Fix RIN inverse for location-scale likelihood parameters --- CHANGELOG.md | 1 + .../models/components/layer_norm_variants.py | 15 +++ .../forecasting/pl_forecasting_module.py | 19 ++- .../forecasting/torch_forecasting_model.py | 15 +++ .../components/test_layer_norm_variants.py | 44 +++++++ .../test_torch_forecasting_model.py | 124 ++++++++++++++++++ darts/utils/likelihood_models/torch.py | 16 +++ 7 files changed, 233 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba2e1d2932..1746b7e361 100644 --- a/CHANGELOG.md +++ b/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** diff --git a/darts/models/components/layer_norm_variants.py b/darts/models/components/layer_norm_variants.py index dbeef19441..cf3bf31666 100644 --- a/darts/models/components/layer_norm_variants.py +++ b/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)) diff --git a/darts/models/forecasting/pl_forecasting_module.py b/darts/models/forecasting/pl_forecasting_module.py index 4e1415ccbd..4bb1cb193d 100644 --- a/darts/models/forecasting/pl_forecasting_module.py +++ b/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 diff --git a/darts/models/forecasting/torch_forecasting_model.py b/darts/models/forecasting/torch_forecasting_model.py index 04347c7dac..acad769242 100644 --- a/darts/models/forecasting/torch_forecasting_model.py +++ b/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""" diff --git a/darts/tests/models/components/test_layer_norm_variants.py b/darts/tests/models/components/test_layer_norm_variants.py index 204fee7210..5d7c780271 100644 --- a/darts/tests/models/components/test_layer_norm_variants.py +++ b/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() diff --git a/darts/tests/models/forecasting/test_torch_forecasting_model.py b/darts/tests/models/forecasting/test_torch_forecasting_model.py index c80ca921d7..8a739cac54 100644 --- a/darts/tests/models/forecasting/test_torch_forecasting_model.py +++ b/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 .""" diff --git a/darts/utils/likelihood_models/torch.py b/darts/utils/likelihood_models/torch.py index c180fb380d..d34ccad4d7 100644 --- a/darts/utils/likelihood_models/torch.py +++ b/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