アクセラレータプログラミングとプログラム変換のためのPythonライブラリである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の画期的な研究が可能になります。