JuliaStats
Better samplers for Bernoulli(p) and Geometric(1/2) random variables
The implementation for generating a Bernoulli(p) random variable right now is rand() <= p. This is totally fine, but there are better ways to handle rational probabilities, and there is a method detailed here which uses less entropy (on expectation, 2 bits).
For rationals the obvious solution is rand(1:denominator) <= numerator, which is faster on my machine, possibly due to the expense of converting the rational prob to a float for comparison.
I also have code for the other sampler, but on my machine it's about 2x as slow as the current implementation, still pretty fast, 6.4ns vs the old method 2.7ns.
In writing that code I needed to sample a Geometric(1/2) distribution, and found that the current implementation is needlessly expensive for this default case. It is faster to generate random bits (sizeof(Int) at a time) until you see a 1, this could go down to generating random Int8's as well if minimising entropy usage was a concern (but the default sampler generates 64 bits at a time and then truncates anyways)
I am interested in implementing these, but want to know whether it is of interest. I also wouldn't be sure how to implement a sampler for a "default" geometric rv that doesn't meddle with the current implementation either.
Posting here to get thoughts, I think the method for rational Bernoulli's should absolutely be implemented as a strict improvement, but otherwise I want to know what people think is important.
Not a maintainer, so don't take my position as canonical.
There's a lot of optimization that can be done with the rand and rand! functions.
Personally, I would also recommend defining explicit rand! methods alongside scalar rand methods. Very few distributions currently do this, but it can yield significant speedups when implemented.
The core reason behind the performance increases is because Random.jl has SIMD implementations for rand! and randn! for both Xoshiro and MersenneTwister. Calling rand(1) 100 times will be significantly slower than calling rand(100) once.
Don't worry too much if you replace an existing method with something faster that yields different results (but is still within expectations for the distribution). Distributions.jl doesn't guarantee that rand(::AbstractRNG, ::Distribution) will yield the same result across versions with the same seed, and neither does Random.jl.
I'm curious about the other method for sampling now; I'm tempted to try an implementation of my own to see if I can get it faster than what we currently have.
Really like the suggestion to implement rand!, I have personally had slow-downs from generating lots of Bernoulli samples in my own work, I'll have to look at what I can do to make this work, but sounds fun.
As far as my own implementation the details are as follows,
benchmarks for my version on a Mac M1 laptop:
Which shows comparable performance (min time), but worse average time (does a lot more operations, which is pretty apparent looking at the built code)
Where this is the code I used:
function my_bernoulli_abbreviated(d::Bernoulli{T})::Bool where T <: Union{Float16, Float32, Float64}
isone(d.p) && return true
significand_bits::UInt = precision(T) - 1
max_coin_flips::UInt = Base.exponent_max(T) + significand_bits
leading_zs::UInt = Base.exponent_max(T) - 1 - Base.unbiased_exponent(d.p)
first_heads_index::UInt = begin
coin_flips::UInt = typemax(UInt)
chunk_size::UInt = 8 * sizeof(UInt)
num_chunks::UInt = cld(max_coin_flips, chunk_size)
for i in 1:num_chunks
geo = leading_zeros(rand(UInt))
if geo != chunk_size
coin_flips = (i-1)*chunk_size + geo
break
end
end
coin_flips
end
first_heads_index == typemax(UInt) && return false # Truncated geometric
first_heads_index < leading_zs && return false # Cannot be 1
first_heads_index == leading_zs && return !iszero(reinterpret(Unsigned, d.p) & Base.exponent_mask(T)) # implicitly 1, for normal floats
first_heads_index > leading_zs + significand_bits && return false # beyond the 1's, implicitly 0
!iszero(reinterpret(Unsigned, d.p) & (1 << (leading_zs + significand_bits - first_heads_index))) # index the i'th value of the significand
end
If you'd like to see a documented version I can point you to this well written implementation which I used as a reference https://docs.rs/opendp/latest/src/opendp/traits/samplers/bernoulli/mod.rs.html#25-137
I found this post which promises to be pretty efficient, but I don't understand the method right now https://discourse.julialang.org/t/faster-bernoulli-sampling/35209
