Resumen Técnico: Yeti – Un Tokenizador Compacto de Estructuras Proteicas para Reconstrucción y Generación Multi-Modal

1. Declaración del Problema

La integración de secuencias proteicas, estructuras 3D y anotaciones funcionales en una representación multi-modal unificada es crítica para avanzar en la biología computacional. Sin embargo, los enfoques actuales enfrentan tres desafíos principales:

1. Heterogeneidad de Datos: Los datos biológicos abarcan diversas modalidades (restricciones biofísicas, texto, imágenes, datos volumétricos) con propiedades estadísticas distintas.
2. Incompletitud y Ruido: Los conjuntos de datos del mundo real a menudo sufren de modalidades faltantes, mediciones ruidosas y cobertura desigual en el espacio estructura-función.
3. Escalabilidad y Acoplamiento: La mayoría de los modelos operan en una sola escala o modalidad, careciendo de mecanismos para transferir información entre niveles (atómico a celular) o para manejar procesos acoplados.

Específicamente, si bien los modelos generativos recientes tratan a las proteínas como objetos 3D utilizando espacios de coordenadas continuas (por ejemplo, difusión SE(3)), existe una alternativa convincente: codificar las estructuras atómicas 3D como tokens discretos. Esto permite un entrenamiento escalable y agnóstico a la modalidad, similar a los modelos de lenguaje. Sin embargo, los tokenizadores de estructuras proteicas existentes (por ejemplo, ESM3, DPLM-2, Kanzi) a menudo priorizan la precisión de reconstrucción sobre las capacidades generativas. Existe una desalineación donde una alta fidelidad de reconstrucción no garantiza una diversidad de tokens de alta calidad ni el soporte para la co-generación incondicional de secuencia y estructura. Además, muchos modelos de última generación dependen de conteos masivos de parámetros o de paradigmas de entrenamiento en dos etapas donde los codificadores se congelan, limitando potencialmente la optimización de extremo a extremo.

2. Metodología

2.1 Arquitectura: Yeti

Yeti (Yielding Encoded Tokens for Intermodality) es un tokenizador compacto de estructura proteica en una sola etapa diseñado para mapear coordenadas atómicas 3D continuas a tokens discretos.

• Entrada: Estructuras proteicas representadas como coordenadas $C_\alpha$ centradas en la media ($x \in \mathbb{R}^{L \times 3}$).
• Codificador: Un codificador basado en Transformer que utiliza atención multi-cabeza con codificación posicional rotatoria (RoPE) mapea las entradas a incrustaciones latentes continuas $Z \in \mathbb{R}^{L \times D}$.
• Cuantización: Las incrustaciones se proyectan en un espacio de cuantización y se procesan mediante un Cuantizador Libre de Búsqueda (LFQ). El espacio LFQ es un producto cartesiano de sub-libros de códigos binarios ($C_i = \{-1, +1\}$). La cuantización se realiza mediante una operación de signo, mapeando vectores a índices discretos sin necesidad de una tabla de búsqueda, lo que produce un tamaño de vocabulario $K = 2^D$.
• Decodificador: Una arquitectura StripedHyena (un modelo híbrido de procesamiento de señales) actúa como decodificador. A diferencia de los decodificadores tradicionales que dependen de componentes invariantes a SE(3) como la Atención de Puntos Invariantes (IPA), el decodificador StripedHyena aprende un campo de velocidad condicional para reconstruir estructuras a partir de coordenadas ruidosas.

2.2 Objetivo de Entrenamiento: Flow Matching

Yeti se entrena de extremo a extremo utilizando un objetivo de Flow Matching, que interpola entre una distribución de ruido gaussiano ($p_0$) y la distribución objetivo de estructura proteica ($p_{data}$).

• Función de Pérdida: El modelo minimiza la diferencia entre el campo de velocidad condicional predicho $v_\theta(x_t, t, c)$ y el campo vectorial objetivo ($x_1 - x_0$), donde $x_1$ es la estructura limpia y $x_0$ es ruido gaussiano.
• Regularización de Entropía: Para evitar el colapso del libro de códigos y asegurar una alta utilización, se añade una penalización de entropía a la función de pérdida: $L = L_{FM} + \lambda L_{entropy}$.
• Inferencia: Se emplea la Guía Libre de Clasificador (CFG) durante la inferencia. Al anular aleatoriamente las incrustaciones de tokens durante el entrenamiento, el modelo aprende tanto campos de velocidad condicionales como incondicionales, permitiendo una decodificación guiada para refinar la precisión de la estructura.

2.3 Co-Generación Multi-Modal

Para validar la utilidad generativa de los tokens Yeti, los autores entrenaron un Modelo de Difusión Enmascarada (MDM) desde cero. Este modelo genera conjuntamente secuencias de aminoácidos y tokens de estructura Yeti sin ninguna inicialización preentrenada (a diferencia de DPLM-2, que arranca en caliente desde un modelo de secuencia, o ESM3, que utiliza pipelines en cascada).

3. Contribuciones Clave

1. Tokenizador Compacto y de Extremo a Extremo: Introducción de Yeti, un tokenizador de una sola etapa entrenado de extremo a extremo con flow matching, evitando el paradigma de dos etapas "congelar codificador" común en trabajos anteriores.
2. Alta Utilización del Libro de Códigos: Yeti logra una superior utilización del libro de códigos y diversidad de tokens en comparación con modelos más grandes (ESM3, DPLM-2) y otros métodos de cuantización (FSQ, VQ-VAE), a pesar de tener significativamente menos parámetros.
3. Eficiencia de Parámetros: Yeti logra una precisión de reconstrucción competitiva con solo 62.5M de parámetros, aproximadamente 10 veces menos que ESM3 (648M) y 2 veces menos que DPLM-2 (118M).
4. Co-Generación de Prueba de Concepto: Demostración de que un modelo multi-modal compacto entrenado desde cero sobre tokens Yeti puede realizar la co-generación incondicional de secuencias y estructuras proteicas, logrando resultados comparables a modelos 10 veces más grandes.
5. Análisis de la Dinámica de Plegamiento: Identificación de una emergencia jerárquica de la estructura secundaria durante la trayectoria de decodificación basada en flow matching, caracterizada por una compactación global temprana seguida de un fenómeno de "compromiso tardío" donde la topología global se consolida en los pasos finales.

4. Resultados

4.1 Análisis del Libro de Códigos

• Utilización: En un libro de códigos de 8,192 entradas, Yeti logra alta perplejidad y entropía, superando significativamente a ESM3 y Kanzi. En el conjunto de datos CATH, Yeti utiliza ~66% de su capacidad (Perplejidad = 5414), con una diversidad intra-estructura ($\bar{U}$) de 98.8%, saturando casi la métrica e indicando asignaciones de tokens altamente distintas para los residuos.
• Comparación: Yeti supera a DPLM-2 en la utilización del libro de códigos a pesar de que ambos utilizan LFQ, atribuido al paradigma de entrenamiento de extremo a extremo frente al enfoque de codificador congelado de DPLM-2.

4.2 Precisión de Reconstrucción

• Métricas: Yeti logra un puntaje TM de 0.96 en CAMEO, 0.97 en CASP14, y 0.95 en CASP15 y CATH.
• Eficiencia: Aunque ESM3 muestra un RMSD ligeramente menor (debido a su escala masiva), los puntajes TM de Yeti son comparables o superiores, indicando que recupera efectivamente la topología global correcta. El decodificador basado en flujo de Yeti produce muestras con variación en las posiciones atómicas locales (penalizando el RMSD) pero con topología global consistente, una característica beneficiosa para la generación de múltiples conformaciones.

4.3 Co-Generación Incondicional

• Rendimiento: Un modelo de 224M de parámetros entrenado desde cero sobre 2.6M de proteínas (utilizando tokens Yeti) logró un puntaje TM de auto-consistencia (scTM) de 0.70 y un pLDDT de 76.12.
• Comparación: Este rendimiento es comparable al de DPLM-2 (scTM 0.87) y significativamente mejor que ESM3-Open (scTM 0.46), a pesar de que DPLM-2 depende de un modelo de secuencia preentrenado de 650M de parámetros y ESM3 utiliza un pipeline en cascada.
• Diseñabilidad: ProteinMPNN recuperó el 30.6% de las secuencias generadas originales a partir de las estructuras generadas, confirmando la consistencia geométrica y la diseñabilidad de los pares co-generados.

4.4 Dinámica de la Trayectoria de Flow

El análisis de la trayectoria de decodificación reveló una emergencia no monótona de características:

• Fase Temprana ($t < 0.6$): Compactación global rápida (disminución del Radio de Giro).
• Fase Tardía ($t > 0.9$): Coalescencia rápida de elementos de estructura secundaria ($\alpha$-hélices y $\beta$-láminas) y consolidación de la topología global.
  Esto sugiere que el tokenizador desacopla efectivamente la densidad a granel de los detalles estructurales de alta frecuencia.

5. Significado y Afirmaciones

El artículo afirma que Yeti sirve como una base compacta y expresiva para entrenar modelos proteicos multi-modales. Su significado principal radica en demostrar que:

1. La Tokenización Discreta es Viable para la Generación: La co-generación incondicional de alta calidad de secuencia y estructura es posible sin preentrenamiento masivo o pipelines en cascada, siempre que el tokenizador sea expresivo y el objetivo de entrenamiento (flow matching) esté alineado con los objetivos generativos.
2. Eficiencia sobre Escala: Un modelo más pequeño cuidadosamente diseñado (62.5M de parámetros) puede superar o igualar a modelos mucho más grandes (600M+) en diversidad de tokens y reconstrucción, desafiando la noción de que la escala es el único camino hacia el rendimiento en el modelado de proteínas.
3. Representación Unificada: La capacidad de entrenar un modelo desde cero sobre datos conjuntos de secuencia-estructura valida el potencial de una representación molecular universal que capture el acoplamiento intrínseco entre secuencias de aminoácidos y estructuras 3D.

Los autores posicionan a Yeti no como una solución final para todas las tareas de diseño de proteínas, sino como una prueba de concepto que establece la viabilidad de su estrategia de tokenización. Reconocen limitaciones, como la dependencia de estructuras de conformación única y la necesidad de mayor escalado para igualar el rendimiento de modelos especializados y altamente ajustados como La-Proteina. Sin embargo, afirman que Yeti proporciona una columna vertebral escalable, eficiente y efectiva para futuros modelos proteicos multi-modales capaces de co-generar secuencias y estructuras altamente plausibles.