Skip to content

Keep Bernoulli training loss and prior KL in logit space - #3305

Open
AHMETHAKANBEZIR1 wants to merge 1 commit into
unit8co:masterfrom
AHMETHAKANBEZIR1:fix/bernoulli-logit-loss
Open

AHMETHAKANBEZIR1 wants to merge 1 commit into
unit8co:masterfrom
AHMETHAKANBEZIR1:fix/bernoulli-logit-loss

Conversation

@AHMETHAKANBEZIR1

Copy link
Copy Markdown

Confident, incorrect Bernoulli predictions can stop learning because sigmoid rounds the network output to a probability before the loss converts it back to logits. For float32 logits [-40, 40] and targets [1, 0], the current loss is capped near 15.94 and both gradients are zero; the Bernoulli NLL should be 40 with nonzero gradients. A linear model initialized with the wrong positive logit also fails to update.

Related to #621. This addresses the Bernoulli case only, not every likelihood discussed in that issue.

Override BernoulliLikelihood.compute_loss to retain the network logits for the NLL. Compute the optional KL(prior || prediction) as logit cross-entropy minus prior entropy, so it also stays finite when the predicted probability rounds to zero or one. Sample weighting still applies to the NLL, not the prior term. Sampling, probability parameter outputs, and the existing parameter-extraction methods are unchanged.

Validation on Windows CPU, Python 3.12.14, PyTorch 2.14.1+cpu:

  • New regression coverage: 9 failures and 2 passing ordinary-value controls on unchanged master; all 11 pass after the fix. It checks float32/float64 losses and analytic gradients, weights, priors, and an optimizer step from a confident wrong prediction.
  • Entire likelihood test directory: 60 passed. The Torch likelihood module alone has 57 passing tests; all 14 executable statements of the added method are covered.
  • Existing Bernoulli RNN training/forecast and likelihood-parameter model tests: 2 passed, 7 dependency/runtime warnings.
  • Separate float16/bfloat16/float32/float64 checks confirm finite loss, analytic gradients with component-wise priors and weights, and unchanged probability outputs.
  • All configured pre-commit hooks pass. Full Darts tests, GPU/distributed execution and performance were not tested.

AI assistance: implementation and tests were prepared and executed autonomously with OpenAI Codex. No independent human review or personal production incident is claimed.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant