This PR makes several needed fixes to augmented-softmax:
The implementation incorrectly constrained logr to be positive.
The Jacobian determinant in the paper was off by a factor of r and is now corrected
The default prior is now better motivated. If y is IID std normal distributed, then logr is approximately normally distributed with a mean in the range (log(N), log(N) + 1/2) and standard deviation in the range (0, 1). We choose the default prior of logr ~ Normal(log(N), 1) to shift the distribution of y to be closer to the origin when x is uniform on the simplex.
I expect these improvements will improve the performance of augmented-softmax in the comparisons.
This PR makes several needed fixes to augmented-softmax:
logrto be positive.rand is now correctedyis IID std normal distributed, thenlogris approximately normally distributed with a mean in the range(log(N), log(N) + 1/2)and standard deviation in the range(0, 1). We choose the default prior oflogr ~ Normal(log(N), 1)to shift the distribution ofyto be closer to the origin whenxis uniform on the simplex.I expect these improvements will improve the performance of augmented-softmax in the comparisons.