ConRetroBert: EMA Stabilized Dual Encoders for Template-Based Single-Step Retrosynthesis

Resumen Técnico de ConRetroBert

Formulación del Problema
La retrosíntesis de un solo paso basada en plantillas tiene como objetivo predecir conjuntos de reactivos para un producto objetivo mediante la selección y aplicación de una plantilla de reacción explícita (codificada como una transformación SMARTS). Aunque este enfoque ofrece interpretabilidad y trazabilidad inherentes, cruciales para la planificación de síntesis, los métodos basados en plantillas han sido históricamente considerados menos competitivos que los modelos libres de plantillas. Los autores argumentan que esta limitación no proviene del uso de las plantillas en sí, sino de la formulación de aprendizaje dominante: la clasificación global sobre una biblioteca de reglas de reacción con distribución de cola larga. Esta formulación genera tres problemas: las plantillas frecuentes dominan la pérdida, el denominador de softmax trata a las plantillas no observadas pero químicamente válidas como etiquetas negativas en competencia (incompletitud de etiquetas) y la capa de salida escala mal con bibliotecas grandes de plantillas.

Metodología: ConRetroBert
El artículo presenta ConRetroBert, un marco de codificador dual que reformula la retrosíntesis basada en plantillas como un problema de recuperación y clasificación densa en lugar de clasificación global. La arquitectura emplea dos codificadores Transformer independientes (uno para los SMILES del producto y otro para las plantillas SMARTS) que mapean las entradas en un espacio de incrustación compartido. El proceso de entrenamiento ocurre en dos etapas:

1. Etapa 1: Preentrenamiento Contrastivo
   El modelo aprende un espacio de representación conjunto producto-plantilla utilizando preentrenamiento contrastivo simétrico (pérdida InfoNCE). Para un minibatch de ejemplos emparejados, el codificador de productos se alinea con su plantilla correspondiente, mientras que otros ejemplos en el batch sirven como negativos. Esto establece un espacio de incrustación buscable donde las plantillas pueden recuperarse mediante búsqueda de vecinos más cercanos.

2. Etapa 2: Clasificación Listwise de Múltiples Positivos
   Esta etapa refina el espacio de recuperación en una política de propuesta de plantillas. En lugar de clasificar contra el vocabulario completo, el modelo clasifica un conjunto candidato $C(x)$ para cada producto. Este conjunto incluye:

   • Positivos Observados: Todas las plantillas registradas con el producto en los datos de entrenamiento.
   • Negativos Difíciles: Plantillas químicamente similares recuperadas de un índice FAISS.
   • Negativos en el Batch: Plantillas asociadas con otros productos en el batch.

   El modelo optimiza un objetivo de verosimilitud negativa logarítmica listwise de múltiples positivos, fomentando la separación de los positivos observados de los negativos difíciles químicamente similares. Esta formulación coincide mejor con el problema de decisión en tiempo de inferencia, donde el modelo debe elegir entre un pequeño conjunto de candidatos recuperados en lugar de toda la biblioteca.

Adaptación de Plantillas Estabilizada por EMA
Un desafío técnico clave abordado es la inestabilidad de la minería de negativos difíciles cuando se actualiza el codificador de plantillas. Si el codificador de plantillas cambia demasiado rápido, el índice de recuperación (utilizado para extraer negativos difíciles) se vuelve obsoleto. ConRetroBert resuelve esto utilizando un mecanismo de Media Móvil Exponencial (EMA):

• Un codificador de plantillas en vivo se actualiza mediante gradientes de la pérdida de clasificación.
• Un codificador sombra EMA de movimiento lento (actualizado con $\alpha = 0.999$) se utiliza para reconstruir el índice de recuperación y extraer negativos difíciles a nivel de época.
  Esto permite que las representaciones de plantillas se adapten al objetivo de clasificación sin desestabilizar el proceso de construcción de candidatos.

Contribuciones Clave

1. Formulación Basada en Recuperación: Los autores reformulan la retrosíntesis basada en plantillas como recuperación densa producto-plantilla seguida de clasificación listwise de conjuntos candidatos, reemplazando la clasificación global con un objetivo que se alinea con las decisiones en tiempo de inferencia.
2. Entrenamiento Estabilizado por EMA: Introducen un codificador de plantillas entrenable estabilizado por una sombra EMA, permitiendo una adaptación controlada de las representaciones de plantillas mientras se mantiene un índice de recuperación estable para la minería de negativos difíciles.
3. Pipeline de Evaluación Riguroso: El artículo proporciona un protocolo de evaluación integral que incluye clasificación de conjuntos de reactivos, aplicación de plantillas, deduplicación y transferencia controlada contra fugas de USPTO-Full a USPTO-50k.
4. Validación Empírica: Mediante análisis locales y de cola larga, los autores demuestran que la clasificación de conjuntos candidatos fortalece sustancialmente la retrosíntesis basada en plantillas, manteniendo al mismo tiempo la trazabilidad explícita de las reglas.

Resultados
Los experimentos se realizaron en el benchmark USPTO-50k (régimen local) y mediante aprendizaje por transferencia desde un punto de control USPTO-Full con control de fugas.

• Rendimiento Local en USPTO-50k:

  • Pasar de la Etapa 1 (recuperación contrastiva) a la Etapa 2 (clasificación con codificador de plantillas congelado) mejoró la precisión de reacción top-1 del 50.5% al 61.4%.
  • Añadir la adaptación de plantillas estabilizada por EMA (Etapa 2 EMA Mid) mejoró aún más la precisión top-1 al 62.4%.
  • El modelo alcanzó precisiones top-3, top-5 y top-10 de 81.6%, 85.3% y 87.8%, respectivamente.
  • La ganancia de rendimiento se atribuye principalmente a la formulación de clasificación de conjuntos candidatos de la Etapa 2, que separa mejor los positivos de los negativos difíciles que la alineación contrastiva por sí sola.

• Aprendizaje por Transferencia:

  • El ajuste fino desde un punto de control USPTO-Full con control de fugas alcanzó una precisión top-1 del 75.4% en USPTO-50k, demostrando la escalabilidad del enfoque con datos de preentrenamiento más grandes.

• Cola Larga y Fiabilidad:

  • La predicción basada en recuperación mostró una fuerza particular en la cola larga de plantillas raras, superando a la puntuación basada en clasificación en clases raras.
  • El análisis reveló que el 42.8% de las predicciones correctas de reactivos top-1 fueron "éxitos secundarios", generados por una plantilla diferente a la etiqueta de verdad fundamental registrada pero químicamente equivalente tras la aplicación y la deduplicación. Esto destaca que el modelo aprende a recuperar alternativas químicamente válidas, no solo la etiqueta específica registrada.

Significado y Afirmaciones
El artículo argumenta que la retrosíntesis basada en plantillas no necesita sacrificar el rendimiento predictivo por la interpretabilidad. Al cambiar de la clasificación global a la clasificación basada en recuperación con minería de negativos difíciles y estabilización EMA, ConRetroBert desafía la suposición de que una alta precisión requiere abandonar las plantillas de reacción explícitas. Los autores afirman que su método produce predicciones que no solo son competitivas con los modelos libres de plantillas y semi-plantillas de última generación, sino que también permanecen químicamente inspeccionables, trazables a reglas de transformación específicas y compatibles con las restricciones de planificación de síntesis aguas abajo. El trabajo sugiere que la "debilidad" de los métodos basados en plantillas es un problema de formulación solucionable mediante técnicas modernas de aprendizaje de recuperación, en lugar de una limitación inherente del uso de reglas explícitas.