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에 맞춥니다.
• 엔드 투 엔드 파이프라인: 데이터 다운로드, 전처리, 학습 및 구조 완화 스크립트가 포함되어 있습니다.