JAX

JAX는 Google Research에서 개발한 고성능 수치 연산 및 머신러닝 라이브러리로, NumPy 스타일의 API를 기반으로 자동 미분(autograd), JIT 컴파일(just-in-time), 벡터화(Vectorization), 병렬 처리 등을 제공한다. XLA(Accelerated Linear Algebra) 컴파일러를 통해 CPU, GPU, TPU에서 모두 효율적으로 실행되며, 대규모 딥러닝 연구와 과학 계산에서 널리 사용된다.

개요

JAX는 기존 NumPy 스타일의 직관성을 유지하면서 다음과 같은 기능을 결합한다:

• 자동 미분 (Autograd)
• JIT 컴파일 (XLA 기반)
• 벡터화(vmap)
• 병렬화(pmap)
• 순수 함수 기반의 함수형 프로그래밍 모델

JAX는 TensorFlow나 PyTorch와 달리 모델 정의용 고수준 프레임워크가 아니라, 수치 연산을 함수형 방식으로 최적화하는 “프레임워크의 기초 도구”에 가깝다. Neural network 라이브러리인 Flax, Haiku, Equinox 등이 JAX 위에서 사용된다.

주요 특징

Autograd (자동 미분)

JAX는 Python 함수를 입력받아 해당 함수의 기울기를 자동 계산한다.

• reverse-mode differentiation
• forward-mode differentiation
• 많은 연산에 대한 강력한 합성(differentiable programming)

예시:

import jax.numpy as jnp
import jax

f = lambda x: x**2 + 3*x
jax.grad(f)(2.0)  # 7.0

JIT 컴파일

XLA 기반의 JIT 컴파일 기능(`jax.jit`)은 Python 코드를 고성능 기계 코드로 변환하여 실행 속도를 크게 높인다.

• 연산 fusion
• device placement 자동 처리
• GPU/TPU에서 높은 효율

@jax.jit
def compute(x):
    return jnp.sin(x) * jnp.cos(x)

벡터화 (vmap)

루프를 Python이 아닌 XLA에서 병렬화하여 대규모 데이터 처리 속도를 높인다.

• Python loop → XLA 병렬 연산 자동 변환

pmap (병렬 처리)

다중 GPU/TPU 장치에서 동일한 함수를 병렬 실행.

• Data parallelism 구현
• TPU Pod 환경에서 핵심 기능

함수형 패러다임

JAX는 순수 함수(pure function)를 강조하며, 상태(state) 없는 계산을 지향한다. 모델 파라미터도 불변성을 유지한 채 함수 인자로 전달하는 방식으로 관리한다.

JAX 생태계

Flax

• JAX 기반의 고수준 neural network 라이브러리. LLM, Vision Transformer 등 대규모 모델 구현에 널리 사용된다.

Haiku

• DeepMind에서 개발한 모듈식 신경망 프레임워크.

Optax

• JAX를 위한 최적화 라이브러리(Adam, RMSProp 등 제공).

Equinox

• PyTorch와 유사한 객체 스타일의 JAX 신경망 라이브러리.

Orbax

• JAX 체크포인트 저장 및 복구 라이브러리.

XLA와의 관계

JAX는 XLA를 기반으로 동작하는 구조를 가지고 있다:

• JIT → XLA로 그래프 변환
• XLA가 CPU/GPU/TPU용 최적화된 기계 코드 생성
• TPU 사용 시 탁월한 효율

XLA는 TensorFlow와 JAX 모두의 공통 백엔드로 사용되지만, JAX는 훨씬 더 자연스럽게 XLA와 통합되도록 설계되었다.

TPU와의 통합

JAX는 TPU를 가장 잘 활용하는 프레임워크 중 하나로 평가된다.

• TPU v2/v3/v4/v5에서 높은 성능
• pmap을 통한 대규모 모델 학습
• Google 내부 연구에서 광범위하게 사용

PaLM, T5, 여러 대형 언어 모델 연구에서 JAX+TPU 조합이 핵심 역할을 한다.

PyTorch 및 TensorFlow와의 비교

항목	JAX	PyTorch	TensorFlow
주요 철학	함수형, 확장성, 연구용	직관적, eager 기반	산업용, 정적 그래프
Autograd	매우 강력함	강력함	강력함
JIT 속도	매우 빠름(XLA)	PyTorch 2.x Inductor로 개선	TensorFlow XLA로 빠름
TPU 지원	최고 수준	제한적	좋음
생태계	Flax/Haiku 중심	가장 큰 생태계	기업 중심

활용 분야

• 대규모 언어 모델(LLM)
• 컴퓨터 비전(ViT, CNN)
• 강화학습
• 과학 계산 및 시뮬레이션
• Google 내부 연구
• TPU 기반 대규모 학습

장점

• 매우 빠른 JIT 성능(XLA 기반)
• TPU/GPU 모두에서 높은 효율
• 함수형 프로그래밍으로 안전한 코드 구조
• 연구·대규모 모델 학습에 최적화
• 자동 병렬화(vmap/pmap) 기능

한계

• 생태계가 PyTorch만큼 크지 않음
• 디버깅 난이도가 높을 수 있음
• 함수형 패턴에 익숙하지 않은 사용자는 진입 장벽 존재
• stateful 연산 및 비순수 함수가 어렵거나 제한적

함께 보기

• PyTorch
• TensorFlow
• XLA
• TPU
• 딥 러닝
• 자동 미분
• Neural Network