PI-Mamba: Linear-Time Protein Backbone Generation via Spectrally Initialized Flow Matching

1. 문제 제기 (Problem)

단백질 백본 (backbone) 생성을 위한 기존 생성 모델들은 다음과 같은 근본적인 딜레마에 직면해 있습니다.

• 기하학적 유효성 vs. 계산 효율성: 대부분의 기존 방법 (확산 모델 등) 은 물리적 제약 (결합 길이, 결합 각도 등) 을 소프트 페널티나 사후 (post-hoc) 정제 과정을 통해 처리합니다. 이로 인해 시퀀스 길이가 길어질수록 결합 위반이 누적되고, 정확한 구조를 얻기 위해 비용이 많이 드는 정제 단계 (예: Rosetta FastRelax) 가 필수적이 되어 효율성이 떨어집니다.
• 확장성 (Scalability) 한계: 기존 트랜스포머 기반의 어텐션 메커니즘은 시퀀스 길이에 대해 이차 (quadratic, $O(L^2)$) 로 복잡도가 증가하여 긴 단백질 (2,000 개 이상의 잔기) 생성 시 메모리 및 계산 비용이 급증합니다.
• 물리적 무결성 부재: 생성된 구조가 물리적으로 불가능한 형태 (예: 비현실적인 결합 길이) 를 가질 수 있으며, 이를 보정하는 과정에서 원래의 설계 의도가 훼손될 수 있습니다.

2. 방법론 (Methodology)

저자들은 PI-Mamba(Physics-Informed Mamba) 라는 새로운 생성 모델을 제안합니다. 이는 물리 법칙을 내재화하고 선형 시간 복잡도를 가진 상태 공간 모델 (SSM) 을 결합한 프레임워크입니다.

• 물리 정보 기반 스펙트럴 초기화 (Physics-Informed Spectral Initialization):
  • Mamba 아키텍처의 상태 전이 행렬을 무작위가 아닌 Rouse 고분자 모델에서 유래한 스펙트럼으로 초기화합니다.
  • 이를 통해 저주파 모드는 전역적 토폴로지를, 고주파 모드는 국소적 변동을 포착하도록 유도 인덕션 (inductive bias) 을 부여합니다. 이는 생성의 안정성을 높이고 'molten-globule'과 같은 기하학적 병리를 방지합니다.
• SE(3) 기반 플로우 매칭 (Geometric Flow Matching):
  • 단백질 백본을 $SE(3)^L$ (리만 다양체) 상의 점으로 모델링합니다.
  • 리 군 (Lie group) 적분을 통해 시간 의존적 벡터 필드를 학습하여 노이즈에서 구조로 이동하는 경로를 생성합니다.
• 생성 과정 내 기하학적 제약 강제 (Constraint Enforcement during Generation):
  • 사후 정제가 아닌 생성 과정 중에 제약 조건을 적용합니다.
  • Kinematic Projection Operator ($\Pi_K$): 생성 단계마다 $C_\alpha$ 궤적을 결정론적으로 재추출 (retraction) 하여 이상적인 결합 길이 (3.80 Å) 와 펩타이드 평면성을 유지합니다.
  • NeRF (Natural Extension Reference Frame): 최종 출력 시 고정된 내부 좌표 (결합 길이, 각도, 평면성) 를 사용하여 원자 좌표 (N, $C_\alpha$, C, O) 를 결정론적으로 재구성합니다. 이를 통해 결합 길이 위반을 0 으로 보장합니다.
• 선형 시간 복잡도:
  • Mamba 의 선택적 상태 공간 (Selective State Space) 아키텍처를 사용하여 시퀀스 길이에 대해 선형 ($O(L)$) 인 추론 시간과 메모리를 달성합니다.

3. 주요 기여 (Key Contributions)

1. 완벽한 국소 기하학 보장: 사후 정제 없이도 0.0% 의 결합 길이 및 각도 위반을 달성합니다. 생성 과정 자체가 물리적으로 유효한 다양체 (manifold) 상에서 이루어집니다.
2. 대규모 단백질 생성의 확장성: 단일 GPU (NVIDIA A5000, 24GB) 에서 2,000 개 이상의 잔기를 가진 단백질 생성이 가능하며, 기존 확산 모델 대비 메모리 사용량이 50~80 배 적고 속도가 수십 배 빠릅니다.
3. 물리 기반 초기화의 효과: Rouse 고분자 모델을 기반으로 한 초기화가 학습된 내부 역학 (잠재 상태의 진동 모드, 이완 시간 등) 에 물리적으로 의미 있는 구조를 부여함을 실험적으로 증명했습니다.
4. 높은 설계 가능성 (Designability): 자기 일관성 TM-스코어 (scTM) 가 $0.91 \pm 0.03$으로, 기존 최첨단 모델들과 경쟁하면서도 선형 시간 효율성을 제공합니다.

4. 실험 결과 (Results)

• 벤치마크 성능 (L=100):
  • scTM: 0.91 (RFdiffusion 0.97, FrameDiff 0.58 대비 매우 우수).
  • 기하학적 위반: 최종 위반률 0.0% (사후 정제 없이도 0.1% 미만의 초기 위반률).
  • 생성 속도: 샘플당 2.3 초 (RFdiffusion 대비 21 배, FrameDiff 대비 37 배 빠름).
• 확장성 테스트 (L=1,000 이상):
  • PI-Mamba 는 L=1,000 에서도 21.3 초 내에 생성 완료 및 0.41 GB 메모리 사용.
  • 반면, FrameDiff, Genie2 등 확산 기반 모델들은 L>500 에서 OOM(메모리 부족) 발생 또는 L=1,000 에서도 10 배 이상의 메모리와 900 초 이상의 시간이 소요됨.
• 물리적 통계 일치:
  • 생성된 구조의 Rouse 모드 진폭 스펙트럼이 이론적 $1/p^2$ 감쇠를 따름.
  • 회전 반경 ($R_g$) 이 $L^{0.48}$로 스케일링되어 고분자 이론과 일치함.
  • 이차 구조 (helix, sheet, loop) 에 따라 학습된 이완 시간 ($\tau$) 이 물리적으로 타당한 계층 구조를 보임.
• 시퀀스 복구 및 다양성:
  • 시퀀스 복구율 23.9% (C$\alpha$만 사용 시), 전체 백본 재구성 시 31.2% 로 RFdiffusion(33.4%) 과 유사한 수준.
  • 새로운 구조 (Novelty Score 0.56) 를 생성하며 학습 데이터와 유사하지 않음.

5. 의의 및 결론 (Significance)

PI-Mamba 는 단백질 설계 분야에서 "학습을 통해 단백질을 모방하는 것"에서 "물리 법칙에 따라 단백질을 구성하는 것"으로의 패러다임 전환을 보여줍니다.

• 효율성과 정확성의 동시 달성: 기존에는 상충관계로 여겨졌던 '계산 효율성 (선형 시간)'과 '기하학적 정확도 (물리 제약)'를 동시에 해결했습니다.
• 실용적 확장성: 단일 GPU 로도 긴 단백질 사슬을 생성할 수 있어, 대규모 단백질 복합체나 합성 생물학 응용에 실용적인 도구가 될 수 있습니다.
• 물리 기반 AI 의 진보: 단순한 데이터 패턴 학습을 넘어, 고분자 물리학 (Rouse 모델) 을 아키텍처 초기화와 제약 조건에 직접 통합함으로써 생성 모델의 신뢰성과 해석 가능성을 높였습니다.

이 연구는 단백질 생성 모델이 더 이상 사후 정제에 의존하지 않고, 물리적으로 유효한 구조를 선형 시간 내에 생성할 수 있음을 입증하여 차세대 단백질 설계 도구의 새로운 기준을 제시합니다.