NumPy와 PyTorch의 난수 함수는 전역 상태를 사용해.
import numpy as np
np.random.seed(42)
a = np.random.normal(size=10) # 어떤 값
b = np.random.normal(size=10) # 또 다른 값 (state 진행됨)
c = np.random.normal(size=10) # 또 다른 값
쓰기에는 편하지만 전역 상태를 읽고 쓰는 순간 함수의 순수성이 깨져. 이런 함수는 jit에서 처음 한 번만 추적되기 때문에 이후 호출에서 같은 값이 반복될 수 있어.
JAX는 난수 상태를 숨기지 않고 명시적인 인자로 꺼내 놨어. 모든 난수 함수가 key를 받는 이유야.
import jax
import jax.numpy as jnp
from jax import random
key = random.PRNGKey(42)
a = random.normal(key, (10,)) # 같은 key 면 항상 같은 값
b = random.normal(key, (10,)) # ← 같은 값! (a == b)
같은 키를 두 번 쓰면 같은 값이 나와. 버그가 아니라 같은 키는 같은 난수를 만든다는 의도된 규칙이야. 이 규칙 덕분에 결과를 정확히 재현할 수 있어.
다른 난수가 필요하면 키를 split:
key, subkey1, subkey2 = random.split(key, 3)
a = random.normal(subkey1, (10,)) # 어떤 값
b = random.normal(subkey2, (10,)) # 다른 값
# subkey1 ≠ subkey2 → a ≠ b
# key 도 다음 split 위해 따로 보존
key, sk = random.split(key)
c = random.normal(sk, (10,))
처음에는 번거로워 보여도 얻는 게 분명해.
- 완벽한 재현성: 같은 시작 키와 같은 연산 순서라면 언제나 같은 값이 나와. CI/CD에서도 학습 결과를 비트 단위로 비교할 수 있어.
- 순수성: 난수 상태가 입력에 드러나므로 jit 안에서도 안전하게 추적할 수 있어.
- 병렬 실행: 장치와 스레드마다 다른 키를 주면 전역 상태가 충돌하지 않아.
- 세밀한 제어: 학습의 어느 부분에 어떤 키를 썼는지 드러나서 디버깅하기 쉬워.
# PyTorch 식: random 이 어디서 흐르는지 볼 수 없음
torch.manual_seed(0)
def step():
noise = torch.randn(...) # 이 noise 는 어디서 왔지?
# JAX 식: 명시적
def step(key):
noise = jax.random.normal(key, ...) # key 가 input
🎲 Pure 난수의 가치
처음엔 "왜 이렇게 까다롭게?" 싶지만, 학습이 발산할 때, 모델 결과가 미묘하게 달라질 때, 그 원인이 난수 시드인지 다른 무엇인지 구분할 수 있는 게 연구에서는 결정적이야. 한 번 익히면 전역 난수가 얼마나 위험한 추상화였는지 보여.
JAX의 PRNG는 ThreeFry나 RBG 같은 카운터 기반 알고리즘을 사용해. 같은 시드에서 같은 방식으로 키를 다루면 CPU, GPU, TPU 어디서든 같은 결과를 얻을 수 있어.