ESM-2 임베딩 실전 — zero-shot으로 변이 효과를 미리 읽는 법

F10에서 몸의 물리로 갔다면, 이제 서열의 확률로 돌아옵니다

F06에서 ESM-2가 ESMFold의 앞단에서 진화 정보를 언어모델 형태로 압축한다고 다뤘습니다. 그런데 그때는 ESM-2를 오직 "구조 예측을 위한 부품"으로만 봤습니다. 사실 ESM-2는 그 자체로 훨씬 다양하게 쓰입니다. 이 편에서는 ESM-2의 임베딩 벡터가 실제로 무엇을 담고 있는지, 그리고 추가 학습(fine-tuning) 없이도 "이 변이가 단백질 기능에 해로울지"를 예측하는 zero-shot 변이 효과 예측을 직접 다룹니다.

원리 — 임베딩은 왜 유용한 정보를 담게 되는가

임베딩 공간: 잔기 하나하나를 벡터로

F06에서 다룬 마스크 언어 모델링으로 학습된 ESM-2는, 서열을 입력하면 각 잔기 위치마다 고차원 벡터(가령 1280차원)를 출력합니다.

$h_i = \text{ESM-2}(x_1, \ldots, x_L)_i \in \mathbb{R}^{d}$

이 벡터 $h_i$ 는 단순히 " $i$ 번째 위치의 아미노산이 무엇이다"라는 정보만 담고 있지 않습니다. 마스크된 잔기를 문맥으로 맞히도록 2억 5천만 서열에서 학습된 결과, 이 벡터에는 "이 위치가 구조적으로 얼마나 보존됐는가", "이 위치가 다른 어떤 위치와 함께 진화했는가", "이 위치가 활성 부위 근처인가" 같은 정보가 암묵적으로 녹아 있습니다. 서열 전체를 하나의 벡터로 요약하고 싶으면 각 잔기 벡터를 평균 내는 방식(mean pooling)을 흔히 씁니다.

$h_{\text{seq}} = \frac{1}{L}\sum_{i=1}^{L} h_i$

이렇게 얻은 서열 임베딩은 분류·회귀 등 다양한 다운스트림 태스크의 입력 특징으로 바로 쓸 수 있습니다. 예를 들어 서열 임베딩 사이의 코사인 유사도가 높을수록 기능적으로 유사한 단백질일 가능성이 높다는 식의 활용이 가능합니다.

Zero-shot 변이 효과 예측: masked marginal 방식

여기서 더 흥미로운 활용이 나옵니다. 추가 학습 데이터 없이, 이미 학습된 마스크 예측 능력만으로 "이 위치에서 야생형 아미노산 대신 다른 아미노산이 왔을 때 얼마나 이상한가"를 점수로 매길 수 있습니다. 방법은 F06의 MLM 손실을 그대로 뒤집어 씁니다.

변이가 일어난 위치 $i$ 를 마스크 토큰으로 가립니다.
모델에게 "이 위치에 무엇이 와야 하는가"를 물어, 20개 아미노산 각각의 확률 $P(x_i = a \mid x_{\setminus i})$ 을 얻습니다.
야생형 아미노산과 변이 아미노산의 로그 확률 차이를 계산합니다.

$\text{score} = \log P(x_i = x_{\text{mut}} \mid x_{\setminus i}) - \log P(x_i = x_{\text{wt}} \mid x_{\setminus i})$

이 점수가 크게 음수라면 모델이 변이 잔기를 야생형보다 덜 그럴듯하게 본 것입니다. 이 접근의 대표 원 논문은 ESM-2가 아니라 ESM-1v를 평가한 Meier et al. (2021)입니다. ESM-2에도 같은 masked-marginal 계산을 적용할 수 있지만, 체크포인트·scoring strategy·벤치마크를 명시하지 않은 채 ESM-1v 논문의 성능을 ESM-2 성능으로 옮겨 적으면 안 됩니다.

손 계산 예제: 로그 확률 차이가 실제로 뜻하는 것

모델이 어떤 위치에서 야생형 아미노산에 확률 $P(x_{\text{wt}}) = 0.6$ , 변이 아미노산에 확률 $P(x_{\text{mut}}) = 0.02$ 를 부여했다고 합시다.

$\text{score} = \log(0.02) - \log(0.6) = -3.912 - (-0.511) = -3.401$

같은 위치에서 조금 덜 극단적인 변이, $P(x_{\text{mut}}) = 0.15$ 인 경우와 비교하면

$\text{score} = \log(0.15) - \log(0.6) = -1.897 - (-0.511) = -1.386$

두 변이 모두 "야생형보다 낮은 확률"이라 음수 점수를 받지만, 첫 번째 변이의 점수(-3.401)가 두 번째(-1.386)보다 훨씬 더 낮습니다. 이는 모델이 첫 번째 변이를 두 번째보다 훨씬 더 "있을 법하지 않다"고 판단했다는 뜻이고, 실제 임상 벤치마크에서는 이런 극단적으로 낮은 점수의 변이일수록 병원성일 확률이 높은 경향이 관찰됩니다. 다만 이는 어디까지나 상관관계이지, 개별 변이 하나하나에 대한 결정론적 판정이 아니라는 점을 다음 절에서 다시 강조합니다.

실습: HuggingFace로 변이 효과 zero-shot 점수 계산

python

# Colab T4에서 실행
!pip install -q transformers torch

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, EsmForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
model = EsmForMaskedLM.from_pretrained("facebook/esm2_t33_650M_UR50D").cuda().eval()

def masked_marginal_score(sequence, position, wt_aa, mut_aa):
    """position은 0-indexed 잔기 위치."""
    seq_list = list(sequence)
    assert seq_list[position] == wt_aa, "야생형 아미노산이 서열과 일치하지 않습니다"

    inputs = tokenizer(sequence, return_tensors="pt")
    # HuggingFace ESM tokenizer는 맨 앞에 <cls>, 맨 뒤에 <eos>를 붙입니다.
    mask_idx = position + 1
    tokens_before = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    assert tokens_before[mask_idx] == wt_aa, (tokens_before, mask_idx, wt_aa)
    inputs["input_ids"][0, mask_idx] = tokenizer.mask_token_id
    inputs = {k: v.cuda() for k, v in inputs.items()}

    with torch.no_grad():
        logits = model(**inputs).logits[0, mask_idx]
        log_probs = F.log_softmax(logits, dim=-1)

    wt_id = tokenizer.convert_tokens_to_ids(wt_aa)
    mut_id = tokenizer.convert_tokens_to_ids(mut_aa)
    return (log_probs[mut_id] - log_probs[wt_id]).item()

sequence = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSPLHSVYVDQWDW"
score = masked_marginal_score(sequence, position=5, wt_aa="I", mut_aa="P")
print(f"변이 효과 zero-shot score: {score:.3f} (더 음수일수록 유해할 가능성)")

위 코드는 마스킹 전에 해당 토큰이 실제 야생형 잔기인지 assertion으로 확인합니다. 다른 tokenizer나 batch padding을 쓰면 position+1을 그대로 재사용하지 말고 residue-to-token mapping을 다시 검증해야 합니다.

CS 매핑

임베딩 공간(Embedding Space): 이산적인 아미노산 서열을 연속적인 고차원 벡터로 사상하는 이 과정은, 자연어의 단어 임베딩(Word2Vec, BERT 임베딩)이 이산 토큰을 연속 벡터 공간에 배치하는 것과 원리적으로 동일합니다.
Zero-shot 추론: 특정 태스크(변이 효과 예측)를 위한 라벨링된 데이터로 별도 학습을 하지 않고, 사전학습 목표(마스크 예측)를 그대로 활용해 새로운 문제에 답하는 것은 대형 언어모델의 zero-shot 프롬프팅과 같은 원리입니다.
로그우도비(Log-likelihood Ratio): 두 가설(야생형 vs 변이)의 로그 확률 차이로 점수를 매기는 방식은 통계학의 우도비 검정(likelihood ratio test)과 형식적으로 동일한 구조입니다.

자주 만나는 결함

zero-shot 점수를 임상 진단으로 오용: 이 점수는 "진화적으로 얼마나 이례적인 변이인가"를 추정할 뿐, 실제 질병 유발 여부를 진단하는 임상 검사가 아닙니다. F14에서 다룰 AlphaMissense 등 전용 모델도 마찬가지로 스크리닝·연구 가설 생성 보조 도구로 봐야 합니다.
토큰 offset 미검증으로 엉뚱한 위치 마스킹: 실습 코드처럼 토크나이저의 특수 토큰 구조를 확인하지 않고 인덱스를 하드코딩하면, 의도한 위치가 아닌 다른 잔기를 마스킹해 완전히 틀린 점수가 나옵니다.
모델 크기에 따른 점수 편차 무시: esm2_t33_650M 같은 중간 규모 모델과 최상위 규모 모델은 zero-shot 성능에 차이가 있습니다. 벤치마크 재현 시 사용한 모델 크기를 명시해야 합니다.

더 깊게 파고 싶다면

본문은 BPD 연구진이 직접 재구성한 서술입니다. 원 논문과 공식 자료로 심화해봅시다.

ESM-1v zero-shot 변이 효과 예측 원 논문: Meier et al. (2021), Language models enable zero-shot prediction of the effects of mutations on protein function, NeurIPS 2021.
ESM-2 원 논문: Lin et al. (2023), Evolutionary-scale prediction of atomic-level protein structure with a language model, Science (F06과 동일 논문, ESM-2 아키텍처 섹션).
HuggingFace ESM-2 모델 카드: huggingface.co/facebook/esm2_t33_650M_UR50D — 모델 크기별 체크포인트 목록.

다음 편 F12에서는 서열-구조-기능을 하나의 모델 안에서 함께 다루는 ESM3로 넘어갑니다.

BioPlayground

ESM-2 임베딩 실전: zero-shot으로 변이 효과를 미리 읽는 법