2025-11-11T14:46:09.738382

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

Información Básica

  • ID del Artículo: 2510.10211
  • Título: Hierarchical Bayesian Flow Networks for Molecular Graph Generation
  • Autores: Yida Xiong, Jiameng Chen, Kun Li, Hongzhi Zhang, Xiantao Cai, Wenbin Hu (Facultad de Informática, Universidad de Wuhan)
  • Clasificación: cs.LG (Aprendizaje Automático)
  • Fecha de Publicación: 11 de octubre de 2025 (preimpresión en arXiv)
  • Enlace del Artículo: https://arxiv.org/abs/2510.10211

Resumen

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.

Antecedentes de Investigación y Motivación

Problema Central

Existe un problema fundamental de inconsistencia entre entrenamiento e inferencia en la generación de grafos moleculares:

  1. 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
  2. Fase de Inferencia: Se requiere convertir valores predichos continuos a categorías discretas mediante redondeo duro
  3. 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

Importancia del Problema

  • 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

Limitaciones de Métodos Existentes

  1. 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
  2. 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
  3. 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

Contribuciones Principales

  1. 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
  2. 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
  3. Propuesta de marco jerárquico de grueso a fino, capturando simultáneamente conectividad atómica local y topología molecular global mediante representaciones multiescala
  4. 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

Explicación Detallada del Método

Definición de la Tarea

Dado un grafo molecular G=(X,A)G = (X, A), donde:

  • X{0,,KX1}DX \in \{0, \ldots, K_X - 1\}^D: Matriz de características de DD átomos, provenientes de KXK_X categorías
  • A{0,,KA1}D×DA \in \{0, \ldots, K_A - 1\}^{D \times D}: Matriz de adyacencia, conteniendo características de KAK_A categorías de enlaces

El objetivo es aprender a generar nuevos grafos moleculares que se ajusten a la distribución molecular real.

Arquitectura del Modelo

1. Marco Jerárquico de Grueso a Fino

  • Representación Multiescala: Utiliza DiffPool para construir LL 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)\phi_1^{(l)} en cada capa convierte la salida de la capa gruesa en condición c(l)c^{(l)} para la capa fina

2. Mapeo de Representación de Grafos

Mapea categorías discretas k{0,,K1}k \in \{0, \ldots, K-1\} al espacio continuo [1,1][-1, 1]:

k_c = (2k + 1)/K - 1  # punto central
k_l = k_c - 1/K       # límite izquierdo  
k_r = k_c + 1/K       # límite derecho

3. Componentes de Red de Flujo Bayesiano

Distribución de Entrada: Modelada mediante distribución gaussiana

p_I(G|θ) = N(G|μ, ρ^{-1}I)

Distribución de Envío: Adición de ruido gaussiano

p_S(Y|G; α) = N(Y|G, α^{-1}I)

Distribución de Salida: Cálculo de probabilidades discretas mediante CDF

p_O^{(d)}(k|θ; t) = F(k_r|μ_x^{(d)}, σ_x^{(d)}) - F(k_l|μ_x^{(d)}, σ_x^{(d)})

Distribución de Recepción:

p_R(Y|θ; t, α) = ∏_{d=1}^D ∑_{k=0}^{K-1} p_O^{(d)}(k|θ; t)N(Y^{(d)}|k_c, α^{-1})

4. Innovación Clave: Mecanismo CDF

Utiliza función de distribución acumulada truncada para conectar distribuciones continuas con categorías discretas:

F(x|μ_x^{(d)}, σ_x^{(d)}) = {
  0,                    si x ≤ -1
  1,                    si x ≥ 1  
  1/2[1 + erf((x-μ_x^{(d)})/(√2σ_x^{(d)}))], en otro caso
}

Puntos de Innovación Técnica

  1. Consistencia Entrenamiento-Inferencia: CDF calcula directamente probabilidades discretas, evitando desajuste entre predicción continua y redondeo discreto
  2. Mapeo de Categorías No Equidistantes: A diferencia de BFN tradicional que asume categorías equidistantes, permite convergencia más rápida y suave
  3. Supervisión Multiescala: El marco jerárquico proporciona información estructural en diferentes granularidades, mejorando la calidad de generación
  4. Optimización Extremo a Extremo: Función de pérdida unificada optimiza simultáneamente la pérdida de generación BFN y la pérdida de agrupamiento

Configuración Experimental

Conjuntos de Datos

  • QM9: Conjunto de datos de química cuántica que contiene 134k moléculas pequeñas
  • ZINC250k: Conjunto de datos de moléculas similares a fármacos que contiene 250k moléculas relativamente grandes

Métricas de Evaluación

  • Validez sin corrección: Proporción de moléculas válidas sin corrección
  • Unicidad: Proporción de moléculas generadas únicas
  • FCD (Distancia de ChemNet de Fréchet): Distancia entre conjunto de entrenamiento y conjunto generado calculada usando características de ChemNet
  • NSPDK MMD: Diferencia de media máxima del núcleo de distancia de pares de subgrafos de vecindario considerando características de átomos y enlaces

Métodos de Comparación

Incluye múltiples líneas base de última generación:

  • Modelos de Flujo: MoFlow
  • Modelos de Difusión: EDP-GNN, GDSS, DiGress, GSDM
  • Coincidencia de Flujo: Dirichlet FM, CatFlow
  • Modelos de Energía: GraphEBM

Detalles de Implementación

  • Pasos de muestreo: GraphBFN utiliza 100×L pasos (L es el número de capas), significativamente menos que los 400-1000 pasos de las líneas base
  • Parámetro de equilibrio de pérdida multiescala: λ₁, λ₂
  • Umbral de tiempo mínimo: t_min = 10⁻⁵

Resultados Experimentales

Resultados Principales

MétodoQM9 Val.↑QM9 Unique↑QM9 FCD↓QM9 NSPDK↓ZINC250k Val.↑ZINC250k Unique↑ZINC250k FCD↓ZINC250k NSPDK↓Pasos de Muestreo
GDSS95.7298.462.5650.003397.1299.6414.0320.01921000
CatFlow99.8199.950.4410.002999.21100.0013.2110.0207-
GraphBFN99.6099.970.2140.000896.00100.005.7430.0069100×L

Hallazgos Clave:

  • Mejora de 51.5% en métrica FCD, mejora de 72.4% en métrica NSPDK
  • Logra rendimiento óptimo con significativamente menos pasos de muestreo
  • Alcanza la máxima unicidad, demostrando excelente diversidad

Estudios de Ablación

GraphBFN vs GraphBFN_w/o (sin supervisión jerárquica):

  • El marco jerárquico mejora en todas las métricas
  • Aunque sacrifica cierta velocidad de muestreo, mejora significativamente la calidad de generación

Análisis de Eficiencia de Muestreo

  • Rendimiento excelente dentro de los primeros 50 pasos
  • En comparación con métodos de línea base que requieren 400-1000 pasos, GraphBFN solo necesita 100 pasos para lograr rendimiento excelente
  • Apropiado para escenarios de aplicación sensibles al tiempo de inferencia

Trabajo Relacionado

Modelos de Generación de Grafos Moleculares

  • Modelos Autorregresivos: Agregan progresivamente nodos y aristas, como la serie GraphRNN
  • Modelos de Una Sola Pasada: Métodos basados en VAE, flujos normalizados, GAN, pero frecuentemente enfrentan problemas como colapso de modo
  • Modelos de Difusión: Dirección principal en años recientes, dividida en dos categorías: discreta y continua

Modelos de Difusión de Grafos

  • Difusión Discreta: Define directamente el proceso de difusión en espacio de estados discreto, como DiGress
  • Difusión Continua: Mapea al espacio continuo aplicando difusión gaussiana, como GDSS, GSDM
  • Desafío Central: Cómo manejar la naturaleza discreta de etiquetas de átomos y enlaces

Redes de Flujo Bayesiano

  • Nuevo tipo de modelo generativo que aprende mapeos entre distribuciones
  • Crea procesos de entrenamiento continuos y diferenciables para datos discretos
  • Este artículo propone un mecanismo de procesamiento de características discretas más simple y efectivo basado en esto

Conclusiones y Discusión

Conclusiones Principales

  1. Resolución Exitosa del Problema de Inconsistencia Entrenamiento-Inferencia: Unifica entrenamiento continuo y muestreo discreto mediante mecanismo CDF
  2. Mejora Significativa de Calidad de Generación: Logra rendimiento de última generación en puntos de referencia estándar
  3. Aumento Dramático de Eficiencia de Muestreo: Reduce pasos de muestreo a 1/4-1/10 de métodos de línea base
  4. Mejora de Diversidad Molecular: Evita sobreajuste a variaciones dentro de clases irrelevantes

Limitaciones

  1. Análisis de Interpretabilidad Insuficiente: Falta análisis profundo sobre cómo la información multiescala optimiza los resultados de generación
  2. Rango de Aplicabilidad Limitado: Validación principalmente en conjuntos de datos de moléculas relativamente pequeñas
  3. Complejidad Computacional: El marco jerárquico añade cierta sobrecarga computacional

Direcciones Futuras

  1. Extensión a dominios de grafos más grandes y complejos
  2. Exploración de aplicaciones de generación condicional
  3. Mejora del análisis de interpretabilidad
  4. Optimización de eficiencia computacional

Evaluación Profunda

Fortalezas

  1. Contribución Teórica Significativa: Identifica y resuelve un problema fundamental de modelos de difusión continua
  2. Innovación Técnica Destacada: El mecanismo CDF conecta ingeniosamente entrenamiento continuo e inferencia discreta
  3. Verificación Experimental Completa: Experimentos comparativos exhaustivos e investigaciones de ablación
  4. Alto Valor Práctico: Mejora significativa de eficiencia, apropiada para aplicaciones prácticas

Insuficiencias

  1. Profundidad de Análisis Teórico: Análisis limitado de propiedades de convergencia y garantías teóricas
  2. Escala Experimental: Validación principalmente en conjuntos de datos de escala pequeña a mediana, falta validación a gran escala
  3. Costo Computacional: Análisis insuficiente del costo computacional adicional del marco jerárquico
  4. Sensibilidad de Hiperparámetros: Análisis insuficiente de sensibilidad a hiperparámetros clave

Impacto

  1. Contribución Académica: Proporciona nuevas perspectivas de solución para tareas de generación discreta
  2. Valor Práctico: Puede acelerar procesos de descubrimiento de fármacos
  3. Reproducibilidad: Descripción clara del método, fácil de reproducir
  4. Potencial de Generalización: El marco es extensible a otras tareas de generación de estructuras discretas

Escenarios de Aplicación

  1. Descubrimiento de Fármacos: Diseño y optimización molecular
  2. Ciencia de Materiales: Generación de estructura de nuevos materiales
  3. Quimioinformática: Expansión de biblioteca de compuestos
  4. Otras Generaciones de Estructuras Discretas: Como secuencias de proteínas, ADN, etc.

Referencias

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.