기술 요약: ENSEMBITS

문제 정의

단백질 언어 모델 (PLM) 과 구조 예측 도구는 주로 정적 구조 표현이나 단순한 아미노산 서열에 의존해 왔습니다. 기존 단백질 구조 토크나이저 (PST) 는 정적 스냅샷의 국소 기하학을 효과적으로 포착하지만, 단백질 앙상블에 내재된 상관된 운동과 대체 입체 구조 상태를 인코딩하지는 못합니다. 생물학적 기능은 종종 단일 구조가 아닌 단백질이 샘플링하는 입체 구조의 분포에 암호화됩니다 (예: 기질 접근을 제어하는 촉매 루프). 연구 분야가 앙상블 생성으로 이동하고 있지만, 트랜스포머 기반 아키텍처에서 이러한 동적 앙상블을 토큰화할 수 있는 이산 어휘는 부재합니다. 또한, 동적 데이터는 종종 희소하며, 기존 방법들은 잔기 간 상관 운동을 공동으로 모델링하지 못하거나 추론 시점에 사용 가능하지 않을 수 있는 전체 궤적을 필요로 합니다.

방법론: ENSEMBITS

ENSEMBITS는 단백질 입체 구조 앙상블을 위해 특별히 설계된 최초의 토크나이저로 소개됩니다. 이는 구조 프레임 (입체 구조체) 의 순서 없는 다중 집합을 잔류 벡터 양자화 변이 오토인코더 (RVQ-VAE) 아키텍처를 통해 이산 토큰으로 매핑합니다.

1. 문제 공식화

이 모델은 앙상블 토크나이징을 다중 집합 - 토큰 압축 문제로 취급합니다.

• 입력: 단일 단백질에 대한 $P$개의 프레임으로 구성된 순서 없는 다중 집합 $P = \{x_1, \dots, x_P\}$이며, 각 프레임은 원자별 좌표를 포함합니다.
• 출력: 앙상블 전체에 걸쳐 잔기의 국소 동적 모티프를 나타내는 코드북 $\mathcal{C}$의 이산 토큰.
• 제약 조건: 토크나이저는 $P$에 구애받지 않아야 하며, 임의의 양의 크기를 가진 앙상블을 수용하고 프레임 순서에 대해 치환 불변이어야 합니다.

2. 기술자 설계

동적 특성을 포착하기 위해 ENSEMBITS 는 프레임 간 SE(3)-불변 기술자를 계산합니다. 저자들은 잔기의 프레임과 그 $K$개의 공간적 가장 가까운 이웃 간의 SE(3) 변환으로 국소 골격 환경을 인코딩하는 ESM3 스타일의 상대 프레임 기술자(생산 환경) 를 활용합니다.

• 동적 모드: 핵심적으로, 이웃 식별자 ($N^p_r$) 는 모든 프레임 $p$에 대해 독립적으로 재계산됩니다. 이는 고정된 이웃 접근법으로는 보이지 않는 신호인 궤적을 따라 이루어지는 접촉 형성과 파괴를 기술자가 자연스럽게 인코딩할 수 있게 합니다.

3. 아키텍처

• 셋 인코더: 잔기별 기술자 다중 집합을 단일 잠재 벡터 $z$로 매핑하는 PerceiverIO 스타일 아키텍처입니다. 공유된 요소별 MLP 와 학습 가능한 쿼리 토큰을 통한 크로스 어텐션을 사용하여 구성상 치환 불변성을 보장합니다.
• 잔류 양자화기: $K=3$ 단계의 다단계 RVQ-VAE 는 잠재 벡터 $z$를 코드북 인덱스 튜플 $(c_1, \dots, c_K)$로 이산화합니다. 최종 토큰은 이러한 단계들의 임베딩 합입니다.
• 디코더 및 손실 함수: 디코더는 기술자 다중 집합을 재구성합니다. 재구성 손실은 고정된 순서를 요구하지 않고 치환 불변성을 보장하기 위해 예측된 기술자와 목표 기술자를 정렬하는 헝가리안 매칭을 사용합니다.

4. 단일 프레임 - 토큰 증류 (SFTD)

동적 데이터의 희소성을 해결하고 단일 정적 구조에 대한 추론을 가능하게 하기 위해 저자들은 SFTD를 제안합니다.

• 메커니즘: 훈련 동안 모델은 동일한 단백질에서 두 개의 앙상블을 처리합니다. 전체 앙상블 ($P_{max}$) 과 무작위로 샘플링된 부분 앙상블 (최소 $P=1$까지) 입니다.
• 목적: 부분 앙상블의 잠재 표현은 전체 앙상블의 잠재 표현 쪽으로 당겨집니다 (전체 앙상블에는 그라디언트 정지 적용).
• 결과: 이를 통해 토크나이저는 테스트 시 단일 예측 구조에서 "동적 토큰"을 예측할 수 있게 되어, 앙상블 수준의 정보를 단일 프레임 표현으로 효과적으로 증류합니다.

주요 기여

1. 공식화 및 파이프라인: SE(3)-불변 기술자와 치환 불변 셋 인코더를 활용하여 단백질 동적 토크나이징을 다중 집합 - 토큰 압축 문제로 처음 종단간 (end-to-end) 공식화했습니다.
2. SFTD 목적 함수: 부분 앙상블 임베딩을 전체 앙상블 대응물과 정렬하는 새로운 훈련 전략으로, 단일 정적 구조만 존재할 때도 모델이 동적 인식 토크나이저로 기능할 수 있게 합니다.
3. 실증적 검증: ENSEMBITS 가 앙상블 수준의 동적 특성을 포착하고 하류 기능 작업으로 효과적으로 전이되며, 종종 훨씬 적은 사전 학습 데이터를 사용함에도 정적 토크나이저보다 성능이 뛰어남을 포괄적으로 검증했습니다.

결과

1. 동적 표현 (RMSF 및 ANOVA)

• RMSF 예측: ENSEMBITS 는 Root Mean Square Fluctuation (RMSF) 예측 벤치마크에서 압도적인 성과를 보입니다.
  • 다중 프레임: ENSEMBITS ($P=full$) 은 mdCATH-div 에서 가장 강력한 베이스라인 (ProtProfileMD) 보다 약 2.5 개의 스피어만 점수를, MISATO 분할에서는 약 11~18 점수를 더 높게 기록하며 우위를 점합니다.
  • 단일 프레임: $P=1$(증류된) 일지라도 ENSEMBITS 는 RMSF 예측에서 AminoAseed 및 ESM3struct 를 포함한 다른 모든 단일 프레임 토크나이저보다 우수한 성능을 보입니다.
• ANOVA 검정: 운동 진폭 ($s_1$) 에 대한 일원 분산 분석 (ANOVA) 결과, ENSEMBITS 토큰은 국소 운동 진폭의 분산 **37.1%**를 설명합니다 ($\eta^2 = 0.371$). 이는 다음으로 가장 좋은 다중 프레임 베이스라인 (Vote_3Di, $\eta^2 \approx 0.128$) 보다 약 3 배 높고, 단일 프레임 정적 토크나이저보다 7 배 높습니다. 이는 토큰이 단순한 정적 접힘 클래스나 위치가 아닌 구별 가능한 국소 동적 특성을 인코딩함을 확인시켜 줍니다.

2. 하류 기능 작업

훨씬 적은 사전 학습 데이터 (PDB 규모 역접힘 사전 학습 없이 MD 앙상블만 사용) 를 사용함에도 불구하고, ENSEMBITS 는 여러 작업에서 정적 토크나이저와 맞먹거나 초과하는 성능을 보입니다.

• 돌연변이 효과 예측: PROTEINGYM 벤치마크에서 ENSEMBITS ($P=1$) 를 ESM2 와 결합하면 ESM2 단독보다 6.9% 향상된 성능을 기록하며, ESM2 와 혼합된 다른 모든 구조 토크나이저를 능가합니다.
• 결합 부위 및 친화도: ENSEMBITS 는 결합 부위 예측 (AUROC 0.750) 및 모든 분할에 걸친 결합 친화도 회귀에서 선두를 달립니다. 특히, 정적 토크나이저가 놓치는 인터페이스 잔기에 대한 국소 입체 구조 유연성이 중요한 신호임을 시사하며, 결합 부위 감지 분야에서 ESM3struct 와 AminoAseed 보다 우수한 성능을 보입니다.
• EC 및 GO 예측: ENSEMBITS 는 효소 위원회 (EC) 및 유전자 온톨로지 (GO) 작업에서 ESM3struct 와 경쟁력을 보이지만, 접힘 클래스와 밀접하게 상관된 작업에서는 역접힘 사전 학습으로 인해 ESM3struct 가 우위를 유지합니다.

의의 및 주장

이 논문은 ENSEMBITS 를 정적 구조 예측에서 앙상블 생성으로의 단백질 모델링 진화에 있어 필수적인 단계로 위치시킵니다.

• 격차 해소: 단백질 동적 특성을 언어 모델링과 디자인으로 가져오기 위해 필요한 이산 어휘를 제공합니다.
• 데이터 효율성: ESM3 과 같은 모델이 사용하는 거대한 서열/구조 코퍼스에 비해 상대적으로 작은 MD 코퍼스로도 고품질의 동적 표현을 학습할 수 있음을 보여주며, 이는 표현이 명시적으로 동적 특성을 위해 설계되었을 때 가능합니다.
• 실용적 유용성: SFTD 메커니즘은 동적 데이터 가용성의 병목 현상을 완화하여, 학습된 동적 사전 지식을 활용하면서도 단일 정적 구조 (예: AlphaFold 예측) 에 대해 모델을 쿼리할 수 있게 합니다.

저자들은 토큰의 품질이 기반이 되는 훈련 궤적의 범위와 충실도에 의해 제한되며, 목표는 모든 단백질 동적 특성을 해결했다고 주장하는 것이 아니라 이산 토큰에 동적 정보를 주입하기 위한 방법론을 제시하는 것임을 겸손하게 언급합니다. 향후 작업은 더 큰 앙상블 코퍼스로 확장하고, 이러한 토큰을 다중 프레임 구조 생성을 위해 단백질 언어 모델에 직접 통합하는 것을 포함합니다.