JAX随机生成器生成的正态数为何呈现有序而非完全随机？

我正从零开始学习jax，尝试生成满足线性关系Y = B1*X + B0 + epsilon的合成数据。令我意外的是，直接使用生成的epsilon时，数据呈现出有序性，极端位置的点误差大于中心位置的点；而将epsilon打乱后，误差则围绕线性关系呈正态分布，效果正常。相关代码如下：

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random
import matplotlib.pyplot as plt

n = 100
p = 2  # number of parameters including linear bias
key = random.PRNGKey(0)

X = random.uniform(key, (n, p - 1))
X = jnp.concatenate([jnp.ones((n, 1)), X], axis=1)
B = random.randint(key, (p, 1), 0, 10)

epsilon = random.normal(key, (n, 1))*0.3

y = X.dot(B) + epsilon
y_s = X.dot(B) + random.shuffle(key, epsilon)
plt.scatter(X[:,1], y, label='Original order')
plt.scatter(X[:,1], y_s, label='Shuffled')
plt.legend()

原因分析 #

核心问题是重复使用了同一个PRNGKey。JAX的随机数生成是确定性的：同一个密钥生成的随机序列完全固定，正确用法是每次调用随机函数前拆分密钥生成子密钥，但你这里所有随机操作都用了初始的key，没有做密钥拆分。

具体细节：

• 生成X时用key生成的均匀分布数据，是按固定顺序输出的序列；
• 生成epsilon时再次用同一个key，得到的正态分布序列和X的序列存在隐性同步相关性——两者的随机序列是完全对应生成的，导致X取值越大的位置，对应的epsilon绝对值也越大，最终呈现出误差随X增大而变大的有序性；
• 而random.shuffle打乱了epsilon的顺序，直接破坏了这种同步关联，误差就回到了正常的正态分布状态。

修正方案 #

每次调用随机函数前，用random.split(key)拆分密钥，生成独立的子密钥供后续操作使用：

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random
import matplotlib.pyplot as plt

n = 100
p = 2  # number of parameters including linear bias
key = random.PRNGKey(0)

# 拆分密钥生成X
key, subkey = random.split(key)
X = random.uniform(subkey, (n, p - 1))
X = jnp.concatenate([jnp.ones((n, 1)), X], axis=1)

# 拆分密钥生成B
key, subkey = random.split(key)
B = random.randint(subkey, (p, 1), 0, 10)

# 拆分密钥生成epsilon
key, subkey = random.split(key)
epsilon = random.normal(subkey, (n, 1))*0.3

y = X.dot(B) + epsilon
y_s = X.dot(B) + random.shuffle(key, epsilon)
plt.scatter(X[:,1], y, label='Original order')
plt.scatter(X[:,1], y_s, label='Shuffled')
plt.legend()

内容的提问来源于stack exchange，提问作者ABaron