Hierarchical Bayesian Flow Networks for Molecular Graph Generation
Xiong, Chen, Li et al.
Molecular graph generation is essentially a classification generation problem, aimed at predicting categories of atoms and bonds. Currently, prevailing paradigms such as continuous diffusion models are trained to predict continuous numerical values, treating the training process as a regression task. However, the final generation necessitates a rounding step to convert these predictions back into discrete classification categories, which is intrinsically a classification operation. Given that the rounding operation is not incorporated during training, there exists a significant discrepancy between the model's training objective and its inference procedure. As a consequence, an excessive emphasis on point-wise precision can lead to overfitting and inefficient learning. This occurs because considerable efforts are devoted to capturing intra-bin variations that are ultimately irrelevant to the discrete nature of the task at hand. Such a flaw results in diminished molecular diversity and constrains the model's generalization capabilities. To address this fundamental limitation, we propose GraphBFN, a novel hierarchical coarse-to-fine framework based on Bayesian Flow Networks that operates on the parameters of distributions. By innovatively introducing Cumulative Distribution Function, GraphBFN is capable of calculating the probability of selecting the correct category, thereby unifying the training objective with the sampling rounding operation. We demonstrate that our method achieves superior performance and faster generation, setting new state-of-the-art results on the QM9 and ZINC250k molecular graph generation benchmarks.
academic
Redes de Flujo Bayesiano Jerárquico para Generación de Grafos Moleculares
La generación de grafos moleculares es esencialmente un problema de generación clasificatoria, cuyo objetivo es predecir las categorías de átomos y enlaces químicos. Los modelos de difusión continua predominantes actualmente tratan el proceso de entrenamiento como una tarea de regresión, prediciendo valores numéricos continuos, pero requieren una operación de redondeo para convertirlos a categorías clasificatorias discretas durante la generación final. Dado que el proceso de redondeo no se incluye durante el entrenamiento, existe una divergencia significativa entre el objetivo de entrenamiento del modelo y el proceso de inferencia, lo que resulta en sobreajuste, baja eficiencia de aprendizaje y diversidad molecular reducida. Para resolver esta limitación fundamental, los autores proponen GraphBFN, un marco jerárquico de grueso a fino basado en redes de flujo bayesiano, que introduce innovadoramente funciones de distribución acumulada para calcular la probabilidad de seleccionar la categoría correcta, unificando así el objetivo de entrenamiento con la operación de redondeo en el muestreo.
Existe un problema fundamental de inconsistencia entre entrenamiento e inferencia en la generación de grafos moleculares:
Fase de Entrenamiento: Los modelos de difusión continua mapean categorías discretas de átomos/enlaces al espacio continuo, optimizando predicciones de valores continuos mediante pérdida de regresión
Fase de Inferencia: Se requiere convertir valores predichos continuos a categorías discretas mediante redondeo duro
Inconsistencia: El modelo no considera reglas de redondeo durante el entrenamiento, causando que se enfoque excesivamente en variaciones dentro de clases mientras ignora la naturaleza discreta
La generación de grafos moleculares es una tecnología clave en el descubrimiento de fármacos, afectando optimización molecular, predicción de afinidad de unión fármaco-objetivo y otras tareas posteriores
La inconsistencia de los métodos existentes resulta en diversidad molecular reducida y capacidad de generalización limitada
Incluso pequeños errores de regresión pueden resultar en resultados de clasificación completamente incorrectos
Modelos de Difusión Discreta: Aunque son apropiados para estructuras de grafos discretos, sacrifican la suavidad de la representación continua y las características de generación dinámica
Modelos de Difusión Continua: El objetivo de entrenamiento se desvincula del proceso de inferencia, siendo propenso al sobreajuste en variaciones dentro de clases irrelevantes
Redes de Flujo Bayesiano Tradicionales: Asumen que todas las categorías están equidistantes en el símplex de probabilidad, resultando en convergencia lenta y mayor ruido
Primera aplicación de redes de flujo bayesiano a generación de grafos moleculares, mejorando los efectos de generación mediante supervisión de representación molecular jerárquica
Introducción innovadora de funciones de distribución acumulada (CDF), calculando probabilidades de categorías en lugar de ajustar valores específicos, unificando el objetivo de entrenamiento con la operación de redondeo en el muestreo
Propuesta de marco jerárquico de grueso a fino, capturando simultáneamente conectividad atómica local y topología molecular global mediante representaciones multiescala
Implementación de entrenamiento y muestreo más rápidos, logrando rendimiento de última generación en los puntos de referencia QM9 y ZINC250k, con reducción significativa en pasos de muestreo
Representación Multiescala: Utiliza DiffPool para construir L capas de agrupamiento, generando una representación piramidal del grafo molecular
Generación de Abajo hacia Arriba: Comienza con generación incondicional desde la capa más gruesa, refinando progresivamente hasta el grafo atómico completo
Transmisión Condicional: El módulo de sobremuestreo ϕ1(l) en cada capa convierte la salida de la capa gruesa en condición c(l) para la capa fina
El artículo cita trabajos importantes en el campo, incluyendo:
Graves et al. (2023): Trabajo original de redes de flujo bayesiano
Vignac et al. (2023): Método de difusión discreta DiGress
Jo, Lee, and Hwang (2022): Modelo de difusión de puntuación GDSS
Ying et al. (2018): Método de agrupamiento jerárquico de grafos DiffPool
Evaluación General: Este es un artículo de investigación de alta calidad que identifica y resuelve exitosamente un problema central en la generación de grafos moleculares. Mediante el innovador mecanismo CDF y marco jerárquico, mejora significativamente el rendimiento práctico mientras mantiene rigor teórico. Aunque hay espacio para mejora en profundidad de análisis teórico y escala experimental, sus contribuciones son suficientes para impulsar el desarrollo del campo.