ChrisHayduk/minAlphaFold2

What it solves

minAlphaFold2 provides a minimal, pedagogical PyTorch implementation of AlphaFold2. It addresses the difficulty of understanding the original system by removing production-grade complexity (like JAX-tuned code or heavy scaffolding) and replacing it with a clear, 1-to-1 mapping between the code and the original 62-page supplement paper.

How it works

The project implements the AlphaFold2 architecture using pure PyTorch primitives (nn.Linear, torch.einsum, etc.) without external ML libraries. It follows the paper's specifications exactly, including:

• Architecture: All 32 supplement algorithms are implemented as distinct modules, from the Evoformer stack to the Structure Module.
• Training: A two-stage training recipe (initial and fine-tune) using gradient accumulation to simulate large-scale TPU batches on a single GPU.
• Data Pipeline: Integration with OpenProteinSet for MSAs and templates, including deterministic filtering for resolution and sequence composition.
• Refinement: A faithful port of the iterative restrained Amber minimisation for relaxing predicted structures.

Who it’s for

• AI researchers and students wanting to learn the inner workings of AlphaFold2 without navigating production codebases.
• Developers looking for a compact, trainable reference implementation to fork and tweak for protein folding research.
• Biologists with a coding background who want a legible entry point into AI-driven structural biology.

Highlights

• Paper-to-Code Mapping: Every module corresponds to a numbered algorithm in the AF2 supplement.
• Pure PyTorch: No custom CUDA kernels or einops, making the code highly readable and standard.
• Single-GPU Trainable: Uses gradient checkpointing and accumulation to fit the 48-block Evoformer on consumer or cloud GPUs.
• End-to-End Pipeline: Includes data downloading, preprocessing, training, and structure relaxation scripts.