AI바라기의 인공지능
protein : 빠른 논문 리뷰 : LEARNING THE LANGUAGE OF PROTEIN STRUCTURE 본문
용어 설명
- de novo generation: 기존에 존재하지 않는 완전히 새로운 단백질 구조를 처음부터 설계하고 생성하는 작업.
- Vector-Quantized Autoencoder: 연속적인(continuous) 3D 데이터를 제한된 개수의 이산적인(discrete) 코드북(codebook) 토큰으로 변환하여 압축하고, 이를 다시 원래 3D 좌표로 복원하는 신경망 구조.
- FSQ (Finite Scalar Quantization): 전통적인 Vector Quantization에서 발생하는 codebook collapse(코드북 내의 일부 토큰만 사용되는 현상)를 해결하기 위해, 연속적인 벡터 값을 제한된 범위 내의 미리 정해진 스칼라 값으로 직접 반올림(rounding)하여 양자화하는 최신 기법.
- MPNN (Message-Passing Neural Network): 단백질을 그래프 형태(노드는 잔기, 엣지는 거리/각도)로 모델링하여 노드와 엣지 간의 정보를 교환하며 특징을 추출하는 인코더 네트워크.
- FAPE (Frame Align Point Error): AlphaFold2에서 도입된 손실 함수로, 예측된 3D 좌표와 실제 3D 좌표를 글로벌 좌표계가 아닌 각각의 로컬 프레임 기준으로 비교하여 오차를 계산함.
- scTM (Self-consistent TM-score) / scRMSD: 생성된 단백질 구조가 생물리학적으로 실제로 접힐 수 있는 현실적인 구조인지(designability) 평가하는 지표. 생성된 구조를 바탕으로 아미노산 sequence를 예측(ProteinMPNN)하고, 이를 다시 3D로 접어보았을 때(ESMFold) 원래 생성했던 구조와 얼마나 일치하는지 측정.
Purpose of the Paper
- 현재 deep learning 패러다임은 discrete sequence 기반의 attention architecture (예: Transformer)를 중심으로 통합되고 있으며, 이미지나 오디오 등 비표준 modality도 sequence-modeling 문제로 취급하여 큰 성공을 거두고 있음.
- 그러나 단백질 구조 모델링은 데이터가 3D이고 연속적(continuous)이기 때문에, 기존 연구들은 표준 Transformer를 직접 적용하지 못하고 SE(3) equivariance를 보장하는 복잡하고 맞춤화된 기하학적 딥러닝(geometric deep learning)이나 diffusion 모델에 주로 의존함.
- 본 연구는 이러한 구조적 한계를 극복하고자, 3D 단백질 구조를 1D 이산 시퀀스(discrete sequential data)로 변환하는 토크나이저를 개발하여, 복잡한 3D 전용 아키텍처 없이도 표준 NLP의 language models (예: GPT)를 단백질 구조 생성에 직접 적용할 수 있는 기반을 마련하고자 함.
Key Contributions
- 고품질 구조 토크나이저 (Vector-Quantized Autoencoder) 개발: MPNN 기반 인코더와 AlphaFold-2의 Structure Module 디코더를 결합하여, 3D 단백질 구조를 4096개~64000개 규모의 코드북을 가진 이산 토큰으로 압축 및 복원하는 파이프라인 완성.
- FSQ 도입을 통한 양자화 한계 극복 (Novelty): 기존 단백질 토크나이저들이 겪었던 codebook collapse와 straight-through estimator의 기울기 오차 문제를 FSQ 기법을 도입하여 원천적으로 해결. 이를 통해 더 큰 vocabulary를 안정적으로 활용하여 정보 병목 현상(information bottleneck)을 줄임.
- 표준 LLM 아키텍처의 3D 생성 능력 입증 (Novelty): 복잡한 3D 생성 전용 모델 대신, 토큰화된 구조 데이터를 바탕으로 아무런 변형이 없는 standard decoder-only GPT를 next-token prediction task로 학습시켰음에도, 매우 현실적인 de novo 단백질 구조 생성이 가능함을 증명함.
Experimental Highlights
- Autoencoder 재구성 성능 (Reconstruction Error):
- Dataset: PDB 데이터셋 (약 31만 개 구조).
- Metrics: RMSD, TM-score.
- 결과: Downsampling ratio 1, Codebook size 64000 기준 평균 RMSD 1.22 옹스트롬, TM-score 0.96을 달성. 이는 실험을 통해 단백질 구조를 규명할 때 발생하는 오차(실험적 해상도 한계)에 근접하는 매우 높은 정확도임.
- De novo 구조 생성 성능 (Structure Generation):
- Baselines: 단백질 구조 생성 전용 diffusion 모델인 FrameDiff 및 SOTA 모델인 RFDiffusion.
- Generation Metrics: scTM, scRMSD (Designability 평가), Novelty, Diversity.
- 결과: 344M 파라미터의 단순한 GPT 모델이 scTM 76.61%를 기록하며, 단백질 전용 모델인 FrameDiff (75.77%)와 대등한 수준의 designability를 보여줌. (단, 대규모 pre-trained 모델인 RFDiffusion의 97.07%에는 미치지 못함).
Limitations and Future Work
- Limitations:
- 생성된 단백질 구조의 Novelty와 Diversity 지표가 diffusion 기반 SOTA 모델들보다 상대적으로 낮음. 이는 단순한 sampling 전략으로 인해 GPT 모델이 학습 데이터셋의 최빈값(modes)과 구조적으로 유사한 도메인을 생성하려는 편향이 발생했기 때문임.
- RFDiffusion과 같이 막대한 컴퓨팅 자원과 사전 학습을 거친 최첨단 맞춤형 3D generative models의 최고 성능을 즉각적으로 뛰어넘지는 못함.
- Future Work:
- 본 연구는 sequence-based 단백질 3D 모델링의 proof-of-concept 역할을 함. NLP 분야가 겪었던 것처럼 데이터 스케일업과 compute-optimal large language model training 기법을 적용하면, 현재의 diffusion 기반 SOTA 방법론들을 능가할 가능성이 매우 높음.
- 단백질 구조를 sequence space로 가져옴으로써, 텍스트, DNA/RNA sequence 등 타 modality와의 끊김 없는 multi-modal 통합(예: text-conditioned protein design)을 위한 완벽한 토대를 제공함.
Overall Summary
이 논문은 연속적인 3D 단백질 구조를 FSQ 기반의 autoencoder를 활용해 1D 형태의 이산 토큰 시퀀스로 고정밀 변환하는 혁신적인 방법론을 제시합니다. 이렇게 토큰화된 데이터를 통해, 복잡한 3D 특화 아키텍처 없이도 범용적인 GPT 모델만으로 매우 현실적이고 디자인 가능한 단백질 구조를 새로 만들어낼 수 있음을 입증했습니다. 이는 단백질 구조 데이터를 언어 모델과 동일한 공간에서 다룰 수 있게 만들었다는 점에서, 향후 생물학 데이터 전반을 통합하는 거대한 multi-modal foundation model 시대를 여는 중요한 이정표가 될 것입니다.
쉬운 설명
기존에는 새로운 단백질을 설계할 때, 찰흙으로 3D 조각상을 정교하게 빚어내듯 공간의 규칙을 모두 이해하는 "3D 전용 특수 기계(Diffusion 등)"가 필수적이었습니다. 하지만 이 논문은 복잡한 3D 단백질 조각상을 일련번호가 적힌 레고 블록 설명서(discrete sequence)로 완벽하게 번역해내는 '통역기'를 만들었습니다.
그 결과, 단순히 글짓기를 잘하는 일반적인 인공지능(GPT)에게 "이 단어들로 새로운 설명서 좀 써봐"라고 학습시켰더니, 특수 기계 없이도 실제로 조립 가능한 그럴싸한 새로운 3D 단백질을 뚝딱 만들어낼 수 있게 된 것입니다. 즉, "3D 건축" 문제를 우리가 아주 잘 푸는 "언어 완성" 문제로 바꿔버린 획기적인 아이디어입니다.
50~512 길이 시퀀스를 가진 단백질 사용
3d 구조 원자별 좌표 데이터를 활용함.
이제 그래프를 만들어야함.
노드는 각 아미노산이 되는거임. 그럼 무슨값을 쓰냐?
내 포지셔널 인코딩 값이 피쳐가됨. 좌표나 이런건 쓰지 않음.
엣지는 어케되냐? 시퀀스 상이 아니라 3D상에서 가까운 50개 하고 연결때려버림
3D좌표는 순전히 거리 계산용도
그래프로 만들면 인베리언트 하기 떄문
그래서 mpnn 인코더에 넣고 처리해서 1d 임베딩 추출
mpnn은 어떻게 동작하냐면
타겟노드, 이웃노드, 엣지정보 이 3개를 이어서 mlp 태움 그래서 나온 값이 메시지인셈인데
내 정보랑 얻은 정보랑 swish가 들어간 레이어를 통과시켜서 새로 정보 갱신 이걸 3번 반복
다운 샘플링, 1,2,4 분의 1로 길이 줄임.
그리고 FSQ를 해버림., 이게 뭐냐ㅕㄴ 소수점인 연속 피쳐를 이산화시키는거임. 토큰 종류는 4096~64000이라고 함.
이걸 다시 늘려버림. ㄷㄷ
그래서 피쳐끼리 짝지어서 아웃터 프로덕트 해버리고 상대적 위치 정보를 섞은 뒤 mlp 태워서 페어와이즈 피처 생성
그래서알파롤드2 태우고
구조 복원완료
그럼 뭐가 학습 되냐? fsq가 학습되는건가