기술 요약: 적응형 단백질 토큰화 (Adaptive Protein Tokenization)

문제 정의

현재의 단백질 구조 토큰화 방식은 단백질 서열을 따라 공간적으로 인접한 잔기(residue)로부터 정보를 풀링(pooling)하여 토큰을 생성하는 **국소성(locality)**에 의존합니다. 이러한 방식은 높은 충실도의 재구성을 달emat지만, 저자들은 두 가지 결정적인 한계를 지적합니다:

1. 생성 시 오차 누적: 국소적 토큰화에 기반한 생성 모델은 오차 전파 문제를 겪습니다. 단 하나의 토큰만 잘못 샘플링되어도 예측된 구조에서 상당한 편차가 발생할 수 있으며, 이로 인해 이산적(discrete) 접근 방식이 연속적 확산 모델(continuous diffusion models)에 비해 생성 지표에서 성능이 떨어지는 원인이 됩니다.
2. 낮은 확장성: 토큰의 수가 단백질 크기에 따라 선형적으로 증가합니다. 이는 거대 단백질 복합체를 모델링할 때 계산 비용을 과도하게 높이며, 단백질을 효과적으로 압축하는 능력을 제한합니다.
3. 표현력의 한계: 국소적 토크나이저는 다운스트림 태스크를 위해 고정된 크기의 표현을 생성할 때 평균 풀링(mean-pooling) 연산을 필요로 하는 경우가 많으며, 이는 중요한 전역적(global) 정보를 유실시키고 단백질 크기에 따라 확장성이 떨어집니다.

방법론

저자들은 **적응형 단백질 토큰화(Adaptive Protein Tokenization, APT)**를 제안합니다. 이는 연속적인 토큰들이 특정 국소 영역에 대응하는 대신, 전역적 표현에 점진적으로 더 높은 수준의 세부 정보를 기여하는 전역적 토큰화 방식입니다.

아키텍처

APT는 이산적 병목(discrete bottleneck)을 가진 **확산 오토인코더(diffusion autoencoder)**로 구현됩니다:

• 입력: 길이 $L$인 원시 $C_\alpha$ 좌표 시퀀스 (zero-mass로 정규화됨).
• 인코더: 양방향 어텐션 트랜스포머가 원시 위치를 잠재 시퀀스 $c \in \mathbb{R}^{L \times d}$로 매핑합니다.
• 이산화: 잠재 시퀀스는 레벨 $(8, 5, 5, 5)$를 가진 **유한 스칼라 양자화(Finite Scalar Quantization, FSQ)**를 사용하여 이산화되며, 유효 코드북 크기는 1,000입니다.
• 적응형 메커니즘: 적응성을 강제하기 위해 모델은 **중첩 드롭아웃(nested dropout)**을 사용합니다. 훈련 중에 상한값 $U$는 $[1, \min(L, k_{max})]$에서 균등하게 샘플링됩니다. $U$를 초과하는 토큰은 드롭됩니다. 이는 모델이 첫 몇 개의 토큰에 핵심적인 전역 정보(저주파 성분)를 배치하고, 후속 토큰에 고주파 세부 사항을 배치하도록 유도합니다.
• 디코더: 플로우 매칭(flow-matching) 목적 함수를 사용하여 훈련된 확산 디코더는 이산 토큰에 조건화되어 원자 좌표를 재구성합니다. 디코더는 확률적 등변성(stochastic equivariance)을 활용하여, 대칭성을 명시적으로 강제하기보다 학습합니다.
• 크기 예측: 단백질 크기는 첫 번째 토큰으로부터 회귀(regression)되며, 이는 모델이 단백질 크기를 조건부 시퀀스 길이로부터 분리할 수 있게 합니다.

훈련 및 추론

• 훈련: 모델은 Foldseek으로 클러스터링된 AFDB 데이터베이스의 약 473,000개 합성 AlphaFold2 예측값으로 훈련됩니다. 훈련은 두 단계로 나뉩니다: (1) 확률적 등변성을 학습하기 위해 무작위 회전을 포함한 오토인코더 훈련, (2) APT 토큰에 대한 자기회귀(GPT 스타일) 모델 훈련.
• 추론 전략:
  • 테일 드롭아웃(Tail Dropout): 생성 중에 모델은 토큰 생성을 조기에 중단할 수 있습니다. 토큰 시퀀스의 어떤 접두사(prefix)라도 유효한 조건화 신호가 됩니다. 꼬리 부분(고주파 세부 사항)은 확산 디코더에 위임됩니다.
  • 엔트로피 기반 중단: 저자들은 생성의 최적 지점을 결정하기 위해 토큰 엔트로피(유한 컷오프, 스플라인 기반 최소 엔트로피)에 기반한 중단 기준을 도입하여, 표현 충실도와 생성 능력 사이의 균형을 맞춥니다.
  • 클래시파이어 어닐링(Classifier Annealing): 실제 데이터 매니폴드에서 벗어남으로써 발생하는 아티팩트를 완화하기 위해, 디코딩 중에 조건부 가이던스 필드와 무조건부 가이던스 필드를 보간하는 클래시파이어 어닐링을 사용합니다.

주요 기여

1. 전역 적응형 토크나이저: 토큰이 국소 이웃이 아닌, 조ละเอียด-to-fine 계층 구조를 가진 전역 기술자(global descriptor)를 나타내는 토크나이저를 도입했습니다.
2. 성능 검증: APT가 생성, 표현 학습 및 재구성 작업에서 최신 기술(SOTA)인 국소 토크나이저(예: DPLM2, ESM3, Kanzi)와 대등하거나 이를 능가함을 입증했습니다.
3. 추론 시간 기법: 샘플 복잡도를 관리하고 오차 노출을 완화하기 위해 엔트로피 기반 샘플링 및 클래시파이어 어닐링을 개발하여 설계 가능성(designability)을 개선했습니다.
4. 제로샷 애플리케이션: 단백질 크기를 조건부 시퀀스 길이로부터 분리함으로써, 단백질 축소(protein shrinking, 구조를 유지하면서 잔기 수를 줄임) 및 추론 시간 스케일링을 통한 **친화도 성숙(affinity maturation)**과 같은 제로샷 애플리케이션을 가능하게 했습니다.

결과

• 재구성: CATH, CAMEO, AFDB 테스트 세트에서 APT는 SOTA 연속 확산 모델과 유사한 RMSD 및 TM-score를 달 achieves 합니다. 특히, 32~64개의 토큰만 사용하여 RMSD < 2 Å를 달성함으로써 단백질 구조의 압축성을 입증했습니다.
• 생성: APT 토큰으로 훈련된 자기회귀 모델은 설계 가능성(designability) 0.871(scRMSD < 2 Å인 샘플의 비율)을 기록하며, DPLM2(0.486)나 ESM3와 같은 이산 확산 모델을 크게 앞질렀습니다. 이는 "Best-of-N" 샘플링 기법 없이 달성되었습니다.
• 표현 학습: CATH 분류 태스크에서 APT 토큰에 대한 비선형(MLP) 프로빙은 ESM3 및 DPLM2의 표현에 대한 프로빙보다 우수한 성능을 보였습니다. 결정적으로, APT는 평균 풀링 없이도 고정된 크기의 전역 표현을 제공하며, 매우 압축된 표현(16개 토큰)조차 더 큰 모델들보다 뛰어난 성능을 보였습니다.
• 애플리케이션:
  • 단백질 축소: 더 적은 잔기 수로 조건화함으로써 전역 및 국소 구조를 유지하면서 단백질(예: 헤모글로빈, 베타 배럴)의 더 작은 버전을 성공적으로 생성했습니다.
  • 친화도 성숙: 잠재 공간 내에서 보상 함수(예: iPAE)에 의해 가이드되는 빔 서치(beam search)를 사용하여 약한 결합체를 강한 결합체로 진화시키는 능력을 입증했습니다.

의의

본 논문은 적응형 단백질 토큰화가 생물학적 태스크를 확장하기 위한 태스크 인지적(task-aware) 접근 방식을 제공한다고 주장합니다. 국소적 토큰화에서 전역적 토큰화로 전환함으로써, 이 방법은 재구성 충실도와 생성 능력 사이의 트레이드오프를 해결합니다. 평균 풀링을 통한 정보 손실 없이 단백질을 고정 크기 벡터로 압축할 수 있는 능력은 거대 단백질 복합체를 모델링하고 추론 시간 스케일링을 수행하는 데 실질적인 경로를 제공합니다. 저자들은 이 연구를 여러 길이 척도(length scales)에 걸쳐 추론할 수 있는 생성 모델을 향한 단계로 규정하며, 이는 프런티어 바이오 엔지니어링의 핵심 과제를 다루고 있습니다.

참고: 저자들은 한계점 또한 인정하는데, 구체적으로 APT 토큰이 전역적 태스크에는 효과적이지만 모티프 스캐폴딩(motif scaffolding)과 같은 국소적 태스크에는 덜 적합하며, 현재의 토크나이저가 특정 기능에 필수적인 사이드체인(sidechain)이나 아미노산 조성을 고려하지 않는다는 점을 명시했습니다.