JAX 101（四）：显式随机数与可复现实验

Code & Development

发布日期: 2026-06-06

文章字数: 2k

阅读时长: 8 分

掌握 JAX 的显式伪随机数设计，正确完成参数初始化、dropout、采样和数据打乱。随机 key 看起来只是 API 差异，实际上它体现了 JAX 对纯函数和并行计算的要求。

1. 为什么没有隐藏的全局随机状态

像 NumPy、Python 内置的 random 模块，内部都维护着一个隐藏的、全局共享的随机数状态（比如 NumPy 底层是一个 Mersenne Twister 状态机）。你每次调用 np.random.rand()，实际上是：读取这个全局状态；用它生成一个”随机”数值；更新（推进）这个全局状态，为下一次调用做准备。关键点：这个状态是”隐藏的”、”可变的”、”全局唯一的”——你在代码里看不到它，但它确实存在，且每次随机调用都会悄悄修改它。

JAX 把随机状态写成显式 key：

import jax
import jax.numpy as jnp

key = jax.random.key(42)
x = jax.random.normal(key, (3,))
print(x)

随机函数是 key 的纯函数。相同 key 与相同参数会产生相同结果。

2. key 不能重复使用

key = jax.random.key(42)
a = jax.random.normal(key, (3,))
b = jax.random.normal(key, (3,))
print(jnp.all(a == b))  # True

JAX 不会自动“推进”原 key。需要新随机数时必须拆分：

key, subkey = jax.random.split(key)
a = jax.random.normal(subkey, (3,))

key, subkey = jax.random.split(key)
b = jax.random.normal(subkey, (3,))

可以把 key 看作随机数资源的根，把 subkey 分配给一次具体用途。已经消费过的 subkey 不应再次使用。

3. 批量拆分

key = jax.random.key(0)
key, *layer_keys = jax.random.split(key, 5)

这里得到一个保留给后续操作的 key，以及 4 个不同的层初始化 key。对于多层模型，推荐初始化函数按树结构逐层拆分，避免不小心让两个权重矩阵完全相同。

4. 参数初始化

def init_embedding(key, vocab_size, dim):
    return 0.02 * jax.random.normal(key, (vocab_size, dim))

def init_model(key, vocab_size, dim, num_layers):
    keys = jax.random.split(key, num_layers + 2)
    return {
        "embedding": init_embedding(keys[0], vocab_size, dim),
        "blocks": [init_block(keys[i + 1], dim) for i in range(num_layers)],
        "lm_head": 0.02 * jax.random.normal(keys[-1], (dim, vocab_size)),
    }

这里的 init_block 在下一课/下一篇完整实现。关键不是某个初始化系数，而是每个随机操作获得独立 key。

5. Dropout 需要 key，也需要训练模式

def dropout(key, x, rate, training):
    if not training or rate == 0.0:
        return x
    keep_prob = 1.0 - rate
    mask = jax.random.bernoulli(key, keep_prob, x.shape)
    return jnp.where(mask, x / keep_prob, 0.0)

dropout 每次调用都需要随机性，如果每一步训练都用同一个 key，dropout 就会失效（变成”假随机”）。除以 keep_prob 是为了让输出的期望尺度保持不变。训练时每一步都应给 dropout 新 key；推理时关闭 dropout。

在 jitted 函数里，Python 分支条件通常应当是静态值，或者改写成适合 JAX 的数组控制流。对初学课程来说，可以分别编译训练函数与推理函数。

6. 自回归采样

6.1 基础采样与温度调节的关系

自回归生成的意思是：模型根据当前序列预测下一个 token，把选出的 token 拼回序列，再用更新后的序列继续预测。每一步选择的结果，都会成为后续预测的上下文。

模型输出的 logits 是词表中各个 token 的原始分数。最基础的随机采样方式，是用 softmax(logits) 得到概率分布，再按这个分布抽取一个 token。分数高的 token 被选中的概率更大，分数稍低的 token 也有机会被选中。这里说的是 token，它可能是一个字、一个词，也可能只是词的一部分。

把 logits 除以温度 T，是生成时常用的概率调节方式，可以按需要启用。 对于 T > 0，带温度的采样分布写成 softmax(logits / T)。当 T = 1 时，logits 保持原样，正好就是前面介绍的基础采样。因此，两种写法可以统一理解：基础采样对应温度为 1，其他正温度则用于调整分布的集中程度。这是解码阶段的设置，不需要修改模型权重或重新训练。

用带温度的随机采样生成文本时，一轮的概念流程可以写成：

将当前序列输入模型，取最后一个位置的 logits，用于预测下一个 token。
根据生成配置，用正温度 T 调整 logits；T 为 1 时保持原分数。
按调整后的 logits 所对应的 softmax 概率分布，随机采样一个 token。
将新 token 拼回序列，继续生成，直到满足停止条件，例如遇到结束标记或达到长度上限。

6.2 温度如何影响选词

当 0 < T < 1 时，分布更集中。 例如，两个 token 的 logits 原本相差 2，除以 0.5 后就相差 4。softmax 会让高分 token 占据更大的概率，采样结果更容易集中在少数高分候选上。在文本生成中，这通常表现为输出更保守、更可预测，也可能更容易重复。

当 T > 1 时，分布更平缓。 同样相差 2 的 logits，除以 2 后只相差 1，低分候选获得的概率会增加。生成结果通常更多样，但温度过高时，也更容易选到不合适的 token，影响连贯性和准确性。多样性增加，并不直接意味着内容更有创意或质量更高。

例如，假设只有三个候选 token，logits 是 [2, 1, 0]，不同温度下的概率大约是：

温度 T	第一个 token	第二个 token	第三个 token
0.5	86.7%	11.7%	1.6%
1.0	66.5%	24.5%	9.0%
2.0	50.6%	30.7%	18.6%

可以看到，正温度不会改变 logits 的高低排序，改变的是各个候选被抽中的概率。在工程实现中，温度还经常与 top-k、top-p 等候选过滤策略配合使用；这里先只讨论温度，便于理解随机采样本身。

6.3 在 JAX 中直接用 logits 采样

def sample_next_token(key, logits, temperature=1.0):
    if temperature <= 0:
        return jnp.argmax(logits, axis=-1)
    return jax.random.categorical(key, logits / temperature, axis=-1)

这里的 jax.random.categorical 接受未归一化 logits，采样分布对应于 softmax(logits / temperature)，所以调用前不用再手写 softmax。应当把调整后的 logits 传给它，不能把 softmax 得到的概率当成 logits 再传入，否则采样分布就变了。 这是接口层面省去了显式计算概率的步骤，概率采样的含义仍然相同，参见 JAX categorical 文档。

代码中的 temperature <= 0 是这个示例约定的贪心分支：直接用 argmax 选择分数最高的 token，不进行随机采样，也不会真的执行除以 0。低正温度仍然属于随机采样；随着温度接近 0，概率会越来越集中到最高分候选上。生产接口可以将负温度判为非法参数，并用 temperature = 0 或独立开关显式选择贪心解码。

生成循环中应当每一步拆分 key：

for _ in range(max_new_tokens):
    key, sample_key = jax.random.split(key)
    logits = model(params, tokens)
    next_token = sample_next_token(sample_key, logits[:, -1, :])
    tokens = jnp.concatenate([tokens, next_token[:, None]], axis=1)