JAX는 Google DeepMind와 Google Research의 연구자들이 개발한 오픈소스 수치 계산 라이브러리이다. NumPy 호환 API에 자동 미분, JIT(just-in-time) 컴파일, GPU/TPU 가속을 제공하여 머신 러닝 연구와 과학 계산의 기반 도구가 되었다. 2018년에 처음 출시된 JAX는 신경망 훈련과 고성능 수치 실험에 널리 채택되었다.
JAX의 핵심 설계는 합성 가능한 함수 변환에 중점을 둔다. 주요 연산으로는 자동 미분을 위한 grad, 가속기 컴파일을 위한 jit, 벡터화를 위한 vmap, 장치 간 병렬화를 위한 pmap이 있다. 이러한 변환은 임의로 중첩될 수 있어 연구자들이 간결한 코드로 복잡한 알고리즘을 표현할 수 있다. JAX는 컴파일 백엔드로 XLA(Accelerated Linear Algebra)를 사용하며, 이는 AMD, Intel, NVIDIA, Google Cloud TPU에서 계산을 최적화한다.
역사와 개발
JAX는 2017년 Google의 연구에서 시작되었으며, autograd와 XLA에 대한 초기 작업을 기반으로 한다. 첫 공개 릴리스는 2018년 12월에 이루어졌다. 프로젝트는 Matthew Johnson, Roy Frostig, Alex Wiltschko를 포함한 연구자들이 주도했으며, Google Brain 팀의 광범위한 기여가 있었다. 2020년에는 JAX가 Flax(신경망 라이브러리), Haiku(DeepMind에서 사용), Trax를 포함한 여러 주목할 만한 라이브러리의 기반이 되었다. 2023년까지 JAX는 Google 내부 ML 인프라의 핵심 구성 요소가 되어 대규모 언어 모델과 트랜스포머 같은 모델을 지원했다.
주요 기능
JAX의 자동 미분은 순방향 및 역방향 모드를 모두 지원하여 임의의 Python 함수에 대한 기울기를 효율적으로 계산할 수 있다. jit 변환은 XLA를 통해 함수를 기계 코드로 컴파일하여 순수 Python보다 상당한 속도 향상을 달성하는 경우가 많다. vmap은 배치 차원에 걸쳐 연산을 자동으로 벡터화하여 수동 루프 언롤링을 제거한다. pmap은 여러 장치에 계산을 분산하여 데이터 병렬 및 모델 병렬 훈련을 용이하게 한다. JAX는 또한 함수형 API를 갖춘 난수 생성기를 포함하여 다양한 하드웨어 구성에서 재현성을 보장한다.
이 라이브러리는 Python 생태계와 원활하게 통합되며, 표준 데이터 구조와 NumPy와의 상호 운용성을 지원한다. 또한 NumPy의 인터페이스를 반영하지만 가속기 장치에서 작동하는 jax.numpy 모듈을 제공한다. JAX의 함수형 프로그래밍 스타일 - 배열이 불변이고 함수에 부작용이 없는 방식 - 은 디버깅을 단순화하고 안전한 병렬 실행을 가능하게 한다.
생태계와 채택
JAX는 풍부한 특수 라이브러리 생태계를 탄생시켰다. Flax와 Haiku는 고수준 신경망 API를 제공하고, Optax는 최적화 알고리즘을 제공한다. 과학 계산을 위해 JAX-MD(분자 역학)와 JAX-COSMO(우주론) 같은 라이브러리가 그 범위를 확장한다. MIT CSAIL, 스탠포드 AI 연구소, 버클리 AI 연구를 포함한 주요 연구 기관은 강화 학습, 확률적 프로그래밍, 미분 가능한 시뮬레이션 프로젝트에 JAX를 사용한다.
산업계에서 JAX는 Google의 프로덕션 시스템을 지원하며, Google Cloud AI 서비스의 일부와 Waymo의 인식 모델을 포함한다. 또한 OpenAI가 일부 연구 프로젝트에 사용하지만, Anthropic은 주로 PyTorch를 사용한다. TPU에서의 성능 덕분에 JAX는 특히 생성형 AI 애플리케이션에서 대규모 모델 훈련에 선호되는 선택지가 되었다.
다른 프레임워크와의 비교
JAX는 TensorFlow 및 PyTorch와 같은 머신 러닝 프레임워크와 경쟁한다. TensorFlow의 정적 그래프 방식과 달리 JAX는 많은 연구자들이 더 직관적이라고 생각하는 함수형, NumPy 스타일을 사용한다. PyTorch와 비교하여 JAX는 컴파일 및 병렬화에 대한 더 명시적인 제어를 제공하지만, 함수형 제약으로 인해 학습 곡선이 더 가파르다. JAX의 jit 컴파일은 종종 PyTorch의 즉시 실행보다 더 빠른 추론을 제공하지만, PyTorch의 동적 그래프는 디버깅에 더 쉽다. 벤치마크에서 JAX는 일반적으로 GPU 워크로드에서 PyTorch와 동등하거나 더 뛰어나며, PyTorch가 기본적으로 지원하지 않는 TPU에서 뚜렷한 이점이 있다.
응용 분야와 향후 방향
JAX는 딥 러닝부터 인공지능 연구까지 다양한 분야에서 사용된다. 미분 가능한 물리 엔진, 베이지안 추론 도구, 최적화 알고리즘을 지원한다. 최근 개발에는 ROCm을 통한 AMD GPU 지원과 향상된 CPU 성능이 포함된다. JAX 팀은 자동 샤딩과 혼합 정밀도 훈련 같은 기능을 계속 개선하고 있다. 2024년 현재 JAX는 활발히 개발 중이며, 성장하는 커뮤니티와 정기적인 릴리스를 보유하고 있다. 합성 가능성, 성능, 재현성이라는 설계 원칙은 차세대 AI 연구의 핵심 도구로 자리매김하게 한다.
같이 보기
참고 문헌
- JAX 공식 문서 및 GitHub 저장소 (2024년 접근)
- Google Research 블로그 게시물 (2018-2023)
- 머신 러닝 및 과학 계산에서 JAX를 인용한 학술 논문 (2020-2024)