AI바라기의 인공지능

Protein : 논문리뷰 : Prot2Token: A Unified Framework for Protein Modeling via Next-Token Prediction 본문

논문리뷰

Protein : 논문리뷰 : Prot2Token: A Unified Framework for Protein Modeling via Next-Token Prediction

AI바라기 2025. 9. 23. 13:27

쉬운 설명 (Easy Explanation)

이 논문의 핵심 아이디어는 단백질에 관한 모든 종류의 질문(예: "이 단백질의 기능은?", "안정성은 얼마?", "3D 구조는?")을 하나의 공통된 '언어'로 번역하는 것과 같습니다. 기존에는 각 질문 유형마다 별도의 '전문 번역기'(specialized model)가 필요했지만, Prot2Token은 일종의 **'만능 번역기'**를 만듭니다.

이 만능 번역기는 먼저 task token이라는 특별한 코드를 통해 "이제부터 '3D 구조'에 대해 번역해줘"라고 지시를 받습니다. 그 후, 어떤 질문이든 "다음 단어는 무엇일까?"를 맞추는 방식으로 정답(토큰 시퀀스)을 순서대로 내놓습니다. 가장 중요한 혁신은 모든 질문의 정답 형태를 이 '단어 맞추기' 게임에 맞도록 표준화하는 규칙(universal tokenization)을 발명한 것입니다.

Prot2Token: A Unified Framework for Protein Modeling via Next-Token Prediction 학습 노트

용어 설명 (Terminology)

  • PLMs (Protein Language Models): 단백질 서열을 자연어처럼 취급하여 패턴과 의미를 학습하는 모델. (예: ESM2, ProtBert)
  • Prot2Token: 본 논문에서 제안하는 통합 프레임워크 이름. 다양한 단백질 예측 task를 '다음 토큰 예측'이라는 단일 문제로 변환.
  • Unified Tokenization: 단백질 예측의 다양한 출력 형태(분류, 회귀, 결합 부위 등)를 표준화된 이산 토큰 시퀀스로 변환하는 핵심 아이디어.
  • Task Token: Decoder에게 어떤 task를 수행해야 하는지 알려주는 특별한 prompt 토큰. Multi-task learning을 가능하게 하는 핵심 요소.
  • Autoregressive Decoder: 이전 토큰들을 기반으로 다음 토큰을 순차적으로 생성하는 GPT 스타일의 모델 구조.
  • Self-supervised Decoder Pre-training: Decoder가 단백질 서열 내 특정 아미노산의 위치를 예측하도록 학습시키는 과정. 별도의 label 없이 데이터 자체로 학습하며, 이를 통해 Decoder에 공간적, 위치적 정보에 대한 inductive bias를 주입.

Purpose of the Paper

  • 기존 PLM들은 특정 task를 위해 별도의 예측 모듈(head)을 추가하고 fine-tuning 해야 하는 'specialized model'의 한계를 가짐. 이는 비효율적이고 확장성이 낮음.
  • 기존의 Autoregressive 단백질 생성 모델(ProGen 등)은 예측 task에 대한 controllability(제어 가능성)가 부족하여 특정 예측을 유도하기 어려웠음.
  • 이 연구의 목적은, 이 두 가지 한계를 극복하기 위해 다양한 종류의 단백질 예측 task(분류, 회귀, 서열-대-서열, 상호작용 등)를 'next-token prediction'이라는 단일하고 통일된 패러다임으로 통합하는 것. 이를 통해 단일 모델로 여러 task를 효율적으로 수행하는 범용 예측 프레임워크를 제안하고자 함.

Key Contributions and Novelty

  • Key Contributions:
    • Unified Framework (Prot2Token): Pre-trained encoder(ESM2)와 autoregressive decoder를 결합하여, 다양한 단백질 예측 task를 통합된 next-token prediction 문제로 해결하는 새로운 architecture를 제안.
    • Multi-task Learning via Task Tokens: 각기 다른 task를 구별하는 task token을 도입하여, 단일 모델이 여러 task를 동시에 학습하고 수행할 수 있도록 함. 관련 task를 함께 학습할 때 시너지를 통해 성능이 향상됨을 보임.
    • Self-supervised Decoder Pre-training: 위치 정보가 중요한 task(예: binding site prediction)의 성능을 높이기 위해, Decoder를 self-supervised 방식으로 pre-train하는 방법을 제안. 이는 Decoder에 필요한 inductive bias를 효과적으로 주입.
    • State-of-the-art Performance & Efficiency: 여러 benchmark에서 기존 specialized model들과 동등하거나 더 높은 성능을 달성했으며, 특히 단백질 3D 구조 예측에서는 AlphaFold2 대비 약 1000배 빠른 추론 속도를 달성.
  • Novelty:
    • Universal Tokenization Protocol: 이 논문의 가장 핵심적인 독창성은 단백질 예측의 이질적인 출력값(classification label, regression 수치, 아미노산 인덱스 리스트 등)을 표준화된 토큰 시퀀스로 변환하는 'universal tokenization' 아이디어를 고안하고 구현한 것. 이것이 통합 프레임워크를 가능하게 한 근본적인 기술.
    • Task Token의 생물학적 의미 학습: Task Token의 embedding을 분석하여, 이들이 단순히 task를 구별하는 ID 역할을 넘어, ligand들의 실제 생화학적 특성에 따라 군집화되는 것을 보여줌. 이는 모델이 task 간의 의미론적 관계를 학습했음을 시사.

Experimental Highlights

  • 3D Structure Prediction Speed-up:
    • CAMEO benchmark에서 384개 잔기 단백질 예측 시, Prot2Token은 1-2초가 걸린 반면, AlphaFold2(w/ MSA)는 18-25분이 소요되어 약 1000배의 속도 향상을 보임 (Table 8).
    • 이는 대규모 단백질 구조 생성이나 실시간 예측 시나리오에서 Prot2Token의 압도적인 유용성을 증명.
  • Multi-task Learning Synergy:
    • ER(Enzyme Reaction) 예측 task에서 단독 학습 대비 다른 task(localization, affinity, stability)와 함께 학습했을 때 성능이 7.5% 향상 (Table 2).
    • Fluorescence 예측 task에서도 multi-task learning을 통해 성능이 5.6%까지 향상됨 (Table 4). 이는 모델이 여러 task를 동시에 배우며 일반화 성능이 강화됨을 시사.
  • Task Token Embedding Analysis:
    • 41개 ligand에 대한 task token embedding을 PCA와 k-means로 clustering한 결과, 실제 ligand들의 생화학적 특성(molecular weight, net charge 등) 기반의 clustering과 높은 상관관계(ARI=0.447)를 보임 (Appendix A.5).
    • 이는 모델이 task token을 통해 ligand 간의 생물학적, 화학적 유사성을 성공적으로 학습했음을 보여주는 매우 중요한 실험 결과.

Limitations and Future Work

  • Limitations:
    • Spatial Understanding: Autoregressive decoder의 본질적 한계로 인해, 초기에는 binding site prediction처럼 정밀한 공간적 이해가 필요한 task에서 성능이 낮았음. 이는 self-supervised pre-training을 통해 완화되었지만, architecture 자체의 근본적인 한계점을 시사.
    • Computational Cost for Universal Model: 논문은 task의 일부 조합에 대한 실험을 수행했으며, 모든 종류의 예측 task를 포함하는 진정한 'universal model'을 학습시키는 것은 막대한 계산 자원을 요구함.
    • Pre-trained Encoder 의존성: Prot2Token의 성능은 입력 embedding을 생성하는 pre-trained protein encoder(본 논문에서는 ESM2)의 성능에 본질적으로 의존함. Encoder의 한계가 그대로 모델의 한계로 이어질 수 있음.
    • Inconsistent Output Length: Sequence-to-sequence task에서 decoder가 입력과 길이가 다른 서열을 생성하는 경우가 있어, 추론 시 end-of-sequence 토큰에 제약을 가하는 후처리가 필요했음.
  • Future Work:
    • Protein Design으로의 확장: 예측과 생성을 하나의 통합 모델 내에서 완벽하게 결합하여, 목표 특성 예측부터 새로운 단백질 후보 생성까지 신약 개발의 여러 단계를 하나의 파이프라인으로 수행하는 연구.
    • Comprehensive Multi-task Learning: 훨씬 더 넓은 범위의 생물학적 task를 포함하는 대규모 multi-task 학습을 통해 모델의 일반화 능력과 시너지 효과를 극대화.
    • Advanced Decoding Strategies: 현재의 greedy decoding을 넘어, top-k sampling, nucleus sampling 등 더 정교한 확률적 샘플링 방법을 탐구하여 예측의 다양성과 제어 가능성을 높이는 연구.

Overall Summary

Prot2Token은 다양한 단백질 예측 task들을 'next-token prediction'이라는 단일 패러다임으로 통합하는 범용 encoder-decoder 프레임워크를 제안한다. 핵심 혁신인 'universal tokenization'과 task token을 통해, 단일 모델이 효율적으로 multi-task learning을 수행하며, 특히 3D 구조 예측에서 기존 대비 1000배 빠른 속도를 달성했다. 이 연구는 단백질 모델링 분야의 파편화된 접근 방식을 통합하고, 고효율, 고성능 예측을 통해 생물학적 발견과 신약 개발을 가속화할 수 있는 중요한 발판을 마련했다.

 

 

Abstract

다양한 성격의 protein prediction task는 전통적으로 특화된 models을 필요로 하여, 광범위하게 적용 가능하고 계산적으로 효율적인 Protein Language Models (PLMs)의 개발을 저해해 왔습니다. 본 연구에서는 sequence-level 속성과 잔기별 속성부터 복잡한 단백질 간 상호작용에 이르기까지 광범위한 단백질 관련 predictions을 표준화된 next-token prediction 형식으로 변환함으로써 이러한 문제들을 극복하는 통합 framework인 Prot2Token을 소개합니다.

Prot2Token의 핵심은 pre-trained된 protein encoders로부터 얻은 embeddings을 조건으로 하고, 학습 가능한 "task tokens"에 의해 유도되는 autoregressive decoder를 사용하여 다양한 predictions을 수행하는 것입니다. 이러한 architecture는 독특하게 multi-task learning을 촉진하여, 단일 model이 향상된 효율성으로 수많은 task를 마스터할 수 있도록 합니다.

우리는 다양한 benchmarks에 걸쳐 광범위한 실험적 검증을 제시하며, 다양한 유형의 protein-prediction tasks에서 Prot2Token의 강력한 predictive power를 입증합니다. 주요 결과로는 상당한 speedups (예: MSA를 사용하는 AlphaFold2 대비 1000배)와 종종 전문적인 접근 방식과 비슷하거나 이를 능가하는 performance가 포함됩니다. 그뿐만 아니라, 우리는 공간적으로 민감한 task performance를 개선하기 위해 보조적인 self-supervised decoder pre-training 접근법을 도입합니다.

따라서 Prot2Token은 protein modeling을 위한 다재다능하고 high-throughput paradigm을 향한 중요한 진전을 제시하며, 생물학적 발견과 새로운 치료법 개발을 가속화할 것을 약속합니다. code는 GitHub에서 확인할 수 있습니다.

 

 

 

1 Introduction

더보기

단백질은 생명의 근본적인 구성 요소로서, 인간의 건강을 유지하는 데 중요한 역할을 합니다. 그러나 그들의 sequences와 structures에 암호화된 복잡한 단백질의 language를 이해하는 것은 연구자들에게 여전히 중요한 과제로 남아 있습니다. 이러한 복잡성은 다양한 생물의학 및 치료 응용 분야를 위해 단백질을 해석, 예측 및 설계하는 우리의 능력을 제한합니다.

Protein function prediction은 방대한 단백질 sequences의 다양성, 구조적 변이, 그리고 labeled data의 제한된 가용성 때문에 특히 어렵습니다. 자연어와 달리, 단백질 sequences는 인간이 이해할 수 있는 명시적인 구문 규칙을 따르지 않기 때문에, models이 광범위한 생물학적 지식 없이는 의미 있는 representations을 학습하기 어렵습니다. Protein language models (PLMs)은 단백질 sequences의 의미 있는 representations을 학습함으로써 혁신적인 해결책을 제공하며, 연구자들이 단백질 데이터를 더 해석 가능한 형식으로 해독하고 번역할 수 있게 합니다. PLMs를 활용함으로써 우리는 원시 단백질 정보와 인간의 이해 사이의 격차를 해소하고, 신약 개발, 질병 메커니즘, 그리고 합성 생물학 연구를 발전시킬 수 있습니다.

PLMs가 protein-prediction tasks를 크게 발전시켰지만, 현재의 models은 pre-training 이후 task-specific 전문화가 필요합니다. 별개의 tasks에 대해 별도의 modules에 의존하는 것은 비효율적인 계산 자원 사용과 제한된 확장성으로 이어집니다. 대부분의 PLMs는 개별 tasks에 대해 특화된 predictor architectures와 post-training 정렬을 거치며, 이는 독립적인 training과 fine-tuning을 요구하는 시간 소모적이고 자원 집약적인 접근 방식입니다. 다양한 protein-prediction tasks를 효율적으로 처리할 수 있는 unified model은 이러한 한계를 극복하고, protein function prediction을 간소화하며 실제 응용 분야에서의 접근성을 향상시킬 것입니다.

우리가 아는 한, 단백질을 위한 foundation models의 등장에도 불구하고, 광범위한 protein prediction tasks에 걸쳐 이들을 체계적으로 정렬할 포괄적인 framework는 존재하지 않습니다. 대신, 연구자들은 종종 특정 응용에 맞게 기존 foundation models을 수정하는데, 예를 들어 맞춤형 기술을 사용하여 sequences로부터 3D 단백질 구조를 예측하는 것과 같습니다. 한 가지 주요 한계점은 대부분의 기존 models이 BERT-style architectures에 기반하고 있다는 것인데, 이는 의미 있는 representation을 제공하는 데 효과적이지만 다양하고 제어 가능한 prediction 능력에 필요한 유연성이 부족합니다. Natural language processing (NLP)에서는 BERT-style models에서 autoregressive GPT-style models로의 전환이 generation process를 제어하기 위한 더 동적이고 인간이 이해할 수 있는 지시(prompts)를 가능하게 했으며, 따라서 NLP domain 내에서 다양한 predictions을 처리할 수 있게 되었습니다. 단백질 연구에서도 유사한 paradigm 전환이 필요하며, 정적인 encoders를 넘어 더 포괄적인 predictive 능력을 제공하는 더 진보된 generative AI 접근 방식으로 나아가야 합니다.

비록 ProGen2, RITA, Ankh와 같은 autoregressive transformer models이 단백질의 language에 대해 탐구되었지만, 특히 protein-prediction tasks에 대해서는 제어 가능성과 task에서 어려움을 겪고 있습니다. 제어 가능하고 해석 가능한 predictions을 위해 prompting mechanisms을 효과적으로 활용하는 NLP의 language models와 달리, autoregressive PLMs는 현재 그들의 출력을 인간이 해석할 수 있는 형식으로 유도할 강력한 방법이 부족합니다. 이러한 격차는 그들의 실용적인 적용 가능성을 저해하며, NLP와는 대조적으로 연구자들이 encoder-style PLMs에 계속해서 크게 의존하도록 만들었고, 종종 특정 protein prediction tasks를 위해 이러한 encoders 주변에 특화된 architectures를 구축하게 했습니다.

이러한 한계점들을 해결하기 위해, 본 연구는 단일의 포괄적인 framework 내에서 다양한 단백질 관련 prediction tasks를 통합하는 중요한 한 걸음을 내딛습니다 (그림 1). 우리는 다양한 protein-prediction tasks를 tokenizing하기 위한 보편적인 프로토콜을 도입하여, 일반적인 autoregressive transformer predictor가 기존의 BERT-style PLMs를 활용할 수 있도록 합니다. next-token prediction loss에 의해 유도되는 이 통합된 autoregressive predictor는 protein-level, residue-level, 그리고 protein-protein interaction-level tasks를 포함한 여러 protein-prediction task 범주에 걸쳐 강력한 generality를 보여줍니다. 우리는 키나아제 인산화 부위 예측, 단백질-리간드 결합 부위 예측, protein 3D structure prediction, 그리고 단백질 돌연변이 안정성 평가와 같은 다양한 예시를 통해 그 다재다능함을 설명합니다. 더욱이, 우리의 framework는 본질적으로 multi-task learning을 지원하며, 관련 tasks가 공동으로 trained될 때 시너지 효과적인 performance 향상을 보여주는 초기 분석을 제공합니다. 결합 부위 예측과 같은 특정 전문화된 tasks의 경우, self-supervised pre-training을 통해 decoder를 초기화하는 것이 performance를 크게 향상시키는 것을 보여줍니다. 특히, protein-ligand binding site prediction에 대해, 우리는 학습된 token representations을 추가로 분석하여, 적게 표현된 리간드에 대한 predictions을 향상시킬 수 있었던 리간드 tokens 간의 의미 있는 관계를 밝혔습니다. 우리는 우리의 접근 방식이 견고하고 유연한 protein prediction tasks를 위해 large language models (LLMs)를 활용하고 업그레이드하는 데 필수적인 단계라고 믿습니다.


1.1 Related work

현재 단백질을 위한 많은 전문화된 또는 foundation models이 존재하지만, generation과 다양한 prediction tasks를 모두 수행할 수 있는 단일의, prompt-controllable interface를 제공하는 것은 없습니다. 따라서 우리는 이전 연구들을 generative protein design, predictive representation learning, 그리고 unified models로 분류합니다.

Generative protein design. Autoregressive language models이 de novo sequence generation을 지배합니다. ProGen은 기능적 태그를 사용하여 제어 가능한 generation을 처음으로 시연했습니다. 이후의 확장—ProtGPT2 1.2b, RITA 1.2b, 그리고 ProGen2 6.4b—은 perplexity와 실험적 성공을 개선했지만, 기능을 조종하기 위해 여전히 task-specific fine-tuning이나 필터링이 필요합니다. 가장 최근에, ProGen3는 이를 상당히 확장함으로써 이러한 추세를 이어가지만, fine-grained generation에 대한 제한된 제어 가능성을 보고합니다.

Predictive representation learning. 병행되는 연구 갈래는 task-specific heads를 구동하는 bidirectional encoders에 초점을 맞춥니다. ESM2-15b와 같은 대규모 masked-language models은 다양한 downstream tasks를 위한 embeddings을 산출하고 심지어 ESMFold로 end-to-end folding을 구동하기도 하지만, folding module은 3-D structure prediction에 특화되어 있습니다. 마찬가지로, AlphaFold2 (AF2)는 EvoFormer encoders를 맞춤형 structure decoder에 결합합니다. 이러한 “wrapper” architectures는 그들의 전용 출력에서 뛰어나지만, 일반적인 predictor를 형성하지는 않습니다. 우리는 단 하나의 cross-task autoregressive 대안인 PTMGPT2만을 발견했는데, 이는 GPT-2를 prompt-based fine-tuning으로 적용하여 단일 model에서 19개 종류의 post-translational modifications (PTMs)를 예측하지만, 여전히 PTMs domain에 국한되어 있습니다.

Unified models. 최근에는 protein design과 prediction을 단일 시스템 내에서 연결하려는 models이 등장했습니다. HelixProtX는 sequence, structure, 그리고 자유 텍스트를 하나의 multimodal autoregressive transformer로 통합하여, 이들 modalities 중 어느 두 가지 사이에서든 번역이 가능하고 sequence로부터 직접 atom-level 3-D structure를 예측할 수 있습니다. ProLLaMA는 단백질 특화 instruction tuning을 통해 LLaMA-2를 적용하여, 하나의 model이 natural-language prompts에 의해 유도되어 안정성, 형광, 결합 친화도, 그리고 원격 상동성 분류와 같은 property-prediction tasks와 함께 제어 가능한 sequence generation을 수행할 수 있도록 합니다. InstructProtein은 knowledge-graph–guided instruction tuning을 통해 protein sequences를 human language와 정렬하여, model이 단백질의 기능을 자유 텍스트로 설명하거나 텍스트 명세를 만족하는 그럴듯한 sequence를 generate할 수 있게 합니다. 비록 이러한 시스템들이 고무적인 modality transfer를 보여주지만, fine-grained control을 위해 여전히 prompt engineering에 의존하며, 본 연구에서 다루는 모든 표준 prediction tasks에 걸쳐 아직 benchmarked되지 않았습니다.

 

1 Introduction - 정리노트 (for AI Researchers)

1. Problem Definition

  • 현존 PLMs의 한계: 현재 Protein Language Models (PLMs)은 task-specific 전문화가 필수적입니다. 각기 다른 protein-prediction task를 위해 별도의 module이나 독립적인 fine-tuning이 요구되어 계산 자원이 비효율적으로 사용되고 확장성이 제한됩니다.
  • Architecture의 한계: 대부분의 PLM은 BERT-style (encoder) architecture에 기반합니다. 이는 representation 생성에는 뛰어나지만, NLP 분야의 autoregressive GPT-style models이 보여주는 동적이고 제어 가능한 prediction 능력은 부족합니다. 기존의 autoregressive PLM들 또한 NLP의 prompting 메커니즘과 같은 강력한 제어 기능이 없어 다양한 prediction task에 적용하기 어렵습니다.

2. Proposed Solution: Unified Autoregressive Framework

  • 핵심 아이디어: 다양한 protein-prediction tasks(protein-level, residue-level, protein-protein interaction)를 **단일 framework**로 통합합니다. 이를 위해 서로 다른 task들을 표준화된 next-token prediction 문제로 변환하는 범용 tokenizing 프로토콜을 제안합니다.
  • Model Architecture: 기존 BERT-style PLM의 embeddings를 입력으로 사용하는 **범용 autoregressive transformer predictor**를 도입합니다. 이 predictor는 next-token prediction loss를 통해 학습되어 다양한 종류의 예측을 수행합니다.

3. Key Contributions & Capabilities

  • 범용성 (Generality): 단일 model이 여러 카테고리의 protein-prediction task 전반에 걸쳐 강력한 성능을 보임을 입증합니다.
  • Multi-task Learning: Framework가 본질적으로 multi-task learning을 지원하며, 관련 task들을 공동으로 학습시킬 때 시너지 효과로 performance가 향상됨을 보여줍니다.
  • 성능 향상 기법: 결합 부위 예측과 같은 특정 task의 경우, self-supervised pre-training을 통해 decoder를 초기화하면 performance가 크게 향상됨을 확인했습니다.
  • 의의: 이 연구는 단백질 연구 분야를 정적인 encoders에서 벗어나, NLP의 LLM처럼 광범위한 prediction task를 유연하고 견고하게 처리할 수 있는 진보된 generative AI 접근 방식으로의 전환을 제시합니다.

쉬운 설명 :

단백질을 연구하는 것은 마치 수많은 전문 도구가 필요한 작업실과 같습니다. 단백질의 3D 구조를 예측하는 model(망치), 단백질의 안정성을 예측하는 model(드라이버), 특정 분자와의 결합을 예측하는 model(렌치)이 모두 따로 필요했습니다. 과학자들은 새로운 문제를 풀 때마다 새로운 전문 도구를 만들거나 기존 도구를 힘들게 개조해야 했죠. 이는 매우 비효율적이고 시간이 많이 걸리는 과정입니다.

이 논문은 **"모든 작업을 할 수 있는 만능 공구 하나를 만들자"**는 아이디어를 제시합니다.

이들이 제안한 Prot2Token이라는 framework는 일종의 AI 만능 공구입니다. 연구자가 "이 단백질의 3D 구조를 알려줘" 또는 "이 단백질이 안정적인지 알려줘" 라고 질문하면, 이 framework는 어떤 질문이든 마치 우리가 GPT에 질문하듯 표준화된 "언어"(tokenizing)로 변환합니다.

그 후, 똑똑한 autoregressive model이 이 변환된 질문을 보고, 마치 문장에서 다음 단어를 예측하듯이 문제의 답에 해당하는 "다음 토큰"을 예측하여 최종 결과를 생성합니다.

결론적으로, 여러 개의 전문화된 model을 따로따로 훈련하고 관리할 필요 없이, **하나의 통합된 model**을 통해 다양한 단백질 관련 문제를 훨씬 빠르고 효율적으로 풀 수 있게 되는 것입니다. 이는 신약 개발이나 질병 연구 같은 분야의 발전을 크게 앞당길 수 있는 중요한 진전입니다.

 

 

 

 

2 Method

더보기

 

2.1 Prot2Token architecture

Prot2Token framework는 encoder-decoder transformers에 기반한 공유 architecture를 사용하여 다양한 단백질 관련 prediction tasks를 통합하도록 설계되었습니다. 핵심 idea는 autoregressive decoder language model을 기존의 encoder-style 단백질 및 선택적 화학 language models과 cross-attention layers를 통해 통합하여, prediction tasks를 통일된 next-token prediction 문제로 변환하는 것입니다.

이 architecture는 pre-trained된 bidirectional transformer (ESM2)를 protein encoder로 사용합니다. 화학 정보(예: 리간드 결합)를 포함하는 tasks의 경우, 선택적 chemical encoder (BARTSmile)를 사용하여 SMILES representations을 처리합니다. 이러한 encoders는 각각의 입력 sequences를 문맥적 embeddings으로 변환합니다:

여기서 $h_{enc} \in \mathbb{R}^{N \times d_{enc}}$는 encoder 출력이고, 은 sequence 길이이며, $d_{enc}$는 encoder의 hidden dimension입니다.

우리는 각 encoder(단백질 및 해당되는 경우 화학)와 decoder에 대해 고유한 embedding tables를 사용하여 architecture에서 그들의 서로 다른 tokenization 방식과 기능적 역할을 반영합니다.

sequence embeddings의 위치 인식을 향상시키기 위해, 우리는 학습 가능한 positional embedding layer $g_{pos}(\cdot)$를 도입하여 증강된 representations을 생성합니다:

여기서 $p \in \mathbb{R}^{N \times d_{enc}}$는 학습 가능한 positional embedding입니다.

encoder 출력을 decoder의 hidden dimension $d_{dec}$와 맞추기 위해, 우리는 linear projection을 적용합니다:

여기서

이 projected representation $h_{proj} \in \mathbb{R}^{N \times d_{dec}}$는 cross-attention을 통해 decoder에 입력됩니다.

decoder는 multi-head self-attention, feed-forward layers, 그리고 GeLU activations과 같은 표준 transformer 구성 요소로 이루어진 causal (autoregressive) transformer입니다. training 속도와 메모리 효율성을 개선하기 위해 FlashAttention-2가 통합되었습니다. 이 연구에서 사용된 특정 architectural configurations에 대해서는 표 10을 참조하십시오.

통합된 training process 내에서 여러 tasks를 지원하기 위해, 우리는 task token을 도입합니다. 각 출력 sequence의 시작 부분에 위치하는 이 tokens는 각 특정 task에 대한 decoder의 행동을 안내하는 prompts 역할을 합니다. task token sequence $t = (T_1, T_2, \ldots, T_m)$는 학습 가능한 embedding function을 통해 embedded됩니다:

decoder는 embedded된 task tokens를 수신하고, 이들과 projected된 encoder 출력 모두에 attention을 적용합니다:

inference 동안, decoder는 autoregressive합니다: 특별한 beginning-of-sequence (<bos>) token과 그 뒤에 오는 task token을 받고, 각 출력 token을 순차적으로 생성합니다.

decoder는 출력 sequence $x = (x_1, x_2, \ldots, x_T)$의 확률을 다음과 같이 인수분해합니다:

training 목표는 negative log-likelihood를 최소화하는 것입니다:

prompt tokens의 역할을 더 잘 관리하기 위해, 우리는 token-specific weights $w_t \in [0, \infty)$를 할당하여 loss에 대한 기여도를 제어합니다. 특히, 우리는 으로 설정하여 prompt (task token)를 loss에서 제외하는 한편, 다른 tokens 는 다르게 가중치를 부여할 수 있도록 합니다:

이 유연한 가중치 부여는 model이 label sequence의 다른 부분에 대한 attention을 tune하는 데 도움이 됩니다. Prot2Token architecture의 개요는 그림 2를, task tokens가 decoder와 상호 작용하는 방식에 대한 자세한 내용은 그림 5를 참조하십시오. Architectural 변형 및 구성 세부 정보는 표 10에 요약되어 있습니다. 다양한 출력을 token sequences로 표현함으로써, 이 설계는 Prot2Token이 단일 decoder 하에서 광범위한 protein prediction tasks를 통합하여 공동 및 독립적인 training 체제를 모두 용이하게 합니다.


2.2 Tokenization

Prot2Token framework는 입력 encoders와 출력 decoder에 대해 별개의 tokenization 전략을 활용합니다. 단백질 아미노산 sequences나 화학 SMILES 문자열과 같은 입력 sequences는 각각의 pre-trained encoders (예: 단백질용 ESM2, 화학물질용 BARTSmiles)의 네이티브 tokenizers에 의해 처리됩니다. 핵심 혁신은 autoregressive decoder에 의해 예측되는 출력 labels에 대한 통합된 tokenization 전략에 있습니다. 이 전략은 광범위한 생물학적 예측 목표를 표준화된 이산 tokens의 sequences로 변환하여 decoder가 일관된 next-token prediction 메커니즘을 통해 다양한 tasks를 처리할 수 있게 하므로 중추적입니다. 모든 tokenized된 출력 sequences는 <bos> token으로 시작하고 <eos> token으로 끝나며, sequence 경계를 명확하게 구분합니다.

그림 3에 묘사된 바와 같이, 이 접근 방식은 이질적인 labels를 균일한 순차 형식으로 변환하여, task-agnostic한 decoding 과정을 용이하게 합니다. 구체적으로, classification tasks의 경우, labels는 고유한 이산 tokens에 매핑되며, multi-label tasks는 일반적으로 이러한 tokens를 (종종 알파벳순으로) 연결합니다. Regression tasks는 문자 구성 요소(예: 부호, 숫자, 소수점)의 세분화된 숫자별 encoding을 통해 연속적인 수치를 나타냅니다. Sequence-to-sequence tasks는 입력 단백질의 각 잔기에 대해 출력 token을 생성하여 직접적인 대응 관계를 유지합니다. Binding site prediction은 상호작용에 참여하는 잔기들의 정렬된 1-기반 인덱스를 tokenizing하는 것을 포함합니다. PTMs와 같은 다른 복잡한 출력 유형도 특정 token sequences로 변환되는데, 예를 들어 잠재적 및 확인된 변형 부위를 특별한 <sep> token으로 분리하여 나열하는 방식입니다. 이 보편적인 tokenization 프로토콜은 Prot2Token이 단일 decoding architecture 내에서 광범위한 protein prediction tasks를 통합하는 능력의 기본입니다. 각 특정 tokenization 방법에 대한 포괄적인 설명은 부록 A.2를 참조하십시오.


2.3 Datasets

이 연구는 PEER, ProteinShake, CATH, AlphaFoldDB, 그리고 ProteinGym과 같은 다른 선별된 소스를 포함한 여러 기존 benchmarks와 저장소에서 가져온 다양한 tasks를 활용합니다. 이러한 datasets는 regression, classification, binding site, 그리고 sequence-to-sequence predictions를 포함한 광범위한 단백질 관련 prediction tasks를 망라합니다. 전처리 단계를 포함한 각 task의 세부 정보는 부록 A.3에 제공됩니다. 이 datasets의 모든 tasks는 섹션 2.2에 설명된 통합 프로토콜에 따라 tokenized됩니다.

 

2 Method - 정리노트 (for AI Researchers)

1. Architecture: Encoder-Decoder Transformer

  • Core Design: Prot2Token은 encoder-decoder transformer architecture를 채택합니다.
    • Encoder: Pre-trained된 단백질 encoder (ESM2)와 선택적으로 화학 encoder (BARTSmile)를 사용하여 입력 sequence로부터 문맥적 embeddings를 추출합니다.
    • Decoder: Causal (autoregressive) transformer model입니다. Encoder의 출력을 cross-attention을 통해 입력받아, 최종 결과물을 token 단위로 순차적으로 생성합니다. FlashAttention-2를 적용하여 학습 효율을 높였습니다.
    • Interface: Encoder와 Decoder 사이의 hidden dimension 차이를 맞추기 위해 linear projection layer를 사용합니다.

2. Key Innovation: Unification Strategy

  • Task Token (Prompting): Multi-task learning을 위해, 출력 sequence의 시작 부분에 학습 가능한 task token을 추가합니다. 이 token은 decoder에게 "어떤 task를 수행해야 하는지" 알려주는 prompt 역할을 하여, 단일 model이 다양한 prediction을 수행하도록 유도합니다.
  • Universal Output Tokenization: 이 방법론의 핵심입니다. 서로 다른 종류의 prediction 목표(예: classification label, regression 수치, sequence, 결합 부위 인덱스)를 모두 표준화된 **이산 token의 sequence**로 변환합니다.
    • 예시:
      • Classification: Label을 고유 token으로 매핑.
      • Regression: -1.25와 같은 연속 값을 -, 1, ., 2, 5 token들의 sequence로 변환.
      • Binding Site: 상호작용하는 아미노산의 인덱스 목록을 token sequence로 변환.
    • 결과: 모든 이질적인 prediction task를 통일된 next-token prediction 문제로 단순화하여, 단일 autoregressive decoder로 모든 문제를 풀 수 있게 합니다.

3. Training Objective

  • Loss Function: 표준적인 autoregressive model의 학습 목표인 Negative Log-Likelihood (Cross-Entropy Loss)를 사용합니다.
  • Loss Weighting: Loss 계산 시 token별로 가중치()를 부여하는 유연한 방식을 도입했습니다. 특히, prompt 역할을 하는 task token의 가중치를 0으로 설정하여, model이 prompt 자체가 아닌 실제 예측 결과에 대해서만 학습하도록 설계했습니다.

4. Datasets

  • Framework의 범용성을 입증하기 위해 PEER, ProteinShake, CATH 등 여러 공신력 있는 benchmarks에서 regression, classification, binding site prediction, sequence-to-sequence 등 다양한 종류의 datasets를 활용했습니다.

쉬운 설명 :

이 논문의 Method 섹션은 자신들이 제안한 Prot2Token이라는 model이 기술적으로 어떻게 작동하는지 설명합니다. 🧠

이 model을 똑똑한 **"단백질 만능 번역가"**라고 생각해보세요.

  1. 1단계: 단백질 정보 읽기 (Encoder)
    • 번역가는 먼저 번역할 단백질의 아미노산 서열을 읽습니다. 이 Encoder 부분은 단백질의 구조나 특징 같은 깊은 의미를 파악해서 풍부한 정보가 담긴 요약본(embeddings)을 만듭니다.
  2. 2단계: 질문의 종류 파악하기 (Task Token)
    • 이제 번역가는 어떤 종류의 "번역"을 해야 할지 알아야 합니다. 이때 "이 단백질이 얼마나 안정적인가요?" 혹은 "어디에 다른 분자가 붙나요?" 와 같은 질문의 종류를 알려주는 특별한 꼬리표(task token)를 받습니다. 이 꼬리표가 번역의 방향을 정해주는 prompt 역할을 합니다.
  3. 3단계: 답변 작성하기 (Decoder & Tokenization)
    • 이것이 이 model의 가장 독창적인 부분입니다. 번역가는 어떤 종류의 질문이든 답변을 "한 글자씩 순서대로 써 내려가는" 방식으로 통일했습니다.
      • 만약 안정성 점수가 "-1.25"라면, "-", "1", ".", "2", "5" 순서로 글자를 예측합니다.
      • 만약 결합 부위가 10번, 52번 아미노산이라면, "10", "52" 순서로 숫자를 예측합니다.
    • 이처럼 모든 복잡한 단백질 문제의 "정답"을 간단한 "다음 글자 맞추기" 게임으로 바꿔버린 것입니다. Decoder는 이 게임의 전문가로서, Encoder가 만든 요약본과 task token 꼬리표를 참고하여 정답 sequence를 한 token씩 정확하게 생성해 나갑니다.

결론적으로, **Prot2Token은 어떤 단백질 관련 질문이든 하나의 표준화된 "다음 글자 예측" 문제로 바꿔서 푸는 매우 효율적이고 똑똑한 architecture**인 셈입니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

주인장 이해

더보기
  1. 아미노산 서열 인코딩
  2. 특수 토큰 지정
  3. 인코더의 컨텍스트와 특수 토큰을 입력 받고
  4. 디코더에서 하나씩 출력
  5. 정답과 비교 후 loss를 줘서 디코더와 인코더 학습