Protein : 논문 리뷰 : BiHiTo: Biomolecular Hierarchy-inspired Tokenization

용어 설명

• Biomolecular Hierarchy: 생체 분자(단백질, RNA 등)가 원자 -> 백본 -> 2차 구조 -> 복합체로 이어지는 자연적이고 다층적인 조립 계층(multi-level organizational structure)을 의미합니다.
• Tokenization: 연속적이고 불규칙한 3D point cloud 데이터를 AI 모델(주로 transformer 계열)이 처리하기 쉬운 순서가 있는 이산적인(discrete) 토큰으로 변환하는 과정입니다.
• VQ-VAE (Vector Quantized Variational Autoencoder): 연속적인 데이터를 이산적인 codebook을 사용하여 압축(Encoding)하고 복원(Decoding)하는 생성 모델 구조입니다.
• FSQ (Finite Scalar Quantization): VQ-VAE의 codebook 매핑 과정에서 생기는 최적화 문제를 피하기 위해, 코드북을 유한한 스칼라 값들의 조합으로 단순화하고 효율적으로 구성하는 최신 양자화 기법입니다.
• Bidirectional Mamba (BML): 긴 시퀀스 데이터를 양방향으로 처리하여 원거리 의존성(long-range dependency)을 효율적으로 포착하면서도 연산 복잡도를 줄인 Mamba 기반의 신경망 레이어입니다.
• C-alpha (Alpha Carbon) / C3': 단백질(C-alpha)과 RNA(C3')의 뼈대(backbone) 구조를 결정짓는 가장 핵심적인 기준 원자들입니다.
• Umeyama-Kabsch algorithm: 두 3D point cloud 집합 간의 회전 및 이동 변환을 계산하여 최적의 구조적 정렬(structural alignment)을 찾아내는 알고리즘으로, 주로 RMSD 오차를 계산하기 직전에 사용됩니다.

---

Purpose of the Paper

• 기존 연구의 한계: 기존의 구조 생성 방법(diffusion, autoregressive 모델 등)은 대규모 생체 분자를 시뮬레이션할 때 계산 복잡도가 기하급수적으로 증가하는 문제가 있습니다. 이를 해결하기 위해 최근 Bio2Token 같은 unified biomolecular VAE가 등장했지만, 이들은 생체 분자를 단순한 '일반 3D point cloud'로만 취급하여 생물학적 고유의 계층 구조(hierarchy priors)를 무시했습니다. 이는 모델 학습을 어렵게 만들고 generalization 능력을 저하시키는 원인이 됩니다.
• 새로운 문제 정의 및 목적: 생체 분자는 단순한 점들의 집합이 아니라 물리화학적 규칙에 따라 조립되는 명확한 hierarchy를 가지고 있습니다. 본 논문은 이러한 생물학적 조립 원리를 3D structure tokenization 과정에 직접 반영하는 것을 목표로 합니다. 원자 단위의 디테일부터 전체적인 거시적 단백질 접힘(global conformation)까지 동시에 포착할 수 있는 다층(multi-level) 구조의 모델인 BiHiTo를 제안하여 기존 방식론의 표현력 한계를 극복하고자 합니다.

---

Key Contributions

• Natural Hierarchy 기반의 Multi-level Tokenizer 설계 (Novelty)
  • 기존 연구들이 모든 원자를 동등한 해상도로 취급한 것과 달리, 생체 분자를 5단계 물리적 계층(Global topology -> Sparse C-alpha/C3' -> Full C-alpha/C3' -> Backbone atoms -> All-atom)으로 나누어 downsampling 및 upsampling을 수행합니다. 이를 통해 생물학적 위상(topology) 정보를 본질적으로 보존하는 참신한 구조를 제안했습니다.
• Multi-codebook Quantization 도입 (Novelty)
  • 단일 codebook에 모든 정보를 압축하던 기존 VQ-VAE 방식에서 벗어나, 5단계의 각 hierarchy level마다 해상도에 맞춰 지수 함수적으로 크기가 증가하는(exponential scaling) 전용 codebook들을 병렬로 할당하는 FSQ 시스템을 구축했습니다. 이는 거대 분자 복합체를 효율적으로 처리하면서도 강력한 generalization 능력을 부여하는 핵심 기여입니다.
• Bidirectional Mamba Architecture 적용
  • Encoder와 Decoder에 Transformer 대신 BML을 적용하여, 연산 병목 현상을 줄이면서도 3D 구조 데이터의 양방향 long-range dependency를 효과적으로 모델링했습니다.

---

Experimental Highlights

• 핵심 실험 설정: 단백질 학습(CATH 4.2), 단백질 테스트(CASP14/15), RNA 테스트(RNA3DB), 다중 형태 OOD 테스트(FastFolding), 소분자 테스트(2DFT). 주요 성능 지표는 RMSD(낮을수록 좋음)와 TM-Score(높을수록 좋음)를 사용하였으며, 가장 강력한 최신 베이스라인인 Bio2Token 및 ESM3 등과 비교했습니다.
• 압도적인 State-of-the-art 달성: CASP14 단백질 벤치마크에서 RMSD 0.515 Angstrom (Bio2Token 대비 17% 감소), TM-Score 0.987을 기록하며 SOTA를 달성했습니다.
• 놀라운 OOD Generalization 성능: 학습 데이터에 없는 Out-of-Distribution 데이터인 FastFolding (다중 형태 단백질) 데이터셋에서, 모델 추가 튜닝 없이 Bio2Token 대비 RMSD 오차를 무려 51% (0.965 -> 0.469 Angstrom) 감소시키는 놀라운 성능을 보였습니다.
• 거대 복합체 및 RNA 모델링 우수성: 단백질-RNA 복합체(4W5O) 복원에서 원자 수준의 정밀도를 유지하며 Bio2Token 대비 RMSD를 46.5% 감소시켰습니다.
• Ablation Study (가설 검증): 각 계층의 중요도를 분석한 결과, Full C-alpha 계층(L3)과 All-atom 계층(L5)이 예측 정확도에 가장 지배적인 기여를 함을 수치적으로 증명하여 구조적 설계의 정당성을 입증했습니다.

---

Limitations and Future Work

• Limitations: BiHiTo는 3D 구조를 토큰화하고 복원하는 데 압도적인 성능을 보였으나, 추출된 고품질의 discrete 토큰들이 실제 downstream tasks (예: 단백질 구조 예측, 신약 물질 디자인 등)에서 얼마나 큰 성능 향상을 가져오는지에 대한 end-to-end 파이프라인 검증은 논문의 범위상 제한적으로만 제시되었습니다.
• Future Work: 저자들은 향후 BiHiTo를 multi-scale biomolecular generation(예: VAR 기반의 자동 회귀적 구조 생성 등) 파이프라인에 적용할 계획입니다.
• 의의: 한계점에도 불구하고, 이 연구가 제안한 tokenization은 대용량 생체 분자 데이터를 Large Language Models이 처리하기 완벽한 형태로 변환해 주었기 때문에, 향후 구조 생물학 분야에서 AI 모델의 연산 효율과 생성 품질을 폭발적으로 증가시킬 잠재력이 있습니다.

---

Overall Summary

BiHiTo는 생체 분자가 원자부터 단백질 뼈대, 그리고 거대 복합체로 이어지는 자연적인 계층 구조를 지녔다는 점에 착안하여 설계된 혁신적인 multi-level 3D tokenization 모델입니다. 기존의 단일 해상도 point cloud 방식과 달리, 5단계의 물리적 샘플링과 multi-codebook 양자화 방식을 결합하여 거시적 구조와 원자 수준의 디테일을 동시에 완벽하게 학습합니다. 그 결과 단백질, RNA, 복합체 등 다양한 벤치마크에서 압도적인 state-of-the-art 복원 성능을 달성했으며, 특히 전혀 학습하지 않은 OOD 구조에 대해 극적인 generalization 능력을 입증하여 향후 대규모 생체 분자 생성 AI를 위한 핵심 기반 기술로 자리 잡을 것입니다.

---

쉬운 설명

BiHiTo의 작동 방식은 **"거대한 도시의 지도를 그리는 방식"**과 매우 유사합니다.

기존의 AI 모델(Bio2Token 등)은 도시를 그릴 때 건물, 나무, 가로등을 가리지 않고 전부 똑같은 중요도를 부여해 한 번에 스캔하려고 했습니다. 그러다 보니 데이터가 너무 무거워지고 숲 전체의 모양(전체 단백질 구조)을 파악하기 어려웠습니다.

반면 BiHiTo는 지도를 그릴 때 단계를 나눕니다(Hierarchy). 먼저 1) 국가의 큰 윤곽을 그리고, 2) 주요 고속도로와 랜드마크(C-alpha 뼈대)의 위치를 핀으로 꽂은 뒤, 마지막에 3) 개별 골목과 건물들(All-atom 디테일)을 세밀하게 채워 넣습니다. 이렇게 각 해상도에 맞는 맞춤형 도화지(Multi-codebook)를 따로 쓰기 때문에, 본 적 없는 완전히 새로운 도시 구조(OOD 데이터)를 만나더라도 전체적인 뼈대부터 디테일까지 훨씬 빠르고 오차 없이 완벽하게 복원해 낼 수 있는 것입니다.

 

 

 

Abstract

생체 분자의 3차원 원자 배열은 생물학적 기능을 규명하는 데 핵심적입니다. 과학 분야의 AI 발전에 힘입어 접근 가능한 구조 데이터가 빠르게 확장됨에 따라, 생물학적 조립 원리에 의해 형성된 high-dimensional systems인 대규모 생체 분자 구조를 효율적으로 modeling해야 하는 중요한 과제가 부각되고 있습니다.

이를 해결하기 위해, 우리는 자연적인 생물학적 조립 계층을 본질적으로 모방하는 multi-level 생체 분자 계층 기반의 Tokenizer인 BiHiTo를 소개합니다.

구체적으로, 우리는 생체 분자 구조의 자연적 계층을 반영하는 multi-codebook quantizer를 설계하여, 원자 motifs부터 전역적인 conformational variations에 이르는 representations를 동시에 포착할 수 있도록 합니다. 이러한 계층적 정렬은 생체 분자 구조의 생물학적 interpretability와 reconstruction fidelity를 크게 향상시킵니다.

광범위한 실험을 통해 BiHiTo가 분자 동역학 trajectories 및 거대 분자 complexes 전반에 걸쳐 state-of-the-art 성능과 견고한 generalization을 제공하며, 구조 generation 및 동적 conformation exploration의 발전을 촉진함을 입증합니다.

CASP14 및 OOD test set의 FastFolding 단백질 multi-conformation data의 reconstruction에서, 우리의 method는 Bio2Token과 비교하여 RMSD를 각각 17% 및 51% 감소시켰습니다.

 

 

 

Introduction

생체 분자 구조는 거대 분자의 conformational landscapes와 기능적 레퍼토리를 결정합니다. X선 결정학과 극저온 전자 현미경으로 대표되는 고해상도 실험 기술은 일상적으로 sub-angstrom 정밀도로 원자 좌표를 해석합니다. 그 결과물인 구조 좌표는 Protein Data Bank (PDB)와 같은 커뮤니티 저장소에 보관됩니다. 보완적으로, AlphaFold2로 대표되는 최근의 발전은 실험적 정확도에 근접한 단백질체 규모의 구조적 models를 생성하여, 접근 가능한 구조적 영역을 수십 배 이상 확장시켰습니다. 이와 수반된 고품질 구조 data의 확장은 이제 거대 분자 상호작용 네트워크에 대한 체계적이고 데이터 기반의 연구를 가능하게 하여, 기초 생물학과 다양한 생명공학 응용 분야 전반의 발전을 촉진하고 있습니다.

그러나 생체 분자에 내재된 구조적 복잡성은 computational modeling에 있어 상당한 과제를 제시합니다. diffusion models 및 autoregressive frameworks를 포함한 현재의 구조 generation methods는 대규모 생체 분자 시스템 내에서 원자 수준의 상호작용을 시뮬레이션할 때 기하급수적으로 증가하는 계산 복잡성 문제와 자주 씨름합니다. 이는 all-atom encoding/decoding techniques 개발의 시급한 필요성을 강조합니다.

최근의 발전은 생체 분자 구조 tokenization에서 큰 진전을 이루었습니다. 일부 methods는 저분자에 초점을 맞춥니다. 단백질의 맥락에서 FoldSeek은 구조적 알파벳을 사용하여 단백질 형태를 1D sequences로 변환하여 빠른 비교를 용이하게 합니다. ESM-3는 단백질 백본을 encode하기 위해 transformer models를 train시킵니다. 이러한 approaches의 대부분은 백본 원자에 초점을 맞추거나 짧은 사슬에서의 training에 국한됩니다. 보다 최근에는 Bio2Token과 같은 통합된 생체 분자 VAEs가 균일성을 달성하기 위해 생체 분자를 순전히 point clouds로 취급합니다. 그러나 이는 생물학적 분자의 고유한 계층적 priors를 무시하여, training을 더 어렵게 만들고 generalization을 방해합니다.

우리는 생체 분자를 일반적인 3D point cloud data로 취급하기보다는, 그것들이 고유한 multi-level 조직 구조를 가지고 있음을 인식합니다. 단백질을 예로 들면, 이들의 3차원 conformations는 수천 개 원자의 계층적 조립에서 비롯되며, 이는 $\alpha$-helices 및 $\beta$-sheets와 같은 2차 구조 요소를 형성하고 나아가 기능적인 초분자 complexes로 조직화됩니다(그림 1 참조). 이러한 자연적 계층은 생체 분자 구조 encoding을 위한 구조적 priors 역할을 할 수 있습니다.

이를 바탕으로, 우리는 자연적인 생물학적 조립 계층을 본질적으로 모방하는 multi-level 생체 분자 계층 기반 Tokenizer인 BiHiTo를 소개합니다. 구체적으로, 우리는 생체 거대 분자를 전역적 topology, $\alpha$-carbon/C3' 간격의 sparse sampling, 전체 $\alpha$-carbon/C3' 백본, 완전한 백본 원자 및 full-atom resolution이라는 5개의 생물학적 구조화된 quantization levels로 분해하는 multi-codebook quantizer를 설계합니다. 이러한 multi-level 설계를 통해, quantizer는 생체 분자의 내재적 구조적 계층을 명시적으로 encode하여, 원자 motifs부터 전역적인 conformational variations에 이르는 representations를 동시에 포착할 수 있도록 합니다. 광범위한 experiments는 제안된 BiHiTo가 reconstruction 정확도 및 입체화학적 타당성을 포함한 모든 metrics에 걸쳐 복잡한 구조에 대해 기존 VQ-VAEs를 크게 outperforms함을 입증합니다. BiHiTo는 또한 modeling 효능과 generalization 능력을 향상시켜, Bio2Token과 비교하여 RNA3DB test set에서 RMSD를 25% 낮추고 out-of-distribution FastFolding 단백질 multi-conformation dataset에서 RMSD를 51% 낮춥니다.

요약하자면, 우리의 주요 contributions는 다음과 같습니다:

• 우리는 생체 분자 구조의 계층적 조립 prior를 기반으로 하는 multi-level 생체 분자 tokenizer를 설계하여, 원자 motifs부터 전역적인 conformational variations에 이르는 representations의 동시 포착을 가능하게 합니다.
• 우리는 생체 분자의 다양한 자연적 계층을 개별적으로 model하기 위해 multi-codebook quantizer를 사용하여, 거대 분자 complexes를 효율적으로 처리하는 동시에 model generalization을 향상시킵니다.
• 광범위한 experiments는 제안된 method가 reconstruction 정확도 및 입체화학적 타당성에서 state-of-the-art performance를 달성함을 입증합니다.

 

 

Introduction 핵심 정리 노트

1. 연구 배경 및 기존 방법론의 한계

• 문제 제기: 대규모 생체 분자 시스템에서 원자 수준의 상호작용을 시뮬레이션할 때 발생하는 기하급수적인 계산 복잡성 문제 해결을 위해 all-atom encoding/decoding techniques 개발이 필수적임.
• 기존 연구의 한계점:
  • 일부 models(예: FoldSeek, ESM-3)는 1D sequences 변환에 그치거나 백본(backbone) 원자 및 짧은 사슬에만 국한됨.
  • 최근의 Bio2Token과 같은 통합 VAEs는 생체 분자를 단순한 3D point clouds로 취급함.
  • 결정적 패착: 생체 분자 고유의 생물학적 계층적 priors를 무시하여, 모델의 training이 어려워지고 generalization 성능이 저하됨.

2. 논문의 핵심 아이디어 (Key Insight)

• 생체 분자를 단순한 3D point cloud data로 보는 대신, 내재된 다단계 조직 구조(multi-level organizational structure)에 주목함.
• 원자 $\rightarrow$ 2차 구조($\alpha$-helices, $\beta$-sheets) $\rightarrow$ 초분자 complexes로 이어지는 자연적인 생물학적 조립 계층(hierarchical assembly) 자체를 구조적 priors로 활용함.

3. 제안 모델: BiHiTo (Biomolecular Hierarchy-inspired Tokenizer)

• 생물학적 조립 계층을 본질적으로 모방하는 multi-level tokenizer.
• Multi-codebook quantizer 설계: 생체 거대 분자를 다음 5가지 생물학적 quantization levels로 분해하여 개별적으로 modeling함.
  1. 전역적 구조 (Global topology)
  2. $\alpha$-carbon/C3' 간격의 sparse sampling
  3. 전체 $\alpha$-carbon/C3' 백본 (Full backbones)
  4. 완전한 백본 원자 (Complete backbone atoms)
  5. Full-atom resolution
• 효과: 원자 수준의 motifs부터 전역적인 conformational variations까지 다양한 범위의 representations를 동시에 포착함.

4. 주요 기여 및 실험 성과

• Generalization 및 효율성 향상: 거대 분자 complexes를 효율적으로 처리하면서도 model generalization 능력을 극대화함.
• SOTA 달성: 기존 VQ-VAEs 대비 reconstruction 정확도 및 입체화학적 타당성(stereochemical validity) 등 모든 metrics에서 압도적 우위.
• 성능 지표: Bio2Token 대비 RNA3DB test set에서 RMSD 25% 감소, OOD FastFolding 단백질 multi-conformation dataset에서 RMSD 51% 감소.

---

쉬운 설명 :

이 논문의 서론을 아주 쉽게 비유하자면, 건물(단백질)을 AI에게 이해시키는 방법의 차이를 설명하고 있습니다.

기존의 AI 모델들(Bio2Token 등)은 단백질을 그저 허공에 떠 있는 수많은 점들의 모음(Point Cloud)으로만 생각했습니다. 이렇게 하면 AI 입장에서는 "이 점이 왜 여기에 있지?"라는 규칙을 처음부터 맨땅에 헤딩하듯 학습해야 해서 훈련도 어렵고 새로운 형태를 예측하는 능력(Generalization)도 떨어졌습니다.

하지만 실제 단백질은 아무렇게나 뭉쳐진 점들이 아닙니다. **벽돌(원자)이 모여서 벽(2차 구조)이 되고, 벽이 모여서 하나의 방(초분자 복합체)이 되는 명확한 건축 단계(계층 구조)**가 있습니다.

그래서 이 논문은 BiHiTo라는 새로운 모델을 만들어, AI에게 단백질을 그냥 점들의 모음으로 주지 않고 이 5단계의 건축 도면(자연적 계층)을 나누어서 가르쳤습니다(Multi-codebook quantizer). 전체적인 뼈대부터 세밀한 원자 위치까지 단계별로 정보를 압축하고 해독하게 만들었더니, AI가 단백질 구조의 핵심(Representation)을 훨씬 더 완벽하게 이해하고 복원(Reconstruction)할 수 있게 되었다는 것이 이 Introduction의 핵심입니다.

 

Related Work

3D Point Clouds Tokenization

Point clouds는 3D data format으로서, 내재된 sparsity, disordered arrangement 및 irregular structure로 인해 deep learning에 독특한 challenges를 제시합니다. 연구자들은 unstructured point cloud data를 ordered token representations로 변환하는 다양한 methods를 개발해 왔습니다. 이러한 tokenization approaches는 neural networks, 특히 transformer-based architectures에 의한 보다 효율적인 processing을 가능하게 합니다.

예를 들어, PointContrast와 DepthContrast는 모두 contrastive learning을 통해 instance discrimination frameworks를 확립하며, 전자는 multi-view observations 전반에 걸쳐 동일한 points의 features를 정렬(align)하는 반면, 후자는 향상된 3D representation learning을 위해 이 패러다임을 depth map augmentations로 확장합니다. 보다 최근에 Point-BERT는 BERT와 유사한 masked pre-training paradigm을 채택하고, tokenization을 위해 farthest point sampling (FPS)을 통해 point clouds를 local patches로 그룹화합니다. Point-MAE는 FPS-based patch partitioning scheme을 유지하지만 MAE-style framework로 전환합니다. Point-M2AE는 global-to-local reconstruction을 향상시키기 위해 hierarchical self-supervised learning을 추가로 도입합니다.

현재의 point cloud tokenization approaches는 biological systems에 내재된 더 강력한 intrinsic order로 인해 biomolecular structures에 대한 applicability가 제한적임을 보여줍니다. 정확한 결합 길이(bond lengths), 각도 제약(angular constraints) 및 비틀림 선호도(torsional preferences)와 같은 생체 분자의 geometric regularities는 명시적인 structural modeling을 요구하며, 이는 point cloud methods가 본질적으로 결여하고 있는 capability입니다.

Biomolecules Structure Tokenization

현재 biomolecular structures에 대한 tokenization methods는 주로 단일 molecule types에 초점을 맞추고 있습니다. 저분자(small molecules)의 경우, UniMol은 pairwise structural features를 사용하여 원자(atoms)를 개별적으로 tokenize합니다. Geo2Seq는 3D geometries를 $SE(3)$-invariant sequences로 encode합니다. 그리고 MolStructTok는 VQ-VAE discretization과 함께 spherical line notation을 사용합니다.

sequential order를 가진 거대한 선형 중합체(linear polymers)로서의 단백질(Proteins)은 structural discretization에 있어 자연적인 이점을 제공합니다. FoldSeek은 '3Di' alphabet을 통해 3D backbone structures를 1D sequences로 변환하는 것을 개척했습니다. ProTokens는 structure prediction과 inverse folding을 통합하는 unsupervised framework를 도입하여 protein backbone structures를 작고 아미노산과 유사한(amino-acid-like) discrete representations로 tokenize합니다. ESM-3는 geometric attention을 갖춘 transformer-based VQ-VAE를 적용하는 반면, InstaDeep과 FoldToken4는 GNN-based vector quantization을 사용합니다. Bio2Token은 $SE(3)$ constraints 없이 all-atom structures를 위한 Mamba-based quantized autoencoders를 추가로 도입합니다.

본 연구에서는 biomolecules의 자연적인 hierarchical organization을 반영하는 multiscale tokens를 통합하도록 structure VQ-VAE를 확장하여, high-fidelity reconstruction 및 biologically interpretable encoding을 가능하게 합니다.

 

 

Related Work 핵심 정리 노트

1. 3D Point Clouds Tokenization의 한계 (The Gap)

• 기존 동향: Point-BERT, Point-MAE 등은 3D point clouds의 sparsity와 irregular structure를 해결하기 위해, farthest point sampling (FPS) 등을 통해 로컬 패치로 묶어 tokenization을 수행함. (주로 contrastive learning이나 masked pre-training 패러다임 활용)
• 생체 분자 적용 시의 치명적 문제: 생체 분자는 정확한 결합 길이(bond lengths), 각도 제약, 비틀림 선호도 등 매우 강한 기하학적 규칙성(geometric regularities)과 내재적 질서를 가짐. 기존의 범용 point cloud methods는 이러한 명시적인 structural modeling을 반영할 수 없어 생체 분자 구조 학습에 부적합함.

2. Biomolecules Structure Tokenization의 현주소

• 저분자 (Small molecules): UniMol, Geo2Seq ($SE(3)$-invariant sequences 활용), MolStructTok 등 주로 개별 원자나 3D geometries를 이산화(discretization)하는 데 집중됨.
• 거대 단백질 (Proteins): 순차적 질서를 가진 선형 중합체라는 특성을 활용한 연구들이 주를 이룸.
  • 1D Sequence 변환: FoldSeek ('3Di' alphabet 활용)
  • VQ-VAE 및 GNN 기반: ESM-3 (transformer-based), InstaDeep, FoldToken4
  • All-atom 접근: Bio2Token (Mamba-based, $SE(3)$ 제약 없음)
• 한계점: 대부분 단일 분자 유형에만 초점을 맞추거나, 복잡한 생체 분자의 입체적 계층 구조를 토큰화하는 데는 한계가 존재함.

3. 본 연구의 차별점 (Core Contribution)

• 기존의 structure VQ-VAE 구조를 확장하여, 생체 분자 고유의 자연적인 계층적 조직(natural hierarchical organization)을 그대로 반영하는 multiscale tokens를 도입함.
• 이를 통해 단순 형태 보존을 넘어선 high-fidelity reconstruction을 달성하고, 모델이 생물학적으로 해석 가능한(biologically interpretable) encoding을 수행하도록 만듦.

---

쉬운 설명 :

이 섹션은 **"기존 AI 모델들은 왜 단백질 구조를 제대로 이해하지 못했나?"**에 대한 배경을 설명하고 있습니다.

• 3D 3D Point Clouds Tokenization 부분: 기존 AI 비전 분야에서는 3D 사물(예: 자동차, 의자)을 수많은 점들의 모음(Point Clouds)으로 보고 덩어리(Patch)로 묶어서 학습했습니다. 하지만 단백질은 단순한 점토 덩어리가 아니라, 원자와 원자 사이의 거리, 꺾이는 각도 등이 매우 엄격한 규칙으로 결합된 정교한 기계와 같습니다. 일반적인 3D 학습 방식을 단백질에 그대로 적용하면 이런 디테일한 '생물학적 규칙'이 다 무시되어 버립니다.
• Biomolecules Structure Tokenization 부분: 그래서 과학자들은 단백질이나 분자 구조에 특화된 번역기(Tokenizer)들을 만들어왔습니다. 하지만 이전 모델들은 작은 분자에만 적용 가능하거나, 복잡한 3D 단백질 구조를 단순한 1차원 끈(Sequence)으로 억지로 펴서 이해하려 했고, 최근 모델들조차 구조의 디테일을 완벽히 살리지는 못했습니다.
• 결론적으로 이 논문이 하고자 하는 것: 이 논문은 기존처럼 단백질을 단순한 점의 모음이나 1차원 선으로 취급하지 않겠습니다. 대신 단백질이 실제로 만들어지는 과정(원자 $\rightarrow$ 뼈대 $\rightarrow$ 전체 구조)이라는 '자연스러운 계층 구조'를 여러 가지 크기의 렌즈(multiscale tokens)로 들여다보며 AI에게 학습시키겠습니다. 이것이 이전 방식들과의 가장 큰 차이점입니다.

 

Method

이 섹션에서는 제안된 method에 대해 자세히 설명하며, 이는 encoder, quantizer, decoder라는 세 가지 기본 구성 요소로 이루어져 있습니다.

다음으로, 효율적이고 강력한 feature extraction을 보장하기 위해 Mamba-based model architecture를 소개하는 것으로 시작합니다.

그런 다음, biomolecular structure의 multilevel tokenization을 달성하기 위해 제안된 natural hierarchical quantization을 개략적으로 설명하며, 이는 복잡한 data patterns를 포착하고 represent하는 model의 능력을 크게 향상시킵니다.

우리의 methodology의 세부 사항은 다음과 같습니다.

---

Model Architecture

encoder에서 $N$개의 heavy atoms를 가진 biomolecular structure는 point cloud $X \in \mathbb{R}^{N \times 3}$으로 표현됩니다. 이 point cloud는 atom-identity-agnostic하며 residue나 atom-type 정보를 전달하지 않습니다. 우리의 encoder와 decoder는 atomic locality를 보존하면서 biomolecular point clouds의 bidirectional structural dependencies를 포착하도록 설계된 neural module인 Bidirectional Mamba Layers (BMLs)를 사용합니다.

각 BML은 bidirectional information flow를 확립하기 위해 flip operations (Flip)와 weight-shared Mamba blocks를 활용하는 대칭적인 residual structure를 사용합니다. 길이 $N$의 input point cloud sequence $X = [x_0, x_1, \dots, x_{N-1}] \in \mathbb{R}^{N \times d}$가 주어지면, processing은 두 가지 경로를 따라 발생합니다:

• Primary Path: $X$는 최종 Add module (residual connection)로 직접 전파됩니다. 이 경로는 원본 structural information을 보존하고 backpropagation 동안 안정적인 gradient flow를 보장합니다.
• Processing Path: $X$는 순차적인 transformations인 Flip $\rightarrow$ Mamba Block $\rightarrow$ Flip을 거칩니다. Mamba Block은 weight sharing을 통해 역전된 sequence $F(X)$를 처리합니다(bidirectional RNNs와 유사함). 이 대칭적인 processing은 model이 forward 및 reverse 방향 모두에서 long-range dependencies를 포착할 수 있게 합니다.

output sequence $Y$는 다음과 같이 공식화됩니다:

$$Y = M(X) + F(M(F(X)))$$

여기서 $X$는 input point cloud sequence를 나타내며 각 구성 요소는 bidirectional feature learning에서 특정 역할을 수행합니다. $F$는 sequence 순서를 역전시키는 flip operation을 나타내어 model이 forward 및 reverse 방향 모두에서 structural dependencies를 처리할 수 있게 합니다. $M$은 Mamba Block의 transformation에 해당합니다.

---

Natural Hierarchical Quantization

그림 2에 설명된 바와 같이 핵심 module인 NHQ component는 세 가지 주요 측면으로 구성됩니다. biomolecular hierarchy priors를 활용하는 multi-level downsampling, computational efficiency와 함께 expressive representations를 가능하게 하는 hierarchical quantization, 그리고 upsampling 및 loss optimization을 통합하는 reconstruction techniques와의 feature fusion입니다.

이 세 부분의 세부 사항은 다음과 같습니다.

Algorithm 1: Hierarchical Sampling Strategy

1: Input:

2: $X \in \mathbb{R}^{N \times d}$: Full atom embeddings

3: $m_{C\alpha} \in \{0, 1\}^N$: $C\alpha/C3'$ mask

4: $m_{bb} \in \{0, 1\}^N$: Backbone mask

5: Output: Multi-scale embeddings $[X_1, \dots, X_5]$

6: L1: Global max-abs feature

7: $s_i \leftarrow \sum_{j=1}^d |X_{i,j}|$

8: $i^* \leftarrow \arg\max_{i \in \{1,\dots,N\}} s_i$

9: $X_1 \leftarrow X[i^*]$

10: L2: Strided $C\alpha/C3'$ sampling

11: $X_{C\alpha} \leftarrow X[m_{C\alpha}]$

12: $X_2 \leftarrow \text{StridedSample}(X_{C\alpha}, \lceil N/2 \rceil)$

13: L3: Full $C\alpha$ set

14: $X_3 \leftarrow X[m_{C\alpha}]$

15: L4: Backbone motifs

16: $X_4 \leftarrow X[m_{bb}]$

17: L5: Atomic resolution

18: $X_5 \leftarrow X$

19: Return $[X_1, X_2, X_3, X_4, X_5]$

Natural Hierarchy. 기존의 biomolecular reconstruction methods는 모든 atoms를 균질화하여 중요한 atoms(예: 그림 1에 시각화된 $\alpha$-carbons)의 topological dominance를 무시합니다. 이를 해결하기 위해, 우리는 physics-informed hierarchical embedding strategy를 제안합니다. hierarchical sampling과 residue-aware encoding을 통해, 이 strategy는 atomic primacy relationships를 명시적으로 models합니다.

$$X_l = H(X, m_l)$$

여기서 $X_l$은 level $l$의 encoder-embedded molecular representation을 나타내고, $H$는 hierarchical sampling function을 나타내며, $m_l$은 각 hierarchy에 대한 sampling mask를 지정합니다.

sampling rule은 layer-specific physical properties와 mask parameters $m_l$ 모두에 동적으로 적응합니다. hierarchical sampling strategy는 Alg. 1로 공식적으로 정의됩니다.

algorithm은 biomolecular representations를 5개의 뚜렷한 resolution levels로 분해합니다:

• Level 1 (L1: Global max-abs feature) $s$는 per-atom activation vector(absolute feature values의 합)를 나타내며, $i^*$는 최대 activation sum을 가진 atom의 index를 나타냅니다. 이 단계는 가장 biophysically significant atom을 topological anchor point로 식별하여 결합된 feature activation이 가장 높은 molecular "hotspot"을 포착합니다. 이는 structural alignment 및 functional site identification을 위한 global reference를 제공합니다.
• Level 2 (L2: Strided $C\alpha/C3'$ sampling) $X_{C\alpha}$는 $C\alpha$ (단백질) 또는 $C3'$ (RNA) atoms의 feature matrix인 반면, $X_2$는 균일하게 subsampled된 $C\alpha/C3'$ atoms를 포함합니다. 이 level은 global 및 $C\alpha/C3'$ features 간의 양적 격차를 줄여 levels 간의 부드러운 전환을 달성합니다.
• Level 3 (L3: Full $C\alpha$ set) $X_3$는 $C\alpha/C3'$ atoms의 완전한 집합을 나타냅니다. 이 단계는 residue-level conformation의 정밀한 modeling을 위해 full backbone representation을 유지합니다. 각 $C\alpha/C3'$는 하나의 residue에 해당하며 dihedral angles ($\phi/\psi$)와 secondary structure motifs를 직접적으로 encoding하여 polypeptide/nucleotide chain의 정확한 reconstruction을 가능하게 합니다.
• Level 4 (L4: Backbone motifs) $X_4$는 전체 backbone atom set (N, $C\alpha$, C, O)를 포함합니다. 이 level은 모든 backbone atoms를 포함하여 peptide bond rigidity constraints를 시행하며, 이는 protein folding에서 structural integrity를 유지하고 stereochemical constraints를 충족시키는 데 중요합니다.
• Level 5 (L5: Atomic resolution) $X_5$는 완전한 atomic embedding matrix입니다. 이 마지막 level은 sidechain conformations, hydrogen-bond networks 및 solvent interaction surfaces의 high-precision reconstruction을 위해 전체 atomic details를 보존합니다. 이는 atomic-level structural refinement를 위한 가장 미세한 granularity를 제공하여, molecular interactions 및 binding sites의 정확한 modeling을 가능하게 합니다.

Hierarchical Quantization. Hierarchical Quantization (HQ) framework 내에서, 우리의 quantization system은 exponential scaling을 가진 composite codebooks를 사용합니다. 이 structure는 낮은 search complexity를 유지하면서 간결한 high-fidelity representations를 가능하게 합니다. 각 hierarchy level $l$에 대해, quantization system은 size vector $s_l = [C_0, C_1, \dots, C_{K_l-1}]$를 통해 구성됩니다(여기서 $K_l$은 hierarchy level $l$에서의 sub-codebook count를 나타내고, $C_k$는 sub-codebook $C_k$의 dimension (learnable prototypes의 수)을 나타냅니다).

system은 Cartesian product를 통해 exponentially scaled quantization space를 구성합니다:

$$\mathcal{C}_l = \bigotimes_{k=0}^{K_l-1} \mathcal{C}_k, \quad \text{dim}(\mathcal{C}_l) = \prod_{k=0}^{K_l-1} C_k$$

각 hierarchy level $l$은 전용 codebook $\mathcal{C}_l$을 유지하며, 여기서 quantization steps의 수 $|\mathcal{C}_l| = K_l$은 configuration file에 미리 정의되어 있습니다. 각 quantization step $k$는 $C_k$개의 고정된 discrete states (learnable prototypes)를 정의합니다. Composite representations는 다음을 통해 multi-level features를 통합합니다:

$$c^{(i)} = \left(c^{(i_0)}_0, c^{(i_1)}_1, \dots, c^{(i_{K_l-1})}_{K_l-1}\right)$$

다음과 같은 index mapping을 가집니다:

$$i = \sum_{k=0}^{K_l-1} i_k \cdot B_k \quad \text{where} \quad B_k = \prod_{m=0}^{k-1} C_m$$

여기서 base $B_k$는 cumulative dimension scaling을 나타냅니다.

Projection and Spatial Reorganization (Input Dimension Adaptation).

$$Z^{res}_l = P^{in}_l (X_l)$$

여기서 $P^{in}_l$은 level-$l$ input projection layer를 나타내며, input $X_l$을 batch size $B$, atomic count $N$, 그리고 quantization steps $K_l$을 가진 residual features $Z^{res}_l \in \mathbb{R}^{B \times N \times K_l}$로 변환합니다. 이 operation은 hierarchical quantization structure와 정렬되도록 linear projection 및 spatial reorganization을 통해 input dimensionality를 조정합니다.

Hierarchical Quantization Core는 다음 두 단계로 나타낼 수 있습니다:

$$Z^{quant}_l = Q_l (Z^{res}_l)$$

$$I_l = E_l(Z^{quant}_l)$$

quantization function $Q_l$은 FSQ와 함께 quantization step마다 작동합니다. index encoding function $E_l$은 quantized projections를 통해 discrete indices를 계산합니다.

$$\hat{Z}_l = P^{out}_l(Z^{quant}_l)$$

output projection layer $P^{out}_l$은 $Z^{quant}_l$을 원래의 feature space로 다시 매핑하여 features를 reconstruct하고, input $X_l$과 일치하는 dimensionality를 가진 $\hat{Z}_l$을 생성합니다. 이 reconstruction은 learnable linear transformations를 통해 feature fidelity를 최적화합니다.

Up-Sampling. feature reconstruction 동안, 우리는 spatial resolution 강화를 위해 직접적인 block-wise repetition strategy를 사용합니다. quantized downsampled feature matrix $\hat{Z}_l \in \mathbb{R}^{s_k \times d}$ ($s_k$: level $k$에서의 sampling points, $d$: feature dimension)가 주어지면, upsampling ratio $r$은 atomic count $N$에 의해 결정됩니다. upsampling operation은 다음과 같이 정의됩니다:

$$U(\hat{Z}_l) = \text{repeat}(\hat{Z}_l, r, \text{dim} = 0), \quad r = \left\lceil \frac{N}{s_k} \right\rceil$$

여기서 $U$는 upsampling operator를 나타내고 $\text{repeat}(\cdot)$은 spatial dimension (dimension 0)을 따라 block-wise repetition을 수행합니다.

Feature Fusion Mechanism. 우리는 간결한 fusion operation을 통해 5개의 모든 hierarchy levels에서 features를 통합합니다:

$$F_{fused} = P^{5d \rightarrow d}_{fused} \left( \text{Concat}_{i=1}^5 U(\hat{Z}_i) \right)$$

여기서 $F_{fused} \in \mathbb{R}^{N \times 5d}$는 fused features를 나타내며, $P^{5d \rightarrow d}_{fused}$는 decoder input을 위해 원래의 feature dimension을 복원하는 linear projection입니다. 이 설계는 hierarchical information을 보존하는 동시에 learned projection parameters를 통한 adaptive weighting을 가능하게 합니다.

Loss Structural Alignment: Ground truth $X$와 reconstructed $\tilde{X}$ point clouds는 Umeyama-Kabsch algorithm을 사용하여 정렬됩니다. Loss functions는 다음과 같습니다:

$$L_{RMSD}(X, \tilde{X}) = \sqrt{\frac{1}{n} \sum_{i=1}^{n} \|x_i - \tilde{x}_i\|^2}$$

$$L_{atom-dist} = \sum_{r} \sum_{i,j \in R_r} \left( \|x_i - x_j\| - \|\tilde{x}_i - \tilde{x}_j\| \right)^2$$

우리는 RMSD Loss와 residue $r$ 내의 각 atom 사이의 ground truth 및 reconstructed pairwise distances 간의 차이를 계산하는 Inter-atomic Distance Loss를 사용합니다. 여기서 $R_r$은 residue $r$ 내의 atomic index set입니다 (저분자의 경우 분자 전체에 대해 계산됨). total loss는 동일한 가중치로 이러한 components를 결합합니다:

$$L_{total} = L_{RMSD} + L_{atom-dist}$$

Method 핵심 정리 노트

1. 기본 Architecture: Bidirectional Mamba Layers (BMLs)

• 구조: Encoder와 Decoder는 point cloud의 양방향 구조적 의존성(bidirectional structural dependencies)을 포착하기 위해 BMLs를 채택함.
• 핵심 메커니즘: 원본 정보를 유지하는 Primary Path와 Flip operation을 활용해 양방향을 처리하는 Processing Path를 결합한 대칭적 잔차 구조(symmetric residual structure).
• 수식: $Y = M(X) + F(M(F(X)))$ (여기서 $F$는 Flip, $M$은 Mamba Block). 이를 통해 계산 효율성을 챙기면서도 forward/reverse 방향의 long-range dependencies를 모두 포착함.

2. 핵심 모듈: Natural Hierarchical Quantization (NHQ)

기존 방식들처럼 모든 원자를 동등하게 취급(homogenize)하지 않고, 생물학적 중요도에 따라 5단계의 resolution levels로 point cloud를 분해(downsampling)하여 개별적으로 임베딩함.

• L1 (Global max-abs feature): 가장 활성화된(activation sum이 높은) 단일 원자를 추출하여 전체 구조의 topological anchor (전역적 기준점)로 활용.
• L2 (Strided $C\alpha/C3'$ sampling): L1과 L3 사이의 부드러운 전환을 위한 중간 단계의 sparse sampling.
• L3 (Full $C\alpha$ set): 잔기(residue) 수준의 conformation 및 2차 구조(secondary structure)를 모델링하는 뼈대.
• L4 (Backbone motifs): 펩타이드 결합의 강성 제약(peptide bond rigidity constraints)을 유지하기 위해 모든 백본 원자(N, $C\alpha$, C, O) 포함.
• L5 (Atomic resolution): 수소 결합 네트워크 등 미세한 상호작용을 모델링하기 위한 full atom 디테일.

3. Hierarchical Quantization (HQ) 시스템

• 전략: 지수적으로 스케일링되는 composite codebooks를 Cartesian product 방식으로 구성.
• 효과: 매우 낮은 search complexity를 유지하면서도, 생체 분자의 복잡한 정보를 잃지 않고 high-fidelity representations로 압축(Quantization) 가능. FSQ(Finite Scalar Quantization) 적용.

4. Feature Fusion 및 Loss function

• Upsampling & Fusion: 압축된 5개 레벨의 feature들을 block-wise repetition 전략으로 원래 길이($N$)로 upsampling한 뒤, 모두 Concat하고 linear projection($P_{fused}^{5d \rightarrow d}$)을 거쳐 최종 융합함.
• Loss: Umeyama-Kabsch algorithm으로 point cloud를 정렬한 후, 전체 구조의 오차를 재는 RMSD Loss와 국소적 원자 간 거리 오차를 재는 Inter-atomic Distance Loss를 1:1 비율로 결합하여 최적화함 ($L_{total} = L_{RMSD} + L_{atom-dist}$).

---

쉬운 설명 :

이 Method 섹션은 AI가 복잡한 단백질 구조를 완벽하게 분해하고 조립하는 3단계 공정을 설명하고 있습니다.

건축에 비유해 보겠습니다. AI가 아주 복잡한 롯데타워 같은 건물을 통째로 스캔해서 기억해야 한다고 상상해 보세요.

1. 눈 (Bidirectional Mamba): 건물을 처음부터 끝까지 한번 훑고, 다시 끝에서 처음으로 거꾸로 훑어봅니다(양방향). 이렇게 하면 앞뒤 문맥을 완벽히 파악하면서도 계산이 아주 빠릅니다.
2. 뇌의 5단계 기억법 (NHQ & 5 Levels): 이것이 이 논문의 진정한 필살기입니다. 건물의 모든 벽돌을 한 번에 외우려다 실패했던 예전 AI들과 달리, 이 모델은 5개의 다른 해상도 렌즈를 낍니다.
   • L1렌즈: 건물의 가장 중심이 되는 '메인 기둥' 하나만 봅니다.
   • L2렌즈: 층별 주요 뼈대만 듬성듬성 봅니다.
   • L3렌즈: 전체 층의 뼈대를 다 봅니다.
   • L4렌즈: 뼈대에 붙은 외벽의 형태까지 봅니다.
   • L5렌즈: 마지막으로 개별 벽돌과 시멘트(모든 원자)까지 세밀하게 봅니다.
3. 압축과 복원 (Quantization & Fusion): 이 5가지 렌즈로 본 정보들을 효율적인 암호장부(Codebook)에 꾹꾹 눌러 담아 압축합니다. 그리고 다시 건물을 지어낼 때(Reconstruction), 이 5단계의 도면을 하나로 쫙 합쳐서(Feature Fusion) 실제 원본 건물과 얼마나 똑같이 지어졌는지, 벽돌 간의 간격은 맞는지(Loss Function) 꼼꼼하게 채점하며 학습하는 방식입니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

3D 원자, 좌표들 x, y,z가 인풋이 됨



task는 토큰화임. 

전체 윤곽 (Global topology): 분자의 가장 큰 특징 하나만 뽑기



듬성듬성 뼈대 (Interval Sampling): 뼈대 중에서도 듬성듬성 포인트만 잡기



핵심 뼈대 (Full Alpha C): 단백질 구조의 핵심이 되는 척추 뼈대 모두 잡기



전체 뼈대 (Backbone): 척추를 이루는 모든 기본 뼈대 원자 잡기



세부 디테일 (All Atom): 곁가지로 붙은 아주 미세한 모든 원자까지 묘사하기



아웃풋, 재구성된 분자 구조.



그니까 약간 VQVAE 같은 느낌이네



인풋 아웃풋 똑같이 맞추고 내부 토큰 디자인







실제 동작은

N개의 원자 좌표가 행렬로 있음 

세로 N줄 가로 3칸 x,y,z



이때 원자 및 서열 정보 완전 배제

이 원자가 C인지 N인지 모름



그대로 인코더로 들어감



양방향으로 처리함  (Bidirectional Mamba Layers)



원자 별로 절대값 합계를 내서 젤 높은놈을 고름 그게 레벨 1



알파탄소, c3쪽 원자만 쏙 고름 그게 레벨2(대신 약간 띄엄띄엄)



레벨2를  띄엄띄엄 없이 그냥쓴게 레벨3



알파 탄소를 포함해 펩타이드 결합을 이루는 핵심 원자들을 뽑음 레벨 4



그냥 레벨 5





그럼 5개의 피쳐가 나올텐데



각 레벨마다 코드북의 숫자는

16 64 256 1024 4096이 됨



일단 각 피쳐를 선형레이어에 넣어 차원을 축소





기본적으로 128차원의 벡터인데

매우 적은 칸수로 줄이고, 그 칸에 들어갈 수 있는 숫자 자체를 제한함. 가질 수 있는 정보의 수가 코드북의 수와 일치 시킴



tanh를 거쳐서 -1~1로 만드네? 그리고 정수 격자 범위로 곱



반올림. 가까운 정수로



그 조합들을 키로 삼아서 하나의 정수 값으로 변환 그게 토큰의 고유 인덱스가 됨.



암튼 고유인덱스 확인햇으면 



다시 128차원으로 각각 맞춰놓고

밑으로 이어 붙임.



그리고 5개 레벨 모두 N개로 복사해서 맞춤.



그럼 N  x 128행렬이 되는데 그걸 그냥 컨캣해서 

디코딩에 활용 후 원래 인풋으로 맞추게끔. 그리고 loss를 줌





이러면 고유인덱스가 정확하게 구분되는 어떤 정보를 담게끔 학습이 된다고 가정을 하는것 같은데  그후 다른 task에 고유인덱스를 쓰는 것 같음.

 

 

 

 

별점 3점 / 5점

도메인 지식은 매우 훌륭함. 큰 뼈대부터 미세한 원자까지 5단계로 쪼갠 아이디어는 직관적이고 훌륭, 하지만 5단계 마스크를 사람이 직접 선택한 것이 약간 아쉬움.