protein : 논문 리뷰 : BIO2TOKEN: ALL-ATOM TOKENIZATION OF ANY BIOMOLECULAR STRUCTURE WITH MAMBA

용어 설명

• Mamba: Transformer의 고질적인 문제인 sequence 길이에 따른 2차 함수적 연산량 증가(O(N^2))를 해결하기 위해 제안된 linear-time sequence modeling 아키텍처(SSM)입니다.
• Quantized Auto-Encoder (QAE): 연속적인 입력 데이터를 압축된 잠재 공간(latent space)으로 인코딩한 뒤, 이를 이산적인(discrete) 코드북 기반의 token으로 양자화(quantization)하고 다시 원래 데이터로 복원(decoding)하는 신경망 구조입니다.
• Finite-Scalar Quantization (FSQ): 기존 Vector Quantization (VQ) 방식이 자주 겪는 codebook collapse(특정 토큰만 과도하게 사용되는 현상) 문제를 해결하기 위해, 입력을 정수 길이의 하이퍼큐브에 투영하여 양자화하는 효율적이고 학습이 안정적인 기법입니다.
• All-atom resolution: 단백질이나 RNA를 아미노산 잔기(residue) 단위나 백본(backbone) 구조로 단순화(coarse-grained)하지 않고, 개별 원자(atom) 단위의 정확한 위치까지 모두 모델링하는 해상도입니다.
• SE(3)-invariance: 3D 구조가 회전하거나 이동해도 모델의 출력이나 특징이 일관되게 유지되는 수학적 특성입니다. 주로 Invariant Point Attention (IPA) 구조에서 쓰이지만, 연산량이 매우 높습니다.
• Umeyama-Kabsch algorithm: 두 개의 3D point cloud 집합 간의 구조적 차이(RMSE)를 최소화하기 위해 최적의 회전 및 이동 변환을 찾아 정렬(alignment)하는 알고리즘입니다.

---

Purpose of the Paper

• 기존 연구의 한계: 기존의 3D biomolecular generation 모델(예: ESM-3, RFDiffusion)은 계산 복잡도 문제로 인해 거대한 분자를 처리할 때 residue 단위의 coarse-grained representation에 의존하거나, backbone과 side-chain을 분리해서 예측해야 했습니다. Transformer나 SE(3)-invariant 아키텍처(IPA 등)는 메모리와 연산량의 한계로 수만 개의 원자를 가진 복합체에 적용하기 불가능했습니다.
• 새로운 문제 정의: 이 논문은 크기에 상관없이 small molecules부터 proteins, RNA, 그리고 이들의 complexes까지 모든 biomolecule을 all-atom 수준에서 하나의 1D discrete sequence로 변환할 수 있는 통합 tokenizer를 개발하고자 했습니다.
• 접근 방식의 차별점: 무겁고 복잡한 Transformer나 SE(3) inductive bias를 완전히 배제하고, 효율적인 long-context modeling이 가능한 Mamba 아키텍처를 도입하여 연산량과 메모리 문제를 해결했습니다.

---

Key Contributions

• 통합 All-Atom Tokenizer (Bio2Token) 개발
  • Novelty: Small molecules, proteins, RNA 각각에 특화된 tokenizer뿐만 아니라, 이들 모두를 아우르는 단일 모델(bio2token)을 구축했습니다. 최대 100,000개에 가까운 원자를 가진 거대 구조도 단일 프레임워크 내에서 처리 가능합니다.
• Mamba 기반 3D Point Cloud Modeling
  • Novelty: 3D 좌표 데이터를 sequence 형태로 변환한 뒤 bidirectional Mamba block을 적용했습니다. 이는 Transformer의 quadratic bottleneck을 극복하여 제한된 GPU 환경에서도 극도로 긴 원자 시퀀스를 처리할 수 있게 한 핵심 혁신입니다.
• SE(3) Inductive Bias의 배제
  • Novelty: 3D 구조 학습 시 필수적이라고 여겨졌던 복잡한 회전 불변성(rotational invariance) 연산을 제거했습니다. FSQ 양자화와 단순화된 아키텍처만으로도 모델 스스로 충분히 공간적 패턴을 학습하여 sub-Angstrom 수준의 정확도를 달성할 수 있음을 증명했습니다.

---

Experimental Highlights

• State-of-the-Art Reconstruction Accuracy: CASP14 단백질 데이터셋에서 all-atom 구조 복원 RMSE 0.58 Angstrom을 달성했습니다. 이는 residue-level에 의존하는 ESM-3의 tokenizer(1.3 Angstrom)나 InstaDeep의 backbone tokenizer(1.89 Angstrom)를 압도적으로 뛰어넘는 수치입니다.
• 거대 분자 확장성 (Scalability): RNA3DB 테스트 셋에 포함된 약 90,000개 원자 규모의 거대 RNA 체인을 0.66 Angstrom의 놀라운 정확도로 복원하는 데 성공했습니다.
• Zero-shot Complex Reconstruction: 학습 과정에서 한 번도 본 적 없는 RNA-protein 복합체(예: PDB 3WBM)나 다중 체인 복합체(multi-chain complexes)를 입력했을 때도 0.77 Angstrom 수준으로 정확하게 구조를 복원해냈습니다.
• 압도적인 연산 효율성: 동일한 GPU 메모리 환경에서 기존 IPA 기반 모델은 최대 batch size 1까지만 학습 가능하고 step당 1.0초가 걸린 반면, Mamba 기반 모델은 batch size 32를 수용하며 step당 0.7초 만에 처리하여 막대한 computational efficiency 향상을 입증했습니다.

---

Limitations and Future Work

• Limitations (화학적 타당성 보장의 한계)
  • RMSE가 매우 낮음에도 불구하고, 복원된 3D 좌표가 완벽한 화학적 결합 구조(chemical validity)를 보장하지는 않습니다. 특히 거대 분자의 외곽 부분에서 원자 간 충돌(steric clashes)이 발생하거나 비정상적인 결합이 형성되는 경우가 관찰되었습니다. 이는 향후 downstream generative modeling 시 물리적으로 불가능한 분자가 생성될 위험을 내포합니다.
• Future Work
  • Post-processing 도입 및 데이터 확장: 물리적 제약 조건을 손실 함수(loss)에 하드코딩하는 대신 더 방대한 데이터를 학습시키거나, AlphaFold 3에서 사용된 것과 같은 물리 기반 휴리스틱 후처리 기법을 도입하여 화학적 유효성을 높이는 연구가 필요합니다.
  • Continuous Latent Space 활용: 양자화(quantization) 단계를 제거하여 압축된 연속적(continuous) 잠재 공간을 만들고, 이를 language model뿐만 아니라 flow matching이나 diffusion 모델과 결합하여 분자 설계의 새로운 파이프라인으로 확장할 계획입니다.

---

Overall Summary
이 논문은 무거운 Transformer와 기하학적 연산(SE(3))을 배제하고 Mamba 아키텍처를 도입하여, 세상의 모든 생체 분자를 원자 단위(all-atom)에서 1D 토큰으로 변환하는 초고효율 tokenizer인 Bio2Token을 제안했습니다. 연산량 한계로 인해 분자 구조를 단순화해야 했던 기존 연구들의 한계를 부수고, 최대 10만 개의 원자를 가진 거대 복합체까지 sub-Angstrom의 정확도로 복원해냈습니다. 이는 향후 Large Language Models를 활용한 고해상도 신약 개발, 단백질 및 RNA 설계 등 다양한 AI 기반 분자 생물학 연구에 범용적으로 활용될 수 있는 강력한 인프라를 제공한다는 점에서 매우 큰 의의를 가집니다.

---

쉬운 설명
거대한 도시(단백질)의 지도를 그릴 때, 기존 AI 모델들은 계산이 너무 복잡해서 모든 건물(원자) 사이의 거리와 각도를 일일이 계산하다가 컴퓨터가 뻗어버렸습니다. 그래서 어쩔 수 없이 큼지막한 '동네 단위(residue)'로 뭉뚱그려 지도를 그렸죠. 하지만 이 논문의 Bio2Token은 복잡한 각도 계산(SE(3))을 과감히 버리고, 아주 빠르고 효율적인 스캐너(Mamba)를 사용합니다. 이 스캐너는 도시 전체를 순식간에 훑고 지나가면서 건물 하나하나의 위치를 아주 짧은 '우편번호(토큰)'로 만들어냅니다. 나중에 이 우편번호들만 쭉 읽어도 건물 하나하나의 위치가 1mm 오차도 없이 완벽하게 원래대로 복원되는 놀라운 마술을 보여준 것입니다.

 

 

 

 

 

 

 

준비물: N개의 중원자 좌표로 이루어진 3D 행렬을 준비하고, 모든 점의 중심을 영점 맞춤

다양한 각도에 모델을 적응시켜야하니까 회전 졸라시킴, 데이터 증강

그 뒤 N개의 3D 좌표 시퀀스가 양방향 맘바 인코더로 들어감 그걸 4번

인코더 나오고 형태 유지해서 잠재 행렬이 만들어짐.

근데 바로는 못쓰고 양자화 해야함. 
FSQ 양자화기에 들어가게 되는데 양자화됨

그럼 특정 정수가 담기게 되는데 약간 해시 테이블 마냥 4096크기의 경우의수가 되게끔 해서 4096으로 1대1 매칭 시켜버림. 어떤 id 개념

그리고 이산화된 토큰시퀀스도ㅓ 똑같이 양방향 맘바에 들어가서 6번 반복
그럼 동일한 구조 형태를 가진 3d 좌표 출력
Umeyama-Kabsch 알고리즘 이용해서 좌표 살짝 보정

그리고 loss 계산 후 역전파



별점 3점 / 5점

all-atom biomolecular tokenization을 꽤 깔끔하게 밀어붙인 좋은 3점 논문. 다만 downstream 생성 효용과 화학적 타당성, atom identity 통합이 비어 있어 3.5까지는 어렵다.