JAX, библиотека Python для ускоренного программирования и преобразования программ, оказывается жизненно важной для разработки белков на основе ИИ, расширяя свое влияние за пределы обучения больших моделей ИИ. Стартап Escalante использует JAX для обучения моделей, которые предсказывают влияние лекарств на уровни экспрессии клеточных белков, демонстрируя функциональный и компонуемый характер JAX. Их долгосрочное видение — проектирование лекарств с нуля, но первоначально они сосредоточены на создании важных биологических наборов данных путем разработки новых лабораторных анализов. Конструирование белков включает многокритериальную оптимизацию, требующую, чтобы белки соответствовали различным критериям, таким как связывание, растворимость и стабильность. JAX упрощает интеграцию многочисленных моделей ИИ, каждая из которых предсказывает различное свойство, в единую функцию потерь. Escalante приняла экосистему JAX, даже переведя модели из других фреймворков, таких как PyTorch. Это позволяет использовать выразительный язык для конструирования белков, где модели могут быть составлены и преобразованы в конечную цель, все оптимизируемые с помощью jax.jit для повышения производительности. Их рабочий процесс переворачивает типичное обучение, оптимизируя входные последовательности с использованием коллекции фиксированных нейронных сетей в качестве сложной дифференцируемой функции потерь. Этот процесс аналогичен DeepDream, где градиенты направляют обновления последовательности к желаемым свойствам. Возможности автоматического дифференцирования и компиляции JAX имеют решающее значение для оптимизации этих сложных функций потерь. Нативная интеграция фреймворка с TPU способствует масштабированию этих рабочих нагрузок, при этом Escalante использует схему запуска и отключения TPU по мере необходимости. Это использование TPU обеспечивает значительную экономическую эффективность по сравнению с GPU для их крупномасштабных задач. Ключевые библиотеки экосистемы JAX, такие как Equinox и Optax, используются для представления модели и гибкости алгоритма оптимизации. Сочетание функционального ядра JAX, его библиотек экосистемы и масштабируемого оборудования TPU обеспечивает новаторские исследования Escalante.