Resumen Técnico: SinAE – Un Autoencoder de Flujo de Arquitectura Única para Sistemas Atómicos de Dominios Cruzados

1. Planteamiento del Problema

La generación de estructuras atómicas 3D es crítica para el descubrimiento de fármacos, el diseño de materiales y la ingeniería de proteínas. Sin embargo, las canalizaciones generativas actuales están fragmentadas en tres dominios distintos: moléculas pequeñas, cristales y proteínas. Cada dominio depende de arquitecturas especializadas adaptadas a sus prioridades estructurales específicas:

• Moléculas: Suelen utilizar redes de paso de mensajes equivariantes para gestionar la conectividad de los enlaces y las reglas de valencia.
• Cristales: Requieren el manejo de condiciones de contorno periódicas e invariancia de la base de la red.
• Proteínas: A menudo dependen de parametrizaciones a nivel de residuo (por ejemplo, marcos SE(3)) o mecanismos de atención de triángulo.

Esta fragmentación exacerba la escasez de datos, ya que los conjuntos de datos por dominio (por ejemplo, QM9 con 134k moléculas, MP-20 con 45k cristales) son pequeños para los estándares modernos. Aunque el entrenamiento de dominios cruzados podría mitigar esto, los enfoques existentes enfrentan dos obstáculos principales:

1. Limitaciones de la Generación Directa: Generar directamente en el espacio de coordenadas 3D obliga a un único modelo a aprender simultáneamente composiciones válidas, geometría precisa y prioridades estructurales heterogéneas (periodicidad, backbones de residuos, valencia), lo que a menudo requiere una tokenización compleja y específica del dominio o una generación por etapas.
2. Limitaciones de los Autoencoders Latentes: Los autoencoders latentes de dominios cruzados existentes sufren problemas de fidelidad de reconstrucción (los decodificadores con pérdida introducen errores geométricos que las prioridades no pueden corregir) y complejidad de entrenamiento (requieren combinaciones de pérdida y ajuste de pesos específicos de cada dominio). Ningún método previo ha logrado un autoencoder de reconstrucción sin pérdidas, arquitectónicamente general y de objetivo unificado para los tres dominios simultáneamente.

2. Metodología: SinAE

Los autores presentan SinAE, un autoencoder de flujo de arquitectura única diseñado para moléculas, cristales y proteínas.

Arquitectura

• Backbone Unificado: SinAE utiliza un Transformer vanilla tanto para el codificador como para el decodificador. Evita explícitamente los operadores equivariantes, el paso de mensajes de grafos o los módulos específicos de dominio (por ejemplo, capas conscientes de la red o atención de triángulo).
• Caracterización de Entrada:
  • Coordenadas: Se utilizan coordenadas cartesianas puras sin descriptores invariantes (por ejemplo, distancias de pares) ni construcción de marcos canónicos.
  • Periodicidad: Para los cristales, los tres vectores de la red se adjuntan como "tokens de átomos virtuales" a la secuencia, lo que permite al modelo procesar la matriz de la red de la misma manera que las coordenadas atómicas.
  • Equivarianza: En lugar de una equivalencia arquitectónica, el modelo depende de la aumentación de rotación SO(3) y traslación aleatoria durante el entrenamiento para asegurar la robustez de la orientación.
• Tokenización: Los átomos se representan mediante la suma de embeddings de elementos aprendidos y proyecciones lineales de sus coordenadas. Para las proteínas, los residuos se tratan como tokens únicos en sus posiciones $C_\alpha$.

Objetivos de Entrenamiento

El modelo se entrena en dos etapas:

1. Autoencoder de Flujo de Coincidencia (Flow-Matching):
   • Codificador: Un Transformer estocástico mapea la estructura de entrada $G=(X, A)$ a un código latente por token $z$.
   • Decodificador: Un Transformer no equivariante reconstruye la estructura desde $z$ utilizando coincidencia de flujo iterativa. A diferencia de los decodificadores VAE de paso único, el decodificador refina la estructura a través de múltiples pasos de ODE.
   • Función de Pérdida: Un objetivo unificado combina el flujo de coincidencia condicional continuo para las coordenadas ($L_{coord}$) y el flujo de coincidencia discreto para los tipos de átomos ($L_{atom}$), además de un ligero regularizador KL ($L_{KL}$). No se requieren ponderaciones de pérdida ni pérdidas auxiliares específicas de dominio.
2. Prior de Coincidencia de Flujo Latente:
   • Un Diffusion Transformer (DiT) vanilla aprende la distribución $p(z)$ sobre las salidas del codificador.
   • Dado que el término KL empuja la posterior hacia una Gaussiana estándar, el prior opera sobre un espacio latente suave y aproximadamente Gaussiano sin restricciones de equivalencia.
   • Durante la generación, se utiliza una etiqueta de dominio como condición (mediante guía libre de clasificador) para dirigir al DiT hacia moléculas, cristales o proteínas.

3. Contribuciones Clave

1. Tokenización de Arquitectura Única: SinAE es la primera arquitectura que maneja moléculas, cristales y backbones de proteínas utilizando un único Transformer vanilla sin operadores equivariantes o específicos de dominio.
2. Reconstrucción Casi Sin Pérdidas: Al desplazar la carga de la reconstrucción hacia un decodificador de coincidencia de flujo iterativo, SinAE logra una RMSD de sub-mili-Ångström en todos los dominios:
   • Moléculas: 0.0002 Å en QM9, 0.0006 Å en GEOM-Drugs.
   • Cristales: 0.0017 Å en MP-20.
   • Proteínas: ~0.013 Å de CA-RMSD del backbone en CASP15.
   • Esto representa una mejora de 10 a 100 veces en la fidelidad de reconstrucción respecto a las líneas base latentes previas, con tasas de coincidencia de estructura del 100% en moléculas y cristales.
3. Transferencia Positiva de Dominios Cruzados: El entrenamiento conjunto en QM9 y MP-20 mejora estrictamente el rendimiento de reconstrucción y generación en ambos dominios en comparación con el entrenamiento de un solo dominio, proporcionando evidencia de que un latente atómico compartido captura regularidades geométricas transferibles (por ejemplo, distribuciones de longitud de enlace, patrones de coordinación).

4. Resultados Experimentales

Los autores evaluaron SinAE en benchmarks estándar (QM9, GEOM-Drugs, MP-20, AFDB-FS/CASP15, QMOF).

• Fidelidad de Reconstrucción: SinAE supera a todas las líneas base (incluyendo GeoLDM, CDVAE, ADiT, UAE-3D, ProteinAE) en RMSD y tasas de coincidencia de estructura. El decodificador iterativo resuelve la colocación atómica tempranamente y refina la geometría de los enlaces/estereoquímica más tarde, logrando una alta precisión sin relajación post-hoc.
• Generación Molecular: En QM9 y GEOM-Drugs, el prior DiT latente logra una alta validez (puntuaciones PB Valid de 0.97 y 0.93, respectivamente), superando o igualando a generadores equivariantes especializados (por ejemplo, EQGAT-diff, SemlaFlow).
• Generación de Cristales: En MP-20, el modelo conjunto SinAE logra las puntuaciones S.U.N. (Estable, Único, Nuevo) y M.S.U.N. más altas entre los métodos comparados, incluyendo aquellos con arquitecturas invariantes de la red. Demuestra una estabilidad termodinámica superior (estructuras metaestables) a pesar de utilizar un Transformer genérico.
• Generación de Proteínas: SinAE logra alta diseñabilidad (0.96 scRMSD $\leq$ 2.0 Å) y diversidad simultáneamente, superando a ProteinAE y LatentDiff. En tareas de co-diseño secuencia-estructura, logra puntuaciones competitivas de scTM y scRMSD en longitudes de 100 y 200.

5. Significado y Reivindicaciones

El artículo afirma que SinAE demuestra que la generación atómica de dominios cruzados puede unificarse bajo una única arquitectura simple sin sacrificar la fidelidad.

• Simplicidad vs. Rendimiento: Los autores argumentan que los operadores complejos específicos de dominio no son estrictamente necesarios para una generación de alta calidad si el decodificador es suficientemente potente (coincidencia de flujo iterativa) y el espacio latente está debidamente regularizado.
• Eficiencia de Datos: Los resultados proporcionan evidencia directa de que los latentes atómicos compartidos permiten que diferentes poblaciones químicas se beneficien mutuamente de sus datos, mejorando el rendimiento en regímenes de pocos datos.
• Desacoplamiento de Geometría y Prioridades: El trabajo sugiere que la complejidad geométrica puede ser manejada por el refinamiento iterativo del decodificador, permitiendo que el codificador y el prior operen sobre un espacio latente simplificado y agnóstico al dominio.

Limitaciones Reconocidas:

• La evaluación de proteínas está actualmente limitada a la generación a nivel de backbone; aún no se aborda el modelado de todos los átomos.
• El modelo actual no soporta la generación condicionada por propiedades, composición o longitud en un único checkpoint a través de los tres dominios simultáneamente.
• Las cadenas de proteínas más largas (>256 residuos) requieren el reentrenamiento con presupuestos de longitud más altos.