伪随机数生成

随机变量

考虑一个函数： $f(x) = x^2 +5$

无论我们调用这个函数多少次，对于相同的输入，它始终会给出相同的输出。

另一方面，考虑一些随机/随机性函数：比如第七次掷骰子的结果、下午3点的气温，或某个粒子的量子态。这些都无法保证每次运行都得到相同的结果。

这类随机函数的输出结果被称为随机变量。几乎很难找到一个完全不依赖随机变量的领域。因此，随机变量在建模几乎所有现实世界系统中都扮演着核心角色。

伪随机数生成（PRNG）

计算机相对于手工操作的一大优势在于其确定性行为：如果我们以相同的条件运行某个算法，它总会产生相同的结果。

然而，这种优势有时也可能变成劣势。计算机科学家早已意识到这一特性，因此设计出一类称为伪随机数生成（Pseudo-Random Number Generation, PRNG） 的算法来模拟随机性。

提出最早期PRNG算法之一的约翰·冯·诺依曼曾有名言警告道：

“任何试图用算术方法生成随机数字的人，当然都处于罪恶之中。” $-$ （1951）

梅森旋转算法（Mersenne Twister）

目前最广泛使用的PRNG之一是梅森旋转算法（Mersenne Twister），其名称来源于所使用的梅森素数，被广泛应用于C++、R、CUDA、NumPy等语言和库中。

尽管应用广泛，梅森旋转算法仍存在一些问题。

Threefry计数器

基于计数器的PRNG（我们即将介绍）旨在解决梅森旋转算法及其他传统PRNG的缺陷。
2011年，D.E. Shaw Research提出了Threefry计数器，该算法正是JAX中PRNG实现的基础。

接下来我们将重点介绍JAX的PRNG实现（后文简称 JAX PRNG）的主要特性。

---

JAX PRNG的优势

很自然地，用户会关心JAX PRNG是如何克服梅森旋转算法的不足之处的。

有状态与无状态生成

JAX PRNG的一个关键优势在于采用了无状态生成机制。对于还不熟悉这些概念的读者，以下是一个简要说明。

有状态生成

每个PRNG都会使用一个状态向量来表示其当前状态，并利用该状态生成下一个随机样本。

梅森旋转算法基于有状态生成，这意味着它具有预定义的状态，每次生成一个随机数后，状态向量会自动更新。

下面的例子将详细说明这一点。我们先回顾几个NumPy随机函数：

• seed(<int>)：初始化PRNG
• get_state()：获取当前状态向量（返回值为包含其他信息的元组）
• uniform()：用于采样一个随机数。我们将在后续课程中深入讨论此函数。

1import numpy as np
2import pandas as pd
3
4np.random.seed(0) 
5
6Twister_State = np.random.get_state()
7print(type(Twister_State))
8print(pd.DataFrame(Twister_State))
9
10# 采样若干数值以观察对状态的影响
11np.random.uniform()
12Twister_State = np.random.get_state()
13print(pd.DataFrame(Twister_State)[1:3]) # 仅显示相关信息
14
15np.random.uniform()
16Twister_State = np.random.get_state()
17print(pd.DataFrame(Twister_State)[1:3])
18
19for i in range(1,101):              # 进一步采样
20    print(np.random.uniform())      # 每次输出值均不同
21
22print("采样2+100次后的状态信息")
23Twister_State = np.random.get_state()
24print(pd.DataFrame(Twister_State)[1:3]) # 如预期所示，步进为204

可以看出，每次采样或生成随机数都会消耗两个状态 —— 即从状态向量中取出两个 uint32 值。

这里的主要问题是：用户对此状态几乎无法安全干预。如果需要手动修改状态，完全由用户自行承担风险。NumPy官方文档明确指出：

如果内部状态被手动更改，用户必须确切知道自己在做什么。

无状态生成

相比之下，Threefry 或任何基于计数器的方法采用了一种更简单且显式的密钥生成机制。无论我们对随机变量采样多少次，其状态始终保持不变 —— 换句话说，它是无状态的。

1import jax.random as random
2
3Threefry_Key = random.PRNGKey(0)
4print(type(Threefry_Key)) # 结构更简洁，体积更小
5print(Threefry_Key)
6
7for i in range(1,101):
8    print(random.uniform(Threefry_Key,(1,1))) # 每次迭代输出值相同
9
10random.uniform(Threefry_Key,(1,1))
11print(Threefry_Key) # 密钥值不会改变，除非显式更新

如上所示，除非我们主动更新密钥，否则输出将不具备随机性。有些人可能已经对引入显式密钥机制的必要性产生怀疑，但请稍等！下一节将解释为何JAX PRNG优于传统方法。

Threefry基于2008年提出的Threefish哈希函数，是Threefish的一个简化版本，“fry”（意为幼鱼）这一名称正源于此。

并行化能力

我们已经看到，传统的PRNG是顺序性的：下一个随机数通过如下简单变换得到：

$s_{n+1} = f(s_n)$

因此，尝试并行生成随机数可能导致混乱。

这个问题常常给数据科学家或工程师带来糟糕的体验，而Threefry的显式生成机制正好解决了这一痛点。

在JAX PRNG或其他基于计数器的PRNG中，我们使用分叉（splitting） 技术，通过生成子密钥（subkeys） 来实现并行化。

下面我们用分叉功能重新实现上述例子。只需使用 split() 函数，该函数接收一个密钥并返回主密钥和子密钥。

1Threefry_Key = random.PRNGKey(0)
2print("原始密钥为：", Threefry_Key)
3
4Threefry_Key, Threefry_Subkey = random.split(Threefry_Key)
5print("分叉后的密钥：", Threefry_Key)
6print("子密钥值：", Threefry_Subkey)

再看一个循环示例：

1Threefry_Key = random.PRNGKey(0)
2
3for i in range(1,101):
4    Threefry_Key, subkey = random.split(Threefry_Key)
5    print(subkey)
6    print(random.uniform(subkey,(1,1)))

整个过程的关键在于：JAX不仅支持可重现的随机数生成，还确保具备以下优势：

• 更小的状态向量，有助于快速初始化PRNG；
• 支持无冲突的并行化，从而能够良好地配合XLA变换使用。