ChrisHayduk/minAlphaFold2

解決する課題

minAlphaFold2は、AlphaFold2の最小限かつ教育的なPyTorch実装を提供します。JAXチューニングされたコードや複雑な足場のようなプロダクションレベルの複雑さを排除し、コードと62ページの補足論文との間に明確な1:1の対応関係を持たせることで、元のシステムを理解する難しさを解消します。

仕組み

このプロジェクトは、外部の機械学習ライブラリを使用せず、純粋なPyTorchプリミティブ（nn.Linear、torch.einsumなど）を使用してAlphaFold2アーキテクチャを実装しています。以下を含む論文の仕様に厳密に従っています。

• アーキテクチャ： 32個の補足アルゴリズムすべてが、EvoformerスタックからStructure Moduleまで、個別のモジュールとして実装されています。
• トレーニング： 勾配累積を使用して単一のGPU上で大規模なTPUバッチをシミュレートする、2段階のトレーニングレシピ（初期および微調整）。
• データパイプライン： MSAおよびテンプレート用のOpenProteinSetとの統合。解像度と配列組成の決定論的フィルタリングを含みます。
• リファインメント： 予測された構造を緩和するための、反復的な拘束付きAmber最小化の忠実な移植。

対象読者

• プロダクションコードベースを読み解くことなく、AlphaFold2の内部動作を学びたいAI研究者や学生。
• タンパク質折り畳み研究のためにフォークして調整できる、コンパクトでトレーニング可能なリファレンス実装を探している開発者。
• AI駆動の構造生物学への読みやすい入り口を求めている、プログラミングの背景を持つ生物学者。

ハイライト

• 論文とコードの対応： すべてのモジュールがAF2補足論文の番号付きアルゴリズムに対応しています。
• 純粋なPyTorch： カスタムCUDAカーネルや einops を使用せず、コードの可読性と標準性を高めています。
• 単一GPUでのトレーニング： 勾配チェックポインティングと累積を使用することで、48ブロックのEvoformerをコンシューマー向けまたはクラウドGPUに収めています。
• エンドツーエンドパイプライン： データのダウンロード、前処理、トレーニング、構造緩和のスクリプトが含まれています。