Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use approximate log_beta() in .fit(), .predict() #2502

Merged
merged 4 commits into from
May 25, 2020
Merged

Conversation

fritzo
Copy link
Member

@fritzo fritzo commented May 24, 2020

Addresses #2426
Blocking #2498 (from which this simpler PR was extracted)

This PR:

  1. Fixes a nan gradient issue in ExtendedBetaBinomial related to nan propagates through backward pass even when not accessed pytorch/pytorch#15506. This may have caused spurious rejections during heuristic initialization.
  2. Starts using the fast log_beta() in CompartmentalModel.fit() and .predict() now that gradients are safe. This reduces likelihood compute cost by about 14% in SuperspredingSEIRModel, and will reduce Add overdispersed models to contrib.epidemiology #2498 more.
  3. Moves safe_log() to pyro.ops.special which seems like a better home. (this change is non-breaking since safe_log() has not yet been released).

Tested

  • added unit tests for nan gradients in ExtendedBinomial and ExtendedBetaBinomial
  • ran locally to verify sane behavior

Comment on lines +48 to +49
n = total_count.clamp(min=0)
k = value.masked_fill(invalid, 0)
Copy link
Member Author

@fritzo fritzo May 24, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is the crux of the NAN gradient fix

@@ -23,7 +23,7 @@ class ExtendedBinomial(Binomial):

def log_prob(self, value):
result = super().log_prob(value)
invalid = ~super().support.check(value)
invalid = (value < 0) | (value > self.total_count)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

curous: why this change?

Copy link
Member Author

@fritzo fritzo May 25, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new version is a little cheaper. The old integer_interval version additionally checks value % 1 == 0, but that is already checked in validation by the above line, and it doesn't affect numerical stability.

@fritzo
Copy link
Member Author

fritzo commented May 25, 2020

Thanks for reviewing!

@martinjankowiak martinjankowiak merged commit aa40beb into dev May 25, 2020
@fritzo fritzo deleted the sir-approx branch June 5, 2020 15:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants