기술 요약: Yeti – 재구성 및 다중 모달 생성을 위한 컴팩트한 단백질 구조 토크나이저

1. 문제 제기

단백질 서열, 3D 구조, 기능적 주석을 통합된 다중 모달 표현으로 통합하는 것은 계산 생물학의 발전에 필수적입니다. 그러나 기존 접근 방식은 세 가지 주요 과제를 직면하고 있습니다:

1. 데이터 이질성: 생물학적 데이터는 고유한 통계적 특성을 가진 다양한 모달리티 (생물물리학적 제약, 텍스트, 영상, 체적 데이터) 에 걸쳐 있습니다.
2. 불완전성과 노이즈: 실제 세계 데이터셋은 종종 누락된 모달리티, 노이즈가 포함된 측정값, 그리고 구조 - 기능 공간 전반에 걸친 불균일한 커버리지로 고통받습니다.
3. 확장성과 결합: 대부분의 모델은 단일 스케일 또는 단일 모달리티에서 작동하며, 원자 수준에서 세포 수준까지의 정보 전달 메커니즘이나 결합된 프로세스를 처리할 수 있는 기능이 부족합니다.

구체적으로, 최근 생성 모델들이 연속 좌표 공간 (예: SE(3) 확산) 을 사용하여 단백질을 3D 객체로 취급하는 반면, 이산 토큰으로 3D 원자 구조를 인코딩하는 것이 매력적인 대안입니다. 이는 언어 모델과 유사하게 확장 가능하고 모달리티에 구애받지 않는 학습을 가능하게 합니다. 그러나 기존 단백질 구조 토크나이저 (예: ESM3, DPLM-2, Kanzi) 는 종종 생성 능력보다 재구성 정확도를 우선시합니다. 높은 재구성 충실도가 높은 품질의 토큰 다양성을 보장하거나 서열과 구조의 무조건적 공동 생성을 지원하지 않는다는 불일치가 존재합니다. furthermore, 많은 최첨단 모델은 방대한 파라미터 수나 인코더를 고정하는 2 단계 학습 패러다임에 의존하여, 엔드 - 투 - 엔드 최적화를 제한할 수 있습니다.

2. 방법론

2.1 아키텍처: Yeti

Yeti(Yielding Encoded Tokens for Intermodality) 는 연속적인 3D 원자 좌표를 이산 토큰으로 매핑하도록 설계된 컴팩트한 단일 단계 단백질 구조 토크나이저입니다.

• 입력: 평균 중심화된 $C_\alpha$ 좌표 ($x \in \mathbb{R}^{L \times 3}$) 로 표현된 단백질 구조.
• 인코더: 회전 위치 인코딩 (RoPE) 을 사용하는 멀티헤드 어텐션을 활용한 Transformer 기반 인코더가 입력을 연속 잠재 임베딩 $Z \in \mathbb{R}^{L \times D}$로 매핑합니다.
• 양자화: 임베딩은 양자화 공간으로 투영되어 **Lookup-Free Quantizer(LFQ)**에 의해 처리됩니다. LFQ 공간은 이진 서브 코드북 ($C_i = \{-1, +1\}$) 의 카르테시안 곱입니다. 양자화는 룩업 테이블 없이 벡터를 이산 인덱스로 매핑하는 부호 연산을 통해 수행되며, 어휘 크기 $K = 2^D$를 생성합니다.
• 디코더: StripedHyena 아키텍처 (하이브리드 신호 처리 모델) 가 디코더로 작용합니다. 불변 포인트 어텐션 (IPA) 과 같은 SE(3)-불변 구성 요소에 의존하는 전통적인 디코더와 달리, StripedHyena 디코더는 노이즈가 포함된 좌표에서 구조를 재구성하기 위해 조건부 속도장을 학습합니다.

2.2 학습 목적: 플로우 매칭

Yeti 는 가우시안 노이즈 분포 ($p_0$) 와 목표 단백질 구조 분포 ($p_{data}$) 간을 보간하는 플로우 매칭 목적 함수를 사용하여 엔드 - 투 - 엔드로 학습됩니다.

• 손실 함수: 모델은 예측된 조건부 속도장 $v_\theta(x_t, t, c)$와 목표 벡터장 ($x_1 - x_0$) 간의 차이를 최소화합니다. 여기서 $x_1$은 깨끗한 구조이고 $x_0$는 가우시안 노이즈입니다.
• 엔트로피 정규화: 코드북 붕괴를 방지하고 높은 활용도를 보장하기 위해 손실 함수에 엔트로피 페널티가 추가됩니다: $L = L_{FM} + \lambda L_{entropy}$.
• 추론: 추론 중 Classifier-Free Guidance(CFG) 가 사용됩니다. 학습 중 토큰 임베딩을 무작위로 0 으로 설정함으로써 모델은 조건부 및 무조건부 속도장을 모두 학습하며, 이는 구조 정확도를 정제하기 위한 유도된 디코딩을 가능하게 합니다.

2.3 다중 모달 공동 생성

Yeti 토큰의 생성적 유효성을 검증하기 위해, 저자들은 **Masked Diffusion Model(MDM)**을 처음부터 학습시켰습니다. 이 모델은 사전 학습된 초기화 없이 (DPLM-2 가 서열 모델에서 워밍업하거나 ESM3 이 캐스케이드 파이프라인을 사용하는 것과 달리) 아미노산 서열과 Yeti 구조 토큰을 공동으로 생성합니다.

3. 주요 기여

1. 컴팩트한 엔드 - 투 - 엔드 토크나이저: 기존 연구에서 흔한 2 단계 "인코더 고정" 패러다임을 피하고 플로우 매칭으로 엔드 - 투 - 엔드 학습된 단일 단계 토크나이저 Yeti 의 도입.
2. 높은 코드북 활용도: Yeti 는 훨씬 적은 파라미터를 가지면서도 더 큰 모델 (ESM3, DPLM-2) 과 다른 양자화 방법 (FSQ, VQ-VAE) 보다 우수한 코드북 활용도와 토큰 다양성을 달성합니다.
3. 파라미터 효율성: Yeti 는 62.5M 파라미터만으로 경쟁력 있는 재구성 정확도를 달성하며, 이는 ESM3(648M) 보다 약 10 배, DPLM-2(118M) 보다 2 배 적은 수치입니다.
4. 개념 증명 공동 생성: Yeti 토큰으로 처음부터 학습된 컴팩트한 다중 모달 모델이 단백질 서열과 구조의 무조건적 공동 생성을 수행할 수 있으며, 10 배 더 큰 모델과 비교 가능한 결과를 달성함을 입증.
5. 접힘 역학 분석: 플로우 매칭 디코딩 궤적 중 이차 구조의 계층적 출현을 확인하였으며, 이는 초기 글로벌 압축을 followed by 최종 단계에서 글로벌 토폴로지가 통합되는 "늦은 헌신 (late commitment)" 현상으로 특징지어집니다.

4. 결과

4.1 코드북 분석

• 활용도: 8,192 개 항목의 코드북에서 Yeti 는 높은 퍼플렉시티와 엔트로피를 달성하여 ESM3 과 Kanzi 를 크게 능가합니다. CATH 데이터셋에서 Yeti 는 약 66% 의 용량을 활용하며 (Perplexity = 5414), 구조 내 다양성 ($\bar{U}$) 은 **98.8%**로, 이 지표를 거의 포화시켜 잔류물에 대한 매우 구별된 토큰 할당을 나타냅니다.
• 비교: Yeti 는 DPLM-2 보다 코드북 활용도에서 우위를 점하는데, 이는 DPLM-2 의 고정된 인코더 접근 방식과 대비되는 엔드 - 투 - 엔드 학습 패러다임에 기인합니다.

4.2 재구성 정확도

• 지표: Yeti 는 CAMEO 에서 TM-score 0.96, CASP14 에서 0.97, CASP15 및 CATH 에서 0.95를 달성합니다.
• 효율성: ESM3 은 그 거대한 규모로 인해 약간 더 낮은 RMSD 를 보이지만, Yeti 의 TM-score 는 비슷하거나 더 우세하여 올바른 글로벌 토폴로지를 효과적으로 복원함을 나타냅니다. Yeti 의 플로우 기반 디코더는 국소 원자 위치의 분산 (RMSD 를 패널티로 작용) 을 가진 샘플을 생성하지만 일관된 글로벌 토폴로지를 유지하며, 이는 다중 컨포메이션 생성에 유익한 특징입니다.

4.3 무조건적 공동 생성

• 성능: 260 만 개의 단백질 (Yeti 토큰 사용) 로 처음부터 학습된 224M 파라미터 모델은 자기 일관성 TM-score(scTM) 0.70과 pLDDT 76.12를 달성했습니다.
• 비교: 이 성능은 DPLM-2(scTM 0.87) 와 비교 가능하며, DPLM-2 가 650M 파라미터 사전 학습 서열 모델에 의존하고 ESM3 이 캐스케이드 파이프라인을 사용하는 반면, ESM3-Open(scTM 0.46) 보다 훨씬 우수합니다.
• 설계 가능성: ProteinMPNN 은 생성된 구조에서 원래 생성된 서열의 **30.6%**를 복구하여 공동 생성된 쌍의 기하학적 일관성과 설계 가능성을 확인했습니다.

4.4 플로우 궤적 역학

디코딩 궤적 분석은 특징의 비단조적 출현을 드러냈습니다:

• 초기 단계 ($t < 0.6$): 빠른 글로벌 압축 (회전 반경 감소).
• 후기 단계 ($t > 0.9$): 이차 구조 요소 ($\alpha$-나선 및 $\beta$-시트) 의 빠른 응집 및 글로벌 토폴로지의 통합.
  이는 토크나이저가 조립된 밀도와 고주파 구조 세부 사항을 효과적으로 분리해낸다는 것을 시사합니다.

5. 중요성과 주장

이 논문은 Yeti가 다중 모달 단백질 모델을 학습하기 위한 컴팩트하고 표현력 있는 기반 역할을 한다고 주장합니다. 그 주요 중요성은 다음을 입증하는 데 있습니다:

1. 생성을 위한 이산 토크나이징의 타당성: 토크나이저가 표현력이 있고 학습 목적 (플로우 매칭) 이 생성 목표와 정렬되어 있다면, 거대한 사전 학습이나 캐스케이드 파이프라인 없이도 서열과 구조의 고품질 무조건적 공동 생성이 가능합니다.
2. 규모보다 효율성: 신중하게 설계된 작은 모델 (62.5M 파라미터) 이 토큰 다양성과 재구성 측면에서 훨씬 큰 모델 (600M+) 을 능가하거나 맞먹을 수 있으며, 이는 단백질 모델링에서 성능으로 가는 유일한 길이 규모라는 관념에 도전합니다.
3. 통합된 표현: 결합된 서열 - 구조 데이터로 처음부터 모델을 학습할 수 있는 능력은 아미노산 서열과 3D 구조 간의 고유한 결합을 포착하는 보편적 분자 표현의 잠재력을 검증합니다.

저자들은 Yeti 를 모든 단백질 설계 작업에 대한 최종 해결책이 아니라, 그들의 토크나이징 전략의 타당성을 확립하는 개념 증명으로 위치시킵니다. 그들은 단일 컨포메이션 구조에 대한 의존성과 La-Proteina 와 같은 전문적으로 세밀하게 조정된 모델의 성능을 맞추기 위한 추가적인 확장 필요성과 같은 한계를 인정합니다. 그러나 Yeti 는 매우 그럴듯한 서열과 구조를 공동 생성할 수 있는 미래의 다중 모달 단백질 모델을 위한 확장 가능하고 효율적이며 효과적인 백본을 제공한다고 주장합니다.