Adaptive Protein Tokenization

Motivation

• Known: 기존 단백질 구조 tokenizer들은 지역(local) 이웃 정보를 pooling하여 residue별 토큰을 생성하며, VQ-VAE류 discrete tokenization과 diffusion/flow matching 기반 생성모델이 단백질 구조 생성 및 멀티모달 시퀀스-구조 모델에 널리 활용되고 있다.
• Gap: Local tokenization 방식은 토큰 수가 단백질 크기에 선형적으로 비례해 확장성이 떨어지고, 하나의 오샘플링된 토큰이 예측 구조에 큰 오차를 유발하는 error accumulation 문제가 있어 생성모델 성능이 비토큰화 모델 대비 크게 저하된다.
• Why: 토큰이 전역 정보를 coarse-to-fine 방식으로 누적 표현하도록 하면 오류 누적을 완화하고, sequence-reduction 없이 임베딩을 얻을 수 있으며, 태스크별로 정보량을 적응적으로 조절할 수 있어 단백질 시퀀스-구조-기능을 통합하는 멀티모달 모델 구축에 중요한 기반이 된다.
• Approach: Diffusion autoencoder 구조에 finite-scalar quantization(FSQ) 기반 discrete bottleneck과 nested dropout을 결합하여, 앞쪽 토큰일수록 전역적이고 중요한 정보를, 뒤쪽 토큰일수록 세부 정보를 담도록 학습하는 adaptive global tokenizer(APT)를 구축한다.

Achievement

1. Adaptive global tokenizer 제안: 지역 이웃이 아닌 전역 서술자를 표현하는 diffusion 기반 adaptive tokenizer를 도입했다.
2. 다중 과제 검증: 생성, 표현학습, 재구성 과제 전반에서 기존 local structure tokenizer 기반 모델 대비 동등하거나 우수한 성능을 입증했으며, 특히 생성 및 표현학습에서 유의미한 개선을 보였다.
3. 정보량 기반 추론 기법 도입: token entropy를 sample complexity의 척도로 활용하는 inference-time 기법을 제안하여, 정보 손실에 기반한 원칙적 샘플링을 가능하게 하고 designability를 향상시켰다.
4. Zero-shot 응용 시연: 조건화(conditioning) 방식을 활용해 protein shrinking과 affinity maturation 같은 zero-shot 응용을 시연했다.

How

• 입력: 길이 L의 raw Cα 좌표(zero-mass로 정규화)
• Bidirectional attention transformer로 latent sequence c ∈ R^{L×d} 생성
• FSQ(levels (8,5,5,5), 코드북 크기 1000)로 discretize하여 토큰 ĉ 생성
• 토큰을 조건으로 flow-matching objective 기반 diffusion decoder로 재구성
• Nested dropout: U(1,...,min(L,kmax))에서 upper cutoff를 균일 샘플링해 적용, 앞쪽 토큰에 전역/중요 정보가 담기도록 유도
• Relative positional encoding 사용, decoder 레이어 전반에 adaLN 파라미터 공유
• 첫 토큰에서 단백질 길이(size)를 projection하여 cross entropy loss(L_size)로 학습, teacher-forcing으로 재구성에 반영
• 전체 loss: L = L_flow + ε_size·L_size (ε_size≈0.01)
• 2단계 학습: (1) random rotation을 통한 stochastic equivariance 학습으로 autoencoder 학습, (2) APT 토큰 위에 GPT-style autoregressive 모델 학습하여 생성 능력 평가
• CATH classification task에서 non-linear probing으로 표현학습 성능 검증

Originality

• Local neighborhood pooling 패러다임을 벗어나 Fourier/wavelet 변환에서 영감을 받은 coarse-to-fine global tokenization 개념을 단백질 구조에 최초로 적용
• FSQ 기반 discrete bottleneck과 nested dropout을 결합해 토큰별 정보량을 adaptive하게 조절하는 diffusion autoencoder 설계
• 컴퓨터 비전 분야의 adaptive tokenization(CAT, ALIT, FlexTok) 아이디어를 단백질 구조 도메인으로 확장 적용
• Token entropy를 sample complexity 척도로 활용한 inference-time 샘플링 기준 제시

Limitation & Further Study

• 재구성 성능이 SOTA 모델과 comparable한 수준에 그쳐, 압축 제어성 대비 절대적 재구성 품질에서의 우위는 제한적일 수 있음
• Diffusion decoder 기반이라 추론 시 ODE/SDE 적분이 필요해 계산 비용이 큼(diffusion autoencoder의 일반적 한계)
• 본문 발췌만으로는 대규모 단백질 복합체나 매우 큰 단백질에 대한 실제 확장성 검증이 충분히 제시되지 않음
• 향후 연구로 protein complex나 멀티체인 구조에 대한 global tokenization의 확장, 그리고 언어·기능 모달리티와의 완전한 멀티모달 통합이 필요함

Evaluation

총평: Local pooling 기반 단백질 구조 tokenization의 근본적 한계를 global coarse-to-fine 접근으로 해결하려는 시도로, 생성·표현학습·재구성 전반에서 실질적 개선을 보인 참신하고 의미 있는 연구이다.