How does JAX PRNG improve on Mersenne Twister beyond splitting and explicit key?

[EngineerKhan]

I have gone through the relevant docs, and although it does answer how JAX's PRNG address it by explicit PRNG and threefry counter to parallelize, I would also like to know how does JAX improve it on the other issues discussed here. Copied below for the ease:

Mersenne Twister

• Fails some statistical tests, with as few as 45,000 numbers. Fails LinearComp Test of the TestU01 Crush and BigCrush batteries.

• Predictable — after 624 outputs, we can completely predict its output.

• Generator state occupies 2504 bytes of RAM — in contrast, an extremely usable generator with a huger-than-anyone-can-ever-use period can fit in 8 bytes of RAM.

• Not particularly fast.

• Not particularly space efficient. The generator uses 20000 bits to store its internal state (20032 bits on 64-bit machines), but has a period of only 219937, a factor of 263 (or 295) fewer than an ideal generator of the same size.

• Uneven in its output; the generator can get into “bad states” that are slow to recover from.

• Seedings that only differ slightly take a long time to diverge from each other; seeding must be done carefully to avoid bad states.

• While jump-ahead is possible, algorithms to do so are slow to compute (i.e., require several seconds) and rarely provided by implementations.

[jakevdp]

For a full discussion of the considerations that went into choosing the ThreeFry RNG, you can see the relevant design doc, as well as the paper that JAX's current PRNG implementation is based on: http://www.thesalmons.org/john/random123/papers/random123sc11.pdf

A couple points on how it improves over the Mersenne Twister within the context that JAX is designed for:

1. Space efficiency: JAX's current implementation of ThreeFry uses a PRNG Key of size 64 bits, while the Mersenne Twister requires over 20000 (you can see this, e.g., by printing np.random.get_state())
2. Parallelism: threefry is well-suited to massively parallel architectures (the MT model of a global state makes this challenging)
3. JAX compatibility: since the state is just a numerical array, threefry is well-suited to representation in XLA and compatible with JAX's existing array-based transforms for batching, AD, and others.

Does that answer your question?

[EngineerKhan]

Almost it does. I did skim that paper but will read in detail now. Does JAX implement the whole paper or skips some features?

[jakevdp]

The paper discusses several classes of counter-based PRNGs; JAX implements ThreeFry 2x32, and there is work underway to extend this (see #7676 and related issues).

[EngineerKhan]

A trivial point regarding point 3 you mentioned above: Does state in the Mersenne Twister (like NumPy implementation) also maintained in the form of a vector? I can understand Threefry's vector coming in handy due to JAX's XLA, but just want to clear this little doubt.

As JAX docs read: "This pseudorandom state vector is automagically updated behind the scenes every time a random number is needed, “consuming” 2 of the uint32s in the Mersenne twister state vector"

[jakevdp]

You can see NumPy's Mersenne Twister state by printing np.random.get_state(), or Python's Merseene Twister state by printing random.getstate(). It contains a vector of integers, as well as some additional metadata.