AlphaFold3 오픈 소스 구현 (Open-Source Implementation)

서론 (Introduction)

이 저장소는 전체 바이오테크 (biotech) 커뮤니티가 자유롭게 사용할 수 있도록, AlphaFold3의 충실하고 완전한 오픈 소스 구현을 향한 진전을 가속화하기 위해 의도되었습니다.

데모 영상 (Demo Video)

우리는 모델 학습 역학 (training dynamics)이 상당히 빠르다는 것을 확인했습니다. 다음 영상은 템플릿 (templates) 없이 8개의 A100 GPU에서 10시간 동안 4,000 스텝 (steps) 동안 학습된 모델의 샘플입니다.

애니메이션 크레딧: Matthew Clark

감사의 글 (Acknowledgments)

이 프로젝트는 다음 프로젝트와 개인들의 기여 없이는 불가능했을 것입니다:

• 혁신적인 연구를 수행하고 핵심 알고리즘을 발표한 Google DeepMind의 AlphaFold3 팀.

• 오픈 소스 단백질 구조 예측 (protein structure prediction)의 토대를 마련한 OpenFold 프로젝트 (https://github.com/aqlaboratory/openfold). 우리는 삼각 어텐션 (triangular attention) 및 곱셈 업데이트 (multiplicative update)와 같은 그들의 많은 핵심 모듈과 데이터 처리 파이프라인 (data processing pipelines)을 재사용합니다.

• ProteinFlow 라이브러리 (https://github.com/adaptyvbio/ProteinFlow), 특히 이 과정 내내 진정한 영웅이었던 ProteinFlow의 설계자 Liza Kozlova (@elkoz). 우리는 단백질 데이터 작업을 위한 깨끗하고 문서화가 잘 된 데이터 파이프라인을 제공하는 ProteinFlow를 사용하여 대부분의 프로토타입 모델을 학습시켰습니다. 우리는 AdaptyvBio와 협력하여 리간드 (ligand) 및 핵산 (nucleic acid)에 대한 완전한 지원을 포함하여 ProteinFlow를 기반으로 한 AlphaFold3의 데이터 파이프라인을 구축했습니다. @elkoz와 @igor-krawczuk은 이러한 데이터 양상 (data modalities)에 대한 완전한 지원을 포함하도록 ProteinFlow의 다음 릴리스를 구축하고 있습니다.

• @alexzhang13님이 Triton에서 구현한 커스텀 MSA 쌍 가중 평균 커널 (custom MSA pair weighted averaging kernel) 덕분에, 긴 시퀀스(sequence)에 대해 PyTorch 구현체보다 10배 이상 메모리 효율이 높아졌습니다. 이는 더 긴 시퀀스를 학습할 때 발생하는 결정적인 확장 병목 현상 (scaling bottleneck)을 제거해 줍니다. 이 프로젝트에 기여해 준 Alex에게 큰 감사를 전합니다!

프로젝트 상태 (Project Status)

이 프로젝트는 초기 단계에 있는 활발한 연구 프로젝트입니다. 저희는 커뮤니티를 위한 안정적인 릴리스 (stable release)를 준비하고 있습니다. 이 작업의 잠재력에 대해 기대하고 있지만, 아직 프로덕션 환경에서 사용할 수 있는 도구 (production-ready tool)는 아니라는 점을 강조하고 싶습니다.

저희는 구현 내용을 테스트하기 위해 단일 사슬 단백질 (single-chain proteins)을 대상으로 AlphaFold3의 버전을 학습시켰으며, 다음 릴리스에는 리간드 (ligand) 및 핵산 (nucleic acid)에 대한 완전한 지원이 포함될 예정입니다.

현재 구현 사항을 테스트하고 피드백을 제공해 줄 소수의 베타 테스터를 모집하고 있습니다. 베타 테스트에 관심이 있으시다면 대기 명단에 등록해 주세요.

AlphaFold3 의사코드와의 불일치 (Discrepancies from AlphaFold3's Pseudocode)

이 프로젝트를 진행하면서, AlphaFold3 보충 정보 (supplementary information)에 기술된 알고리즘의 몇 가지 특성이 주변의 딥러닝 (deep learning) 문헌들과 일치하지 않는다는 점을 발견했습니다.
이에 대해 아래에서 논의하겠습니다:

• MSA 모듈 순서 (MSA Module Order): 보충 정보 (Supplementary Information)에 따르면, MSA 모듈 통신 단계가 MSA 스택 (MSA stack)보다 먼저 발생합니다. 이로 인해 모든 정보가 쌍 표현 (pair representation)을 통해 흘러나가기 때문에, 마지막 블록의 MSA 스택은 구조 예측 (structure prediction)에 기여하지 못하게 됩니다. 의사 코드 (pseudocode)에 명시된 순서대로라면, 마지막 블록의 MSA 스택은 쌍 표현을 업데이트할 기회를 갖지 못합니다. 저희는 모든 블록이 구조 예측에 기여할 수 있도록 OuterProductMean 연산과 MSA 스택의 순서를 변경했습니다. 이러한 수정 사항이 AlphaFold2의 ExtraMSAStack에 있는 연산 순서와 일치한다는 점은 중요합니다. DeepMind는 이 MSA 모듈 블록들이 "동질적 (homogeneous)"이라고 언급했습니다. 이것이 가중치 공유 (shared weights)를 의미하는지, 아니면 블록 전체에 걸쳐 동일한 아키텍처 (architecture)를 의미하는지는 불분명합니다. 만약 레이어 (layers)가 공유된다면, 그래디언트 (gradients)는 모든 레이어를 통해 흐르겠지만 MSA 스택의 최종 계산은 유휴 상태가 됩니다. 이 경우 (의사 코드에는 언급되지 않았지만) 안전하게 건너뛸 수 있습니다. 저희는 DeepMind의 답변을 참고하여 이 모호함을 해결할 예정입니다.

• 손실 스케일링 (Loss scaling): 보충 정보 (Supplementary Information)에 기술된 손실 스케일링 계수는 초기화 시점에 단위 손실 (unit-loss)을 제공하지 않습니다. 초기화 시의 단위 손실은 Karras et al. (2022)이 확산 모델 (diffusion models)을 학습할 때 손실 함수가 갖추어야 할 바람직한 속성 중 하나로 설정한 특성 중 하나이며, Max Jaderberg는 이번 강연 여기에서 그들이 Karras et al.의 프레임워크를 선택한 이유 중 하나로 이 속성을 언급했습니다. 저희는 이것이 보충 정보의 단순한 오타라고 생각하며, 덧셈이 곱셈으로 잘못 기재된 것으로 보입니다. 저희의 구현에서는 Karras et al. (2022)과 일치하는 손실 스케일링 계수를 사용합니다. 저희의 측정 결과에 따르면, 이 방식은 초기화 시점에 단위 MSE 손실을 제공하는 반면, 보충 정보의 스케일링은 초기화 시점에 2~3 자릿수(orders of magnitude) 더 큽니다. 또한, 논문의 손실 스케일링 계수는 t = 16.0에서 국소 최솟값 (local minimum)을 갖지만, 노이즈 수준이 증가함에 따라 다시 증가합니다. 이는 높은 노이즈 수준에서 손실의 가중치를 낮추는 것의 중요성을 강조한 Karras et al. (2022)이 제안한 손실 함수의 속성과 일치하지 않습니다. 저희는 실험 내용을 보여주는 Jupyter notebook을 저장소에 추가했습니다.

• DiT 블록 설계 (DiT block design): AttentionPairBias 및 DiffusionTransformer 블록의 설계는 Peebles & Xie (2022)가 여기에서 소개한 DiT 블록 설계를 밀접하게 따르는 것으로 보입니다. 하지만 잔차 연결 (residual connections)이 누락되어 있습니다. DeepMind가 왜 이를 생략하기로 선택했는지는 논문에 설명되어 있지 않습니다. 저희는 두 방식 모두를 실험해 보았으며, (저희가 모델을 학습시킨 스텝 범위 내에서) 잔차 연결이 있는 DiT 블록이 훨씬 더 빠른 수렴 (convergence)과 네트워크를 통한 더 나은 그래디언트 흐름 (gradient flow)을 제공한다는 것을 발견했습니다. 참고로, 이 부분은 저희가 가장 확신하지 못하는 불일치 사항이며, 만약 원본 구현에서 잔차 연결을 사용하지 않는다면 저희 코드의 몇 줄만 수정하여 변경할 수 있습니다.

이러한 사항들은 투명성을 위해 여기에 기록되었으며, 이를 해결하기 위한 최선의 접근 방식에 대해 커뮤니티의 의견을 구하고자 합니다.

모델 효율성 (Model Efficiency)

이번 구현의 주요 초점은 속도와 메모리 효율성을 위해 모델 구성 요소를 최적화하는 데 있었습니다. AlphaFold3에는 많은 트랜스포머(Transformer) 유사 구성 요소가 포함되어 있지만, AlphaFold3의 쌍 편향 (pair biasing)으로 인해 FlashAttention2와 같이 하드웨어 인지적인 효율적인 어텐션 (attention) 구현체들이 이러한 모듈과 즉시 통합되지는 않습니다. 모든 어텐션 연산은 쌍별 표현 (pairwise representation)으로부터 쌍 편향을 투영하며, 이는 키-쿼리 (key-query) 내적 이후에 더해집니다. 또한 이 편향은 역전파 (backpropagation)를 위한 그래디언트 (gradient)를 필요로 합니다. 이는 FlashAttention2의 범위를 벗어나는 문제는 아닙니다. 왜냐하면 편향 그래디언트가 스케일링된 QK^T 내적과 동일한 그래디언트를 갖기 때문이지만, 현재 구현체는 이를 지원하지 않습니다. FlexAttention과 같은 더 최신의 어텐션 구현체들은 매우 유망하지만, 이들 역시 현재로서는 편향 그래디언트를 지원하지 않습니다. 이는 순전파 (forward pass) 동안 편향 텐서의 브로드캐스팅 (broadcasting) 연산이 역전파 (backward pass) 동안 리덕션 (reduction) 연산이 되기 때문이며, 이러한 기능은 FlexAttention의 첫 번째 릴리스에 구현되어 있지 않습니다.

• 저희는 가능한 곳마다 OpenFold 프로젝트의 TriangularAttention 및 TriangularMultiplicativeUpdate와 같이 검증된 구성 요소들을 재사용합니다. OpenFold 프로젝트의 모듈식 설계 덕분에 이러한 모듈들을 저희 코드베이스로 쉽게 가져올 수 있습니다. 저희는 Triton을 사용하여 이러한 모듈의 효율성을 개선하고 있으며, 연산을 퓨징 (fusing)하여 성능을 높이고 중간 텐서 할당을 줄이는 작업을 진행 중입니다.

• PyTorch에서의 Diffusion Module (확산 모듈)을 단순하게 구현할 경우, 배치당 Diffusion Module이 48회 복제되기 때문에 빈번하게 메모리 부족 (Out of Memory) 현상이 발생하는 것을 관찰했습니다. 이 문제를 해결하기 위해, Deepspeed4Science의 MSARowAttentionWithPairBias 커널을 재용도화하여 Diffusion Module의 메모리 효율적인 버전을 구현하였으며, 서로 다른 노이즈 레벨을 가진 배치 복제본들을 추가적인 배치 차원 (batch dimension)으로 취급하도록 했습니다. AtomAttentionEncoder 및 AtomAttentionDecoder 모듈의 경우, 메모리 사용량을 이차 복잡도 (quadratic)에서 선형 복잡도 (linear)로 줄이기 위해 커스텀 PyTorch 네이티브 구현을 실험해 보았으나, AttentionPairBias 커널을 단순하게 재용도화하는 것에 비해 이점이 그리 크지 않았습니다. 저장소에는 두 가지 구현을 모두 포함하고 있지만, 복잡성을 줄이기 위해 단순한 구현을 사용하고 있습니다.
  이러한 최적화에도 불구하고, 프로파일링 (profiling) 실험 결과 모델 연산의 60% 이상이 메모리 대역폭에 제한되는 메모리 바운드 (memory-bound) 상태임을 보여줍니다. 저희는 ScaleFold의 아이디어를 사용하여 훨씬 더 효율적이고 확장 가능한 구현을 작업 중이며, 이를 통해 원본 AlphaFold3의 학습 규모 (training scale)에 도달할 수 있을 것입니다.

MSA Pair Averaging 효율성

@alexzhang13은 Triton을 사용하여 빠르고 메모리 효율적인 커스텀 MSA pair 가중 평균 (weighted averaging) 커널을 구현했습니다.

• 저희는 주요 메모리 병목 현상 중 하나가 MSA pair 가중 평균 (weighted averaging) 연산이라는 점을 관찰했습니다. AlphaFold3 논문에서는 이 연산이 MSARowAttentionWithPairBias 연산을 더 "저렴한" pair 가중 평균으로 대체한다고 명시하고 있지만, 이를 PyTorch에서 단순하게 (naively) 구현할 경우 Deepspeed4Science의 MSARowAttentionWithPairBias 커널과 비교했을 때 메모리 사용량이 4배 증가합니다. 저희는 이것이 Deepspeed4Science MSARowAttentionWithPairBias 커널에도 포함되어 있는 FlashAttention의 타일링 (tiling) 및 재계산 (recomputation) 기법을 통한 메모리 효율성 이득 때문이라고 가설을 세웠습니다.
• pair 가중 평균을 단순하게 구현하면 (*, N_seq, N_res, N_res, heads, c_hidden) 크기의 중간 텐서 (intermediate tensor)를 할당하게 되는데, 이는 중간 길이의 서열에 대해서도 GPU 메모리에 담기에는 너무 큽니다.
• Alex의 커널은 단일 GPU에서 수천 개의 토큰까지 네트워크를 확장할 수 있게 해줍니다!

PyTorch vs. Triton 커널 메모리 사용량

PyTorch vs. Triton 커널 실행 시간

시작하기

리간드-단백질 (ligand-protein) 및 핵산 (nucleic acid) 예측 기능은 아직 학습되지 않았기 때문에, 아직 샘플링 코드는 제공하지 않습니다. 체크포인트 가중치는 실험 및 모델 수술 (model surgery)을 위해 PyTorch Lightning의 체크포인트 로딩 기능을 통해 불러올 수 있습니다. 현재 모델은 오리지널 AlphaFold2와 동일한 기능인 단일 사슬 단백질 (single-chain proteins)만을 예측합니다. 모델 구성 요소들은 연구자들이 자신의 프로젝트에 쉽게 통합할 수 있도록 재사용 가능하고 모듈화되도록 작성되었습니다.
리간드-단백질 및 핵산 예측의 베타 테스트를 원하시면: 대기 명단에 참여하세요

사용법

현재 이 저장소의 주요 용도는 연구 및 개발입니다. 리간드-단백질 및 핵산 예측 기능이 준비되는 대로 향후 더 많은 사용자용 기능을 포함할 예정입니다.

기여하기

커뮤니티의 기여를 환영합니다! 저희 코드에는 수많은 버그와 미세한 구현 오류가 존재할 가능성이 높습니다. 딥러닝 (Deep learning) 학습은 종종 오류가 발생하더라도 네트워크가 수렴(converge)은 되어 성능이 약간 저하될 뿐, 오류가 겉으로 드러나지 않고 조용히 실패하는 경우가 많습니다. 기여에 관심이 있으시다면, 버그 설명을 포함한 Github 이슈 (issue)를 제기하거나, 저장소 (repository)를 포크 (fork)하여 수정 사항이 담긴 새 브랜치 (branch)를 만든 후 변경 사항에 대한 명확한 설명과 함께 풀 리퀘스트 (pull request)를 제출해 주세요.

그 외의 의견이나 제안 사항은 이메일(alphafold3@ligo.bio)을 통해 저희에게 연락해 주시기 바랍니다.

인용 (Citations)

본 코드를 연구에 사용하실 경우, 다음 논문을 인용해 주세요:

@article{Abramson2024-fj,
  title    = "Accurate structure prediction of biomolecular interactions with
              {AlphaFold} 3",
...