Add neural score estimation (NSE) #7
Conversation
| phi=(x, *self.parameters()), | ||
| ), | ||
| base=DiagNormal(self.zeros, self.ones).expand(batch_shape), | ||
| ) |
There was a problem hiding this comment.
I am not sure if there is any definite API yet but I find a bit inconsistent that NPE defines a sample method based on its own self.flow and here there is no sample method implemented.
There was a problem hiding this comment.
Yes, I dislike it as well, but I am not sure of the best way to sample from the trained estimator yet. From my toy experiments sampling from the probability flow ODE gives good samples and is pretty fast, but Song et al. (2021) put forward some evidence that ancestral and reverse SDE sampling yields more qualitative samples.
So maybe the sample method could be another sampling method than just return self.flow(x).sample()?
There was a problem hiding this comment.
Yeah, I was more worried about the lack of consistency between the interfaces of the classes than about the particular implementation of the sample interface.
|
Most comments have been addressed in the last commits. The description of NSE was mostly re-written. |
| phi=(x, *self.parameters()), | ||
| ), | ||
| base=DiagNormal(self.zeros, self.ones).expand(batch_shape), | ||
| ) |
There was a problem hiding this comment.
After experimenting with this part of the code, it would be nice to have a unified API such that nse.flow can be passed directly to the coverage utils. This is currently not convenient because nse.flow.log_prob is a noisy estimate of the log-density.
|
Thanks for the changes! I have experimented with the codebase and confirm that it works fine -- it produces good results on SLCP, and the coverage is reasonable. +1 for merging in its current state |
Targets #6