Pytree가 뭐야?

pytree는 JAX의 거의 모든 함수가 받는 자료 구조야. 이름은 거창하지만 배열을 리프로 삼는 중첩된 Python 컨테이너 트리일 뿐이야.

import jax
import jax.numpy as jnp

# 단순 array 도 pytree
a = jnp.array([1, 2, 3])

# dict — pytree
params = {
    "W": jnp.zeros((3, 4)),
    "b": jnp.zeros((4,)),
}

# list — pytree
weights = [jnp.zeros((3, 4)), jnp.zeros((4, 2))]

# tuple — pytree
state = (jnp.zeros(()), jnp.zeros((10,)))

# nested — 다 pytree
deep = {
    "encoder": [jnp.zeros((3, 8)), jnp.zeros((8, 4))],
    "decoder": [jnp.zeros((4, 8)), jnp.zeros((8, 3))],
    "stats": (jnp.array(0), jnp.array(0.0)),
}

JAX가 pytree를 다룰 수 있는 컨테이너:

list, tuple, dict
collections.OrderedDict, defaultdict
NamedTuple
@dataclass (등록 후)
None (리프가 없는 빈 자리)

"리프"는 pytree에서 더 이상 나눌 수 없는 단위로, 보통 jnp.array나 Python 스칼라야.

주요 함수, jax.tree.map

params = {
    "W": jnp.zeros((3, 4)),
    "b": jnp.zeros((4,)),
}

# 모든 leaf 에 함수 적용
doubled = jax.tree.map(lambda x: x * 2, params)
# {"W": (3,4) of 0, "b": (4,) of 0}  — 0 이라 다 0 이지만, 모양 보존

zeros = jax.tree.map(jnp.zeros_like, params)
# 같은 모양의 새 pytree, 모두 0

ones = jax.tree.map(jnp.ones_like, params)

여러 pytree를 동시에

grads = {
    "W": jnp.ones((3, 4)),
    "b": jnp.ones((4,)),
}

# params 와 grads 를 zip 해서 update
new_params = jax.tree.map(
    lambda p, g: p - 0.01 * g,
    params,
    grads,
)
# 두 pytree 의 구조가 같아야 함 (key, shape, ...)

이게 모든 JAX 옵티마이저와 학습 코드의 핵심 패턴이야. 매개변수가 단순 배열, dict, NamedTuple, 깊게 중첩된 구조 가운데 어떤 모양이든 같은 코드야.

🌳 pytree의 의미

JAX가 PyTorch와 가장 다른 점 중 하나, 모델 매개변수를 dict / list / dataclass로 자유롭게 표현해. nn.Module 클래스 상속 안 해도 돼. jit, grad, vmap 같은 모든 변환이 pytree를 자동으로 순회해. 모델이 더 이상 마법 같은 객체가 아니라, 그냥 중첩 데이터.

실용 예, Transformer 블록의 매개변수:

block_params = {
    "attention": {
        "W_q": jnp.zeros((d, d)),
        "W_k": jnp.zeros((d, d)),
        "W_v": jnp.zeros((d, d)),
        "W_o": jnp.zeros((d, d)),
    },
    "mlp": {
        "W1": jnp.zeros((d, 4*d)),
        "W2": jnp.zeros((4*d, d)),
    },
    "ln1": {"gamma": jnp.ones(d), "beta": jnp.zeros(d)},
    "ln2": {"gamma": jnp.ones(d), "beta": jnp.zeros(d)},
}

# zero gradient
zero_grads = jax.tree.map(jnp.zeros_like, block_params)

# parameter 개수 세기
n_params = sum(x.size for x in jax.tree.leaves(block_params))
print(f"params: {n_params:,}")

모든 게 중첩 딕셔너리라, 자유롭게 살펴보고 조작할 수 있어. 이게 함수형 모델 표현의 자유도야.

🌳 pytree의 의미

Code

External links

Exercise

Progress

댓글 0