Resumen Técnico: Tokenización Adaptativa de Proteínas

Planteamiento del Problema

Los enfoques actuales para la tokenización de la estructura de proteínas se basan en la localidad, donde los tokens se crean mediante el agrupamiento (pooling) de información de residuos espacialmente vecinos a lo largo de una secuencia proteica. Si bien estos métodos logran una reconstrucción de alta fidelidad, los autores identifican dos limitaciones críticas:

1. Acumulación de Errores en la Generación: Los modelos generativos basados en la tokenización local sufren de propagación de errores. Un solo token mal muestreado puede provocar desviaciones significativas en la estructura predicha, causando que los enfoques discretos tengan un rendimiento inferior en comparación con los modelos de difusión continua en métricas generativas.
2. Escalabilidad Deficiente: El número de tokens escala linealmente con el tamaño de la proteína. Esto hace que el modelado de complejos proteicos grandes sea computacionalmente costoso y limita la capacidad de comprimir proteínas de manera efectiva.
3. Limitaciones de Representación: Los tokenizadores locales a menudo requieren operaciones de mean-pooling para generar representaciones de tamaño fijo para tareas posteriores, lo que puede descartar información global crítica y escalar mal con el tamaño de la proteína.

Metodología

Los autores proponen la Tokenización Adaptativa de Proteínas (APT), un método de tokenización global donde los tokens sucesivos contribuyen con niveles crecientes de detalle a una representación global, en lugar de corresponder a vecindarios locales específicos.

Arquitectura

APT se implementa como un autoencoder de difusión con un cuello de botella discreto:

• Entrada: Una secuencia de coordenadas $C_\alpha$ puras (normalizadas a masa cero) de longitud $L$.
• Codificador (Encoder): Un transformador de atención bidireccional mapea las posiciones puras a una secuencia latente $c \in \mathbb{R}^{L \times d}$.
• Discretización: La secuencia latente se discretiza utilizando Cuantización Escalar Finita (FSQ) con niveles $(8, 5, 5, 5)$, lo que resulta en un tamaño de código efectivo de 1,000.
• Mecanismo de Adaptatividad: Para imponer la adaptatividad, el modelo emplea un dropout anidado. Durante el entrenamiento, un límite superior $U$ se muestrea uniformemente de $[1, \min(L, k_{max})]$. Los tokens más allá de este límite se descartan. Esto fomenta que el modelo coloque la información global crítica (componentes de baja frecuencia) en los primeros tokens y los detalles de alta frecuencia en los tokens subsiguientes.
• Decodificador (Decoder): Un decodificador de difusión, entrenado mediante un objetivo de ajuste de flujo (flow-matching), reconstruye las coordenadas atómicas condicionadas a los tokens discretos. El decodificador utiliza equivalencia estocástica, aprendiendo simetrías en lugar de imponerlas explícitamente.
• Predicción de Tamaño: El tamaño de la proteína se regresa desde el primer token mediante una pérdida de entropía cruzada, lo que permite al modelo desacoplar el tamaño de la proteína de la longitud de la secuencia de condicionamiento.

Entrenamiento e Inferencia

• Entrenamiento: El modelo se entrena con aproximadamente 473,000 predicciones sintéticas de AlphaFold2 de la base de datos AFDB agrupada por Foldseek. El entrenamiento consta de dos etapas: (1) entrenamiento del autoencoder con rotaciones aleatorias para aprender la equivalencia estocástica, y (2) entrenamiento de un modelo autorregresivo (estilo GPT) sobre los tokens de APT.
• Estrategias de Inferencia:
  • Tail Dropout (Descarte de la Cola): Durante la generación, el modelo puede detener la generación de tokens de forma temprana. Cualquier prefijo de la secuencia de tokens es una señal de condicionamiento válida. La cola (detalles de alta frecuencia) se delega al decodificador de difusión.
  • Parada Basada en Entropía: Los autores introducen criterios de parada basados en la entropía de los tokens (límite finito, entropía mínima basada en splines) para determinar el punto óptimo para detener la generación, equilibrando la fidelidad de la representación y la capacidad generativa.
  • Recocido de Clasificador (Classifier Annealing): Para mitigar los artefactos causados por el alejamiento de la variedad (manifold) real de los datos, los autores emplean el recocido de clasificador, interpolando entre los campos de guía condicionales e incondicionales durante la decodificación.

Contribuciones Clave

1. Tokenizador Adaptativo Global: Introducción de un tokenizador donde los tokens representan descriptores globales en una jerarquía de grueso a fino, en lugar de vecindarios locales.
2. Validación de Rendimiento: Demostración de que el modelo iguala o supera a los tokenizadores locales de vanguardia (SOTA) (ej. DPLM2, ESM3, Kanzi) en tareas generativas, de aprendizaje de representación y de reconstrucción.
3. Técnicas de Inferencia: Desarrollo de muestreo basado en entropía y recocido de clasificador para gestionar la complejidad de la muestra y mitigar la exposición al error, mejorando la diseñabilidad.
4. Aplicaciones Zero-Shot: Demostración del desacoplamiento del tamaño de la proteína del condicionamiento, permitiendo aplicaciones zero-shot tales como el encogimiento de proteínas (reducción del recuento de residuos manteniendo la estructura) y la maduración de afinidad mediante el escalado en el tiempo de inferencia.

Resultados

• Reconstrucción: En los conjuntos de prueba CATH, CAMEO y AFDB, APT logra RMSD y puntuaciones TM comparables a los modelos de difusión continua SOTA. Notablemente, usando solo 32–64 tokens, los RMSD son $< 2$ Å, afirmando la compresión de las estructuras proteicas.
• Generación: Un modelo autorregresivo entrenado en tokens de APT logra una diseñabilidad de 0.871 (fracción de muestras con scRMSD $< 2$ Å), superando significativamente a los modelos de difusión discreta como DPLM2 (0.486) y ESM3. Esto se logra sin técnicas de muestreo "best-of-N".
• Aprendizaje de Representación: En la tarea de clasificación CATH, el sondeo no lineal (MLP) en los tokens de APT supera el sondeo en las representaciones de ESM3 y DPLM2. Crucialmente, APT proporciona representaciones globales de tamaño fijo sin mean-pooling, e incluso las representaciones altamente comprimidas (16 tokens) superan a modelos más grandes.
• Aplicaciones:
  • Encogimiento de Proteínas: Generó con éxito versiones más pequeñas de proteínas (ej. hemoglobina, beta-barrel) preservando la estructura global y local al condicionar el proceso de difusión en menos residuos.
  • Maduración de Afinidad: Demostró la capacidad de evolucionar uniones débiles en uniones fuertes utilizando búsqueda de haz (beam search) guiada por funciones de recompensa (ej. iPAE) en el espacio latente.

Significancia

El artículo afirma que la Tokenización Adaptativa de Proteínas proporciona un enfoque consciente de la tarea para escalar tareas biológicas. Al pasar de una tokenización local a una global, el método resuelve el compromiso entre la fidelidad de la reconstrucción y la capacidad generativa. La capacidad de comprimir proteínas en vectores de tamaño fijo sin pérdida de información mediante mean-pooling ofrece una vía práctica para el modelado de grandes complejos proteicos y para realizar el escalado en el tiempo de inferencia. Los autores posicionan este trabajo como un paso hacia modelos generativos capaces de razonar a través de múltiples escalas de longitud, abordando un desafío clave en la bioingeniería de frontera.

Nota: Los autores reconocen limitaciones, específicamente que los tokens de APT son efectivos para tareas globales pero menos adecuados para tareas locales como el andamiaje de motivos (motif scaffolding), y que el tokenizador actual no tiene en cuenta las cadenas laterales o la composición de aminoácidos, que son críticos para ciertas funciones.