JAX는 가속기 프로그래밍 및 프로그램 변환을 위한 파이썬 라이브러리로, 대규모 AI 모델 훈련을 넘어 AI 기반 단백질 공학 분야에서 중요한 역할을 하고 있습니다. 스타트업인 Escalante는 JAX를 사용하여 세포 단백질 발현 수준에 대한 약물 효과를 예측하는 모델을 훈련하고 있으며, 이는 JAX의 함수적이고 조합 가능한 특성을 잘 보여줍니다. 그들의 장기적인 비전은 처음부터 약물을 설계하는 것이지만, 초기에는 새로운 실험실 분석법을 개발하여 중요한 생물학적 데이터 세트를 생성하는 데 집중하고 있습니다. 단백질 공학은 결합력, 용해도, 안정성과 같은 다양한 기준을 충족해야 하는 다중 목표 최적화를 포함합니다. JAX는 각기 다른 속성을 예측하는 수많은 AI 모델을 통합하여 통합 손실 함수로 만드는 것을 단순화합니다. Escalante는 JAX 생태계를 적극적으로 활용하며, PyTorch와 같은 다른 프레임워크의 모델도 변환했습니다. 이를 통해 모델을 조합하고 최종 목표로 변환할 수 있는 표현력 있는 언어를 단백질 설계에 사용할 수 있으며, jax.jit를 사용하여 성능을 최적화할 수 있습니다. 그들의 워크플로우는 고정된 신경망 집합을 복잡하고 미분 가능한 손실 함수로 사용하여 입력 시퀀스를 최적화함으로써 일반적인 훈련 방식을 뒤집습니다. 이 프로세스는 DeepDream과 유사하며, 기울기가 원하는 속성을 향해 시퀀스 업데이트를 안내합니다. JAX의 자동 미분 및 컴파일 기능은 이러한 정교한 손실 함수를 최적화하는 데 매우 중요합니다. 이 프레임워크는 TPU와의 네이티브 통합을 통해 이러한 워크로드를 확장할 수 있으며, Escalante는 필요에 따라 TPU를 가동하고 종료하는 패턴을 사용합니다. 이러한 TPU 채택은 대규모 작업에 대해 GPU보다 상당한 비용 효율성을 제공합니다. Equinox 및 Optax와 같은 주요 JAX 생태계 라이브러리는 모델 표현 및 최적화 알고리즘 유연성을 위해 사용됩니다. JAX의 함수적 핵심, 생태계 라이브러리, 그리고 확장 가능한 TPU 하드웨어의 조합은 Escalante의 획기적인 연구를 가능하게 합니다.