JAX가 과학 계산에서 빛나는 이유

JAX는 머신러닝 프레임워크로도 이미 강력해. 하지만 진짜 차별점은 머신러닝 밖의 영역에서 드러나. JAX가 차별화되는 이유는 수치 코드를 자유롭게 미분하고 가속할 수 있기 때문이야.

전통적인 과학 계산 도구:

• NumPy / SciPy, 주로 CPU에서 실행되고 일반 수치 코드의 자동 미분은 지원하지 않아.
• MATLAB, 도구 상자가 풍부하지만 폐쇄형이고 그래디언트는 따로 작성해야 해.
• Fortran / C++, 빠르지만 그래디언트를 직접 작성해야 해.
• Numba는 JIT을 지원하지만 자동 미분은 지원하지 않아.

JAX가 합치는 두 가지:

1. 임의 수치 코드의 자동 미분: ODE 풀이, 시뮬레이션, 샘플링, 수치 코드를 미분할 수 있어.
2. 가속기 활용: 같은 Python 코드를 GPU와 TPU에서 실행할 수 있어.

이 둘이 합쳐지면 새 카테고리의 알고리즘이 가능해져.

전형적 사용 사례

1. 미분 가능한 물리 시뮬레이션: 시뮬레이터를 통해 그래디언트를 흘려 물리 시스템을 학습해.

def simulate_pendulum(initial_angle, length, dt, n_steps):
    '''단순한 진자 시뮬레이션'''
    theta, omega = initial_angle, 0.0
    for _ in range(n_steps):
        omega += -9.81 / length * jnp.sin(theta) * dt
        theta += omega * dt
    return theta

# 진자가 특정 각도에 도달하도록 길이를 학습
target_angle = 0.5
def loss(length):
    final = simulate_pendulum(0.1, length, 0.01, 100)
    return (final - target_angle) ** 2

# 자동 미분 — 시뮬레이션 통해 gradient 흐름
optimal_length = optimize_with_grad(loss)

2. 미분방정식: Diffrax는 미분 가능한 ODE와 SDE solver를 제공해.

from diffrax import diffeqsolve, Tsit5, ODETerm

def lorenz(t, y, args):
    x, y_, z = y
    return jnp.array([
        10 * (y_ - x),
        x * (28 - z) - y_,
        x * y_ - 8/3 * z,
    ])

solution = diffeqsolve(
    ODETerm(lorenz),
    Tsit5(),
    t0=0.0, t1=10.0, dt0=0.01,
    y0=jnp.array([1., 1., 1.]),
)

3. 베이지안 추론: NumPyro는 JAX 위에서 확률적 프로그래밍을 제공해.

import numpyro
import numpyro.distributions as dist

def model(data):
    mu = numpyro.sample("mu", dist.Normal(0, 1))
    sigma = numpyro.sample("sigma", dist.HalfNormal(1))
    with numpyro.plate("data", len(data)):
        numpyro.sample("obs", dist.Normal(mu, sigma), obs=data)

# NUTS sampler — JAX-jit 으로 자동 가속
mcmc = numpyro.infer.MCMC(numpyro.infer.NUTS(model), num_samples=1000)
mcmc.run(rng_key, data)

4. 최적화: JAXopt와 Lineax 같은 라이브러리를 사용할 수 있어.

import jaxopt

# 비선형 least squares — 자동 미분 + GPU
solver = jaxopt.LevenbergMarquardt(residual_fun=residuals)
result = solver.run(init_params, data=observations)

5. 확률적 프로그래밍, TensorFlow Probability JAX 백엔드

6. 강화학습, Brax (미분할 수 있는 물리 시뮬레이션), MJX (MuJoCo)

7. 양자 컴퓨팅, qujax, PennyLane JAX 백엔드

8. 계산화학, JAX-MD (molecular dynamics)

이어지는 레슨에서는 Diffrax의 ODE, NumPyro의 베이지안 추론, Brax와 JAX-MD의 시뮬레이션을 통해 JAX 생태계의 각 영역을 살펴봐.