Understanding protein function with a multimodal retrieval-augmented foundation model

이 논문은 시퀀스 컨텍스트 순서에 대한 불변성을 가진 계층적 트랜스포머 인코더와 구조 조건부 및 검색 증강 기능을 갖춘 다중 모달 기반 모델인 PoET-2 를 제안하여, 제로샷 변이 효과 예측 및 소규모 데이터셋에서의 단백질 기능 예측 성능을 획기적으로 개선했음을 보여줍니다.

핵심

이 논문은 **'PoET-2'**라는 새로운 인공지능 모델을 소개합니다. 이 모델은 단백질을 이해하고, 새로운 단백질을 설계하는 데 도움을 주는 '초지능 비서'라고 생각하시면 됩니다.

단백질은 우리 몸의 모든 일을 처리하는 작은 '작업자'들이고, 이 작업자들의 행동은 아미노산이라는 레고 블록들이 어떻게 쌓여 있느냐에 따라 결정됩니다. PoET-2 는 이 레고 블록들의 언어를 배우고, 더 나아가 새로운 구조를 만들어내는 능력을 가졌습니다.

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

1. 기존 모델의 한계: "외운 것만 아는 학생" vs "PoET-2: "맥락을 읽는 천재"

기존의 인공지능 모델들은 방대한 양의 단백질 데이터를 외워서 학습했습니다. 하지만 이 모델들은 몇 가지 큰 문제가 있었습니다.

• 한 글자만 고치는 것만 알았어요: "단백질 A 의 5 번째 글자를 B 로 바꾸면 어떻게 될까?"는 물어보면 답을 잘 했지만, "글자를 2 개나 3 개를 동시에 바꾸거나, 아예 글자를 끼워 넣거나 빼면 (삽입/삭제) 어떻게 될까?"는 질문에는 엉뚱한 답을 하거나 아예 못 했습니다. 마치 문법 공부는 했지만, 문장을 길게 늘이거나 줄이는 창의적인 작문은 못 하는 학생 같죠.
• 데이터가 부족하면 망했어요: 실험실 데이터를 아주 조금만 주면 (소량의 데이터), 새로운 일을 배우는 데 매우 서툴렀습니다.
• 크기만 키우는 건 비효율적: 성능을 올리려고 모델의 크기 (파라미터) 를 무작정 키우니, 컴퓨터가 너무 무거워지고 비용이 많이 들었습니다.

2. PoET-2 의 핵심 아이디어: "가족 앨범"과 "구조도"를 함께 보는 비법

PoET-2 는 이 문제들을 해결하기 위해 세 가지 마법 같은 기술을 섞었습니다.

① 검색 증강 (Retrieval-Augmentation): "가족 앨범을 펼쳐보는 것"

PoET-2 는 새로운 단백질을 볼 때, 혼자서 외운 지식만 믿지 않습니다. 대신 **"이 단백질과 비슷한 친척 (동족) 들은 어떤 특징을 가지고 있을까?"**라고 물어보며, 관련 있는 단백질들의 '가족 앨범'을 실시간으로 찾아봅니다.

• 비유: 새로운 친구를 소개받았을 때, 그 친구 혼자만 보는 게 아니라 그 친구의 가족 사진들을 보며 "아, 이 친구는 가족처럼 성실할 거야"라고 추측하는 것과 같습니다. 이 덕분에 거대한 모델을 만들지 않아도, 필요한 정보를 그때그때 찾아서 똑똑하게 판단할 수 있습니다.

② 멀티모달 (Multimodal): "글자"와 "그림"을 동시에 읽는 눈

기존 모델은 단백질의 아미노산 순서 (글자) 만 봤다면, PoET-2 는 **글자 순서와 3 차원 구조 (그림)**를 동시에 봅니다.

• 비유: 요리 레시피 (글자) 만 보고 요리를 하는 게 아니라, 완성된 요리의 모양 (구조) 도 함께 보며 "이 재료를 이렇게 섞어야 이 모양이 나오겠구나"라고 이해하는 것입니다. 특히 단백질이 어떻게 접히는지 (구조) 를 알면, 단백질이 얼마나 튼튼한지 (안정성) 를 훨씬 잘 예측할 수 있습니다.

③ 두 가지 뇌 (Dual Decoders): "창작가"와 "분석가"의 협업

PoET-2 는 두 개의 다른 두뇌를 가지고 있습니다.

1. 창작가 (생성 모델): 새로운 단백질 시퀀스를 만들어내거나, 특정 변이가 생겼을 때 "이게 얼마나 잘 작동할까?" 점수를 매기는 역할을 합니다. (예: "이 글자를 바꾸면 단백질이 망가질까?")
2. 분석가 (임베딩 모델): 단백질의 전체적인 특징을 파악하여, "이 단백질은 어떤 일을 잘할까?"를 분석하는 역할을 합니다.

3. PoET-2 가 보여준 놀라운 성과

이 모델은 실제로 많은 분야에서 기존 최고 수준 (State-of-the-art) 을 깨뜨렸습니다.

• 복잡한 변이 예측: 글자를 여러 개 동시에 바꾸거나, 글자를 끼워 넣는/빼는 (삽입/삭제) 복잡한 변이에서도 기존 모델보다 훨씬 정확하게 예측했습니다. (기존 모델이 20% 더 잘하는 수준)
• 적은 데이터로도 천재: 실험실 데이터가 아주 적을 때 (소량의 데이터) 도, 다른 모델들보다 훨씬 빠르게 새로운 기능을 학습했습니다. 마치 독서량이 적은 학생도 맥락을 잘 파악해서 시험을 잘 보는 것과 같습니다.
• 가볍고 빠름: 성능은 최고인데, 모델 크기는 1 억 8 천만 개 파라미터로 매우 가볍습니다. 거대한 슈퍼컴퓨터가 아니라도 일반 그래픽 카드 (GPU) 하나로도 충분히 돌릴 수 있어 접근성이 좋습니다.

4. 결론: 왜 이것이 중요한가요?

PoET-2 는 단백질 연구자들에게 **"더 적은 비용으로, 더 적은 데이터로, 더 복잡한 문제까지 해결할 수 있는 도구"**를 제공했습니다.

• 질병 치료: 유전병을 일으키는 돌연변이가 왜 나쁜지 정확히 예측할 수 있습니다.
• 새로운 약물 개발: 바이러스를 막거나 암을 치료하는 새로운 단백질을 빠르게 설계할 수 있습니다.
• 환경 보호: 플라스틱을 분해하는 효소 등을 디자인하여 환경 문제를 해결하는 데 기여할 수 있습니다.

요약하자면, PoET-2 는 단백질이라는 복잡한 언어를 단순히 외우는 게 아니라, 맥락 (가족 관계) 과 구조 (모양) 를 이해하여 창의적으로 새로운 단백질을 설계하고, 그 기능을 정확히 예측하는 차세대 AI입니다.

기술 요약

이 논문은 PoET-2라는 새로운 다중 모달 (multimodal) 검색 증강 (retrieval-augmented) 단백질 기초 모델 (foundation model) 을 소개하며, 단백질 기능 이해와 설계 능력을 획기적으로 향상시켰음을 보고합니다. 기존 단백질 언어 모델 (PLM) 의 한계를 극복하고, 제로샷 (zero-shot) 변이 효과 예측 및 소량 데이터에서의 지도 학습 (supervised learning) 성능을 모두 개선한 핵심 기술들을 상세히 설명합니다.

1. 문제 정의 (Problem)

기존의 단백질 언어 모델 (PLM) 은 자연 단백질 시퀀스 수억 개를 학습하여 단백질 구조 예측에는 큰 성과를 거두었으나, 다음과 같은 중요한 한계점이 존재했습니다:

• 단일 돌연변이 예측의 한계: 대부분의 PLM 은 마스킹 언어 모델 (MLM) 기반 접근법을 사용하여 단일 치환 (substitution) 돌연변이만 예측할 수 있습니다. 삽입/삭제 (Indels) 나 고차원 돌연변이 (여러 위치가 동시에 변이된 경우) 의 효과를 예측하는 데는 실패합니다.
• 지도 학습의 데이터 비효율성: 단백질 엔지니어링 현장에서는 제한된 실험 데이터 (few-shot) 로부터 직접 학습하여 특정 기능을 최적화해야 합니다. 기존 모델들은 소량 데이터에서 일반화 능력이 부족했습니다.
• 모델 확장성의 역설: 모델 파라미터 수를 늘리는 것 (Scaling) 이 구조 예측에는 도움이 되지만, 적합도 (fitness) 모델링이나 기능 예측에는 오히려 부정적이거나 중립적인 영향을 미칠 수 있으며, 계산 비용이 급증합니다.
• 정보 통합 부족: 기존 연구들은 다중 모달 (시퀀스 + 구조) 접근법이나 검색 증강 (retrieval-augmentation) 중 하나만 활용했을 뿐, 두 가지를 동시에 통합한 모델은 부재했습니다.

2. 방법론 (Methodology)

PoET-2 는 시퀀스, 구조, 진화적 맥락을 통합하여 학습하는 다중 모달 검색 증강 기반의 생성 모델입니다. 주요 아키텍처 및 학습 전략은 다음과 같습니다.

가. 다중 모달 및 검색 증강 (Multimodality & Retrieval-Augmentation)

• 컨텍스트 기반 학습 (In-context Learning): 모델은 사용자가 입력한 '프롬프트' (관련 단백질 가족의 시퀀스 집합) 를 기반으로 진화적 제약 조건을 추론합니다. 이는 수백억 파라미터의 거대 모델 없이도 새로운 시퀀스에 대한 컨텍스트 학습을 가능하게 합니다.
• 계층적 어텐션 (Hierarchical Attention): 입력된 단백질 집합의 순서에 무관한 (equivariant) 계층적 어텐션 아키텍처를 사용하여, 단백질 간의 진화적 관계를 효과적으로 포착합니다.
• 구조 조건부 생성: 아미노산 시퀀스뿐만 아니라 백본 원자 (N, Cα, C) 의 3D 구조 정보도 입력으로 받아, 구조에 조건부 (conditioned) 로 시퀀스를 생성하거나 평가할 수 있습니다.

나. 듀얼 디코더 아키텍처 (Dual Decoder Architecture)

PoET-2 는 하나의 인코더와 두 개의 디코더로 구성됩니다:

1. 자기회귀 디코더 (Autoregressive Decoder, CLM):
   • 목적: 시퀀스 생성 및 확률적 점수 매기기.
   • 기능: 시퀀스의 결합 확률 분포 $P(\text{sequence} | \text{prompt})$ 를 모델링하여, 변이 (변형) 시퀀스의 적합도 (fitness) 를 정확히 계산합니다. 이 방식은 길이가 다른 Indel 변이와 고차원 돌연변이의 상호작용 (epistasis) 을 자연스럽게 처리합니다.
2. 양방향 디코더 (Bidirectional Decoder, MLM):
   • 목적: 표현 학습 (Representation Learning).
   • 기능: 전체 시퀀스 컨텍스트를 고려한 임베딩을 생성하여, 구조 및 기능 예측과 같은 하위 작업에 강력한 특징 벡터를 제공합니다.

다. 학습 목표 (Training Objectives)

• CLM 손실: 시퀀스 생성 및 변이 점수 매기기를 위한 자기회귀적 학습.
• MLM 손실: 양방향 이해를 위한 마스킹 언어 모델 학습.
• 구조 기반 어텐션 편향 (Structure-based Attention Bias): 시퀀스 내 아미노산 쌍 간의 3D 거리 (Cα-Cα 거리) 를 어텐션 점수에 편향 (bias) 으로 추가하여 구조적 정보를 통합합니다.

3. 주요 기여 및 결과 (Key Contributions & Results)

가. 제로샷 변이 효과 예측 (Zero-shot Variant Effect Prediction)

ProteinGym 벤치마크에서 PoET-2 는 기존 최첨단 모델들을 능가하는 성능을 보였습니다.

• Indel 및 고차원 돌연변이 예측: 기존 모델들이 처리하지 못했던 삽입/삭제 (Indel) 변이와 3 개 이상의 돌연변이가 포함된 고차원 변이에 대해 Spearman 상관계수 (ρ) 에서 약 20% 이상 개선된 성능을 기록했습니다. (예: Indel 예측에서 PoET-1 대비 $\Delta\rho \approx 0.05$ 향상).
• 임상 변이 예측: 인간 질병 관련 변이 (pathogenicity) 예측에서도 기존 모델 대비 AUROC 가 크게 향상되었습니다.
• 효율성: 1.82 억 (182M) 파라미터라는 경량 모델임에도 불구하고, 650M~64B 파라미터 규모의 기존 모델들보다 뛰어난 성능을 발휘했습니다.

나. 지도 학습 및 소량 데이터 효율성 (Supervised Learning & Data Efficiency)

제한된 실험 데이터 (few-shot) 를 활용한 단백질 기능 학습에서 PoET-2 는 압도적인 성능을 보였습니다.

• 소량 데이터에서의 우위: 학습 데이터가 100 개 미만인 극단적인 Few-shot 상황에서도 기존 최첨단 모델 (Kermut 등) 이 최대 데이터셋으로 학습했을 때의 성능을 따라잡거나 능가했습니다.
• 일반화 능력: 무작위 (Random), 모듈로 (Modulo), 연속 (Contiguous) 등 다양한 교차 검증 (Cross-Validation) 방식에서 모두 SOTA 성능을 기록했습니다. 특히 연속 구간을 테스트하는 'Contiguous' 분할에서 기존 모델들의 일반화 실패를 극복했습니다.

다. 구조 조건부의 역할 분석

• 제로샷 예측: 구조 정보를 프롬프트에 포함하면 (특히 안정성 관련 assay) 성능이 향상됩니다.
• 지도 학습: 소량 데이터 학습에서는 구조 정보를 명시적으로 추가하는 것이 큰 이점을 주지 않았습니다. 이는 PoET-2 의 임베딩 자체가 이미 구조 정보를 암묵적으로 잘 인코딩하고 있음을 시사합니다.

4. 의의 및 결론 (Significance & Conclusion)

• 새로운 패러다임 제시: 거대 모델의 단순한 확장 (Scaling) 대신, **검색 증강 (Retrieval-Augmentation)**과 **다중 모달 (Multimodality)**을 결합하여 데이터 효율성과 일반화 능력을 동시에 개선하는 새로운 방향을 제시했습니다.
• 실용적 가치: Indel 및 고차원 돌연변이 예측이 가능해짐으로써, 기존에 불가능했던 복잡한 단백질 설계 및 변이 분석이 가능해졌습니다. 또한, 적은 실험 데이터로도 높은 정확도의 예측이 가능해져 단백질 엔지니어링의 비용과 시간을 획기적으로 줄일 수 있습니다.
• 경량화: 182M 파라미터라는 작은 모델 크기로 SOTA 성능을 달성하여, GPU 리소스가 제한된 환경에서도 널리 활용될 수 있는 가능성을 열었습니다.

요약하자면, PoET-2 는 단백질의 진화적 맥락과 구조적 정보를 통합하여, 기존 모델들이 해결하지 못했던 복잡한 변이 (Indel, 고차원) 예측과 데이터가 부족한 상황에서의 학습 문제를 성공적으로 해결한 차세대 단백질 기초 모델입니다.