Consider Fast Loaded Dice Roller for `WeightedIndex`

[vks]

It seems to outperform the alias method (and others), being close to the theoretical optimum of entropy use. According to the paper, it is faster than the alias method (if the latter is modified to produce exact samples), while using several orders of magnitudes less preprocessing time. Maybe it can replace our current algorithms in Rand?

There is a reference implementation in C and Python.

[vks]

I played a bit with the algorithm, but I could not make it perform as well as our alias-method implementation.

The algorithm consumes the random numbers bit-by-bit. On the one hand, this avoids throwing away random data, on the other hand, the approach does not buy much, because generating random bits is cheap with modern generators. I suspect the results in the paper are with a slow RNG like the Mersenne Twister.

There is also a similar Fast Dice Roller algorithm for uniform sampling, which would also save a lot of random bits for small ranges. However, in practice it seems to be much slower with a fast RNG.

[peteroupc]

Vks wrote:

I played a bit with the algorithm, but I could not make it perform as well as our alias-method implementation.

The algorithm consumes the random numbers bit-by-bit. On the one hand, this avoids throwing away random data, on the other hand, the approach does not buy much, because generating random bits is cheap with modern generators. I suspect the results in the paper are with a slow RNG like the Mersenne Twister.

There is also a similar Fast Dice Roller algorithm for uniform sampling, which would also save a lot of random bits for small ranges. However, in practice it seems to be much slower with a fast RNG.

@fsaad, can you please comment?

[fsaad]

Thanks for trying out FLDR.

I played a bit with the algorithm, but I could not make it perform as well as our alias-method implementation.

@vks Does your alias implementation use exact sampling or floating-point approximations of real numbers? The runtime comparison, both theoretical and practical, between FLDR and alias sampling is with respect to error-free sampling. If the alias implementation uses floating-point approximations of real numbers then there is not much that can analyzed in terms of performance with FLDR, since an inexact alias implementation typically trade-offs error with runtime, whereas FLDR is designed to always be error free.

If you have an error-free alias sampler, I would be really interested to explore its performance as compared to FLDR, and I would be quite surprised if FLDR is slower (especially on low to medium entropy distributions), both for preprocessing and sampling.

The algorithm consumes the random numbers bit-by-bit. On the one hand, this avoids throwing away random data, on the other hand, the approach does not buy much, because generating random bits is cheap with modern generators.

While the FLDR paper uses a simple (naive) serial implementation of the sampler (Alg 5), there are many other ways to implement that leverage alternate data structures, which include linear encodings of the DDG sampling tree (Algs 4--7 Saad et al. 2020), bit-slicing for constant time generation (Sec 3, Karmakar et al. 2018), and batch generating with bit recycling (Alg 3, Devroye & Gravel 2018), among many others.

The choice of implementation depends on the constraints of the application. For example, lazy DDG sampling is typically far more efficient as compared to floating-point arithmetic when implementing the sampler directly on hardware (e.g., an FPGA). While FLDR is characterized by its theoretical near-optimal entropy use, many of these techniques can be used in practice to push down the constant factors of runtime (or memory).

There is also a similar Fast Dice Roller algorithm for uniform sampling, which would also save a lot of random bits for small ranges. However, in practice it seems to be much slower with a fast RNG.

The FDR of Lumbroso 2013 is an exact uniform sampler. Unlike FLDR, the FDR is theoretically optimal (for uniform sampling), so it should be superior to any other exact uniform sampler, though I agree that non-exact samplers can be faster in practice.

Overall, I think it is very valuable to compare the performance of exact and approximate sampling, as your ticket may point toward. The status quo of most existing software libraries is inexact sampling with floating-point. FLDR is an initial step toward showing that exact sampling can be practical, both theoretically and for applications, for general discrete distributions. There is of course a lot more work to do to explore the trade-offs with inexact samplers.

[dhardy]

The code for our weighted alias implementation is here. IIRC it uses exact sampling, with the requirement that the sum of weights can be represented by the weight type. Weight type may be an integer or floating-point or a user-defined type; the latter implies that rounding can occur, depending on the user's choice of types.

In many benchmarks however I found the alias method a worse choice than the binary search method due to the set-up time, hence the latter is our default (and also exact if used with integer types).

[vks]

Does your alias implementation use exact sampling or floating-point approximations of real numbers?

I don't think we use floating-point when the weights are integers. I did not look into the details, but I doubt we use the alias method documented in the paper:

(vi) exact alias sampler (Walker, 1977), using entropy-optimal uniform and Bernoulli sampling (Lumbroso, 2013) and the one-table implementation (Vose, 1991).

In particular, we are not using the FDR for uniform sampling, which would be twice as slow.

Here are the benchmark results:

weighted_u32              :      12,472 ns/iter (+/- 308)
weighted_large_set        :      72,129 ns/iter (+/- 2,319)
weighted_new              :      39,283 ns/iter (+/- 876)

weighted_alias_u32        :      11,029 ns/iter (+/- 327)
weighted_alias_large_set  :      12,318 ns/iter (+/- 366)
weighted_alias_new        :     261,343 ns/iter (+/- 7,985)

weighted_fldr_i32         :      17,420 ns/iter (+/- 365)
weighted_fldr_large_set   :      29,980 ns/iter (+/- 671)
weighted_fldr_new         :     935,315 ns/iter (+/- 26,667)

I also implemented the FDR uniform sampling. It's a bit awkward to implement with our API, so I'm not sure it's optimized as well as it could be.

uniform     :       2,032 ns/iter (+/- 119) = 492 MB/s
uniform_fdr :       4,373 ns/iter (+/- 212) = 228 MB/s

The code is here.

[peteroupc]

For your information there are numerous data structures and algorithms for weighted sampling with replacement, besides the Fast Loaded Dice Roller, binary search, and the alias method, and I list a great number of them. Perhaps it would be a good idea to compare their performance.

[fsaad]

@vks Thanks a lot for the code. I'm not an expert in Rust, but it seems that there are many ways that this FLDR implementation can be optimized. One immediate thought is that, in your current implementation, the PRNG is always called once per sample, i.e.,

rand/rand_distr/src/weighted_fldr.rs

Line 85 in e2b4b12

let r = rng.next_u64();

In our reference implementation, we separately store a "global" buffer of 31 random bits (in C, rand() produces a positive 32 bit signed integer, so only 31 of the bits are random), which is the function "flip()" in the following file:

https://github.com/probcomp/fast-loaded-dice-roller/blob/master/src/c/flip.c

The main FLDR sampler then calls "flip()" to obtain a fresh bit:

https://github.com/probcomp/fast-loaded-dice-roller/blob/2fb32e94a04c744c86aa4cb54b3ee0259894e777/src/c/fldr.c#L93

which may or not may not involve a call to the PRNG, depending on whether the buffer is empty. This approach ensures that we do not spuriously generate random bits. If FLDR is implemented to use a separate buffer that is shared across calls to "sample", instead of calling the PRNG once per sample, then I do expect the runtime should improve in many regimes.

[vks]

@fsaad Thanks for having a look!

I'm not an expert in Rust, but it seems that there are many ways that this FLDR implementation can be optimized. One immediate thought is that, in your current implementation, the PRNG is always called once per sample, i.e.,

Fair enough, that is a very good point! Unfortunately, implementing this is a bit tricky, because our API requires sample to not modify the distribution object. There are the following solutions:

1. Add some API for sampling with modifying the distribution object.
2. Use Cell. This means that the distribution object can no longer be shared by threads.

Going with 2. seems to have an adverse effects on the performance for large sets, while improving it for small sets:

weighted_fldr_i32         :      12,733 ns/iter (+/- 53)
weighted_fldr_large_set   :      36,303 ns/iter (+/- 536)
weighted_fldr_new         :     911,047 ns/iter (+/- 17,077)

[newpavlov]

Personally I don't like option 2. Can we implement Distribution<T> only for &mut reference of a type?

[dhardy]

Distribution::sample is immutable with respect to the distribution for good reasons: (1) API simplicity, (2) basically all distributions have fast (enough) immutable implementations and (3) because it places fewer restrictions on usage.

It's going to take quite a good reason to convince me that we should change that. Especially since the benchmarks so far have the alias-method winning everywhere (especially in set-up, where I already thought the alias method was slow).

[fsaad]

Distribution::sample is immutable with respect to the distribution for good reasons

I agree. Does your API For Rng support efficient and lazy production of random bits one at a time, for example by storing an internal buffer (using a similar method to the "global" buffer approach)? That way, sample can leverage mutability with respect to rng.

the benchmarks so far have the alias-method winning everywhere (especially in set-up, where I already thought the alias method was slow).

Do the benchmarks use test distributions with varying entropy, i.e., for a distribution with n outcomes, entropies from 0 to log_2(n)?

The average runtime of sampling in FLDR is a direct function of the entropy. The runtime of preprocessing also varies with the values of probabilities in the test distribution. The point about set-up is surprising. Do you have a pointer to the benchmarks and the measurements of setup runtime?

[dhardy]

Does your API For Rng support efficient and lazy production of random bits one at a time

Good question. No. Theoretically it could, even using different approaches depending on the RNG, but we chose not to do that. For simpler PRNGs we use no buffer at all and for block-based PRNGs we use a block which we read off one word (32-bit) at a time.

Another factor to consider (for both mutable distribution and buffered RNGs) is reproducible behaviour. If we allowed out-of-order usage of random bits/bytes it becomes much harder to reason about and especially document the behaviour, thus making reproducibility hard. If we don't allow this, then every read has to either be compatible with misaligned buffers (slow) or check and clear up to the next alignment point (also slow). At this point we're no longer optimising one RNG+distribution pair but expected usage of the RNG across a number of distributions.

I suppose we could side-step that problem by having a special RNG optimised for bit- or byte-sized reads, though that's another step in complexity.

@vks which RNG did you use for these benchmarks?

[vks]

Another factor to consider (for both mutable distribution and buffered RNGs) is reproducible behaviour. If we allowed out-of-order usage of random bits/bytes it becomes much harder to reason about and especially document the behaviour, thus making reproducibility hard.

Aren't we having this problem already when mixing fill_bytes and next_u64?

I suppose we could side-step that problem by having a special RNG optimised for bit- or byte-sized reads, though that's another step in complexity.

In principle, we could have a wrapper type implementing the buffering. It can then be used as needed.

For example, we could just add a bit: u8 field to BlockRng and a gen_bool method that uses it.

@vks which RNG did you use for these benchmarks?

It's just our default benchmarks, they use Pcg64Mcg.

Do you have a pointer to the benchmarks and the measurements of setup runtime?

See here.

[newpavlov]

For example, we could just add a bit: u8 field to BlockRng and a gen_bool method that uses it.

I think we can use a single counter for used bits, which would be incremented by 8 for read byte, 32 for u32, etc. We would need to perform additional computations to correctly align read point (e.g. if index is equal to 9, we would have to align it to 16 for reading byte and to 64 for reading u32). It would degrade performance a bit, but I am not sure if it will be noticeable.

As an advantage this approach will be less wasteful (and thus more performant) for users using primarily fill_bytes. But it would change behavior of BlockRng between releases, which may create problems for users who rely on reproducibility.