Toward Better Geometric Representations for Molecule Generative Models

Resumen Técnico: LENSEs – Mejora Latente para Codificaciones Estructurales No Suaves

Planteamiento del Problema
La generación condicionada a representaciones geométricas (GeoRCG) ha surgido como un paradigma eficiente para la generación de moléculas 3D, desacoplando el modelado de representaciones moleculares de la generación de estructuras 3D. En este marco, un codificador molecular preentrenado (por ejemplo, UniMol, Frad) define un espacio latente, un Modelo de Difusión de Representación (RDM) muestrea un vector a nivel de molécula desde este espacio, y un generador reconstruye la estructura 3D condicionada a este vector.

Sin embargo, la eficacia de GeoRCG está fundamentalmente limitada por la calidad del espacio latente proporcionado por los codificadores preentrenados. Los autores identifican dos cuellos de botella críticos:

1. Variedades de Representación No Suaves: Los codificadores preentrenados, optimizados principalmente para la predicción de propiedades mediante eliminación de ruido de coordenadas, producen representaciones con constantes de Lipschitz altas. Pequeñas perturbaciones en la conformación molecular conducen a desplazamientos desproporcionadamente grandes en el espacio latente. Esta no suavidad hace que la distribución sea difícil de modelar para el RDM, a menudo requiriendo mecanismos compensatorios como la inyección de ruido ajustada manualmente.
2. Semántica Multinivel Subexplotada: Los codificadores preentrenados exhiben características jerárquicas donde las capas superficiales capturan motivos estructurales locales (por ejemplo, grupos funcionales) y las capas más profundas codifican semánticas globales. Las pipelines existentes de GeoRCG suelen utilizar únicamente la representación de la capa final, descartando la rica información semántica preservada en las capas anteriores. En consecuencia, el generador debe reaprender estas semánticas desde cero, aumentando la dificultad de optimización.

Metodología: LENSEs
Para abordar estos problemas, los autores proponen LENSEs (Mejora Latente para Codificaciones Estructurales No Suaves), un marco que refina activamente el espacio latente de condicionamiento durante el entrenamiento generativo. En lugar de depender pasivamente de un codificador congelado, LENSEs introduce tres mecanismos complementarios:

1. Cabeza de Representación con Agrupación de Capas Aprendible:
   En lugar de utilizar una única salida determinista del codificador, LENSEs agrega características de múltiples capas. Emplea un mecanismo de agrupación de capas aprendible para ponderar representaciones de diferentes profundidades del codificador ($H^{(l)}$). Estas características agrupadas se mapean a los parámetros de una posterior gaussiana diagonal ($\mu, \sigma$) mediante cabezas de proyección. El código latente $z$ se muestrea utilizando el truco de reparametrización. Esta cabeza se entrena conjuntamente con el generador para extraer representaciones que sean tanto semánticamente ricas como favorables a la generación, regularizadas por una pérdida de divergencia KL para garantizar un espacio latente compacto.

2. Pérdida de Percepción Molecular ($L_{PERC}$):
   Inspirada en las pérdidas de percepción en visión por computadora, esta pérdida alinea la representación semántica de alto nivel de la molécula desruidada (generada) con la de la molécula limpia de verdad fundamental dentro del espacio de características del codificador congelado. Al pasar tanto la molécula generada como la limpia a través del codificador y minimizar la distancia entre sus representaciones agrupadas a nivel de grafo, se incentiva al generador a preservar patrones estructurales significativos (por ejemplo, grupos funcionales) en lugar de simplemente minimizar errores de coordenadas de bajo nivel. Se aplica una operación de stop-gradient a la representación de la verdad fundamental.

3. Pérdida de Alineación de Representación (REPA) ($L_{REPA}$):
   Motivada por la observación de que los estados ocultos del generador se alinean naturalmente con las características del codificador, esta pérdida alinea explícitamente las características a nivel de nodo del generador con las del codificador preentrenado. Un proyector aprendible mapea las características del generador al espacio de características del codificador, y se maximiza la similitud del coseno entre los nodos correspondientes. Esto reduce la brecha semántica entre el preentrenamiento y la generación, acelerando el entrenamiento y mejorando la adquisición de conocimiento químicamente significativo.

El objetivo total de entrenamiento combina la pérdida de generación base ($L_{gen}$) con la pérdida KL, la pérdida perceptual y la pérdida REPA.

Resultados Clave
Los autores evalúan LENSEs en los conjuntos de datos GEOM-DRUG y QM9, comparándolo con líneas base de última generación, incluyendo GeoRCG, CanonFlow y SemlaFlow.

• Calidad de Generación: En el desafiante conjunto de datos GEOM-DRUG, LENSEs logra un rendimiento de última generación. Alcanza una validez del 97.28% y una estabilidad molecular del 98.51%, superando al mejor GeoRCG anterior (95.3% de validez, 97.20% de estabilidad) y a CanonFlow (95.9% de validez). También logra las métricas de energía y tensión más bajas, indicando una plausibilidad física superior.
• Suavidad de Representación: El análisis de la constante de Lipschitz revela que el espacio latente refinado es significativamente más suave que el espacio del codificador crudo. En GEOM-DRUG, la constante de Lipschitz empírica se reduce en 4.6 veces, confirmando que la variedad de representación es más apta para el modelado de difusión.
• Riqueza de Representación: La rango efectivo de la representación aumenta de 65.32 (codificador crudo) a 88.42 (LENSEs), lo que indica una reducción del colapso dimensional y un espacio latente más rico y estructurado.
• Utilidad en Tareas Posteriores: Cuando se utiliza para tareas de regresión posteriores en propiedades de QM9 con un codificador y una cabeza de representación congelados, LENSEs supera al codificador crudo en todas las propiedades, demostrando que el propio proceso de entrenamiento generativo actúa como un paradigma de preentrenamiento efectivo para refinar representaciones moleculares.
• Eficiencia de Muestreo: LENSEs demuestra una mayor eficiencia de muestreo, logrando alta validez y estabilidad con menos pasos de generación en comparación con la línea base.

Significado y Afirmaciones
El artículo afirma que la generación de moléculas 3D de alta calidad no depende solo de la arquitectura del generador, sino críticamente de la calidad de la representación de condicionamiento. LENSEs demuestra que el espacio latente de los codificadores preentrenados puede mejorarse sistemáticamente para ser más suave, más informativo semánticamente y mejor alineado con los objetivos generativos.

Los autores postulan que su trabajo establece una relación bidireccional entre la calidad de la generación y la calidad de la representación. Específicamente, sugieren que el entrenamiento generativo con objetivos de alineación puede servir como un paradigma potencial de preentrenamiento para codificadores moleculares, refinando representaciones para que sean más predictivas y útiles para tareas posteriores. El trabajo destaca que las representaciones moleculares útiles no son inherentemente favorables a la generación y requieren un refinamiento explícito para desbloquear su pleno potencial en el modelado generativo.