NAF reported log-likelihood values are excesively high
Description
Now that I am revisiting my experience with Zuko, I have remembered that we did not use NAF in our work since, during testing, we highly suspected that there was a bug in NAF and did not feel confident using it.
This is how we discovered it (back in Zuko 0.2.0, maybe it is fixed by now): we had different synthetic experiments for which we knew the ground-truth data-generating process, and we could compute the real log-likelihood of the data. When we fitted a NAF model using data from this model, the reported log-likelihoods reported by NAF were orders of magnitude larger than the ground-truth (which cannot happen since the KL is always non-negative). This behaviour was not observed with MAF or NSF, where they reported values slightly lower than the true log-likelihoods.
(I will update the issue when I find time, but I guessed it is useful to raise the issue despite being incomplete.)
Reproduce
TBD, but it can be tested by substituting MAF by NAF in any of the synthetic experiments of this repository.
Expected behavior
The log-likelihoods computed by the flow should be lower (given enough data) than that of the model that generated the data.
Causes and solution
I am uncertain of what could be the reason, but it should be somewhere in the log-likelihood computation.
Environment
- Zuko version: 0.2.0
- PyTorch version: 1.12.0
- Python version: 3.8
- OS: Ubuntu 22.04
Hello @adrianjav, thank you for the bug report.
NAF and UNAF are prone to numerical collapses. Notably, it is possible that the monotonic network learns a function that is not striclty monotonic, which would break almost everything... A common solution is to use weight decay during training.
That said, this could also be an implementation bug, which would definitely require attention. The implementation of NAF (and UNAF) has changed a lot since Zuko 0.2.0, so it would be great if you could provide a minimal (non-)working example with the latest version.
Yes, that sounds great. I will try to reproduce it again while working in my new project and update the issue.
As I said, I just thought of leaving this here so that I don't forget :)
On Tue, Jul 23, 2024, 11:41 AM François Rozet @.***> wrote:
Hello @adrianjav https://github.com/adrianjav, thank you for the bug report.
NAF and UNAF are prone to numerical collapses. Notably, it is possible that the monotonic network learns a function that is not striclty monotonic, which would break almost everything... A common solution is to use weight decay during training, but that does not fixes everything.
That said, this could also be an implementation bug, which would definitely require attention. The implementation of NAF (and UNAF) has changed a lot since Zuko 0.2.0, so it would be great if you could provide a minimal (non-)working example with the latest version.
— Reply to this email directly, view it on GitHub https://github.com/probabilists/zuko/issues/55#issuecomment-2244736296, or unsubscribe https://github.com/notifications/unsubscribe-auth/AGEFZJY4WSGKIZHEBSBCNATZNYQLLAVCNFSM6AAAAABLJ7CT66VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENBUG4ZTMMRZGY . You are receiving this because you were mentioned.Message ID: @.***>
I had run some experiments with the new version of NAF, and I am not able to reproduce the results. Now the fit looks quite reasonable, so I guess I will close the issue for now. Sorry for the bother!
Thank you for digging in!