`dist/categorical` sometimes produces out-of-bounds error when total weight is very small

[bzinberg]

Summary

The following expression sometimes produces an out of bounds error during evaluation:

(dist/categorical '(1.0E-323 0.0))

@zane and I looked into this, and it seems to be that there is an underflow-like issue on this line: https://github.com/probcomp/metaprob/blob/a399f820ef3e9a324d50f03a5fdee30dd55389c2/src/metaprob/distributions.cljc#L39-L42

When sum is very small, the < condition will never be met, so the loop will run too long. This phenomenon is illustrated by this example:

user=> (def eps 1.0E-323)
#'user/eps
user=> (< (* 0.9 eps) eps)
false

Fix

The best fix is to never convert log-probabilities ("scores") into probabilities ("weights"), as this is a lossy conversion. Instead, pass the log-probabilities directly to log-categorical. Inside the body of log-categorical, the scores are normalized in log space before being exponentiated, which avoids the underflow issue.

The plain (non-log) categorical distribution should only be used when higher-precision log-probabilities are not available. Still, there could be cases in which there is underflow and log-probabilities are not available. We can't do much to get "correct" behavior here, but we can at least have the program not crash. As of https://github.com/probcomp/metaprob/pull/159 (merged after this bug was created), the crash is gone; in cases where the pre-https://github.com/probcomp/metaprob/pull/159 looping logic would have produced an out of bounds error, we now instead return n (the index of the last element of probs).

[bzinberg]

• In log-categorical, adjust the weights in log space to prevent underflow. Probably a decent way to do this is by subtracting (max scores) from each entry of scores (here scores are log-probabilities)

Oh wait, the code already does that: https://github.com/probcomp/metaprob/blob/0c2839632890dd425797fc30b19afd42c445b796/src/metaprob/distributions.cljc#L56-L58