Design-CP: Context Parallelism for Design of Protein Nanoparticles

기술 요약: Design-CP

문제 정의

RFD3와 같은 전원자(all-atom) 생성 단백질 모델은 모든 체인을 공동으로 모델링함으로써 대규모 다량체 복합체를 설계할 수 있는 잠재력을 가지고 있습니다. 그러나 기본 아키텍처는 토큰 수($I$)와 원자 수($L$)에 따라 제곱($O(I^2)$ 및 $O(L^2)$)으로 증가하는 쌍 표현(pairwise representations; 토큰-토큰 및 원자-원자)에 의존합니다. 대상 어셈블리의 크기가 커짐에 따라, 이러한 이차적인 피처 텐서와 중간 활성화 값은 단일 GPU의 메모리 용량을 빠르게 초과하게 됩니다. 이러한 "단일 장치 한계(single-device ceiling)"로 인해 현재의 파이프라인은 대규모 대칭 어셈블리(예: 정이십면체 또는 정팔면체 나노입자)를 위한 엔드 투 엔드 공동 생성 대신, 독립적으로 설계된 올리고머를 강체 도킹(rigid-body docking)하는 방식에 의존해야 합니다. LLM과 같은 인접 분야에서는 멀티 GPU 클러스터가 표준이지만, 구조 모델에 이와 유사한 병렬화 전략을 적용하는 것은 쌍형 어텐션(pairwise attention) 메커즘의 특유한 메모리 및 통신 패턴으로 인해 제한되어 왔습니다.

방법론

저자들은 RFD3를 위한 두 가지 컨텍스트 병렬(context-parallel, CP) 추론 전략을 구현하여, 사전 학습된 가중치를 보존하면서(즉, 재학습이나 미세 조정을 요구하지 않음) 다중 GPU 메시에 이차적인 활성화 값을 분산시키는 프레임워크인 Design-CP를 소개합니다.

1. 추론 전략

본 논문은 두 가지 샤딩(sharding) 기법을 구현하고 비교합니다:

• 1D 행 샤딩(1D Row-Sharding): 쌍 표현($Z$ 및 $P$)을 $P$개의 GPU에 대해 단일 축(쿼리)을 따라 줄무늬 형태로 배치하는 경량화된 전략입니다. 키(key)와 값(value)은 모든 GPU에 복제됩니다. 이 방식은 셀프 어텐션을 크로스 어텐션으로 재구성하며, 각 어텐션 블록 이후에 단일 트랙 활성화를 한 번의 ALLGATHER로 수행하여 모든 장치가 다음 단계를 위해 동일한 키와 값을 갖도록 보장합니다.
• 2D 그리드 샤딩(2D Grid Sharding): Fold-CP에서 채택된 전략으로, 쌍 텐서를 $\sqrt{P} \times \sqrt{P}$ 장치 그리드 상에 타일링합니다. 이 접근 방식은 링 어텐션(Ring Attention)(Liu et al., 2023)을 활용하여 온더플라이(on the fly)로 어텐션 점수를 계산합니다. 쿼리는 행을 따라 샤딩되는 반면, 키와 값은 열을 따라 샤딩되어 링 시프트(ring shifts)를 통해 순환합니다. 이를 통해 장치들은 전체 $I \times I$ 어텐션 행렬을 실제로 생성하지 않고도 로컬 쿼리 행과 들어오는 키/값 샤드 사이의 어텐션을 계산할 수 있습니다. 2D 방식은 집합 통신을 $\sqrt{P}$ 크기의 서브그룹으로 제한하여, 1D 샤딩에 비해 통신 볼륨과 일시적 메모리 오버헤드를 줄입니다.

2. 대칭 제약 조건

방법론의 핵심 요소는 샘플링 중 강력한 점군 대칭(point-group symmetry) 제약 조건을 적용하는 것입니다. 고도로 대칭적인 대상(예: 정이십면체, 정팔면체)의 경우, 확산 모델은 전체 복합체에 대해 실행되지만 각 디노이징 단계마다 예측값을 재대칭화(resymmetrizes)합니다. 단일 비대칭 단위(Asymmetric Subunit, ASU)의 좌표만 유지되며, 나머지 단위들은 군 연산(group operations)을 통해 생성된 결정론적 복사본으로 덮어씌워집니다. 이는 설계 공간을 효과적으로 축소하여 자유도를 단일 ASU로 줄이면서도, 전체 어셈블리의 상호작용을 여전히 모델링할 수 있게 합니다. 이 메커니즘을 통해 모델은 고유한 훈련 크롭(training crop) 제한을 초과하지 않으면서도 전체 어셈블리 문맥과 자기 일관성(self-consistency)을 유지할 수 있습니다.

3. 구현 세부 사항

• 메모리 최적화: 두 방식 모두 희소 k-최근접 이웃(kNN) 인덱스에서 원자 쌍 텐서 $P$를 온더플라이로 구축하는 "청크형 쌍 임베더(chunked pairwise embedder, 저메모리 모드)"를 자동 활성화하여 밀집된 실체화를 방지합니다.
• 결정론적 동작: 저자들은 수치적 동등성을 보장하기 위해 특정 통신 가드(예: 랭크 0에서 집계된 토큰 피처를 브로드캐스팅)를 구현하여, 원자 부동 소수점 리덕션에서의 비결정성을 해결했습니다.
• 하드웨어: 본 방법론은 고성능 데이터 센터 GPU(NVIDIA HG200, 95 GB)와 워크스테이션급 GPU(NVIDIA RTX A4000, 16 GB) 모두에서 테스트되었습니다.

주요 기여

1. Design-CP 프레임워크: 모델 가중치를 수정하지 않고도 여러 GPU에 걸쳐 대규모 다량체 단백질의 공동 디노이징을 가능하게 하는 RFD3용 두 가지 컨텍스트 병렬 추서 전략(1D 행 샤딩 및 2D 그리드 샤딩)을 도입했습니다.
2. 확장성 특성 분석: 최대 가능한 비대칭 단위(ASU) 크기가 GPU 개수의 제곱근($\sqrt{P}$)에 비례하여 확장된다는 분석을 제시하였으며, 이는 이차적인 메모리 비용의 분포와 일치합니다. 2D 방식은 계산과 통신을 중첩함으로써 더 우수한 벽시계 시간(wall-clock) 스케일링을 보여주었습니다.
3. 대칭 기반 설계: 강력한 점군 대칭 제약 조건이 정이십면체 나노입자의 엔드 투 엔드 전원자 설계를 위해 컨텍스트 병렬 처리를 "즉시 사용 가능한(usable out of the box)" 상태로 만든다는 것을 입증했습니다. 이 접근 방식은 샘플링 시 네이티브 훈련 범위를 벗어날 때 발생하는 분포 변화(distribution shift) 문제를 완화합니다.
4. 설계의 민주화: 소규모 클러스터의 워크스테이션급 GPU(16 GB 메모리)에서 정팔면체 케이지를 성공적으로 설계함으로써, 대규모 어셈블리 단백질 설계가 고성능 데이터 센터 하드웨어 없이도 가능하다는 것을 증명했습니다.

결과

• 정이십면체 설계: 4-GPU 클러스터(2x2 그리드)에서 2D 샤딩을 사용하여 60개 체인(총 12,600개 잔기)을 가진 정이십면체 나노입자를 생성했습니다. 각 체인의 백본 품질(체인 끊김, 충돌, 2차 구조)은 단일 GPU RFD3 단량체와 대체로 일치했으나, 대칭 어셈블리에서 약간 높은 기하학적 변형(예: 더 넓은 C$\alpha$–C$\alpha$ 결합 길이 편차)이 관찰되었습니다. 대칭 인식 인터페이스 지표(충돌, 접촉 수)는 자연적인 기능적 기준선(Lumazine Synthase, PDB: 1NQX)과 잘 일치하여 물리적으로 타당한 단위 간 인터페이스를 나타냈습니다.
• 스케일링 성능: 2D 샤딩 방식은 통신 지연 시간을 줄임으로써 동일한 GPU 개수에서 1D 방식보다 더 나은 추론 시간을 달acking하며, 종종 약 절반의 시간 내에 추론을 완료했습니다. OOM(Out-Of-Memory) 오류가 발생하기 전의 최대 ASU 길이는 단일 GPU의 약 58개 잔기에서 16개 GPU 사용 시 약 237개 잔기로 증가했습니다.
• 보급형 하드웨어에서의 정팔면체 설계: 저자들은 16개의 NVIDIA RTX A4000 GPU(각 16 GB)를 사용하여 정팔면체 케이지(24개 체인, 약 4,272개 잔기)를 성공적으로 설계했습니다. 결과물은 건강한 백본 지표와 잘 형성된 대칭 인터페이스를 보여주었으며, 이는 적절한 하드웨어에서도 이 접근 방식이 유효함을 확인시켜 줍니다.
• 대칭에 대한 절제 연구(Ablation): 대규모 시스템(예: 10,800개 잔기)에서 대칭 제약을 제거할 경우 퇴보하고 단백질 같지 않은 구조가 생성됨을 시각적 및 정량적 분석을 통해 확인했으며, 이는 대규모 규모에서 샘플 품질을 유지하는 데 대칭이 필수적임을 강조합니다.

의의 및 주장

본 논문은 Design-CP가 대규모 단백질 어셈블리에 대해 기존에 강체 도킹을 필요로 했던 단일 GPU 메모리 장벽을 제거한다고 주장합니다. 모든 원자를 단일 궤적 내에서 공동 디노이징할 수 있게 함으로써, 설계 패러다임을 "도킹 후 설계(dock-and-design)"에서 "엔드 투 엔드 전원자 생성"으로 전환합니다.

저자들은 이 연구를 대규모 어셈블리 단백질 설계의 민주화를 향한 단계로 규정합니다. 워크스테이션급 GPU에서도 고품질의 대칭 나노입자 설계가 가능하다는 것을 입증함으로써, 하드웨어 진입 장벽을 낮추고 자원이 풍부한 그룹 외의 참여를 넓힐 수 있습니다. 저자들은 자신들의 기여가 추론 및 병렬화 계층에 있으며, 새로운 모델 가중치나 훈련 데이터를 도입한 것이 아님을 명시합니다. 또한, 대칭 제약 없이 네이티브 훈련 크프를 넘어설 경우 비대칭 대상에 대한 출력이 저하될 수 있다는 한계를 인정하며, 설계된 어셈블리에 대한 실험적 검증이 향-후 필요한 과제임을 밝히고 있습니다.