Resumen Técnico: Beneficios Sinérgicos de la Generación Conjunta de Moléculas y la Predicción de Propiedades

Formulación del Problema

El artículo aborda un desafío de larga data en el aprendizaje automático: desarrollar modelos que sobresalgan simultáneamente en tareas generativas y predictivas. En el ámbito del diseño de fármacos, el éxito depende de equilibrar la generación de moléculas novedosas a partir de espacios químicos inexplorados con una predicción de propiedades robusta que pueda extrapolarse a estas nuevas moléculas. Actualmente, la generación de moléculas y la predicción de propiedades se abordan predominantemente por separado. Si bien los modelos basados en transformadores son el estado del arte para ambas tareas individualmente, el entrenamiento conjunto plantea desafíos arquitectónicos y de optimización significativos. Combinar las partes generativas y predictivas en un único modelo a menudo conduce a una sobre-regularización o a interferencia de gradientes entre los objetivos generativos y predictivos, lo que hace que los modelos renuncien a los beneficios sinérgicos potenciales del aprendizaje conjunto.

Metodología: Hyformer

Para abordar estos desafíos, los autores proponen Hyformer, un modelo conjunto basado en transformadores que unifica las funcionalidades generativas y predictivas dentro de un único conjunto de parámetros compartidos.

Arquitectura

Hyformer combina un decodificador transformador autoregresivo con un codificador transformador bidireccional. La innovación central reside en un mecanismo de atención alterno controlado por un token de tarea $[TASK] \in \{[LM], [PRED], [MLM]\}$:

• Modo Generativo ($[LM]$): El modelo actúa como un decodificador utilizando una máscara de atención auto-referencial causal ($M_\rightarrow$), restringiendo que cada token atienda solo a su contexto izquierdo.
• Modo Predictivo/Pre-entrenamiento ($[PRED], [MLM]$): El modelo actúa como un codificador utilizando una máscara de atención auto-referencial bidireccional ($M_\leftrightarrow$), permitiendo el acceso completo al contexto.

Este esquema alterno alivia la interferencia de gradientes típica en los modelos conjuntos. El Jacobiano de la función softmax bajo una máscara causal altera los gradientes del módulo de atención de manera diferente a como lo hace bajo una máscara bidireccional, evitando que los objetivos entren en conflicto durante las actualizaciones.

Estrategia de Entrenamiento

Hyformer emplea un proceso de entrenamiento en dos etapas:

1. Pre-entrenamiento Conjunto: El modelo se pre-entrena con datos no etiquetados utilizando una función de pérdida combinada:
   $$\ell_{Hyformer} = \ell_{LM} + \mu\ell_{MLM} + \eta\ell_{PRED}$$
   Aquí, $\ell_{LM}$ es la pérdida de modelado de lenguaje, $\ell_{MLM}$ es la pérdida de modelado de lenguaje enmascarado y $\ell_{PRED}$ es una pérdida predictiva basada en descriptores moleculares computables analíticamente (por ejemplo, peso molecular, hidrofobicidad). El equilibrio entre tareas se controla mediante un vector de probabilidades de tarea $p_{[TASK]}$.
2. Ajuste Fino: El modelo se ajusta fino en tareas posteriores etiquetadas utilizando una pérdida conjunta:
   $$\ell_{Hyformer} = \ell_{LM} + \lambda\ell_{pred}$$
   El modelo alterna entre tareas generativas y predictivas durante el ajuste fino basándose en las probabilidades de tarea, omitiendo la pérdida de modelado de lenguaje enmascarado para centrarse en la tarea posterior mientras mantiene las capacidades generativas.

Muestreo

Hyformer soporta tanto muestreo incondicional como condicional:

• Incondicional: Muestra moléculas $x \sim p_\theta(x)$ utilizando la parte autoregresiva.
• Condicional: Para generar moléculas que satisfagan una condición de propiedad $Y$, el modelo muestrea $K$ candidatos de $p_\theta(x)$ y acepta aquellos donde el predictor entrenado conjuntamente $p_\theta(y|x)$ los clasifica como poseedores de la propiedad $Y$. Esto se enmarca como una variante del muestreo mejor-de-$K$, aprovechando el espacio de representación unificado para una alineación más estrecha entre la generación y el control.

Contribuciones Clave

1. Nuevo Modelo Conjunto: Introducción de Hyformer, que unifica tareas generativas y predictivas en un único conjunto de parámetros utilizando un mecanismo de atención alterno.
2. Demostración de Beneficios Sinérgicos: Evidencia empírica que muestra que Hyformer supera a las líneas base en:
   • Generación condicional de moléculas.
   • Predicción de propiedades fuera de distribución (OOD).
   • Aprendizaje de representaciones moleculares (mediante sondeo).
3. Paridad de Rendimiento: Demostrar que Hyformer rivaliza con el rendimiento de modelos puramente generativos y puramente predictivos de última generación, probando que el entrenamiento conjunto no requiere un sacrificio en el rendimiento de tareas individuales.
4. Aplicación del Mundo Real: Un caso de uso en el descubrimiento de nuevos péptidos antimicrobianos (AMP), demostrando la utilidad del modelado conjunto en el diseño de fármacos.

Resultados Experimentales

Generación Condicional de Moléculas

Evaluado en el conjunto de datos GuacaMol, el Hyformer ajustado fino conjuntamente logró la menor Desviación Absoluta Media (MAD) y Desviación Estándar (SD) en las propiedades objetivo (QED, SA, logP) mientras mantenía una alta validez. Superó significativamente a las variantes no conjuntas y a líneas base como MolGPT y GraphGPT, confirmando que el ajuste fino conjunto mejora la controlabilidad.

Predicción de Propiedades Fuera de Distribución (OOD)

En la tarea de Identificación de Éxitos del benchmark Lo-Hi, Hyformer logró la mayor AUPRC media en todos los conjuntos de datos (DRD2, HIV, KDR, Sol), superando a las líneas base basadas en huellas dactilares (k-NN, GB, SVM, MLP) y a modelos de aprendizaje profundo (Chemformer, Graphformer). Esto demuestra la capacidad del modelo para generalizar a andamios moleculares no vistos durante el entrenamiento.

Aprendizaje de Representaciones

Utilizando un nuevo protocolo de sondeo en el benchmark MoleculeNet, las representaciones extraídas de Hyformer pre-entrenado conjuntamente fueron las más predictivas tanto en el sondeo lineal como en el de KNN. Hyformer logró el mejor rendimiento en 4 de 10 conjuntos de datos para el sondeo lineal y en 5 de 10 para el sondeo KNN, superando a modelos predictivos de última generación como Uni-Mol y otros modelos conjuntos.

Generación Incondicional y Diseño de AMP

• Generación Incondicional: En los benchmarks GuacaMol y MOSES, Hyformer se mantuvo consistentemente en la frontera de Pareto, equilibrando la fidelidad distribucional (FCD, divergencia KL), validez y unicidad, rivalizando con modelos basados en SMILES y basados en grafos de última generación.
• Diseño de Péptidos Antimicrobianos: En una tarea real de diseño de fármacos, Hyformer generó péptidos con alta aptitud y probabilidad de AMP, superando significativamente a las líneas base (PepCVAE, AMPGAN, HydrAMP, AMP-Diffusion) mientras mantenía una perplejidad y diversidad competitivas. El análisis de atención reveló que el modelo priorizó aminoácidos biológicamente relevantes (Arginina, Lisina, Triptófano, Prolina) asociados con la actividad antimicrobiana.

Significado y Afirmaciones

El artículo afirma que Hyformer unifica con éxito la generación molecular y la predicción de propiedades para representaciones basadas en SMILES, abriendo vías para su integración en pipelines de descubrimiento de fármacos del mundo real. Los autores enfatizan que los beneficios sinérgicos del aprendizaje conjunto, específicamente un mejor control sobre el proceso generativo, robustez ante datos OOD y aprendizaje de representaciones predictivas de alto nivel, son alcanzables sin sacrificar el rendimiento de las tareas individuales.

Los autores reconocen compensaciones inherentes, señalando que los parámetros compartidos pueden limitar la especialización específica de la tarea. Sugieren que trabajos futuros podrían explorar arquitecturas de atención dinámicas o modulares para asignar capacidad de manera más flexible. Además, señalan que la medida en que estos beneficios sinérgicos se aplican a otras modalidades, como estructuras 3D o transcriptómica, sigue siendo una pregunta abierta. El trabajo se presenta como un paso hacia herramientas que facilitan el diseño de nuevos medicamentos potenciales, en lugar de una solución para aplicaciones maliciosas.