From 7e05e26bc8f36652429ed6bc84fd9eb95cf5f24f Mon Sep 17 00:00:00 2001 From: AHMETHAKANBEZIR1 Date: Mon, 26 Oct 2026 09:56:10 +0300 Subject: [PATCH] Keep Bernoulli training loss and prior KL in logit space --- CHANGELOG.md | 1 + .../test_torch_likelihoods.py | 66 +++++++++++++++++++ darts/utils/likelihood_models/torch.py | 29 ++++++++ 3 files changed, 96 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba2e1d2932..147a5bd479 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co **Fixed** +- Fixed `BernoulliLikelihood` losing training gradients for confident, incorrect predictions by computing its loss and optional prior regularization directly from logits. Related to [#621](https://github.com/unit8co/darts/issues/621). - 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). **Dependencies** diff --git a/darts/tests/utils/test_likelihood_models/test_torch_likelihoods.py b/darts/tests/utils/test_likelihood_models/test_torch_likelihoods.py index e93c94eea9..4036d05abb 100644 --- a/darts/tests/utils/test_likelihood_models/test_torch_likelihoods.py +++ b/darts/tests/utils/test_likelihood_models/test_torch_likelihoods.py @@ -185,6 +185,72 @@ def test_negative_binomial_mean_and_variance(self): assert torch.isfinite(distr.log_prob(target)).all() +class TestBernoulliLogitLoss: + @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) + @pytest.mark.parametrize("weighted", [False, True]) + @pytest.mark.parametrize("prior_p", [None, 0.25]) + def test_confident_incorrect_predictions(self, dtype, weighted, prior_p): + likelihood = BernoulliLikelihood(prior_p=prior_p, prior_strength=0.3) + output = torch.tensor([-40.0, 40.0], dtype=dtype).reshape(1, 1, 2, 1) + output.requires_grad_() + target = torch.tensor([1.0, 0.0], dtype=dtype).reshape(1, 1, 2) + weights = torch.tensor([0.5, 2.0], dtype=dtype).reshape(1, 1, 2) + sample_weight = weights if weighted else None + + loss = likelihood.compute_loss(output, target, sample_weight) + gradient = torch.autograd.grad(loss, output)[0].squeeze(-1) + logits = output.detach().squeeze(-1) + expected_loss = torch.nn.functional.softplus((1 - 2 * target) * logits) + expected_gradient = logits.sigmoid() - target + if weighted: + expected_loss = expected_loss * weights + expected_gradient = expected_gradient * weights + expected_loss = expected_loss.mean() + expected_gradient = expected_gradient / target.numel() + if prior_p is not None: + prior = torch.tensor(prior_p, dtype=dtype) + neg_entropy = prior.xlogy(prior) + (1 - prior).xlogy(1 - prior) + cross_entropy = torch.nn.functional.softplus(logits) - prior * logits + expected_loss += 0.3 * (cross_entropy + neg_entropy).mean() + expected_gradient += 0.3 * (logits.sigmoid() - prior) / target.numel() + torch.testing.assert_close(loss, expected_loss) + torch.testing.assert_close(gradient, expected_gradient) + + @pytest.mark.parametrize("prior_p", [None, 0.25]) + def test_ordinary_predictions_match_probability_loss(self, prior_p): + likelihood = BernoulliLikelihood(prior_p=prior_p, prior_strength=0.3) + output = torch.tensor([-2.0, 0.0, 3.0], dtype=torch.float64) + output = output.reshape(1, 1, 3, 1).requires_grad_() + target = torch.tensor([1.0, 0.0, 1.0], dtype=output.dtype).reshape(1, 1, 3) + weights = torch.tensor([0.5, 1.0, 2.0], dtype=output.dtype).reshape(1, 1, 3) + reference = super(BernoulliLikelihood, likelihood).compute_loss( + output, target, weights + ) + actual = likelihood.compute_loss(output, target, weights) + expected_grad = torch.autograd.grad(reference, output, retain_graph=True)[0] + actual_grad = torch.autograd.grad(actual, output)[0] + torch.testing.assert_close(actual, reference) + torch.testing.assert_close(actual_grad, expected_grad) + + def test_training_step_recovers_from_confident_wrong_prediction(self): + model = torch.nn.Linear(1, 1) + with torch.no_grad(): + model.weight.zero_() + model.bias.fill_(20) + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + likelihood = BernoulliLikelihood() + inputs = torch.ones((2, 1)) + target = torch.zeros((2, 1, 1)) + before = likelihood.compute_loss( + model(inputs).reshape(2, 1, 1, 1), target, None + ) + optimizer.zero_grad() + before.backward() + optimizer.step() + after = likelihood.compute_loss(model(inputs).reshape(2, 1, 1, 1), target, None) + assert after < before + + class TestTorchLikelihoodInputValidation: def test_gaussian_negative_prior_mu(self): with pytest.raises(ValueError, match="strictly positive"): diff --git a/darts/utils/likelihood_models/torch.py b/darts/utils/likelihood_models/torch.py index c180fb380d..7aa8388beb 100644 --- a/darts/utils/likelihood_models/torch.py +++ b/darts/utils/likelihood_models/torch.py @@ -485,6 +485,35 @@ def __init__(self, prior_p=None, prior_strength=1.0): prior_strength=prior_strength, ) + def compute_loss( + self, + model_output: torch.Tensor, + target: torch.Tensor, + sample_weight: torch.Tensor, + ) -> torch.Tensor: + """Compute the NLL and optional prior KL without rounding probabilities.""" + # Keep logits for the loss: converting to probabilities first can round + # confident predictions to 0 or 1 and lose their training gradients. + out_distr = _Bernoulli(logits=model_output.squeeze(dim=-1)) + loss = -out_distr.log_prob(target) + if sample_weight is not None: + loss = loss * sample_weight + loss = loss.mean() + if self.prior_p is not None: + prior_p = torch.tensor(self.prior_p).to(model_output.device) + dtype = torch.promote_types(model_output.dtype, prior_p.dtype) + prior_p = prior_p.to(dtype=dtype) + prior_distr = _Bernoulli(probs=prior_p) + # KL(prior || output) is cross-entropy minus prior entropy. Keep + # logits here too, since output probabilities can round to 0 or 1. + cross_entropy = nn.functional.binary_cross_entropy_with_logits( + out_distr.logits.to(dtype=dtype), + prior_p.expand_as(out_distr.logits), + reduction="none", + ) + loss += self.prior_strength * (cross_entropy - prior_distr.entropy()).mean() + return loss + @property def _prior_params(self) -> tuple[Any, ...] | None: return (self.prior_p,)