JAX, eine Python-Bibliothek für Accelerator-Programmierung und Programmtransformation, erweist sich als entscheidend für die KI-gestützte Protein-Engineering und erweitert seine Wirkung über das Training von KI-Modellen im großen Maßstab hinaus. Escalante, ein Startup, nutzt JAX zum Trainieren von Modellen, die die Auswirkungen von Medikamenten auf die zelluläre Proteinexpression vorhersagen, was die funktionale und komponierbare Natur von JAX verdeutlicht. Ihre langfristige Vision ist es, Medikamente von Grund auf neu zu entwickeln, aber zunächst konzentrieren sie sich auf die Generierung wichtiger biologischer Datensätze durch die Entwicklung neuer Laborassays. Protein-Engineering beinhaltet eine Multi-Objektiv-Optimierung, die erfordert, dass Proteine verschiedene Kriterien wie Bindung, Löslichkeit und Stabilität erfüllen. JAX vereinfacht die Integration zahlreicher KI-Modelle, von denen jedes eine andere Eigenschaft vorhersagt, in eine einheitliche Verlustfunktion. Escalante hat das JAX-Ökosystem übernommen und Modelle sogar aus anderen Frameworks wie PyTorch übersetzt. Dies ermöglicht eine ausdrucksstarke Sprache für das Proteindesign, bei der Modelle komponiert und in ein endgültiges Ziel transformiert werden können, das alles mit jax.jit für die Leistung optimierbar ist. Ihr Workflow kehrt das typische Training um, indem er Eingabesequenzen mithilfe einer Sammlung fester neuronaler Netze als komplexe, differenzierbare Verlustfunktion optimiert. Dieser Prozess ist analog zu DeepDream, bei dem Gradienten Sequenzaktualisierungen in Richtung gewünschter Eigenschaften steuern. Die automatische Differenzierung und Kompilierungsfähigkeiten von JAX sind entscheidend für die Optimierung dieser hochentwickelten Verlustfunktionen. Die native Integration des Frameworks mit TPUs erleichtert die Skalierung dieser Workloads, wobei Escalante ein Muster des Hochfahrens und Herunterfahrens von TPUs nach Bedarf verwendet. Diese TPU-Nutzung bietet eine erhebliche Kosteneffizienz gegenüber GPUs für ihre groß angelegten Aufträge. Wichtige JAX-Ökosystembibliotheken wie Equinox und Optax werden für die Modellrepräsentation und die Flexibilität der Optimierungsalgorithmen genutzt. Die Kombination aus dem funktionalen Kern von JAX, seinen Ökosystembibliotheken und der skalierbaren TPU-Hardware ermöglicht die bahnbrechende Forschung von Escalante.