JAX, une bibliothèque Python pour la programmation d'accélérateurs et la transformation de programmes, s'avère essentielle pour l'ingénierie des protéines pilotée par l'IA, étendant son impact au-delà de l'entraînement de modèles d'IA à grande échelle. Escalante, une startup, utilise JAX pour entraîner des modèles qui prédisent les effets des médicaments sur les niveaux d'expression des protéines cellulaires, illustrant la nature fonctionnelle et composable de JAX. Leur vision à long terme est de concevoir des médicaments à partir de zéro, mais ils se concentrent initialement sur la génération de jeux de données biologiques cruciaux en développant de nouveaux essais en laboratoire. L'ingénierie des protéines implique une optimisation multi-objectifs, exigeant que les protéines répondent à divers critères tels que la liaison, la solubilité et la stabilité. JAX simplifie l'intégration de nombreux modèles d'IA, chacun prédisant une propriété différente, dans une fonction de perte unifiée. Escalante a adopté l'écosystème JAX, traduisant même des modèles d'autres frameworks comme PyTorch. Cela permet un langage expressif pour la conception de protéines, où les modèles peuvent être composés et transformés en un objectif final, le tout optimisable avec jax.jit pour la performance. Leur flux de travail inverse l'entraînement typique en optimisant les séquences d'entrée à l'aide d'une collection de réseaux neuronaux fixes comme fonction de perte complexe et différentiable. Ce processus est analogue à DeepDream, où les gradients guident les mises à jour de séquence vers les propriétés souhaitées. Les capacités de différenciation automatique et de compilation de JAX sont cruciales pour optimiser ces fonctions de perte sophistiquées. L'intégration native du framework avec les TPU facilite la mise à l'échelle de ces charges de travail, Escalante utilisant un schéma de démarrage et d'arrêt des TPU selon les besoins. Cette adoption des TPU offre une efficacité de coût significative par rapport aux GPU pour leurs travaux à grande échelle. Des bibliothèques clés de l'écosystème JAX comme Equinox et Optax sont utilisées pour la représentation des modèles et la flexibilité des algorithmes d'optimisation. La combinaison du cœur fonctionnel de JAX, de ses bibliothèques d'écosystème et du matériel TPU évolutif permet la recherche révolutionnaire d'Escalante.