伪随机数

伪随机数#

在本节中，我们专注于 jax.random 和伪随机数生成（PRNG）；即，通过算法生成数字序列的过程，这些序列的属性近似于从适当分布中采样的随机数序列的属性。

PRNG 生成的序列并不是真正的随机，因为它们实际上是由它们的初始值决定的，这个初始值通常被称为 种子，并且每次随机抽样都是某个 状态 的确定性函数，这个状态从一次抽样传递到下一次。

伪随机数生成是任何机器学习或科学计算框架中的一个重要组成部分。通常，JAX 力求与 NumPy 兼容，但伪随机数生成是一个显著的例外。

为了更好地理解JAX和NumPy在随机数生成方面的方法差异，我们将在本节中讨论这两种方法。

NumPy 中的随机数#

伪随机数生成在 NumPy 中由 numpy.random 模块原生支持。在 NumPy 中，伪随机数生成基于一个全局 状态，可以使用 numpy.random.seed() 将其设置为确定性的初始条件。

import numpy as np
np.random.seed(0)

您可以使用以下命令检查状态的内容。

def print_truncated_random_state():
  """To avoid spamming the outputs, print only part of the state."""
  full_random_state = np.random.get_state()
  print(str(full_random_state)[:460], '...')

print_truncated_random_state()

('MT19937', array([         0,          1, 1812433255, 1900727105, 1208447044,
       2481403966, 4042607538,  337614300, 3232553940, 1018809052,
       3202401494, 1775180719, 3192392114,  594215549,  184016991,
        829906058,  610491522, 3879932251, 3139825610,  297902587,
       4075895579, 2943625357, 3530655617, 1423771745, 2135928312,
       2891506774, 1066338622,  135451537,  933040465, 2759011858,
       2273819758, 3545703099, 2516396728, 127 ...

state 通过每次调用随机函数进行更新：

np.random.seed(0)
print_truncated_random_state()

('MT19937', array([         0,          1, 1812433255, 1900727105, 1208447044,
       2481403966, 4042607538,  337614300, 3232553940, 1018809052,
       3202401494, 1775180719, 3192392114,  594215549,  184016991,
        829906058,  610491522, 3879932251, 3139825610,  297902587,
       4075895579, 2943625357, 3530655617, 1423771745, 2135928312,
       2891506774, 1066338622,  135451537,  933040465, 2759011858,
       2273819758, 3545703099, 2516396728, 127 ...

_ = np.random.uniform()
print_truncated_random_state()

('MT19937', array([2443250962, 1093594115, 1878467924, 2709361018, 1101979660,
       3904844661,  676747479, 2085143622, 1056793272, 3812477442,
       2168787041,  275552121, 2696932952, 3432054210, 1657102335,
       3518946594,  962584079, 1051271004, 3806145045, 1414436097,
       2032348584, 1661738718, 1116708477, 2562755208, 3176189976,
        696824676, 2399811678, 3992505346,  569184356, 2626558620,
        136797809, 4273176064,  296167901, 343 ...

NumPy 允许你通过单个函数调用对单个数字或整个数字向量进行采样。例如，你可以通过以下方式从均匀分布中采样一个包含3个标量的向量：

np.random.seed(0)
print(np.random.uniform(size=3))

[0.5488135  0.71518937 0.60276338]

NumPy 提供了 顺序等价保证，这意味着逐个采样 N 个数字或采样一个包含 N 个数字的向量会产生相同的伪随机序列：

np.random.seed(0)
print("individually:", np.stack([np.random.uniform() for _ in range(3)]))

np.random.seed(0)
print("all at once: ", np.random.uniform(size=3))

individually: [0.5488135  0.71518937 0.60276338]
all at once:  [0.5488135  0.71518937 0.60276338]

JAX 中的随机数#

JAX 的随机数生成与 NumPy 的有重要区别，因为 NumPy 的 PRNG 设计使得同时保证多个理想特性变得困难。具体来说，在 JAX 中我们希望 PRNG 生成能够：

  1. 可复现的，

  2. 可并行化,

  3. 可向量化。

我们将在下面讨论原因。首先，我们将关注基于全局状态的PRNG设计的含义。考虑以下代码：

import numpy as np

np.random.seed(0)

def bar(): return np.random.uniform()
def baz(): return np.random.uniform()

def foo(): return bar() + 2 * baz()

print(foo())

1.9791922366721637

函数 foo 将从一个均匀分布中采样的两个标量相加。

如果我们假设 bar() 和 baz() 的执行顺序是可预测的，这段代码的输出只能满足需求 #1。在 NumPy 中，这不是一个问题，因为它总是按照 Python 解释器定义的顺序来执行代码。然而，在 JAX 中，这问题更大：为了高效执行，我们希望 JIT 编译器能够自由地重新排序、省略和融合我们定义的函数中的各种操作。此外，在多设备环境中执行时，每个进程需要同步全局状态，这会妨碍执行效率。

显式随机状态#

为了避免这个问题，JAX 避免了隐式的全局随机状态，而是通过一个随机的 key 显式地跟踪状态：

from jax import random

key = random.key(42)
print(key)

Array((), dtype=key<fry>) overlaying:
[ 0 42]

备注

本节使用由 jax.random.key() 生成的新样式类型化 PRNG 键，而不是由 jax.random.PRNGKey() 生成的旧样式原始 PRNG 键。详情请参见 JEP 9263: 类型化键 & 可插拔随机数生成器。

密钥是一个具有特殊dtype的数组，对应于所使用的特定PRNG实现；在默认实现中，每个密钥由一对 uint32 值支持。

键实际上是 NumPy 隐藏状态对象的替代品，但我们将其显式传递给 jax.random() 函数。重要的是，随机函数会消耗键，但不会修改它：将相同的键对象传递给随机函数将始终生成相同的样本。

print(random.normal(key))
print(random.normal(key))

-0.18471177
-0.18471177

即使使用不同的 random API，重复使用相同的键也可能导致输出相关，这通常是不希望的。

经验法则是：永远不要重复使用键（除非你想要相同的输出）。

为了生成不同的独立样本，你必须在将键传递给随机函数之前显式地 split() 它：

for i in range(3):
  new_key, subkey = random.split(key)
  del key  # The old key is consumed by split() -- we must never use it again.

  val = random.normal(subkey)
  del subkey  # The subkey is consumed by normal().

  print(f"draw {i}: {val}")
  key = new_key  # new_key is safe to use in the next iteration.

draw 0: 1.369469404220581
draw 1: -0.19947023689746857
draw 2: -2.298278331756592

（在这里调用 del 不是必需的，但我们这样做是为了强调一旦使用后不应再重复使用该键。）

jax.random.split() 是一个确定性函数，它将一个 key 转换为几个独立的（在伪随机性意义上）键。我们将其中一个输出作为 new_key，并且可以安全地将唯一的额外键（称为 subkey）作为随机函数的输入，然后永远丢弃它。如果你想从正态分布中获取另一个样本，你会再次分割 key，依此类推：关键点是你永远不会使用同一个键两次。

调用 split(key) 输出的哪一部分称为 key，哪一部分称为 subkey，这并不重要。它们都是具有同等地位的独立键。键/子键命名约定是一种典型的使用模式，有助于跟踪键的消耗方式：子键用于随机函数的即时消耗，而键则保留以供以后生成更多随机性。

通常，上述示例会简洁地写成

key, subkey = random.split(key)

它会自动丢弃旧的密钥。值得注意的是，split() 可以创建你需要的任意数量的密钥，而不仅仅是2个：

key, *forty_two_subkeys = random.split(key, num=43)

缺乏顺序等价性#

NumPy 和 JAX 的随机模块之间的另一个区别与上述提到的顺序等价保证有关。

与 NumPy 一样，JAX 的随机模块也允许对数字向量进行采样。然而，JAX 不提供顺序等价保证，因为这样做会干扰 SIMD 硬件上的向量化（上述要求 #3）。

在下面的例子中，从正态分布中分别使用三个子键采样3个值，与使用单个键并指定 shape=(3,) 得到的结果不同：

key = random.key(42)
subkeys = random.split(key, 3)
sequence = np.stack([random.normal(subkey) for subkey in subkeys])
print("individually:", sequence)

key = random.key(42)
print("all at once: ", random.normal(key, shape=(3,)))

individually: [-0.04838832  0.10796154 -1.2226542 ]
all at once:  [ 0.18693547 -1.2806505  -1.5593132 ]

缺乏顺序等价性使我们能够更高效地编写代码；例如，我们可以使用 jax.vmap() 以向量化方式计算相同的结果，而不是通过顺序循环生成上述 sequence：

import jax
print("vectorized:", jax.vmap(random.normal)(subkeys))

vectorized: [-0.04838832  0.10796154 -1.2226542 ]

下一步#

关于 JAX 随机数的更多信息，请参阅 jax.random 模块的文档。如果你对 JAX 随机数生成器的设计细节感兴趣，请参见 JAX PRNG 设计。

目录