기술 요약: Emyx – 빠르고 효율적인 전원자(All-Atom) 단백질 생성

문제 정의

계산적 효소 설계는 촉매 잔기와 리간드를 정밀하게 배치하는 단백질 스캐폴드를 생성해야 합니다. 이 작업은 기하학적 정확성과 구조적 다양성을 모두 요구합니다. RFdiffusion3 및 Proteína-Complexa와 같은 기존의 전원자 생성 모델들은 단백질 구조 예측(예: 조밀한 쌍 표현, 무거운 임베딩 스택)에서 유래한 복잡한 아키텍처를 상속받습니다. 이러한 아키텍처는 높은 훈련 비용, 제한된 샘플 다양성을 초래하며, 고품질 서열을 생성하기 위해 종종 비용이 많이 드는 역접힘(inverse folding) 단계를 필요로 합니다. 또한, 기존의 평가 지표(예: 중원자 자기 일관성 RMSD)는 전역 폴드 회복(global fold recovery)을 검증하지 못함으로써 성공을 과장할 수 있으며, 이는 국부적인 모티프 제약 조건은 만족하지만 안정적이고 접힘 가능한 단백질을 형성하지 못하는 설계로 이어질 수 있습니다.

방법론

모델 아키텍처

Emyx는 전원자 단백질 생성을 위해 설계된 1억 4천만 파라미터 규모의 조건부 플로 매칭(conditional flow matching) 모델입니다. 조밀한 쌍 표현과 도메인 특화 임베딩 레이어에 의존하는 선행 모델들과 달리, Emyx는 표준 트랜스포머 블록 내에 용량을 집중시킵니다.

• 표현(Representation): 모델은 통합된 2단계 표현을 사용합니다. 원자 수준에서 잔기와 리간드는 고정된 크기의 "Rep14" 행렬(14개의 원자 슬롯)로 매핑되며, 사용되지 않는 위치는 백본 좌표를 상속받는 "고스트 원자(ghost atoms)"로 채워집니다. 토큰 수준에서 잔기와 HETATM은 단일 토큰을 구성합니다.
• 연결성(Connectivity): 조밀한 쌍 표현 대신, Emyx는 고정된 엣지 예산을 가진 희소 연결 그래프를 활용합니다. 엣지는 서열 이웃, 화학 결합, 리간드 k-NN 및 모티프 연결성을 기반으로 구축됩니다.
• 트랜스포머 설계: 아키텍처는 "패치화-처리-언패치화(patchify–process–unpatchify)" 파이프라인을 채택합니다:
  1. 원자 인코더(Atom Encoder): 경량 트랜스포머 블록이 Rep14 표현 내의 국부적 원자 세부 사항을 인코딩합니다.
  2. 다운캐스트(Downcast): 게이트형 교차 주의(gated cross-attention)를 통해 원자 표현을 토큰 수준 임베딩으로 풀링합니다.
  3. 토큰 트렁크(Token Trunk): 심층 트랜스포머 트렁크(18개 레이어)가 학습된 쌍 편향(learned pair biases)을 가진 희소 셀프 어텐션을 사용하여 장거리 구조적 맥락을 처리합니다.
  4. 업캐스트(Upcast): 두 번째 교차 주의가 토큰 맥락을 다시 원자로 브로드캐스트합니다.
  5. 원자 디코더(Atom Decoder): 최종 출력 속도 헤드(velocity head) 전 단계에서 마지막 트랜스포머 블록이 국부적 원자 세부 사항을 정제합니다.
• 수정 사항: 표준 DiT(Diffusion Transformer) 블록은 조건화를 위한 병목 투영(bottleneck projections), 잔차 업데이트를 제한하기 위한 시그모이드 게이팅(sigmoid gating), 그리고 정규화를 위한 스토캐스틱 뎁스(stochastic depth, drop-path)를 적용하여 수정되었습니다.

훈련 및 목적 함수

Emyx는 카테시안 좌표에 대해 조건부 플로 매칭을 사용하여 훈련됩니다.

• 목적 함수: 모델은 노이즈 분포에서 데이터 분포로 샘플을 이동시키기 위해 시간 의존적 속도장 $v_\theta(x_t, t, c, G)$를 학습합니다. 훈련 손실은 모티프가 아닌 원자에 대한 $\ell_2$ 회귀와 구조적 품질을 강제하기 위한 보조 미분 가능 로컬 거리 차이 테스트(lDDT) 손실을 결합합니다.
• 조건화(Conditioning): 모델은 희소 기하학적 제약(모티프 기하학, 리간드 위치)과 전역 특징(2차 구조 조성, 회전 반경)에 조건화됩니다. 이러한 특징들은 훈련 중에 확률적으로 마스킹되어 조건부 및 무조건부 생성을 모두 가능하게 합니다.
• 효율성: 훈련에는 노이즈와 데이터 사이의 선형 보간법을 사용하여, 분산 보존 확산 스케줄(variance-preserving diffusion schedules)에 비해 더 낮은 분산의 그래디언트와 빠른 수렴을 제공합니다. 모델은 약 8개의 NVIDIA H200 GPU(681.6 GPU-시간)에서 약 3.55일 동안 실험적 PDB 데이터를 사용하여 훈련되었습니다.

샘플링 및 재매개변수화

주요 방법론적 기여는 플로 매칭 보간법을 EDM(Elucidating the Design Space of Diffusion-based Generative Models) 노이즈 레벨 프레임워크로 변환하는 **정확한 재매개변수화(exact reparametrisation)**의 유도입니다.

• 이 매핑을 통해 플로 매칭 모델은 재훈련이나 아키텍처 변경 없이 Karras 노이즈 스케줄과 확률적 쳐너(stochastic churn) 메커니즘을 활용할 수 있습니다.
• 이 재매개변수화는 플로 매칭의 훈련 효율성과 최첨단 확산 모델용 샘플링 방법을 연결하여, 표준 SDE Euler–Maruyama 적분보다 성공률을 크게 향 향상시킵니다.

주요 기여

1. 효율적인 전원자 생성기: Emyx는 예측 모델에서 상속된 무거운 아키텍처를 제거함으로써, 훨씬 작은 파라미터 수(140M)와 낮은 훈련 비용(RFdiffusion3보다 4배 적음)으로 최첨단 성능을 달고합니다.
2. 엄격한 평가 프로토콜: 저자들은 전역 폴드 회복(backbone RMSD < 2.0 Å)과 국부 촉매 기하학 정확도(motif tip-atom RMSD < 1.5 Å)를 모두 요구하는 "엄격한 sc-RMSD" 프로토콜을 제안합니다. 그들은 기존의 지표(heavy-atom sc-RMSD)가 활성 부위는 회복하지만 전역 폴드는 회복하지 못하는 설계를 처벌하지 못함으로써 성공률을 약 2배 과장한다는 것을 입증했습니다.
3. EDM 재매개변수화: 본 논문은 플로 매칭 속도 모델을 EDM 프레임워크로 변환하는 모델 불가지론적(model-agnostic) 방법을 도입하여, 재훈련 없이 우수한 샘플링 전략(Karras 스케줄, stochastic churn)을 사용할 수 있게 합니다.
4. 스펙트럼 분석: 저자들은 "생성기-예측기 간극(generator–predictor gap)"을 드러내는 스펙트럼 분석을 제공합니다. 풍부한 MSA/템플릿 신호로 초기화된 예측기는 더 높은 유효 랭크(effective rank)와 낮은 과적합을 보이는 반면, 노이즈가 섞인 좌표로부터 초기화된 생성기는 더 낮은 랭크 영역에서 작동합니다. Emyx는 아키텍처적 병목 현상과 스토캐스틱 뎁스 덕분에 생성기 중 가장 낮은 과적합 비율을 보여줍니다.

결과

AME(Atom-level Motif Enzyme) 벤치마크(41개 촉매 활성 부위)를 대상으로 평가한 결과:

• 성공률: 엄격한 sc-RMSD 프로토콜 하에서, Emyx는 **13.4%**의 성공률(41개 타겟 중 39개 해결)을 달성하여 Proteína-Complexa (8.8%)와 RFdiffusion3 (6.7%)를 능가했습니다.
• 구조적 신규성: Emyx는 가장 신규한 스캐폴드를 생성하며, 가장 가까운 PDB 히트와의 중앙값 TM-score는 0.475(낮을수록 더 신규함)로, Proteína-Complexa(0.499) 및 RFdiffusion3(0.512)보다 낮았습니다.
• 다양성: Emyx는 성공적인 설계 중 가장 많은 수의 고유 구조 클러스터(441개)를 생성했습니다.
• 기하학적 타당성: Emyx는 설계의 82.2%에서 기하학적 타당성 검사(결합 길이/각도)를 통과하여, RFdiffusion3 (0.8%) 및 Proteína-Complexa (49.9%)를 크게 앞질렀습니다. 저자들은 RFdiffusion3의 낮은 기하학 점수가 입력 모티프 좌표가 생성된 좌표를 덮어써서 비물리적인 결합을 만드는 후처리 아티팩트 때문이라고 언급했습니다.
• 효율성: Emyx는 약 682 GPU-시간 동안 훈련되며(RFdiffusion3보다 4배 빠름), 모든 체인 길이에 대해 단일 A10G GPU에서 추론 속도가 2배 이상 빠릅니다.

의의 및 주장

본 논문은 아키텍처의 단순성이 단백질 생성기에 유익하다고 주장합니다. 풍부한 공진화 신호(MSA, 템플릿)를 사용하여 복잡한 쌍 표현을 초기화해야 하는 예측기와 달리, 생성기는 희소한 기하학적 제약에 조건화됩니다. 따라서 구조 예측에서 유래한 무거운 아키텍처는 불필요하며 오히려 효율성과 다양성에 해롭습니다.

Emyx는 표준 트랜스포머 블록에 용량을 집중하고 효율적인 플로 매칭 훈련과 EDM 스타일의 샘플링을 결합하면, 훨씬 적은 계산 자원으로 효소 설계 작업에서 우수한 성능을 달성할 수 있음을 보여줍니다. 또한, 이 연구는 국부적으로는 정확하지만 전역적으로는 불안정한 설계와 진정으로 실행 가능한 효소 스캐폴드를 구분하기 위해 엄격한 평가 프로토콜(strict sc-RMSD)이 중요함을 강조합니다. 마지막으로, 제안된 EDM 재매개参数화는 다른 플로 매칭 모델들이 아키텍처의 전면적인 개편 없이 고급 샘플링 기법을 채택할 수 있는 경로를 제공합니다.