Resumen Técnico: GraViti – Autoencoders Variacionales a Nivel de Grafo con Invarianza de Permutación Relajada

Enunciado del Problema

La autoencodificación a nivel de grafo busca mapear grafos completos en vectores latentes compactos para facilitar tareas como la interpolación, la búsqueda guiada por propiedades y el modelado generativo. Aunque los Autoencoders Variacionales (VAE) a nivel de nodo están bien establecidos, sufren de una limitación fundamental: el espacio latente se define a nivel de nodo y depende de un ordenamiento de nodos arbitrario. En consecuencia, dos grafos isomorfos pueden mapearse a conjuntos de vectores latentes completamente diferentes, lo que impide la definición de distancias significativas entre grafos o la captura de variedades a nivel de grafo.

Los enfoques existentes a nivel de grafo enfrentan una compensación. Métodos como PIGVAE y GRALE intentan imponer una invarianza de permutación estricta, a menudo dependiendo de procedimientos de emparejamiento computacionalmente costosos (por ejemplo, emparejamiento de grafos, algoritmos de Sinkhorn) para alinear los grafos de entrada y reconstruidos. Estas etapas de alineación introducen una sobrecarga computacional significativa (típicamente complejidad cúbica) y pueden limitar la escalabilidad en conjuntos de datos grandes. Además, muchos decodificadores a nivel de grafo dependen de estrategias de relleno que imponen un tamaño máximo fijo de grafo o introducen clases de "nodos ficticios", lo que restringe la generalización a grafos más grandes y sesga las tareas de clasificación.

Metodología: La Arquitectura GraViti

GraViti es un VAE a nivel de grafo basado en transformadores diseñado para codificar grafos completos en vectores latentes únicos, relajando el requisito estricto de reconstrucción invariante a permutaciones cuando existe un ordenamiento canónico de nodos consistente.

1. Representación del Grafo

El modelo opera sobre grafos con clases de nodos y aristas discretas (por ejemplo, tipos de átomos, tipos de enlaces).

• Entrada: Una matriz de clases de nodos $X$, un tensor de clases de aristas $E$ y un vector de características globales $y$ que codifica estadísticas a nivel de grafo (conteos de tipos de nodos/aristas y tamaño del grafo).
• Codificaciones Posicionales: Para inyectar información estructural, el modelo utiliza vectores propios del Laplaciano y características de caminata aleatoria como codificaciones posicionales ($P$).
• Conocimiento del Dominio: Para dominios específicos como moléculas, se incorporan características adicionales específicas de nodos ($K$) como la masa atómica o la electronegatividad.
• Agrupación (Batching): Los grafos se rellenan hasta el tamaño del grafo más grande en un lote, pero el modelo está diseñado para ignorar completamente estos elementos rellenos, evitando la necesidad de clases de nodos ficticios.

2. Codificador

El codificador mapea la tripleta de entrada $(X, E, y)$ a los parámetros de la distribución latente $(z_\mu, z_\sigma)$.

• Incrustación (Embedding): Las clases de nodos, las clases de aristas y las características globales se proyectan en incrustaciones.
• Pila de Transformadores de Grafo: Las incrustaciones se procesan a través de una pila de capas de transformadores de grafo (basadas en las arquitecturas DiGress y CatFlow). Esto permite actualizaciones simultáneas de características a nivel de nodo, arista y grafo, asegurando la equivarianza de permutación durante la fase de codificación.
• Lectura (Readout): A diferencia del agrupamiento por suma/media tradicional, GraViti emplea una lectura basada en transformadores de secuencia. Dos transformadores decodificadores de secuencia paralelos procesan las representaciones de nodos y aristas respectivamente, utilizando el vector de características globales $y_{rep}$ como objetivo para inyectar contexto global. Esto asegura que la salida sea independiente del ordenamiento de nodos en la etapa de lectura.

3. Predicción del Tamaño del Grafo

Una innovación crucial es el desentrelazamiento explícito del tamaño del grafo del resto de la representación latente.

• El modelo fuerza a que las primeras $d_s$ dimensiones del espacio latente codifiquen el log-tamaño del grafo.
• Un MLP predice el log-tamaño $\hat{s}$ a partir de estas dimensiones.
• Durante la decodificación, el log-tamaño real $s$ se vuelve a codificar y se reinyecta en la representación latente para guiar la generación del número correcto de nodos, asegurando que el decodificador conozca el tamaño del grafo objetivo sin depender de enmascaramiento.

4. Decodificador

El decodificador reconstruye el grafo a partir del vector latente $z$.

• Difusión (Broadcasting): El vector latente se duplica $N$ veces (donde $N$ es el tamaño predicho o real) y se aumenta con codificaciones posicionales sinusoidales para romper la simetría.
• Transformador de Secuencia: Un codificador de transformador de secuencia procesa esta secuencia, dividiéndola en secuencias orientadas a nodos y orientadas a aristas.
• Refinamiento: Una segunda pila de transformadores de grafo refina las representaciones de nodos y aristas.
• Cabezas de Predicción: MLPs finales clasifican los tipos de nodos y los tipos de aristas. El modelo predice el tamaño del grafo directamente, eliminando la necesidad de una clase ficticia de "sin arista".

5. Función Objetivo de Entrenamiento

El modelo se entrena utilizando un objetivo VAE estándar:

• Pérdida de Reconstrucción: Una combinación de Focal Loss para la clasificación de nodos y aristas (para manejar el desequilibrio de clases) y una pérdida L1 suave para la predicción del tamaño del grafo.
• Pérdida Variacional: Un término de divergencia KL regulariza el espacio latente hacia una priori normal estándar.
• Invarianza de Permutación Relajada: Crucialmente, la pérdida de reconstrucción se aplica basándose en el orden de nodos de entrada. Los autores argumentan que en dominios con un ordenamiento canónico fiable (por ejemplo, moléculas, redes bayesianas), imponer una invarianza de permutación estricta mediante emparejamiento es perjudicial. Al depender de un ordenamiento consistente, el modelo logra una complejidad cuadrática y una mayor fidelidad de reconstrucción.

Contribuciones Clave

1. Arquitectura Escalable: GraViti logra una precisión de reconstrucción state-of-the-art en grandes conjuntos de datos de grafos (PubChem16, PubChem32) al evitar procedimientos costosos de emparejamiento invariante a permutaciones.
2. Espacio Latente Estructurado: Se demuestra que el espacio latente aprendido está bien estructurado, apoyando la interpolación suave, la edición controlada y la optimización impulsada por propiedades. En puntos de referencia moleculares, el modelo recupera principios químicos establecidos (por ejemplo, la relación entre la longitud de la cadena alquílica y la hidrofobicidad) directamente desde el espacio latente.
3. Invarianza de Permutación Relajada: El artículo demuestra empíricamente que la invarianza de permutación estricta no siempre es necesaria. Cuando está disponible un ordenamiento canónico de nodos consistente, relajar esta restricción mejora significativamente la calidad de la reconstrucción y la escalabilidad en comparación con los métodos que imponen invarianza a través de emparejamientos complejos.

Resultados Experimentales

Los autores evaluaron GraViti en varios puntos de referencia:

• Reconstrucción: En conjuntos de datos moleculares (QM9, PubChem16, PubChem32), GraViti supera a las líneas base (GRALE, PIGVAE) en hasta un 25% en precisión de reconstrucción, particularmente en grafos más grandes donde los métodos basados en emparejamiento tienen dificultades.
• Estudios de Ablación: Eliminar el codificador transformador o la lectura basada en transformador degrada significativamente el rendimiento, confirmando la importancia de estos componentes.
• Denoising: El modelo restaura eficazmente estructuras químicamente válidas a partir de entradas ruidosas, mostrando la variante VAE una robustez superior en comparación con la variante AE.
• Optimización de Propiedades: Utilizando CMA-ES en el espacio latente, el modelo maximiza con éxito el logP (lipofilicidad), generando moléculas con cadenas de carbono progresivamente más largas o moieties hidrofóbicas específicas, alineándose con tendencias químicas conocidas.
• Interpolación: La interpolación lineal en el espacio latente produce transiciones suaves entre estructuras moleculares.
• Tareas de Descenso: Los vectores latentes entrenados en PubChem predicen con alta precisión propiedades moleculares (momento dipolar, brecha HOMO-LUMO) y puntuaciones del Criterio de Información Bayesiano (BIC) para redes bayesianas.
• Generación: Si bien la validez de la generación es menor en PubChem (aproximadamente 50-70% dependiendo de la representación aromática) en comparación con QM9, el modelo produce moléculas diversas y novedosas. Los autores atribuyen la menor validez en conjuntos de datos más grandes a la compacidad global del espacio latente en lugar de fallos estructurales locales.

Significado y Afirmaciones

El artículo afirma que GraViti proporciona una base práctica para el aprendizaje de representaciones a nivel de grafo. Su significado principal radica en demostrar que la invarianza de permutación estricta no es un prerrequisito para una autoencodificación a nivel de grafo efectiva cuando existe un ordenamiento de nodos estable. Al relajar esta restricción, GraViti evita los cuellos de botella computacionales del emparejamiento de grafos, permitiendo un entrenamiento escalable en conjuntos de datos masivos mientras mantiene una alta calidad de reconstrucción.

Los autores posicionan a GraViti como una alternativa ligera a pipelines de generación complejos, ofreciendo un mecanismo de decodificación de un solo paso que equilibra el rendimiento generativo con la calidad práctica de las muestras. El trabajo sugiere que para dominios como grafos moleculares y redes bayesianas, donde los ordenamientos canónicos son inherentes o pueden construirse, los modelos a nivel de grafo pueden lograr un rendimiento superior sin la sobrecarga de pérdidas basadas en alineación. Se propone trabajo futuro para integrar restricciones conscientes del dominio para mejorar aún más la validez de la generación y explorar la transferencia a otros dominios de grafos como circuitos y proteínas.