Optax: 합성할 수 있는 그래디언트 Transformation

Optax는 JAX의 사실상 표준 옵티마이저 라이브러리. 핵심 아이디어는 옵티마이저가 monolithic 클래스가 아니라, 작은 transformation 들의 합성해.

pip install optax

가장 단순한 사용

import optax
import jax
import jax.numpy as jnp

# AdamW optimizer
optimizer = optax.adamw(learning_rate=1e-3, weight_decay=1e-4)

# state 초기화
params = {"w": jnp.zeros(10), "b": jnp.zeros(())}
opt_state = optimizer.init(params)

# 학습 step 안에서
@jax.jit
def step(params, opt_state, x, y):
    grads = jax.grad(loss_fn)(params, x, y)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state

세 단계로 진행해. update로 그래디언트를 변환하고 apply_updates로 매개변수에 더해. PyTorch의 optimizer.step() 한 줄과 같은 의미야.

합성의 힘

Optax의 매력, optax.chain으로 transformation 합성:

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),       # gradient clip
    optax.add_decayed_weights(1e-4),      # weight decay
    optax.scale_by_adam(),                 # Adam moments
    optax.scale_by_schedule(schedule),     # learning rate schedule
    optax.scale(-1.0),                     # 부호 반전 (descent)
)

이 5줄이 AdamW + grad clip + scheduled lr의 표준 학습기의 옵티마이저. PyTorch에서는 Adam 클래스, 매 스텝의 그래디언트 자르기, 스케줄 콜백을 따로 다뤄야 하지만 JAX와 Optax에서는 하나의 chain으로 표현해.

주요 transformation

이름	역할
`scale(c)`	모든 grad에 c 곱하기 (보통 -lr)
`scale_by_adam()`	Adam moments (1차, 2차)
`scale_by_belief()`	AdaBelief의 variance estimate
`scale_by_rms()`	RMSProp의 second moment
`scale_by_schedule(s)`	스케줄에 따라 lr 변화
`add_decayed_weights(wd)`	L2 weight decay
`clip(max)`	per-element clip
`clip_by_global_norm(max)`	전체 grad norm clip
`ema(decay)`	exponential moving average
`masked(t, mask)`	일부 리프에만 변환 적용하기

실전 예, bias와 norm은 weight decay 안 함

def make_mask(params):
    '''W 면 True, b/gamma/beta 면 False'''
    def is_weight(path, value):
        return value.ndim > 1   # 1D 이상이면 weight matrix
    return jax.tree_util.tree_map_with_path(is_weight, params)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.masked(
        optax.add_decayed_weights(1e-4),
        mask=make_mask(params),
    ),
    optax.scale_by_adam(),
    optax.scale_by_schedule(cosine_schedule),
    optax.scale(-1.0),
)

built-in 옵티마이저

# 자주 쓰는 표준
optax.sgd(lr, momentum=0.9)
optax.adam(lr)
optax.adamw(lr, weight_decay=1e-4)
optax.adamax(lr)
optax.rmsprop(lr)
optax.lamb(lr)
optax.lion(lr, weight_decay=1e-4)   # 2023 의 Google
optax.amsgrad(lr)

각 함수는 위 변환을 미리 엮어 둔 chain이야. 무엇을 연결하는지 Optax 소스에서 한 번 보면 구조를 익힐 수 있어.

🧬 합성함으로써 옵티마이저

Optax의 정신은 jit, grad, vmap을 합성하는 JAX의 정신과 같아. 하나의 거대한 클래스 대신 작은 변환을 합성하고, 기존 변환의 순서를 달리 연결해 새 옵티마이저를 만들어. AdamW가 Adam과 decay를 연결한 chain인 것과 같은 원리야. 이 프레임워크 안에서, 학술 논문의 새 옵티마이저를 한 시간 안에 구현할 수 있어.

Track 11-3에서 스케줄, 11-4에서 전체 학습 루프. 11-5에서 scan + 체크포인트를 다뤄.

Code

import optax

# Common optimizers — each is a gradient transformation
optimizer = optax.adam(learning_rate=1e-3)
optimizer = optax.adamw(learning_rate=1e-3, weight_decay=0.01)
optimizer = optax.sgd(learning_rate=0.1, momentum=0.9)
optimizer = optax.lion(learning_rate=1e-4)  # newer optimizer

# But the real power is composition
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),  # gradient clipping
    optax.adam(learning_rate=1e-3),   # Adam optimizer
)

# Or even more custom:
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),      # clip gradients
    optax.scale_by_adam(),                # Adam scaling (no LR)
    optax.add_decayed_weights(0.01),     # L2 regularization
    optax.scale(-1e-3),                  # apply learning rate
)

import jax
import jax.numpy as jnp
import optax

# 1. Create optimizer
optimizer = optax.adamw(learning_rate=1e-3)

# 2. Initialize optimizer state from params
params = {'w': jnp.ones((3, 4)), 'b': jnp.zeros(4)}
opt_state = optimizer.init(params)

# 3. Compute gradients (however you like)
grads = jax.grad(loss_fn)(params, x, y)

# 4. Get updates from optimizer
updates, new_opt_state = optimizer.update(grads, opt_state, params)

# 5. Apply updates to params
new_params = optax.apply_updates(params, updates)

# The full cycle:
# grads → optimizer.update(grads, opt_state, params) → updates, new_opt_state
#                                                      ↓
# new_params = optax.apply_updates(params, updates)

# PyTorch                           # JAX + Optax
# optimizer = Adam(model.params())  # optimizer = optax.adam(1e-3)
#                                   # opt_state = optimizer.init(params)
# optimizer.zero_grad()             # (not needed — grads are values)
# loss.backward()                   # grads = jax.grad(loss_fn)(params)
# optimizer.step()                  # updates, opt_state = optimizer.update(...)
#                                   # params = optax.apply_updates(params, updates)

🧬 합성함으로써 옵티마이저

Code

External links

Exercise

Progress

댓글 0