단백질 : 논문 리뷰 : FoldToken4: Consistent & Hierarchical Fold Language

쉬운 설명:

FoldToken4는 단백질 구조를 다양한 "해상도" (coarse-to-fine)로 표현하는 여러 언어를 하나의 똑똑한 번역기로 만드는 것과 같습니다. 기존에는 각 해상도마다 별도의 번역기를 만들어야 해서 비효율적이었지만, FoldToken4는 Code Adapter라는 장치를 이용해 모든 해상도의 언어가 같은 의미 체계를 공유하게 만듭니다. 또한, 학습된 transition matrix를 이용하면 가장 자세한 지도(고해상도)만 저장해두고 필요에 따라 더 넓은 범위의 지도(저해상도)로 모델 재실행 없이 바로 변환할 수 있어 효율적입니다. 즉, 하나의 모델로 여러 수준의 구조 정보를 일관되게 다루고 그 관계까지 파악할 수 있게 해주는 기술입니다.

 

 

 

FoldToken4: Consistent & Hierarchical Fold Language 학습 노트

Purpose of the Paper:

• 기존 연구의 한계:
  • FoldToken1/2/3, ESM3 등 기존 단백질 구조 언어 모델들은 각기 다른 scale (codebook size) 에서 개별적으로 학습되어야 함 (e.g., 2^8, 2^12, 2^16). 이는 비효율적이고 중복적인 학습 노력을 요구함.
  • 이렇게 독립적으로 학습된 여러 scale의 언어들 간의 관계 (consistency, hierarchy)가 불분명하여, 서로 다른 해상도의 구조 정보를 통합적으로 이해하고 활용하기 어려움.
• 새로운 문제 정의 및 접근:
  • "하나의 모델로 여러 scale의 fold language를 생성할 수 있을까?" 라는 질문 제기.
  • 단일 모델 내에서 다양한 scale (e.g., 2^5 ~ 2^12)의 fold language를 일관성 (consistent) 있게 학습하고, 이들 간의 계층 구조 (hierarchy) 를 발견하는 것을 목표로 함.
  • 기존의 단일 scale 표현 능력의 한계를 넘어, multiscale fold language 를 통합적으로 학습하고 분석할 수 있는 새로운 프레임워크 제시.

Key Contributions & Novelty:

• Single Model for Multiscale Fold Languages:
  • Contribution: 단 하나의 모델 학습으로 2^5 부터 2^12 까지 다양한 codebook size를 갖는 여러 fold language를 생성함.
  • Novelty: 개별 scale마다 별도의 모델을 학습시킬 필요 없이 통합된 학습 및 언어 생성이 가능해져 효율성을 크게 높임. VQ 연구 및 단백질 구조 언어 분야에서 multiscale token consistency와 hierarchy를 학습하려는 최초의 시도.
• Scale-aware Code Adapter:
  • Contribution: 서로 다른 scale의 이진 코드 (masked-binary code)를 공유된 semantic space (encoder의 output space)로 projection하는 MLP 기반 adapter 도입 (Eq. 6).
  • Novelty: 각 scale의 code embedding이 encoder의 의미 공간과 정렬(align)되도록 하여, scale 간의 일관성 있는 표현 학습을 가능하게 함.
• Multiscale Mix-Token Training:
  • Contribution: 학습 시 다양한 scale을 무작위로 샘플링하고 해당 scale의 code embedding을 생성하여 학습에 사용하는 방식 (Eq. 10). Encoder, Quantifier, Decoder를 공유함.
  • Novelty: 단일 모델이 여러 scale에 대한 지식을 동시에 학습하고 공유하도록 유도하며, scale 간의 knowledge consensus를 구축함.
• Consistency & Hierarchy Analysis:
  • Contribution: 학습된 code embedding 간의 유사도 계산을 통해 scale 간 token mapping 관계를 나타내는 transition matrix (Mm→m') 를 정의하고 분석함 (Eq. 7).
  • Novelty: 가장 상세한 scale의 code를 transition matrix를 이용해 다른 coarse scale의 code로 모델 재실행 없이 변환 가능. 이를 통해 명시적으로 fold language의 계층 구조를 밝히고 시각화/분석 가능.
• Efficient & Stable Training Techniques:
  • Contribution: Binary Stochastic Quantifier (BSQ), Teacher Guidance (adaptive VQ skipping), Stable Gradient (partial gradient trick) 등 도입 (Sec 2.4, 2.5).
  • Novelty: 미분 불가능한 quantization 과정을 안정적이고 효율적으로 학습시키고, 특히 temperature annealing 시 gradient 폭발 문제를 해결하여 multiscale 학습을 안정화.

Experimental Highlights:

• Datasets: PDB (training), T493, T116, N128 (single-chain evaluation), SAbDab/M1031 (multi-chain evaluation).
• Metrics: TMScore (↑), RMSD (↓), Code Vector Similarity (Sim ↓).
• Single Unified Model Performance: FoldToken4 단일 모델이 다양한 scale에서 개별적으로 학습된 FoldToken3 모델과 비교 가능한 (comparable) 재구성 성능(RMSD, TMScore)을 달성함 (Table 1). 이는 multiscale 통합 학습의 성공을 보여줌.
• High Compression: 최소 32개 code (2^5) 라는 매우 작은 codebook size에서도 합리적인 수준의 단백질 구조 재구성이 가능함을 보임 (Table 1, Fig 2). 성능은 128 또는 256 codes에서 안정화(plateau)되는 경향.
• Multi-Chain Reconstruction: 복잡한 multi-chain 단백질 구조도 compact한 codebook (e.g., 256)으로 효과적으로 재구성 가능함을 입증 (Table 2, Fig 4).
• Consistency & Hierarchy Verified: 생성된 multiscale 언어의 일관성 (Fig 5) 및 학습된 transition matrix를 통한 계층적 변환 능력 (Fig 7)을 실험적으로 확인함. 모델 재실행 없이 scale 간 변환이 가능.
• Code Diversity: 기존 FoldToken 모델들보다 더 나은 code diversity (낮은 code vector similarity)를 달성하여, 더 구별력 있는 언어를 학습했음을 시사함 (Table 3).
• Comparison with ESM3: 훨씬 적은 파라미터와 PDB라는 상대적으로 작은 데이터셋 학습만으로도 ESM3와 유사한 재구성 성능을 보임 (Sec 3.1).

Limitations and Future Work:

• Limitations:
  • Slight Performance Degradation: 단일 scale에 특화되어 학습된 FoldToken3 모델보다는 재구성 성능이 약간 저하됨. 이는 단일 모델 내에서 여러 scale 간의 trade-off 때문으로 보임.
  • Edge Case Handling: 매우 짧은 단백질 (<30 residues) 재구성에 약점을 보임 (training data filtering 영향). 매우 긴 단백질의 경우 작은 codebook size에서는 미세 구조 재현이 어려울 수 있음 (Fig 2 discussion).
• Future Work:
  • 학습된 consistent & hierarchical fold language를 protein structure alignment, generation, representation learning, similarity search 등 다양한 downstream task에 적용.
  • 서로 다른 scale의 언어를 융합(fuse)하여 구조의 복잡도에 따라 적응적으로(adaptively) 다른 해상도의 token을 사용하는 방식 연구 (Fig 6 관련 논의).
  • 매우 짧거나 긴 단백질 등 edge case에 대한 성능 개선 연구.

Overall Summary:

FoldToken4는 기존 단백질 구조 언어 모델들의 비효율성과 scale 간 관계 부재 문제를 해결하기 위해 제안되었습니다. 이 모델은 Scale-aware Code Adapter와 Multiscale Mix-Token Training 기법을 사용하여, 단 하나의 모델로 2^5부터 2^12까지 다양한 scale의 fold language를 일관성(consistent) 있고 계층적(hierarchical)으로 학습합니다. 실험 결과, 단일 모델임에도 기존 개별 학습 모델들과 유사한 성능을 보이면서 높은 압축률(최소 32 codes)을 달성하고, 학습된 transition matrix를 통해 scale 간 변환이 가능함을 입증했습니다. 이 연구는 multiscale fold language를 통합적으로 이해하고 활용하는 새로운 방향을 제시하며, 단백질 구조 분석 및 관련 응용 연구의 효율성과 깊이를 더할 것으로 기대됩니다.

 

 

 

 

Abstract

단백질 structure language를 만드는 것은 단백질 sequence와 structure의 modality를 통합하는 데 있어 increasingly 많은 관심을 받고 있습니다.

FoldToken1&2&3와 같은 최근 연구들이 이 방향에서 큰 진전을 이루었지만, 다른 scales에서 다른 models에 의해 생성된 languages 간의 관계는 여전히 불분명합니다.

더욱이, 여러 scales(예: $2^5, 2^6, \dots, 2^{12}$와 같은 다른 code space size)의 models는 개별적으로 trained되어야 하므로 중복된 노력이 발생합니다.

우리는 다음과 같은 질문을 제기합니다: 단일 model이 multiscale fold languages를 생성할 수 있을까?

본 논문에서는 multiscale fold languages의 consistency와 hierarchy를 학습하기 위해 FoldToken4를 제안합니다.

multiscale code adapters와 code mixing techniques를 도입함으로써, FoldToken4는 동일한 model에서 multiscale languages를 generate할 수 있으며, scales에 걸친 hierarchical token-mapping relationships를 발견할 수 있습니다.

우리가 아는 한, FoldToken4는 VQ research에서 multi-scale token consistency와 hierarchy를 학습하려는 첫 번째 시도입니다; 또한, 이는 단백질 structure language learning 분야에서 더욱 새로운 시도일 것입니다.

 

 

1 Introduction

"SE-(3) structure는 특별하고 어려워서는 안 됩니다. 장벽을 낮춥시다." – Our Goal

단백질 structure language를 만드는 것은 sequence와 structure의 modeling paradigm을 통합하는 데 있어 increasingly 많은 관심을 받고 있습니다. FoldToken1, FoldToken2, FoldToken3와 같은 최근 연구들은 이 방향에서 큰 진전을 이루었습니다: FoldToken2는 FoldToken1에 비해 reconstruction RMSD를 80% 향상시켰고, FoldToken3는 FoldToken2의 code space의 0.39%를 사용하여 compression ratio를 더욱 향상시켰습니다. 일반적으로 compression과 reconstruction 사이에는 균형이 있습니다: 더 작은 code space를 가진 model은 더 높은 compression ability로 거친 geometry를 포착하고, 그 반대도 마찬가지입니다. fold language의 multi-scale consistency와 hierarchy는 단백질 structure를 이해하고 분석하는 데 가치가 있을 것입니다. 불행히도, 현재의 vector quantization 연구들은 이를 간과해 왔습니다. multiscale languages를 연관시키고 그것들의 내재적(intristic) 법칙을 발견하는 방법은 여전히 open question으로 남아 있습니다.

기존의 fold tokenization 방법들은 single-scale expression ability에 제한됩니다. 예를 들어, FoldToken2, ESM3, FoldToken3는 각각 65536, 4096, 256의 codebook size를 가집니다. codebook size를 점진적으로 줄임으로써 더 나은 compression ratio를 얻을 수 있지만, reconstruction quality는 저하되어 coarse-grained hierarchical representation을 초래합니다. 다른 scales의 models이 별도로 trained되어 knowledge sharing을 위한 unified representation space가 부족하기 때문에, 여러 다른 scales에 걸친 fold language의 consistency와 hierarchy는 알려지지 않은 상태로 남아 있습니다. 결과적으로, scales 간에 tokens를 translate하기 위한 또 다른 model을 구축하지 않고서는 language의 consistency와 hierarchy를 분석할 방법이 없습니다. 또한, 다른 scales에서 models를 training하는 것은 time-consuming하고 resource-intensive하며, 각 scale에 대한 codes를 saving하는 것은 compression principle에 위배됩니다. model training과 code saving 모두에서 redundance 없이 semantically consistent한 multiscale tokens를 generate하는 단일 model을 design할 수 있을까요?

우리는 scales에 걸쳐 consistent하고 hierarchical한 fold language를 만들기 위해 FoldToken4를 제안합니다. 첫째, 우리는 multi-scale models이 동일한 encoder를 공유하여 consistent한 semantic space를 학습하도록 만듭니다. multi-scale tokens를 encoder의 output과 align하기 위해, 우리는 다른 scales에서 code embeddings를 generate하고 scales 간에 codes를 mix하는 code adapters를 도입합니다. 둘째, 우리는 token embedding similarity를 compute하여 scales 간의 translation relationships를 reveal할 수 있으며, 이는 fold language의 hierarchy를 반영하는 transition matrix에 기록됩니다. 이 hierarchy는 coarse-to-fine 방식으로 protein structures를 이해하고 visualize하는 데 도움이 됩니다. 마지막으로, FoldToken4는 training과 code saving에서의 redundant 노력을 제거할 수 있습니다: 모든 scales에 대해 단 하나의 model만 train하고 가장 미세한 scale의 code(finest-scale code)만 save하면 되며, 이 code는 model을 rerunning하지 않고 transition matrix를 사용하여 다른 어떤 scale로든 translated될 수 있습니다.

우리는 single-chain 및 multi-chain reconstruction tasks 모두에서 FoldToken4를 evaluate하며, FoldToken3와 비교하여 competitive한 performance를 달성함을 보여줍니다. FoldToken4는 FoldToken3보다 competitive하지만 약간 더 낮은 reconstruction quality를 달성합니다. 그러나 FoldToken4는 code space를 32까지 더욱 compress하며, 32, 64, 128, 256, · · · , 4096과 같은 multi-scale languages가 동일한 model에서 generated될 수 있습니다. 요약하자면, FoldToken4는 FoldToken2&3의 invariant, compact, generative 장점 외에도 consistency, hierarchy, efficiency 측면에서 advantages를 제공합니다. 우리는 FoldToken4가 coarse-to-fine perspective에서 fold languages를 analyzing하는 데 가치가 있으며 광범위한 protein structure 관련 tasks에 benefit이 될 것이라고 믿습니다.

 

 

 

정리노트 (AI 연구자 대상)

핵심 문제 (Problem):

• 기존 Protein Structure Language (fold tokenization) 연구 (e.g., FoldToken1/2/3, ESM3)는 single-scale (codebook size 고정) 표현에 국한됨.
• 다른 scale (다른 codebook size)로 trained된 model들이 생성하는 languages 간의 consistency (일관성) 및 hierarchy (계층 관계)가 불명확함.
  • 예: codebook size 32로 학습된 token과 4096으로 학습된 token 간의 의미론적 연결고리를 모름.
• 각 scale마다 별도의 model을 training하고 code를 저장하는 것은 비효율적 (redundance in training & code saving).

제안 방법 (FoldToken4):

• 목표: 단일 model을 사용하여 consistent하고 hierarchical한 multiscale fold language 생성.
• 핵심 아이디어:
  1. Shared Encoder: 모든 scale이 동일한 encoder를 공유하여 consistent한 semantic space 학습.
  2. Code Adapters & Code Mixing: Encoder output을 각 scale에 맞는 code embedding으로 변환하고(adapters), scale 간 정보를 혼합(mixing)하여 alignment 수행.
  3. Transition Matrix: 학습된 token embedding 간의 similarity를 계산하여 scale 간 translation relationships 및 hierarchy를 포착하는 transition matrix 생성.
• 결과: 단일 model training으로 32, 64, ..., 4096 등 다양한 scale의 language 생성 가능.

주요 장점 (Advantages):

• Consistency & Hierarchy: Scale 간 language의 관계를 명시적으로 학습하고 제공하여 coarse-to-fine 분석 가능.
• Efficiency: 단일 model training. Finest-scale code만 저장 후 transition matrix로 다른 scale 변환 가능 (추가 rerunning 불필요).
• Performance: Reconstruction quality는 single-scale SOTA (FoldToken3) 대비 약간 낮지만, multiscale 지원 및 더 작은 code space (e.g., 32) 압축을 단일 model에서 달성. 기존 FoldToken2&3의 invariant, compact, generative 장점 유지.

---

쉬운 설명:

이 섹션 뭐 하는 얘기냐면요?

단백질 3D 구조를 컴퓨터가 이해하기 쉬운 "단어"(token)들의 나열, 즉 "단백질 구조 언어"(protein structure language)로 바꾸는 연구들이 있었습니다. 그런데 기존 방법들은 문제가 좀 있었어요.

기존 방법의 문제점:

1. 하나의 설명 수준만 가능: 마치 어떤 사물을 설명할 때, "엄청 자세하게" 설명하는 방법(예: 4096 단어 사용) 또는 "매우 간단하게" 설명하는 방법(예: 256 단어 사용) 중 하나만 골라서 배워야 했어요. 이걸 single-scale이라고 불러요.
2. 수준 간 관계 모름: "간단 설명"에 쓰인 단어랑 "자세한 설명"에 쓰인 단어가 서로 어떻게 연결되는지 알 수 없었어요. 예를 들어, 간단 설명의 '덩어리'라는 단어가 자세한 설명의 어떤 구체적인 부분들에 해당하는지 모르는 거죠. (consistency, hierarchy 부족)
3. 비효율적: 설명 수준(scale)마다 별도의 모델을 처음부터 학습시켜야 해서 시간과 자원이 많이 들었어요.

그래서 이 논문이 제안하는 것 (FoldToken4):

"하나의 똑똑한 모델(FoldToken4)을 만들자!" 이 모델은 하나의 학습 과정만으로 여러 설명 수준(예: 32단어, 64단어, ..., 4096단어)의 "단백질 구조 언어"를 모두 만들어낼 수 있어요 (multiscale).

어떻게?

• 모든 설명 수준이 같은 기본 이해 방식(shared encoder)을 공유해요.
• 각 설명 수준에 맞게 단어를 바꿔주는 번역기(code adapter)를 달고, 서로 정보를 섞어줘요(code mixing).
• 학습하고 나면, 간단한 단어와 자세한 단어 사이에 어떤 관계가 있는지 알려주는 관계도(transition matrix)를 만들 수 있어요.

결과적으로 좋은 점:

• 효율적: 한 번만 학습하면 끝! 저장도 제일 자세한 설명만 해두면 나머지는 관계도를 통해 얻을 수 있어요.
• 관계 파악: 간단한 설명과 자세한 설명이 어떻게 연결되는지(hierarchy, consistency) 알 수 있어요. 덕분에 단백질 구조를 대략적으로 봤다가 점점 상세하게 뜯어보는 식(coarse-to-fine)의 분석이 쉬워져요.
• 성능: 기존 최고 성능 모델(FoldToken3)보다는 약간 못하지만, 훨씬 적은 단어(32개)로도 구조를 표현할 수 있게 되고, 무엇보다 하나의 모델로 이 모든 걸 할 수 있다는 게 큰 장점이에요.

 

2 Method

2.1 Overall Framework

Fig.1에 표시된 바와 같이, 전체 framework는 encoder, quantifier, decoder를 포함하여 FoldToken1, FoldToken2, FoldToken3와 동일하게 유지됩니다.

Brief Overview Step1: Encoding. 단백질 G가 주어지면, BlockGAT를 사용하여 residue features를 학습합니다: [ [ f_1, f_2, \cdots , f_n ] \leftarrow \text{BlockGATs}(G) ] 여기서 fs는 s번째 residue의 학습된 features입니다.

Step2: Quantization. F = [f1, f2, · · · , fn]이 주어지면, binary stochastic quantifier (BSQ)를 사용하여 embeddings를 discrete tokens (VQ-IDs) Z = [z1, z2, · · · , zn]으로 quantize합니다. quantifier (Q : f → z)와 de-quantifier (Q⁻¹ : z → ˆf )는 continuous embedding f를 discrete latent code z로 변환하고, 그 반대도 마찬가지입니다: [ \begin{cases}     z_s = Q(f_s) \     \hat{f}_s = Q^{-1}(z_s) \end{cases} ] codebook size를 2m으로 설정함으로써, discrete tokens는 [0,2m−1]로 제한됩니다. 단백질 sequence를 설명하는 20개의 아미노산 타입과 유사하게, VQ-IDs는 단백질 structures를 represent하는 데 사용됩니다. quantized sequence Z = [z1, z2, · · · , zn]는 fold language라고 불립니다.

Step3: Decoding. SE-(3) BlockGAT Decoder를 적용하여 단백질 structures를 복구합니다: [ x_1, x_2, \cdots, x_n = \text{SE3BlockGAT}^{(l)}(z_1, z_2, \cdots, z_n) ] 여기서 xs​∈R4×3는 s번째 residue의 backbone atoms의 coordinates입니다.

Contribution. FoldToken3에서 FoldToken4로 오면서 다음과 같은 개선 사항을 적용했습니다:

1. Scale-aware Code Adapter: 다른 scales에서 code embeddings를 generate하기 위해 code adapters를 도입했으며, 이들 모두는 encoder의 output embedding space와 aligned됩니다.
2. Multiscale Mix-Token Training: 25,26,⋯,212의 code book size를 가진 Multiscale models이 동시에 trained되어, model training에서의 redundant 노력을 피합니다. 공유된 encoder, quantifier, decoder는 또한 scales에 걸쳐 knowledge consensus를 제공합니다.
3. Consistency & Hierarchy Analysis: token similarity를 computing하여 multi-scale hierarchy를 reveal하기 위해 token transition matrix를 결정합니다. transition matrix는 model을 rerun하지 않고 multi-scale fold language를 translate하는 데 사용될 수 있습니다.

2.2 Invariant Graph Encoder

rotation 및 translation equivariant 특성으로 인해 동일한 단백질이 다른 coordinate 기록을 가질 수 있으며, 이는 compact invariant representations를 learning하는 데 어려움을 제기합니다. 이전 연구들에서는 invariant featurizer가 유익한 structure patterns를 추출할 수 있음을 보여주었으며, 우리도 동일한 방식을 따릅니다: 단백질 structures를 invariant node 및 edge features로 구성된 graph로 representing합니다. 그런 다음 BlockGAT를 사용하여 high-level representations를 학습합니다.

Protein Block Graph. n개의 blocks를 포함하는 단백질 M = {Bs​}s=1n​이 주어지면 (각 block은 아미노산을 나타냄), kNN algorithm을 사용하여 block graph G({Bs​}s=1n​, E)를 구축합니다. block graph에서 s번째 node는 Bs​=(Ts​,fs​)로 표현되고, (s, t) 사이의 edge는 Bst​=(Tst​,fst​)로 표현됩니다. Ts​=(Rs​,ts​)와 Tst​=Ts−1​∘Tt​는 각각 s번째의 local frames와 s번째와 t번째 blocks 사이의 relative transform입니다. fs​와 fst​는 node 및 edge features입니다. Rs​와 ts​는 s번째 block의 rotation 및 translation이며, local frame parameters 역할을 합니다.

BlockGAT Encoder. BlockGAT layer fθ​를 사용하여 block-level representations를 학습합니다: [ f_s^{(l+1)}, f_{st}^{(l+1)} \leftarrow \text{BlockGATs}(f_s^{(l)}, f_{st}^{(l)} | T_s, T_{st}, \mathcal{E}) ] 여기서 fs(l)​와 fst(l)​는 l번째 layer의 input node 및 edge features를 나타냅니다. Ts​=(Rs​,ts​)는 s번째 block의 local frame이고, Tst​=Ts−1​∘Tt​=(Rst​,tst​)는 s번째와 t번째 blocks 사이의 relative transform입니다. Ts​,Tst​,fs(0)​,fst(0)​는 UniIF에서 제안된 invariant featurizer를 사용하여 ground truth structures로부터 initialized됩니다. encoder의 output을 F = [f1​,f2​,⋯,fn​]으로 표기하며, 여기서 fs​는 s번째 residue의 embedding이고, encoding space를 semantic space라고 부릅니다.

2.3 Multiscale Code Generator

Multiscale Code. scales에 걸쳐 unified representation 역할을 하도록 masked-binary code를 도입합니다. codebook size가 2m이고 padding length가 M일 때, decimal integer z의 binary code는 다음과 같습니다: [ \mathbf{b}{z|m} = [b_0, b_1, \cdots , b{M-1}], b_i \in {-1, 0, 1} ] 여기서 bi​는 binary code의 i번째 bit입니다. i<M−m이면 bi​=−1이고, 그렇지 않으면 b(M−m):M−1​은 z의 binary form입니다. 예를 들어, M = 4이면, b0∣2​=[−1,−1,0,0], b1∣2​=[−1,−1,0,1], b3∣2​=[−1,−1,1,1], b7∣3​=[−1,1,1,1], ⋯, b15∣4​=[1,1,1,1]입니다.

Code Embedding via Adapter. MLP, 즉, CodeAdapter$_m : \mathbb{R}^M \rightarrow \mathbb{R}^d$는 각 bit 위치에서의 correlations를 고려하기 위해 bi∣m​∈RM을 code vectors vi​∈Rd로 project하는 데 사용됩니다. 공식적으로 다음과 같이 씁니다: [ \begin{cases} b_i = \text{Bit}{\log_2(m)}(z_i) \ b{i|m} = \text{Pad}(b_i, (M - m, 0), -1) \ v_{i|m} = \text{CodeAdapter}m(b{i|m}) \end{cases} ] 여기서 Bit(·)는 decimal numbers를 binary form으로 convert하고, bi∣m​은 masked binary vector이며, vj​는 scale m에서의 j번째 code embedding입니다. 만약 m = 10이면, codebook size는 210이고, MLP는 1024개의 binary vectors를 1024개의 d-차원 code embeddings로 projects합니다.

Code Consistency & Hierarchy. multiscale models을 동시에 training함으로써, CodeAdapter$_m$은 masked-binary code를 동일한 semantic space로 projects하여 공유 encoder의 output과 aligning합니다. 따라서, semantic space에서 similarity를 computing하여 다른 scales 간의 code mapping relationship을 discover할 수 있습니다. m>m′인 scale m과 m'에 대해, code transition relationship을 record하기 위해 transition matrix Mm→m′를 정의합니다: [ M_{i,j} = \arg \max_{k} (v_{i|m}^T v_{k|m'}) \text{ and } \sum_j M_{i,j} = 1 ] 만약 Mi,j​=1이면, scale m의 i번째 code는 scale m'의 j번째 code와 associate됩니다. transition matrix는 model을 rerunning하지 않고 scales 간에 code embeddings를 translate하는 데 사용될 수 있습니다.

2.4 Quantifier: Code Alignment, Optimization, and Mixing

FoldToken3에서와 같이, embeddings를 quantize하기 위해 Binary Stochastic Quantifier (BSQ)라는 새로운 quantifier를 적용합니다. 핵심 문제는 다음과 같습니다: latent embedding fi​를 가장 similar한 token embedding vj​로 differential한 방식으로 어떻게 대체할 것인가?

Find Neighbor. vanilla vector quantization에서는 가장 가까운 code vector를 찾기 위해 "argmax" operation을 사용하는데, 이는 non-differential합니다. 본 논문에서는 selection process를 scale m에서 multi-class distribution zi∣m​∼Mult(pi∣m​)로부터 sampling하는 것으로 간주합니다: [ \begin{cases} p_{i|m} = \text{SoftMax}([f_i^T v_{0|m} /T, f_i^T v_{1|m} /T, \cdots , f_i^T v_{m-1|m} /T]) \ z_{i|m} \sim \text{Mult}(p_{i|m}) \end{cases} ] 그런 다음, pi∣m​을 differential한 방식으로 optimize합니다. temperature parameter T는 attention query operation의 softness를 제어합니다. T가 크면 attention weights는 uniform에 가까워지고, 그렇지 않으면 attention weights는 one-hot에 가까워집니다.

Optimize Neighbor. Eq. 8에서 sampling operation은 non-differentiable하므로, pi∣m​을 optimize하기 위해 reparameterization trick을 사용합니다: [ \begin{cases}     c_{i|m} = z_{i|m} \odot (1 - p_{i|m}) + (1 - z_{i|m}) \odot (- p_{i|m}) \     \hat{z}{i|m} = p{i|m} + \text{sg}(c_{i|m}) \     \hat{f}{i|m} = \hat{z}{i|m}^T V + f_{i|m} - \text{sg}(f_{i|m}) \end{cases} ] 여기서 sg(·)는 stop gradient operation, ⊙는 element-wise multiplication이고, zi∣m​∈Rm은 scale m에서의 zi​의 onehot version입니다. 첫 두 방정식은 pi∣m​이 gradient를 받을 수 있도록 zi∣m​을 z^i∣m​으로 reparameterize하며, 이는 이전 연구에서 영감을 받았습니다; 세 번째 방정식은 hi∣m​(f의 오타로 추정)이 encoder를 optimizing하기 위한 direct gradient를 받도록 합니다. z^i∣mT​V→vzi∣m​​ operation은 onehot z^i∣m​을 사용하여 zi​번째 code vector를 selects합니다.

Multiscale Code Mixing. Training 중에 scale m을 uniformly sample하고 f^​i∣m​으로 f^​i​ 역할을 하는 code embeddings를 generate합니다: [ \begin{cases}     m \sim U(5, 12) \     \hat{f}{i} \leftarrow \hat{f}{i|m} \end{cases} ] 여기서 U(a, b)는 a와 b 사이의 uniform distribution입니다. average code embedding은 scales에 걸친 codes의 mixture입니다: [ \mathbb{E}[\hat{f}i] = \sum_m \hat{f}{i|m} p(m) = \frac{1}{8} \sum_5^{12} \hat{f}_{i|m} ]

Teacher Guidance. Training convergence를 accelerate하기 위해, 확률 T로 encoder output fi​를 decoder input f^​i​로 randomly copy합니다: [ \hat{f}{i} =  \begin{cases}     \hat{z}{i}^{T}V + f_{i} - sg(f_{i}) & \text{ if } p > T;\     h_{i} (\text{아마도 } f_i) & \text{ else}. \end{cases} ] 여기서 p∼U(0,1)은 uniform distribution에서 sampled됩니다. T = 1.0일 때, vector quantization module은 skipped되어 encoder-decoder가 쉽게 optimized될 수 있습니다. T = 0.0일 때, vector quantization module은 fully used됩니다. 0.0 < T < 1.0 값에 대해, shortcut은 vector quantization model이 encoder의 inputs과 aligning되는 code vectors를 learn하도록 guides합니다. adaptive temperature scheduler는 loss L에 dependent합니다: [ \beta = \begin{cases} 1.0, & \text{ if } 0.1 < \mathcal{L} \ 0.05, & \text{ if } 0.05 < \mathcal{L} \leq 0.1 \ 0.01, & \text{ if } 0.02 < \mathcal{L} \leq 0.05 \ 0.001, & \text{ if } \mathcal{L} \leq 0.02 \end{cases} ]

Stable Gradient. Eq. 9의 scaled softmax operation은 continual model (T > 0)과 discrete vector quantization (T = 0)을 연결합니다; 따라서 VVQ에서의 gradient mismatch 대신 precise gradient computation을 가능하게 합니다. Training 동안 temperature를 1.0에서 1e-8로 점진적으로 anneal합니다; 그러나 scaled softmax의 gradient는 T가 극도로 작을 때 explode하는 경향이 있습니다: [ \begin{cases} [ p_1, p_2, \cdots, p_k ] = \text{Softmax}( a_1 /T, a_2 /T, \cdots, a_k /T ) \ \frac{\partial p_i}{\partial a_j} = \frac{1}{T} p_i \cdot (\mathbb{1}(i = j) - p_j) \end{cases} ] unstable gradient는 CodeAdapter$_m$과 encoder parameters가 큰 gradient를 한 번 updating한 후 collapse하기 때문에 representation 및 codebook collapses로 이어질 수 있습니다. 이 issue를 overcome하기 위해, 'partial gradient' trick을 도입합니다: [ \frac{1}{T} h_i^T v_j \leftarrow (\frac{1}{T} - 1)\text{sg}(h_i^T v_j) + h_i^T v_j ] 여기서 첫 번째 항은 stoped gradient이고 오직 두 번째 항만이 gradient computation에 contribute합니다. 명백하게, Eq.15는 Eq.9와 동일한 forward behavior를 가지면서 gradient는 stable하고 T의 극도로 작은 값에 affected되지 않습니다.

2.5 Equivariant Graph Decoder

invariant representations에 conditioned된 단백질 structures를 Generating하는 것은 computing efficiency에서 상당한 challenges를 제기합니다. 예를 들어, 잘 알려진 AlphaFold2를 scratch에서 training하는 데는 11일 동안 128개의 TPUv3 cores가 필요합니다; OpenFold는 training에 50000 GPU 시간이 걸립니다. 이 연구에서는 structure prediction을 위해 어떤 GNN layer에도 추가할 수 있는 efficiency plug-and-play SE(3)-layer를 제안합니다. SE(3)-layer의 simplified module과 sparse graph attention을 사용하는 BlockGAT 덕분에, 8개의 NVIDIA-A100s를 사용하여 1일 만에 전체 PDB dataset에서 model을 train할 수 있습니다.

SE-(3) Frame Passing Layer. 우리는 frame-level message passing을 도입하며, 이는 s번째 block의 local frame을 그 neighbors로부터의 relative rotation Rs​와 translation ts​를 aggregating하여 updates합니다: [ \begin{cases} vec(R_s) = \sum_{j \in \mathcal{N}s} a{sj}^r \text{ vec}(R_{sj}) \ R_s \leftarrow \text{Quat2Rot} \circ \text{Norm} \circ \text{MLP}^{9 \rightarrow 4}(\text{vec}(R_s)) \ t_s = \sum_{j \in \mathcal{N}s} a{sj}^t t_{sj} \end{cases} ] 여기서 asjr​과 asjt​는 rotation 및 translation weights이고, Ns​는 s번째 block의 neighbors입니다. vec(·)는 3×3 matrix를 9차원 vector로 펼칩니다. MLP$^{9 \rightarrow 4}$(·)는 9차원 rotation matrix를 4차원 quaternion으로 maps하고, Norm(·)은 quaternion을 normalize하여 유효한 rotation을 represents하도록 보장합니다. Quat2Rot(·)은 quaternion to rotation function입니다. 세부 사항을 다음과 같이 추가합니다: [ \begin{cases}     w_{st}^r, w_{st}^t = \sigma(\text{MLP}(f_{st})) \     \text{vec}(R_{st}) \leftarrow w_{st}^r \text{vec}(R_{st}) + (1 - w_{st}^r) \text{MLP}^{d \rightarrow 9}(f_{st}) \     t_{st} \leftarrow w_{st}^t t_{st} + (1 - w_{st}^t) \text{MLP}^{d \rightarrow 3}(f_{st}) \     a_{st}^r, a_{st}^t = \text{Softmax}(\text{MLP}^{d \rightarrow 1}(f_{st})) \end{cases} ] 여기서 wstr​과 wstt​는 rotation과 translation을 위한 updating weights이고, astr​과 astt​는 attention weights입니다. 제안된 SE-(3) layer는 local frame updating을 위해 어떤 GNN에도 추가될 수 있습니다.

Iterative Refinement 우리는 SE-(3) layer를 BlockGAT에 adding하여 SE-(3) BlockGAT라는 새로운 module을 제안합니다. structures를 iteratively refine하기 위해 multi-layer SE-(3) BlockGAT를 stack합니다: [ \begin{cases}     \mathcal{E}^{(l)} = \text{kNN}(T_s^{(l)}) \     T_{st}^{(l)} = T_s^{(l)^{-1}} \circ T_t^{(l)} \     f_s^{(l+1)}, f_{st}^{(l+1)} = \text{BlockGAT}^{(l)}(f_s^{(l)}, f_{st}^{(l)} | T_s^{(l)}, T_{st}^{(l)}, \mathcal{E}^{(l)}) \     T_s^{(l+1)} = \text{SE3Layer}(\text{sg}(T_{st}^{(l)}), f_{st}^{(l+1)}) \end{cases} ] 여기서 sg(·)는 stop-gradient operation이고, SE3Layer(·)는 Eq.17에서 설명된 SE-(3) layer입니다. 예측된 local frame Ts(l)​이 주어지면, 다음과 같이 3D coordinates를 얻을 수 있습니다: [ \begin{cases}     h_s = \text{MLP}(f_s^{(l)}) \     x_s = T_s^{(l)} \circ h_s \end{cases} ]

2.6 Reconstruction Loss

Chroma에서 영감을 받아, model을 train하기 위해 multiple loss functions를 사용합니다. overall loss는 다음과 같습니다: [ \mathcal{L} = \mathcal{L}{global} + \mathcal{L}{fragment} + \mathcal{L}{pair} + \mathcal{L}{neighbor} + \mathcal{L}_{distance} ] loss terms를 설명하기 위해, ground truth 3D coordinates X∈Rn×3와 예측된 3D coordinates X^={x1​,x2​,x3​,⋯,xn​}∈Rn×3가 주어졌을 때, aligned RMSD loss를 Lalign​(X^,X)=∥Align(X^,X)−X∥로 정의합니다. global, fragment, pair loss는 aligned MSE loss로 정의되지만, 다른 input data shape를 가집니다:

• Global Loss: shape [n,4,3]인 X. global structure의 RMSD.
• Fragment Loss: shape [n,c,4,3]인 X. 각 residue에 대한 c개 neighbors의 RMSD.
• Pair Loss: shape [n,K,c⋅2,4,3]인 X. 각 kNN pair에 대한 c개 neighbors의 RMSD.
• Neighbor Loss: shape [n,K,4,3]인 X. 각 residue에 대한 K개 neighbors의 RMSD. 여기서 n은 residues의 수, c = 7은 fragments의 수, K = 30은 kNN의 수, 4는 네 개의 backbone atoms {N, CA, C, O}를 고려함을 의미하고, 3은 3D coordinates를 의미합니다. distance loss는 예측된 pairwise distances와 ground truth pairwise distances 간의 MSE loss로 정의됩니다: [ \mathcal{L}_{distance} = \left| \text{Dist} ( \hat{X} ) - \text{Dist} ( X ) \right| ] 여기서 Dist(X) ∈Rn×n는 3D coordinates X의 pairwise distance matrix입니다. 각 decoder layer에 loss를 apply하며, final loss는 평균값이며, 이는 좋은 performance에 crucial합니다.

 

 

 

 

정리노트 (AI 연구자 대상)

Overall Architecture: 표준적인 Encoder-Quantifier-Decoder 구조 유지 (2.1).

핵심 Contribution & Mechanisms:

1. Multiscale Code Generation & Consistency (2.3):
   • Unified Representation: Scale (m=5..12, codebook size 2m) 통합 위해 Masked-binary code 사용.
   • Shared Encoder + Scale-aware Code Adapters: SE(3)-invariant features를 추출하는 공유 BlockGAT Encoder (2.2)의 output(semantic space)을 각 scale m별 전용 MLP (CodeAdapter_m)가 해당 scale의 code embedding vi∣m​으로 project. ⇒ Scale 간 consistency 확보.
2. Multiscale Training & Optimization (2.4):
   • Mix-Token Training: Training step마다 scale m을 U(5,12)에서 랜덤 sample하여 해당 scale의 quantization loss로 학습 진행 (Multiscale Code Mixing). 하나의 model로 모든 scale 동시 학습.
   • Quantifier (BSQ 기반): Softmax 기반 확률적 token 선택 (Find Neighbor) + Reparameterization Trick (Optimize Neighbor) 사용.
   • Stable Gradient Trick: BSQ의 Softmax에서 낮은 temperature (T→0) 사용 시 gradient 폭발 문제를 해결하기 위해 partial gradient 기법 도입 (Eq. 15). ⇒ Training 안정성 확보.
   • Teacher Guidance 및 Adaptive Temperature Scheduler 등 training 보조 기법 사용.
3. Hierarchy Analysis & Translation (2.3):
   • Semantic space 상에서 다른 scale(m,m′) 간 code embedding (vi∣m​,vk∣m′​) similarity를 계산하여 Transition Matrix (Mm→m′) 구축.
   • 이 matrix는 scale 간 hierarchical relationship을 정량적으로 나타내며, model 재실행 없이 finest-scale code를 다른 scale로 translate 가능하게 함.
4. Efficient Equivariant Decoder (2.5):
   • Structure generation을 위해 새로운 SE(3)-equivariant layer (SE-(3) Frame Passing Layer) 제안.
   • 이를 BlockGAT과 결합한 SE-(3) BlockGAT module을 stack하여 iterative refinement 수행.
   • 기존 방법(e.g., AlphaFold2) 대비 높은 computing efficiency 강조 (8x A100으로 1일 학습).
5. Loss Function (2.6):
   • Aligned RMSD 기반의 다양한 loss (global, fragment, pair, neighbor)와 pairwise distance loss 조합 사용.
   • 모든 decoder layer에 loss를 적용하여 performance 향상 도모.

---

쉬운 설명:

이 섹션은 FoldToken4 모델이 실제로 어떻게 작동하는지 설명하는 부분입니다.

전체 흐름: 단백질 3D 구조 정보를 받아서 (Encoder) -> 컴퓨터가 다루기 쉬운 "구조 단어"(token)로 바꾸고 (Quantizer) -> 이 단어들로 다시 3D 구조를 만들어내는 (Decoder) 과정은 기존과 비슷해요.

FoldToken4만의 특별한 점:

1. 여러 수준의 단어 동시 생성 (Multiscale Generator):
   • 단어 표현 통일: 단백질 구조를 설명하는 단어의 개수(수준, scale)가 32개든 4096개든, 일단 통일된 방식("가면 쓴 이진 코드")으로 표현해요.
   • 수준별 맞춤 번역기 (Code Adapter): 단백질 구조를 읽는 부분(Encoder)은 하나로 공유하되, 각 단어 수준(scale)마다 "전용 번역기"(Adapter)를 둬서, 같은 구조 정보라도 해당 수준에 맞는 단어(code embedding)로 딱 맞게 바꿔줘요. 이렇게 하면 다른 수준의 단어라도 일관된 의미를 가지게 되죠.
2. 똑똑한 학습 방법 (Quantifier & Training):
   • 단어 선택: 구조 정보와 가장 비슷한 단어를 고를 때, 딱 하나만 고르는 게 아니라 확률적으로 부드럽게(Softmax) 고르고, 이걸 수학적 트릭(BSQ, reparameterization)으로 잘 학습시켜요.
   • 섞어서 학습: 학습할 때마다 무작위로 단어 수준(scale)을 골라서 "오늘은 64개 단어 수준으로 학습!", "다음엔 1024개 단어 수준으로 학습!" 하는 식으로 섞어서 한 번에 학습시켜요. 모델 하나로 모든 수준을 커버하는 거죠.
   • 학습 안정화: 학습이 불안정해질 수 있는 부분을 안정시키는 기술(stable gradient trick)도 추가했어요.
3. 단어 수준 간 관계 파악 (Hierarchy Analysis):
   • 학습된 단어들이 얼마나 서로 의미가 비슷한지 계산해서, 예를 들어 "32개 단어 수준의 이 단어는 1024개 단어 수준의 저 단어들과 가장 관련이 깊다" 같은 "번역표"(Transition Matrix)를 만들어요. 이걸로 단어 수준 간의 계층 관계를 알 수 있고, 가장 자세한 단어 정보만 있으면 다른 수준의 단어로 쉽게 바꿀 수도 있어요.
4. 빠르게 구조 그리기 (Decoder):
   • 만들어진 구조 단어들을 보고 원래 3D 구조를 복원할 때, 새롭고 더 효율적인 방식(SE-(3) BlockGAT)을 사용해서 더 빠르게 구조를 그릴 수 있게 만들었어요.
5. 평가 기준 (Loss):
   • 모델이 구조를 얼마나 잘 그렸는지 평가할 때, 전체적인 모양, 부분적인 모양, 원자 간 거리 등 여러 기준을 종합적으로 사용해서 꼼꼼하게 평가하고 학습시켜요.

 

 

 

가장 처음 인풋은 PDB (Protein Data Bank) 등에서 얻은 단백질의 3차원 원자 좌표 정보

그걸 바탕으로 Graph형태로 변환을 한다. 각 아미노산이 node가 되고 가까운 node 끼리 edge로 연결 된다.

이때 node와 edge는 invariant features를 초기값으로 가진다.

 

그 Graph를 BlockGAT를 통과 시켜서 아미노산 단위로 f1...fn를 나타내는 f를 얻게 된다.

 

거기서 특정 아미노산의 정보를 담고 있는 f_s 를 특정 scale m에 해당하는 특정 token z_s|m으로 변환 해야함.
근데 바로 f_s를 사용하지 않음. 왜냐면 코드북 사이즈가 유동적이기 때문 그래서 f_s를  M길이를 가진 masked binary code로 변환해서 사용

 

만약 현재 목표 scale이 m이고 token id가 z라면 M-m개가 마스크 되고 m개의 비트는 z의 실제 이진수 값이 됨 

 

그걸 다시 codeadapter 라는 mlp를 통과해서  d차원의 code embedding 벡터가 나오게 됨.
이게 바로 m에서의 token z를 나타내는 벡터 표현

 

 

즉 f_s를 v_0|m ... v_...|m m에 속하는 모든 가능한 code embedding과 비교

 

유사도는 내적

 

그, 후 소프트 맥스를 사용해서 확률을 얻음

 

그래서 token z값을 얻게 됨

 

그걸로 f햇을 얻고 디코딩을 해서 복원 기반으로 학습

SE-(3) Frame Passing Layer라는 모듈을 통해서 

 

디코딩

 

 

전체 구조에서 좌표차이, 즉 전체적인 모양이 얼마나 비슷한지

작은 조각 fragment가 얼마나 비슷한지

아미노산 쌍 주변의 조각들을 비교해서 얼마나 정확한지

가까운 이웃들의 좌표가 얼마나 비슷한지

아미노산 쌍 사이의 거리들이 얼마나 비슷한지

 

5개의 loss를 가지고 인코더 코드북 디코더 없데이트