Si en F10 pasamos a la física del cuerpo, ahora volvemos a las probabilidades de la secuencia

En F06 analizamos cómo ESM-2 comprime la información evolutiva en forma de modelo de lenguaje en la etapa inicial de ESMFold. Sin embargo, en ese momento solo consideramos a ESM-2 como un "componente" para la predicción de estructuras. En realidad, ESM-2 tiene muchos más usos. En este episodio abordaremos directamente qué contiene realmente el vector de incrustación de ESM-2 y cómo realizar una predicción del efecto de variantes zero-shot (sin necesidad de entrenamiento adicional) para determinar si una variante será perjudicial para la función de la proteína.

Principio: ¿por qué los embeddings terminan conteniendo información útil?

Espacio de incrustación: cada residuo como un vector

ESM-2, entrenado mediante el modelado de lenguaje enmascarado (MLM) que se trató en F06, produce un vector de alta dimensión (por ejemplo, de 1280 dimensiones) para cada posición de residuo al recibir una secuencia.

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

Este vector $h_i$ no solo contiene la información de "cuál es el aminoácido en la posición $i$". Como resultado del entrenamiento en 250 millones de secuencias para predecir el residuo enmascarado dado el contexto, este vector incluye implícitamente información sobre "qué tan conservada estructuralmente está esta posición", "con qué otras posiciones ha evolucionado conjuntamente esta posición" y "si esta posición está cerca del sitio activo". Si se desea resumir toda la secuencia en un solo vector, se utiliza comúnmente el método de promediado (mean pooling) de los vectores de cada residuo.

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

El embedding de secuencia obtenido de esta manera se puede utilizar directamente como característica de entrada para diversas tareas posteriores, como la clasificación y la regresión. Por ejemplo, se puede aprovechar el hecho de que una mayor similitud del coseno entre los embeddings de secuencias indica una mayor probabilidad de que las proteínas sean funcionalmente similares.

Predicción del efecto de variantes zero-shot: método marginal enmascarado

Aquí surge un uso aún más interesante. Sin datos de entrenamiento adicionales, es posible asignar una puntuación a "qué tan extraño sería si en esta posición apareciera otro aminoácido en lugar del aminoácido de tipo salvaje" utilizando únicamente la capacidad de predicción enmascarada ya aprendida. El método consiste simplemente en invertir la pérdida MLM de F06.

1. Ocultar la posición $i$ donde ocurre la variante con un token enmascarado.
2. Preguntar al modelo "qué debería haber en esta posición" para obtener la probabilidad $P(x_i = a \mid x_{\setminus i})$ de cada uno de los 20 aminoácidos.
3. Calcular la diferencia de log-probabilidades entre el aminoácido de tipo salvaje y el aminoácido variante.

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

Si esta puntuación es significativamente negativa, indica que el modelo considera que la variante es menos probable que la forma silvestre. El artículo de referencia principal para este enfoque es ESM-1v, evaluado por Meier et al. (2021), y no ESM-2. Si bien el mismo cálculo marginal enmascarado se puede aplicar a ESM-2, no se deben atribuir los resultados del artículo de ESM-1v como el rendimiento de ESM-2 sin especificar explícitamente los puntos de control, la estrategia de puntuación y los conjuntos de datos de referencia.

Ejemplo de cálculo manual: qué significa realmente la diferencia de log-probabilidades

Supongamos que el modelo asigna una probabilidad $P(x_{\text{wt}}) = 0.6$ al aminoácido de la forma silvestre en una posición determinada y una probabilidad $P(x_{\text{mut}}) = 0.02$ al aminoácido de la variante.

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

Si lo comparamos con una variante menos extrema en la misma posición, donde $P(x_{\text{mut}}) = 0.15$:

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

Ambas variantes reciben una puntuación negativa porque tienen una probabilidad menor que la forma silvestre, pero la puntuación de la primera variante (-3.401) es mucho más baja que la de la segunda (-1.386). Esto indica que el modelo considera que la primera variante es mucho menos probable que la segunda. En los conjuntos de datos de referencia clínicos reales, se observa que las variantes con puntuaciones extremadamente bajas tienden a tener una mayor probabilidad de ser patógenas. Sin embargo, como se enfatizará nuevamente en la siguiente sección, esto es solo una correlación y no constituye un juicio determinista para cada variante individual.

Práctica: Cálculo de puntuaciones zero-shot del efecto de variantes con Hugging Face

# Ejecutar en 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):
    \"\"\"La posición es un índice basado en 0 para la ubicación del residuo.\"\"\"
    seq_list = list(sequence)
    assert seq_list[position] == wt_aa, "El aminoácido de tipo silvestre no coincide con la secuencia"

    inputs = tokenizer(sequence, return_tensors="pt")
    # El tokenizador ESM de HuggingFace añade <cls> al principio y <eos> al final.
    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\"Puntuación zero-shot del efecto de la mutación: {score:.3f} (más negativo indica mayor probabilidad de ser perjudicial)\")

El código verifica mediante una aserción si el token correspondiente es realmente un residuo de tipo salvaje antes del enmascaramiento. Si utiliza otro tokenizador o relleno por lotes, no reutilice position+1 tal cual; debe volver a verificar el mapeo de residuos a tokens.

Mapeo CS

• Espacio de incrustaciones: Este proceso, que proyecta secuencias discretas de aminoácidos en vectores continuos de alta dimensión, es conceptualmente idéntico a cómo los embeddings de palabras del lenguaje natural (Word2Vec, embeddings de BERT) colocan tokens discretos en un espacio vectorial continuo.
• Inferencia zero-shot: Responder a nuevos problemas utilizando directamente el objetivo de aprendizaje preentrenado (predicción de máscaras) sin entrenar adicionalmente con datos etiquetados para una tarea específica (predicción del efecto de mutaciones) sigue el mismo principio que el prompting zero-shot de los grandes modelos de lenguaje.
• Razón de verosimilitud logarítmica: La forma de puntuar basada en la diferencia de probabilidad logarítmica entre dos hipótesis (tipo salvaje vs. mutante) tiene una estructura formalmente idéntica a la prueba de razón de verosimilitud en estadística.

Defectos comunes

• Uso indebido de puntuaciones zero-shot para diagnóstico clínico: Esta puntuación estima solo "qué tan inusual es evolutivamente una mutación", no es un examen clínico que diagnostique si realmente causa una enfermedad. Los modelos dedicados como AlphaMissense, tratados en F14, también deben considerarse herramientas auxiliares para cribado y generación de hipótesis de investigación.
• Enmascaramiento en la posición incorrecta por no verificar el offset del token: Si se codifican los índices sin comprobar la estructura de tokens especiales del tokenizador, como en el código de práctica, se enmascarán residuos distintos a los deseados, obteniendo puntuaciones completamente erróneas.
• Ignorar las diferencias de puntuación según el tamaño del modelo: Los modelos de tamaño intermedio como esm2_t33_650M y los modelos de mayor escala presentan diferencias en el rendimiento zero-shot. Al reproducir benchmarks, debe especificarse el tamaño del modelo utilizado.

Para profundizar más

El texto ha sido reconstruido directamente por el equipo de investigación de BPD. Profundice consultando el artículo original y los recursos oficiales.

• Artículo original sobre la predicción zero-shot de efectos de mutaciones con ESM-1v: Meier et al. (2021), Language models enable zero-shot prediction of the effects of mutations on protein function, NeurIPS 2021.
• Artículo original de ESM-2: Lin et al. (2023), Evolutionary-scale prediction of atomic-level protein structure with a language model, Science (el mismo artículo que F06, sección de la arquitectura ESM-2).
• Ficha del modelo ESM-2 en Hugging Face: huggingface.co/facebook/esm2_t33_650M_UR50D — lista de puntos de control según el tamaño del modelo.

En la siguiente entrega, F12, pasaremos a ESM3, que aborda la secuencia, la estructura y la función conjuntamente dentro de un único modelo.