ProteinZero: Self-Improving Protein Generation via Online Reinforcement Learning

이 논문은 ESMFold 구조 가이드와 자체 개발된 ddG 예측기를 결합한 보상 파이프라인과 임베딩 다양성 정규화를 통해 단일 GPU 노드에서 3 일 내 실행 가능한 온라인 강화학습 프레임워크 'ProteinZero'를 제안하며, 기존 최첨단 모델 대비 설계 실패율을 36-48% 줄이고 90% 이상의 성공률을 달성하여 라벨 없는 데이터로 단백질 생성 모델의 지속적인 자기 개선과 다목적 최적화를 가능하게 합니다.

핵심

🧬 단백질 설계의 새로운 혁신: 'ProteinZero'란 무엇인가요?

이 논문은 인공지능이 스스로 단백질을 더 잘 설계하도록 가르치는 새로운 방법을 소개합니다. 기존 방법들의 한계를 뛰어넘어, AI 가 실험실 없이도 스스로 배우고 발전할 수 있게 한 획기적인 기술입니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 문제: "레시피북"만으로는 부족해요 📚

지금까지 과학자들은 자연계에 존재하는 단백질 (약 200 만 종) 의 데이터만 보고 AI 를 훈련시켰습니다.

• 비유: 마치 오직 기존에 나온 요리책 (레시피북) 만 보고 요리사를 훈련시키는 것과 같습니다.
• 한계: 요리책에 없는 새로운 요리를 만들 수 없고, 책에 실린 레시피가 완벽하지 않아 실패할 확률이 높습니다. 또한, "맛있어야 한다 (안정성)"와 "새로운 맛이어야 한다 (다양성)"를 동시에 만족시키기 어렵습니다.

2. 해결책: "ProteinZero" - 스스로 배우는 요리 학교 🎓

연구팀은 ProteinZero라는 새로운 시스템을 만들었습니다. 이는 AI 모델이 스스로 생성한 단백질들을 평가하고, 그 결과를 바탕으로 실시간으로 스스로를 업그레이드하는 방식입니다.

• 비유: 이제 AI 는 요리책만 보는 게 아니라, 스스로 요리를 만들어보고 ("생성"), 그 요리를 맛보고 ("평가"), 실패하면 다음에 더 잘 만들려고 노력하는 ("학습") 과정을 반복합니다.

3. 핵심 기술 3 가지: 어떻게 작동할까요?

① 빠른 점수판: "ESMFold"와 "가상 에너지 계산기" ⚡

단백질이 제대로 접히는지, 튼튼한지 확인하려면 보통 물리 실험이나 복잡한 시뮬레이션을 해야 하는데, 이는 한 번에 몇 시간씩 걸려 실시간 학습이 불가능했습니다.

• ProteinZero 의 해결책:
  • 디자인 능력 점수 (ESMFold): AI 가 만든 단백질이 원래 모양과 얼마나 닮았는지 초고속으로 확인합니다. (기존보다 25~100 배 빠름)
  • 안정성 점수 (Fast-ddG): 단백질이 녹아내리지 않고 튼튼한지 가상 에너지 계산기로 빠르게 측정합니다.
• 효과: 이제 AI 는 하루에 수천 번의 실험을 할 수 있게 되어, 3 일 만에 기존에 몇 달 걸리던 학습을 끝낼 수 있습니다.

② "모드 붕괴" 방지: "다양성 리마인더" 🎨

AI 가 점수만 따기 위해 똑같은 요리만 반복해서 만들어내는 현상 (모드 붕괴) 을 막아야 합니다.

• 비유: AI 가 "매운맛"만 최고라고 생각해서 매운 요리만 계속 만든다면? 우리는 **다양한 맛 (신맛, 단맛, 짠맛)**을 경험해야 합니다.
• ProteinZero 의 해결책: AI 가 만들어낸 단백질들의 **내부 특징 (임베딩)**을 비교합니다. 너무 비슷하면 점수를 깎아주고, 서로 다른 특징을 가진 단백질을 만들도록 유도합니다.
• 결과: 실패 확률을 36~48% 줄이면서도, 다양한 새로운 단백질을 계속 찾아낼 수 있게 됩니다.

③ 온라인 강화 학습: "실전 연습" 🏆

기존 방법은 미리 준비된 정답지 (데이터) 만으로 학습했습니다. 하지만 ProteinZero 는 AI 가 스스로 만든 문제와 정답을 실시간으로 만들어가며 학습합니다.

• 비유: 시험 문제집만 푸는 게 아니라, 스스로 문제를 내고 풀고, 틀린 부분을 바로 고치는 '스스로 교정' 시스템입니다.

---

4. 어떤 성과를 냈나요? 🏆

이 시스템을 통해 AI 는 놀라운 성과를 거두었습니다.

• 성공률 90% 이상: 기존 최고의 AI 들보다 훨씬 더 높은 확률로 원하는 모양의 단백질을 성공적으로 설계했습니다.
• 실패율 대폭 감소: 실패하는 경우를 절반 가까이 줄였습니다.
• 다양성과 안정성 동시 달성: 튼튼하면서도 (안정성), 전혀 새로운 모양 (다양성) 을 가진 단백질을 만들어냈습니다.
• 비용 절감: 고가의 슈퍼컴퓨터나 실험실 장비 없이, 일반적인 GPU 서버 1 대에서 3 일 만에 학습이 완료됩니다.

5. 결론: 미래는 어떻게 바뀔까요? 🚀

ProteinZero 는 인공지능이 인간의 지시 없이도 스스로 진화할 수 있는 가능성을 보여줍니다.

• 의약품 개발: 암을 치료하는 새로운 약이나 백신을 훨씬 빠르게 설계할 수 있습니다.
• 환경 문제: 플라스틱을 분해하는 효소나 친환경 연료를 만드는 단백질을 찾아낼 수 있습니다.

한 줄 요약:

"ProteinZero 는 AI 가 스스로 요리를 만들고 맛을 보며, 실패를 두려워하지 않고 다양한 새로운 레시피를 개발하는 혁신적인 시스템으로, 단백질 설계의 미래를 바꿀 것입니다."

이 기술은 이제 막 시작되었지만, 앞으로 우리가 상상하지 못했던 새로운 생명 공학의 시대를 열어줄 것으로 기대됩니다! 🌱✨

기술 요약

1. 문제 정의 (Problem)

단백질 생성 모델은 단백질 설계 분야에서 큰 잠재력을 보이지만, 다음과 같은 근본적인 한계로 인해 성공률이 제한받고 있습니다.

• 데이터 의존성 및 편향: 기존 모델들은 단백질 데이터베이스 (PDB) 의 정제된 시퀀스 - 구조 쌍 데이터에 의존하여 학습됩니다. 이는 자연계 단백질 공간의 극히 일부만을 반영하며, 모델이 새로운 설계를 탐색하는 능력을 제한합니다.
• 목표 불일치: 지도 학습 (Supervised Learning) 기반의 역접힘 (Inverse Folding) 작업은 주로 자연 시퀀스 복원 (Sequence Recovery) 에 최적화되어 있지만, 실제 단백질 설계 목표는 설계 가능성 (Designability), 열역학적 안정성 (Stability), 그리고 **다양성 (Diversity)**을 동시에 만족하는 것입니다.
• 온라인 RL 의 난제: 기존 강화학습 (RL) 기반 접근법은 계산 비용이 매우 높은 물리 기반 평가 도구 (FoldX, Rosetta 등) 에 의존하거나, 모드 붕괴 (Mode Collapse, 즉 모델이 Reward 를 극대화하기 위해 제한된 시퀀스만 반복 생성하는 현상) 로 인해 실효성이 떨어집니다.

2. 방법론 (Methodology)

저자들은 ProteinZero를 제안합니다. 이는 역접힘 모델을 위한 온라인 강화학습 (Online RL) 프레임워크로, 정제된 선호 데이터셋 없이도 생성된 출력물로부터 지속적으로 자기 개선 (Self-improvement) 을 가능하게 합니다.

핵심 구성 요소:

1. 효율적인 보상 파이프라인 (Fast Proxy Rewards):

   • 설계 가능성 (Designability): AlphaFold 와 같은 MSA 기반 모델 대신, ESMFold를 사용하여 단일 패스로 구조를 예측하고, US-Align 을 통해 TM-score 를 계산합니다. 이는 기존 방법보다 25~100 배 빠른 평가 속도를 제공합니다.
   • 열역학적 안정성 (Thermal Stability): 물리 기반 시뮬레이션 (FoldX 등) 의 고비용을 피하기 위해, **자가 유도 $\Delta\Delta G$ 예측기 (Self-derived $\Delta\Delta G$ predictor)**를 개발했습니다. 이는 역접힘 모델의 백본 조건부 확률 (backbone-conditioned likelihood) 을 무조건부 사전 확률 (unconditional prior) 로 정규화하여 계산합니다.
   • 다중 목표 보상: 위 두 가지 지표를 Min-Max 정규화 후 가중 합하여 최종 보상으로 사용합니다.

2. 임베딩 수준 다양성 정규화 (Embedding-Level Diversity Regularizer):

   • 온라인 RL 에서 발생하는 모드 붕괴를 방지하기 위해 제안된 핵심 기술입니다.
   • 단순한 시퀀스 수준 (해밍 거리) 이 아닌, **모델의 임베딩 공간 (Embedding Space)**에서 생성된 시퀀스들의 코사인 유사도를 계산하여 다양성 점수를 도출합니다.
   • 이를 손실 함수의 정규화 항 ($L_{Div}$) 으로 추가하여, 모델이 Reward 를 극대화하더라도 기능적으로 의미 있는 다양한 시퀀스를 생성하도록 유도합니다.

3. RL 알고리즘 적용:

   • 제안된 프레임워크를 **RAFT (Reward-ranked Fine-tuning)**와 GRPO (Group Relative Policy Optimization) 알고리즘에 적용하여 최적화합니다.
   • 전체 목적 함수는 다음과 같이 구성됩니다:
     $$\mathcal{L}(\theta) = \mathcal{L}_{RL}(\theta) + \alpha_{KL} \cdot KL(p_\theta || p_{ref}) + \mathcal{L}_{Div}(\theta)$$
     (여기서 $KL$은 참조 모델과의 편차를 제어하고,$\mathcal{L}_{Div}$는 다양성을 유지합니다.)

3. 주요 기여 (Key Contributions)

1. ProteinZero 프레임워크: 정제된 선호 데이터셋 없이, 생성된 시퀀스로부터의 피드백을 통해 단백질 설계 모델이 다중 목표 (설계 가능성, 안정성, 다양성) 를 균형 있게 최적화하며 지속적으로 개선되는 온라인 RL 파이프라인을 제시했습니다.
2. 계산 효율성 확보: ESMFold 기반 설계 가능성 보상과 자가 유도 $\Delta\Delta G$ 예측기를 도입하여, 물리 기반 방법의 prohibitive 한 비용을 해결하고 온라인 RL 훈련을 실용적으로 만들었습니다.
3. 모드 붕괴 해결: 단백질 임베딩 공간에서 작동하는 새로운 다양성 정규화 기법을 개발하여, 온라인 RL 중 발생하는 모드 붕괴를 방지하면서도 기능적 일관성을 유지했습니다.
4. RL 파인튜닝 설계 공간 규명: 다양한 RL 알고리즘 (GRPO, RAFT, DPO 등) 과 보상/정규화 전략을 비교 분석하여, 단백질 설계에 최적화된 구성을 도출했습니다.

4. 실험 결과 (Results)

CATH-4.3 벤치마크 (0-150 잔기 및 150-300 잔기) 에서 ProteinZero 는 기존 최첨단 (SOTA) 모델들을 압도적으로 능가했습니다.

• 성능 향상:
  • 성공률 (Success Rate): ProteinMPNN, ESM-IF, InstructPLM 등 기존 모델 대비 90% 이상의 성공률을 달성했습니다 (단백질 크기에 따라 90.13% ~ 91.19%).
  • 실패율 감소: ProteinMPNN 대비 설계 실패율을 36~48% 감소시켰습니다.
  • 안정성: FoldX ddG 기준으로 안정성이 크게 향상되었으며 (예: 0-150 잔기 기준 -20.878 → -24.924 kcal/mol), 실험 데이터 (Ssym 벤치마크) 와의 상관관계도 물리 기반 도구 (FoldX, Rosetta) 와 유사한 수준 (PCC 0.60~0.62) 을 보였습니다.
• 다양성 유지:
  • 기존 오프라인 RL 방법 (DPO 등) 은 성공률은 높였으나 시퀀스 다양성을 감소시키는 경향이 있었으나, ProteinZero 는 성공률과 다양성을 동시에 향상시켰습니다 (예: 0-150 잔기 기준 다양성 0.281 → 0.306).
• 계산 효율성:
  • 단일 8×GPU 노드에서 보상 계산 및 데이터 생성을 포함해 3 일 이내에 완전한 RL 훈련을 완료할 수 있음을 입증했습니다.
• 검증:
  • 훈련 시 ESMFold 를 사용했음에도, 독립적인 평가 도구인 AlphaFold3와 FoldX에서도 일관된 성능 향상을 보여, 모델이 특정 예측기의 아티팩트가 아닌 일반적인 생리학적 원리를 학습했음을 증명했습니다.

5. 의의 및 결론 (Significance)

ProteinZero 는 단백질 생성 모델의 발전에 중요한 이정표가 됩니다.

• 지속 가능한 자기 진화: 레이블이 없는 생성 데이터로부터 학습하여 모델을 지속적으로 개선하는 패러다임을 제시함으로써, 제한된 실험 데이터에 의존하는 기존 방식의 한계를 극복했습니다.
• 실용적인 RL 적용: 계산 비용이 높은 물리 기반 평가 없이도 효율적인 온라인 RL 이 가능함을 보여주어, 대규모 단백질 설계 공간 탐색을 가속화할 수 있는 길을 열었습니다.
• 다양성 보존: 모드 붕괴를 방지하는 임베딩 기반 정규화 기법은 단백질 설계뿐만 아니라 다른 생성 모델의 RLHF 적용에도 중요한 통찰을 제공합니다.

이 연구는 계산적 설계와 실험적 검증 사이의 간극을 좁히고, 신약 개발, 효소 공학, 합성 생물학 등 다양한 분야에서 새로운 단백질 설계를 가능하게 할 잠재력을 가지고 있습니다.