Interpreting GFlowNets for Drug Discovery: What probes can and cannot show

Resumen Técnico: Interpretación de GFlowNets para el Descubrimiento de Fármacos

Planteamiento del Problema
Las Redes de Flujo Generativo (GFlowNets), particularmente aquellas entrenadas para el diseño molecular como SynFlowNet, ofrecen un marco poderoso para explorar el espacio químico adhiriéndose a restricciones sintéticas. Sin embargo, sus políticas de decisión internas permanecen opacas, lo que limita su adopción en el descubrimiento de fármacos, donde los químicos medicinales requieren razonamientos interpretables para las estructuras propuestas. Los métodos de interpretabilidad existentes suelen estar diseñados para espacios latentes continuos o modelos basados en la verosimilitud (likelihood), no para las políticas discretas, secuenciales y basadas en grafos de las GFlowNets. Además, existe una necesidad crítica de distinguir entre el conocimiento químico genuinamente aprendido por la política y la información que está disponible de forma trivial a través de la arquitectura del grafo y la featurización de átomos.

Metodología
Los autores introducen un marco de interpretabilidad multiescala aplicado a SynFlowNet (entrenada con una función de recompensa QED) utilizando embeddings congelados de un checkpoint específico. El marco integra tres enfoques complementarios:

1. Saliencia Basada en Gradientes y Contrafactuales: Se aplican Gradientes Integrados (IG) a la log-probabilidad de la acción "Stop" para generar mapas de saliencia a nivel de átomo. Los átomos de alta saliencia se agrupan en motivos (componentes conectados y sistemas de anillos). Posteriormente, se aplican perturbaciones contrafactuales a estos motivos utilizando reglas de transformación de RDKit (por ejemplo, intercambio de halógenos, modificación de amidas) para medir los cambios en QED ($\Delta$QED).
2. Autoencoders Dispersos (SAEs):
   • Subcompleto (MLP-SAE): Se utiliza un autoencoder compresivo (256 $\to$ 128) para identificar factores latentes dominantes y dispersos, y para probar su correlación con propiedades fisicoquímicas (por ejemplo, polaridad, lipofilicidad).
   • Sobrecompleto (BatchTopK SAE): Se emplea un autoencoder sobrecompleto (256 $\to$ 1024) con un mecanismo de dispersión BatchTopK ($k=64$) para recuperar detectores de características de grano fino, incluyendo aquellos en superposición, sin regularización $\ell_1$.
3. Sondeo y Controles: Se entrenan clasificadores de alimentación hacia adelante (feedforward) superficiales y sondas estrictamente lineales (regresión Ridge/logística) sobre los embeddings para predecir la presencia de motivos químicos (grupos funcionales, halógenos) y propiedades fisicoquímicas. Crucialmente, el estudio emplea controles rigurosos:
   • Líneas Base de Descriptores: Comparación del rendimiento de las sondas contra descriptores de RDKit y huellas dactilares (fingerprints) ECFP4.
   • Controles de Etiquetas Aleatorias: Asegurar que las sondas no memoricen ruido.
   • Divisiones Disjuntas de Andamiaje (Scaffold-Disjoint Splits): Evitar la fuga de andamiaje en las divisiones de entrenamiento/prueba.
   • Control de Red No Entrenada: Ejecución de una red idéntica con la misma arquitectura, pero con inicialización aleatoria, para aislar la contribución del entrenamiento frente a la arquitectura.
   • Ablación de Características: Poner a cero características individuales para probar la especificidad en la decodificación descendente.

Resultados Clave

• Saliencia y Contrafactuales: Los Gradientes Integrados resaltan regiones químicamente significativas (sustituyentes polares, fragmentos aromáticos). Las ediciones contrafactuales a estos motivos producen cambios pequeños pero medibles en QED (por ejemplo, +0.015 a +0.051), aunque los tamaños del efecto son pequeños en relación con el rango de QED, y la comparación controlada contra ediciones aleatorias sigue siendo un trabajo futuro.
• Estructura Latente: El SAE subcompleto identifica factores fuertemente correlacionados con el tamaño y la polaridad, pero un predictor MLP sobre estos factores tiene un desempeño pobre en QED compuesto ($R^2=0.25$) comparado con propiedades individuales.
• Reconstrucción y Dispersión: El SAE sobrecompleto BatchTopK reconstruye los embeddings con un NMSE bajo (0.0015) y mantiene una estricta dispersión de media de lote de 64 características activas. Supera a una línea base lineal emparejada (NMSE 0.0026) con niveles de actividad similares.
• Decodificabilidad Lineal: Los embeddings decodifican bien las propiedades fisicoquímicas (por ejemplo, QED $R^2 \approx 0.59$, tamaño $R^2 \approx 0.996$). Sin embargo, estas puntuaciones no son significativamente mejores que las logradas por una línea base de descriptores RDKit "libre de fugas" bajo divisiones aleatorias.
• El Control de la Red No Entrenada (Hallazgo Decisivo): Una red con inicialización aleatoria de la misma arquitectura decodifica cada propiedad sondeada esencialmente tan bien como la política entrenada. Para la propiedad principal (similitud con fármacos/QED), el entrenamiento proporciona una ganancia insignificante de solo +0.005 (0.581 no entrenada vs. 0.586 entrenada), lo cual cae dentro de la varianza de las semillas aleatorias. En cuanto al tamaño molecular, la red no entrenada es marginalmente mejor.
• Especificidad de Características: En el diccionario sobrecompleto, ejecuciones individuales revelan detectores químicamente enriquecidos (por ejemplo, características separadas para yodo, cloro, bromo, flúor y boro). La ablación de una sola característica muestra efectos específicos de la propiedad (ratios de especificidad de 8–23$\times$), pero los índices de las características no son reproducibles entre semillas; solo un pequeño núcleo de direcciones es estable.

Significancia y Reivindicaciones
El resultado más trascendental del artículo es un control negativo: la decodificabilidad lineal de las propiedades fisicoquímicas de los embeddings de GFlowNet refleja la arquitectura del grafo y la featurización de átomos, no el conocimiento químico aprendido durante el entrenamiento.

• Lo que NO se apoya: Las altas puntuaciones de las sondas (incluso con AUROCs cercanos a la perfección) no son evidencia de comprensión química adquirida o de razonamiento de relación estructura-actividad aprendido. La capacidad de decodificar propiedades es en gran medida un artefacto de la representación de entrada y la arquitectura de la red, como lo demuestra el hecho de que la red no entrenada iguala el desempeño de la política entrenada.
• Lo que SÍ se apoya:
  1. El diccionario sobrecompleto proporciona una descomposición válida de dispersión exacta de los embeddings.
  2. Las ejecuciones individuales contienen detectores de subestructuras químicamente enriquecidas (por ejemplo, detectores por halógeno) que son estables a través de un pequeño subconjunto de semillas.
  3. Estas características tienen efectos intervencionistas de propiedades específicas en la decodificación de la sonda.
  4. La similitud con fármacos (QED) se vuelve más accesible linealmente en el espacio latente disperso que en el embedding denso.

Conclusión
Los autores concluyen que los métodos de interpretabilidad deben controlarse rigurosamente contra las líneas base arquitectónicas. Para SynFlowNet, el marco identifica con éxito detectores de características de grano fino y valida la robustez de la representación dispersa. Sin embargo, el estudio advierte explícitamente contra la interpretación de altas puntuaciones de decodificabilidad como prueba de química aprendida. El artículo limita sus afirmaciones a la existencia de un núcleo reproducible de características estables y químicamente enriquecidas dentro del diccionario, mientras afirma que la capacidad del modelo para predecir propiedades fisicoquímicas es principalmente una función de su arquitectura y no de su entrenamiento. El trabajo futuro debería explorar si estos detectores persisten cuando los SAE se entrenan sobre embeddings de redes no entrenadas y extender el marco a entornos de múltiples objetivos.