PRNG 키와 Splitting

JAX 난수 처리의 핵심은 PRNGKey와 split이야. 이 두 함수부터 익히면 돼.

import jax
from jax import random

# 시작 key 만들기
key = random.PRNGKey(0)
print(key)   # uint32 array of shape (2,) — 64-bit seed

# split — 한 key 에서 N 개의 새 key
key, subkey = random.split(key)         # 2 개 (default)
keys = random.split(key, 4)              # 4 개
key1, key2, key3 = random.split(key, 3)  # 3 개 unpack

구조를 보면 키는 shape이 (2,)인 uint32 배열일 뿐이야. split의 결과도 같은 모양:

key = random.PRNGKey(42)
print(key.shape)         # (2,)

new_keys = random.split(key, 5)
print(new_keys.shape)    # (5, 2)

키 사용의 표준 패턴

관용적으로 split 결과의 첫 번째를 "next 키"로 보존하고 나머지를 사용해:

key = random.PRNGKey(0)

# 단계 1
key, subkey = random.split(key)
a = random.normal(subkey, (10,))

# 단계 2
key, subkey = random.split(key)
b = random.uniform(subkey, (10,))

# 단계 3
key, subkey = random.split(key)
c = random.bernoulli(subkey, p=0.5, shape=(10,))

이 패턴은 JAX 난수 코드 어디서나 보여. key, subkey = random.split(key)가 거의 mantra.

여러 난수가 동시에 필요할 때

key = random.PRNGKey(0)

# 한 번에 여러 split
keys = random.split(key, 5)
key, k_init, k_dropout, k_noise, k_perm = keys

# 또는 unpack
key, k1, k2, k3 = random.split(key, 4)

트리 모양의 키

큰 모델이면 매개변수의 각 layer에 다른 키 필요해. 한 번에 split:

key = random.PRNGKey(0)
n_layers = 5
layer_keys = random.split(key, n_layers)

for i, lk in enumerate(layer_keys):
    layer_params[i] = init_layer(lk, in_dim, out_dim)

💡 split은 cheap해

split이 비싸지 않을까 걱정? 안 그럼. split은 결정적인 해시 함수 비슷해서, 매우 빨라. 키를 함수 인자로 자유롭게 통과시키고, 필요할 때마다 split 해도 성능 영향 거의 없어.

fold_in은 키를 정수 식별자에 따라 deterministic하게 변형해:

# key 를 step 번호와 결합
def step_key(base_key, step):
    return random.fold_in(base_key, step)

base = random.PRNGKey(0)
for step in range(100):
    key_at_step = step_key(base, step)
    # 매 step 의 key 가 deterministic + 독립적

fold_in은 split보다 "정수 → 키" 매핑이라는 의도가 더 명확해. 학습 스텝마다 키를 만들 때 흔히 사용해.

Code

import jax
import jax.numpy as jnp

# Create a starting key from a seed
key = jax.random.key(0)
print(key)  # a special key array

# WRONG: reusing the same key gives the same numbers
print(jax.random.normal(key))  # 0.18784384
print(jax.random.normal(key))  # 0.18784384 — identical!

# RIGHT: split the key to get new, independent keys
key, subkey = jax.random.split(key)
print(jax.random.normal(subkey))  # -1.2515389

key, subkey = jax.random.split(key)
print(jax.random.normal(subkey))  # -0.5841975 — different!

# Split into many keys at once
key = jax.random.key(42)
keys = jax.random.split(key, num=5)
print(keys.shape)  # (5,) — five independent keys

# Use each key for a different purpose
samples = jax.vmap(jax.random.normal)(keys)
print(samples)  # 5 independent random numbers

# Modern API (recommended)
key = jax.random.key(0)

# Legacy API (still works)
key_legacy = jax.random.PRNGKey(0)

# Both produce valid keys for all jax.random functions

Exercise

PRNGKey(0)에서 시작하는 키 트리를 만들어. 키를 네 개로 나누고 각 키를 다시 네 개로 나눠 독립적인 정규분포 표본 16개를 생성해. 표본 사이의 상관계수가 0에 가까운지 검증하고, 분기형 키 모델이 난수 처리를 합성 가능하게 만드는 이유를 적어.

💡 split은 cheap해

Code

External links

Exercise

Progress

댓글 0