키 Management 패턴

실전에서 키를 어떻게 전달하느냐가 학습 코드의 가독성을 좌우해. 흔한 패턴 네 가지를 살펴보자.

패턴 1: 키를 전달하고 값과 함께 반환하기

def sample_noise(key, shape):
    return random.normal(key, shape)

def model_with_dropout(params, x, key, p=0.5):
    key, subkey = random.split(key)
    mask = random.bernoulli(subkey, 1 - p, x.shape)
    return x * mask / (1 - p), key   # 새 key 도 반환

# 호출
key = random.PRNGKey(0)
y, key = model_with_dropout(params, x, key)
y2, key = model_with_dropout(params, y, key)

장점은 흐름이 명시적이라는 점이고, 단점은 모든 함수가 키를 입력과 출력으로 다뤄야 해서 장황해진다는 점이야.

패턴 2: 미리 split, list로 전달

def init_model(key, n_layers):
    keys = random.split(key, n_layers)
    return [init_layer(k) for k in keys]

장점은 코드가 깔끔하다는 점이고, 단점은 필요한 split 횟수를 미리 알아야 한다는 점이야.

패턴 3: 함수 안에서만 split

def sample_things(key):
    k1, k2, k3 = random.split(key, 3)
    a = random.normal(k1, ...)
    b = random.uniform(k2, ...)
    c = random.bernoulli(k3, ...)
    return a, b, c

# 호출 측은 한 key 만 주면 됨
key, sub = random.split(key)
a, b, c = sample_things(sub)

장점은 함수 시그니처가 깔끔하다는 점이고, 단점은 함수 내부의 난수 흐름이 호출자에게 보이지 않는다는 점이야.

패턴 4: fold_in으로 스텝 별 키

@jax.jit
def train_step(params, x, y, base_key, step):
    '''매 step 마다 다른 random — fold_in 으로 deterministic'''
    step_key = random.fold_in(base_key, step)
    k_dropout, k_noise = random.split(step_key)

    # ... random 사용

장점: 학습 루프에서 추가 상태 불필요해. 스텝 번호만 있으면 OK. JIT 친화. 단점: fold_in의 통계적 독립성이 split만큼 강하진 않다는 (이론적) 우려가 있어, 실제론 문제 안 돼.

JAX 표준, train 상태 안에 키 보존

@dataclass
class TrainState:
    params: Any
    opt_state: Any
    step: int
    key: Any   # ← key 도 state 의 일부

@jax.jit
def train_step(state, batch):
    key, subkey = random.split(state.key)
    # subkey 로 random
    ...
    return TrainState(
        params=new_params,
        opt_state=new_opt_state,
        step=state.step + 1,
        key=key,   # 다음 step 위해
    )

이 패턴이 Flax / Equinox / 표준 학습기 코드 어디서나. 키는 상태를 이루는 리프 가운데 하나야.

📐 어떤 패턴 쓸지

(1) 단순 함수, 패턴 3 (함수 안에서 split). (2) 학습 루프, 패턴 4 (fold_in) 또는 상태 안에 키 보존. (3) 큰 모델 초기화, 패턴 2 (미리 list). 가장 중요한 원칙은 같은 코드 안에서 한 패턴을 일관되게 쓰는 거야. 패턴을 뒤섞으면 어느 코드가 난수를 소비했는지 추적할 수 없어.

흔한 함정: 키 재사용해. 같은 키를 두 번 쓰면 같은 난수가 나와. 일부러 그러는 거 (test)이면 OK, 아니면 split 빠뜨린 버그.

Code

import jax
import jax.numpy as jnp

def init_layer(key, in_dim, out_dim):
    """Initialize a single layer with split keys."""
    k1, k2 = jax.random.split(key)
    weights = jax.random.normal(k1, (in_dim, out_dim)) * 0.01
    biases = jax.random.normal(k2, (out_dim,)) * 0.01
    return {'w': weights, 'b': biases}

def init_network(key, layer_sizes):
    """Initialize a full network, splitting keys for each layer."""
    params = []
    for i in range(len(layer_sizes) - 1):
        key, subkey = jax.random.split(key)
        params.append(init_layer(subkey, layer_sizes[i], layer_sizes[i+1]))
    return params

key = jax.random.key(42)
params = init_network(key, [784, 256, 128, 10])
print(f"Layer 0 weights shape: {params[0]['w'].shape}")  # (784, 256)

# WRONG: global key — leads to reuse bugs
# global_key = jax.random.key(0)  # Don't do this!

# RIGHT: pass key through the call chain
def dropout(x, key, rate=0.5):
    """Apply dropout with an explicit key."""
    mask = jax.random.bernoulli(key, 1.0 - rate, x.shape)
    return jnp.where(mask, x / (1.0 - rate), 0.0)

def forward(params, x, key):
    """Forward pass with dropout — key passed explicitly."""
    k1, k2 = jax.random.split(key)
    h = jnp.tanh(x @ params[0]['w'] + params[0]['b'])
    h = dropout(h, k1, rate=0.3)
    h = jnp.tanh(h @ params[1]['w'] + params[1]['b'])
    h = dropout(h, k2, rate=0.3)
    return h @ params[2]['w'] + params[2]['b']

# Option A: fold_in the step number (deterministic, no key threading)
base_key = jax.random.key(0)
for step in range(1000):
    step_key = jax.random.fold_in(base_key, step)
    # step_key is unique for each step, reproducible from base_key + step

# Option B: split at each step (standard threading)
key = jax.random.key(0)
for step in range(1000):
    key, subkey = jax.random.split(key)
    # use subkey for this step

Exercise

난수 표본 세 개가 필요한 함수를 리팩터링해. (1) 함수 안에서 키를 나누기, (2) 키 목록을 받기, (3) 키 하나를 받아 안에서 나누기라는 세 패턴을 모두 시도해. 사용감과 정확성의 절충을 비교하고 프로덕션 코드에 쓸 방식을 골라 이유를 적어.

📐 어떤 패턴 쓸지

Code

External links

Exercise

Progress

댓글 0