Aprendizaje Automático de Grafos @ ICML 2023

Aprendizaje Automático de Grafos @ ICML 2023

¿Qué hay de nuevo en Graph ML?

Avances recientes y tendencias más destacadas, edición de agosto de 2023

¡Ni las magníficas playas ni los paisajes tropicales de Hawái 🌴 han alejado a los valientes científicos de asistir a la Conferencia Internacional sobre Aprendizaje Automático en Honolulu y presentar sus trabajos recientes! Veamos qué hay de nuevo en nuestra área favorita de Aprendizaje Automático en Grafos.

Gracias a Santiago Miret por revisar el artículo.

Para que el artículo no sea tan aburrido y se centre solo en los documentos, tomé algunas fotos alrededor de Honolulu 📷

Tabla de contenidos (con enlaces):

1. Transformadores de Grafos: Más esparsos, más rápidos y dirigidos
2. Teoría: Dimensión VC de GNN, inmersión profunda en sobre-aplastamiento
3. Nuevas arquitecturas de GNN: retrasos y saltos medios
4. Modelos generativos: difusión estable para moléculas, difusión discreta
5. Aprendizaje geométrico: WL geométrico, Álgebras de Clifford
6. Moléculas: preentrenamiento 2D-3D, Estimación de incertidumbre en MD
7. Materiales y proteínas: CLIP para proteínas, Pasaje de mensajes de Ewald, Aumentos equivariantes
8. Aplicaciones interesantes: razonamiento algorítmico, completado inductivo de KG, GNNs para espectros de masas
9. La parte de meme conclusiva

Transformadores de Grafos: Más esparsos, más rápidos y dirigidos

Hace aproximadamente un año presentamos GraphGPS y es gratificante ver que muchos documentos de ICML están construyendo sobre nuestro marco de trabajo y ampliando aún más las capacidades de GT.

• Explorando el lenguaje de programación Julia MongoDB
• Diferenciabilidad de una Función Dada su Gráfica
• Cómo yo, como estudiante, utilizo ChatGPT para aumentar mi productividad en un 10x

➡️ Exphormer de Shirzad, Velingker, Venkatachalam et al agrega una pieza faltante de atención esparsa motivada por grafos a los GT: en lugar de BigBird o Performer (diseñados originalmente para secuencias), la atención de Exphormer se basa en aristas de 1 salto, nodos virtuales (conectados a todos los nodos en un grafo) y una idea ingeniosa de aristas expansoras. Los grafos expansores tienen un grado constante y se ha demostrado que aproximan los grafos completamente conectados. En conjunto, todos los componentes hacen que la atención cueste O(V+E) en lugar de O(V²). Esto permite que Exphormer supere a GraphGPS en casi todos los aspectos y se pueda escalar a grafos realmente grandes de hasta 160k nodos. Un trabajo increíble y todas las posibilidades de convertir a Exphormer en el mecanismo estándar de atención esparsa en GTs 👏.

➡️ Al mismo tiempo que los transformadores de grafos, los grafos expansores ya se pueden utilizar para mejorar el rendimiento de cualquier arquitectura MPNN, como se muestra en Propagación de Grafos Expansores por Deac, Lackenby y Veličković.

En una línea similar, Cai et al demuestran que las MPNN con nodos virtuales pueden aproximar una atención lineal similar a Performer, de modo que incluso el clásico GCN y GatedGCN imbuidos con nodos virtuales muestran un rendimiento casi estatal del arte en tareas de grafos de largo alcance (lanzamos el banco de pruebas LGRB el año pasado específicamente para medir las capacidades de largo alcance de GNNs y GTs).

➡️ Algunos enfoques de submuestreo basados en parches para GTs inspirados en modelos de visión: “Una generalización de ViT/MLP-Mixer para grafos” de He et al divide la entrada en varios parches, codifica cada parche con un GNN en un token y ejecuta un transformador sobre esos tokens.

En GOAT de Kong et al, las características de los nodos se proyectan en un libro de códigos de K clústeres con K-Means, y un vecindario de 3 saltos muestreado de cada nodo se aplica al libro de códigos. GOAT es un modelo de 1 capa y es escalable para grafos de millones de nodos.

➡️ Los grafos dirigidos también recibieron amor de los transformers 💗. “Transformers Meet Directed Graphs” de Geisler et al introduce el Laplaciano Magnético, una generalización del Laplaciano para grafos dirigidos con una matriz de adyacencia no simétrica. Los autovectores del Laplaciano Magnético emparejados con caminatas aleatorias dirigidas son características de entrada sólidas para el transformer que permiten establecer un nuevo SOTA en el conjunto de datos de predicción de propiedades de grafos OGB Code2 con una buena diferencia.

🏅 Por último, pero no menos importante, tenemos un nuevo GT SOTA en el conjunto de datos estándar de la comunidad ZINC: GRIT de Ma, Lin, et al incorpora la matriz de caminata aleatoria d-dimensional completa, denominada probabilidades relativas de caminata aleatoria (RRWP), como características de borde para el cálculo de atención (en comparación, las características RWSE populares son solo los elementos diagonales de esta matriz). Se ha demostrado que RRWP es más potente que las características de distancia de camino más corto y establece un MAE récord de 0.059 en ZINC (en comparación con 0.070 de GraphGPS). GRIT a menudo supera a GPS en otros puntos de referencia también 👏. En una línea similar, Eliasof et al proponen una idea ingeniosa para combinar características aleatorias y espectrales como codificaciones posicionales que superan a RWSE pero no se probaron con GTs.

Teoría: Dimensión VC de las GNNs, en profundidad sobre el exceso de compresión

➡️ La dimensión VC mide la capacidad y expresividad del modelo. Se ha estudiado bien para algoritmos de ML clásicos, pero sorprendentemente nunca se ha aplicado para estudiar las GNNs. En “WL meet VC” de Morris et al, finalmente se descubre la conexión entre la prueba WL y la dimensión VC: resulta que la dimensión VC puede ser limitada por la longitud de bits de los pesos de las GNNs, es decir, los pesos float32 implicarían una dimensión VC de 32. Además, la dimensión VC depende logarítmicamente del número de colores WL únicos en la tarea dada y polinómicamente de la profundidad y el número de capas. Este es un gran resultado teórico y ¡te animo a que lo veas!

🍊🖐️ El efecto de exceso de compresión, la pérdida de información cuando intentas rellenar mensajes de demasiados nodos vecinos, es otro problema común de las MPNNs, y no entendemos completamente cómo lidiar adecuadamente con él. Este año, se dedicaron 3 documentos a este tema. Quizás el más fundamental sea el trabajo de Di Giovanni et al que explica cómo el ancho, la profundidad y la topología del grafo de las MPNNs afectan el exceso de compresión.

Resulta que el ancho podría ayudar (pero con problemas de generalización), la profundidad realmente no ayuda y la topología del grafo (caracterizada por el tiempo de conmutación entre nodos) juega el papel más importante. Podemos reducir el tiempo de conmutación mediante diversas estrategias de reestructuración del grafo (agregando y eliminando aristas basadas en propiedades espaciales o espectrales), y hay muchas de ellas (es posible que hayas oído hablar de la reestructuración basada en el flujo de Ricci que se llevó el premio al Mejor Artículo en ICLR 2022). De hecho, hay un trabajo de seguimiento de este estudio que profundiza aún más y deriva algunas afirmaciones de imposibilidad con respecto al exceso de compresión y algunas propiedades de las MPNNs, ¡te animo a que lo leas también!

➡️ La resistencia efectiva es un ejemplo de estrategias de reconexión espacial, y Black et al lo estudian en detalle. La reconexión basada en el flujo de Ricci trabaja con la curvatura del grafo y se estudia más a fondo en el trabajo de Nguyen et al.

➡️ Las GNN de subgrafo continúan siendo el centro de atención: dos trabajos (Zhang, Feng, Du, et al y Zhou, Wang, Zhang) derivan de manera concurrente las jerarquías de expresividad de las GNN de subgrafo propuestas recientemente y su relación con las pruebas de WL de orden 1 y superior.

Nuevas arquitecturas de GNN: Retardos y medios saltos

Si estás cansado de otra variación de GCN o GAT, aquí tienes algunas ideas frescas que pueden funcionar con cualquier GNN de tu elección:

⏳ Como sabemos por la sección de Teoría, la reconexión ayuda a combatir el exceso de compresión. Gutteridge et al presentan “DRew: Mensaje dinámico reconectado con retardo”, que densifica gradualmente el grafo en capas posteriores de GNN para que los nodos a larga distancia vean los estados originales de los nodos anteriores (la versión original de DRew) o se agreguen conexiones saltadas basadas en el retardo, dependiendo de la distancia entre dos nodos (la versión vDRew). Por ejemplo ( 🖼️👇), en el paso de mensajes con retardo vDRew, un nodo inicial de la capa 0 mostrará su estado a vecinos de 2 saltos en la capa 1, y mostrará su estado a un vecino de 3 saltos en la capa 2. DRew mejora significativamente la capacidad de las GNN ‘vanilla’ para realizar tareas a largo plazo — de hecho, un GCN habilitado con DRew es el estado del arte actual en el conjunto de datos Peptides-func del Long Range Graph Benchmark 👀

🦘 Otra idea interesante de Azabou et al es ralentizar el paso de mensajes mediante la inserción de nuevos nodos lentos en cada arista con un patrón de conectividad especial: solo una conexión entrante desde el nodo inicial y una arista simétrica con el nodo de destino. Los nodos lentos mejoran el rendimiento de las GNN ‘vanilla’ en pruebas heterofílicas en gran medida, y también es posible utilizar nodos lentos para el aprendizaje auto-supervisado mediante la creación de vistas con diferentes ubicaciones de nodos lentos para el mismo grafo original. HalfHop es un componente de SSL que no puede faltar y que impulsa el rendimiento, y debería estar en el conjunto de herramientas estándar de muchas bibliotecas de GNN 👍.

Modelos generativos: Difusión estable para moléculas, difusión discreta

➡️ Los modelos de difusión pueden funcionar en el espacio de características (por ejemplo, el espacio de píxeles en la generación de imágenes, como el DDPM original) o en el espacio latente (como la Difusión Estable). En el espacio de características, debes diseñar el proceso de ruido para respetar las simetrías y equivariancias de tu espacio de características. En el espacio latente, solo tienes que agregar ruido gaussiano a las características producidas por un codificador (preentrenado). La mayoría de los modelos de generación de moléculas 3D trabajan en el espacio de características (como el EDM pionero), y el nuevo modelo GeoLDM de Xu et al (autores del prominente GeoDiff) es el primero en definir la difusión latente para la generación de moléculas 3D. Es decir, después de entrenar un autoencoder EGNN, GeoLDM se entrena en el objetivo de desruido donde el ruido se muestrea de una distribución gaussiana estándar. GeoLDM aporta mejoras significativas sobre EDM y otros enfoques de difusión no latente 👏.

➡️ En el ámbito de los grafos no geométricos (solo con una matriz de adyacencia y tal vez características categóricas de los nodos), la difusión discreta de grafos pionera por DiGress (ICLR’23) parece ser la opción más aplicable. Chen et al proponen EDGE, un modelo de difusión discreta guiado por la distribución de grados de los nodos. A diferencia de DiGress, el grafo objetivo final en EDGE es un grafo desconectado sin aristas, un modelo de ruido hacia adelante elimina aristas mediante una distribución de Bernoulli, y un proceso inverso agrega aristas a los nodos activos más recientes (activos son los nodos cuyos grados cambiaron en el paso anterior). ¡Gracias a la esparcidad introducida por la guía de grados, EDGE puede generar grafos bastante grandes de hasta 4k nodos y 40k aristas!

➡️ Por último, “Modelos de difusión estructurados gráficamente” de Weilbach et al une la brecha entre modelos generativos continuos y modelos gráficos probabilísticos que inducen una cierta estructura en el problema de interés, a menudo estos problemas tienen una naturaleza combinatoria. La idea central es codificar la estructura del problema como una máscara de atención que respeta las invarianzas de permutación y utilizar esta máscara en el cálculo de atención en el codificador Transformer (que, por definición, es equivariante a la permutación de los tokens de entrada a menos que se utilicen incrustaciones posicionales). GSDM puede abordar la factorización de matrices continuas binarias, circuitos booleanos, generar sudokus y realizar ordenamientos. Particularmente disfrutable es un toque de ironía con el que el artículo está escrito 🙃.

Aprendizaje Geométrico: Geométrico WL, Álgebras de Clifford

¡El Aprendizaje Profundo Geométrico está prosperando! Hubo tantos artículos interesantes presentados que ocuparían prácticamente todo el post, así que solo destacaré algunos.

➡️ Geométrico WL finalmente ha llegado en el trabajo de Joshi, Bodnar, et al. Geométrico WL extiende la noción de la prueba WL con características geométricas (por ejemplo, coordenadas o velocidad) y deriva la jerarquía de expresividad hasta GWL de orden k. Puntos clave: 1️⃣ los modelos equivariantes son más expresivos que los invariantes (con una nota de que en grafos completamente conectados la diferencia desaparece), 2️⃣ el orden tensorial de las características mejora la expresividad, 3️⃣ el orden corporal de las características mejora la expresividad (ver la imagen 👇). Es decir, esférico > cartesiano > escalares, e interacciones de muchos cuerpos > solo distancias. El artículo también presenta la increíble fuente de aprendizaje Geometric GNN Dojo, donde puedes derivar e implementar la mayoría de los modelos SOTA desde los principios básicos.

➡️ Yendo más allá de los vectores, Ruhe et al derivan Redes de Álgebra de Clifford Geométricas (GCANs). Las álgebras de Clifford admiten naturalmente interacciones de orden superior mediante bivectores, trivectores y (en general) multivectores. La idea clave es el teorema de Cartan-Dieudonné que afirma que toda transformación ortogonal se puede descomponer en reflejos en hiperplanos, y las álgebras geométricas representan los datos como elementos del grupo Pin(p,q,r). Las GCANs introducen una noción de capas lineales, normalizaciones, no linealidades y cómo pueden ser parametrizadas con redes neuronales. Los experimentos incluyen la modelización de la dinámica de fluidos y las ecuaciones de Navier-Stokes.

De hecho, ya existe un trabajo de seguimiento que presenta redes neuronales Clifford equivariantes. Puedes aprender más sobre las bases del álgebra de Clifford y los artículos más recientes sobre CliffordLayers respaldados por Microsoft Research.

💊 Equivariant GNN (EGNN) es la aspirina del DL geométrico que se aplica a casi todas las tareas y ha experimentado una serie de mejoras. Eijkelboom et al unen EGNN con redes simpliciales que operan en estructuras de orden superior (específicamente, complejos simpliciales) en EMPSN. Este es uno de los primeros ejemplos que combina características geométricas y topológicas ¡y tiene un gran potencial de mejora! Finalmente, Passaro y Zitnick descubren un truco ingenioso para reducir las convoluciones SO(3) a SO(2), reduciendo la complejidad de O(L⁶) a O(L³) pero con garantías de equivalencia matemática 👀. Este hallazgo permite escalar modelos geométricos en conjuntos de datos más grandes como OpenCatalyst y ya se ha implementado en Equiformer V2, ¡pronto en muchas otras bibliotecas para modelos geométricos 😉

Moleculas: Preentrenamiento 2D-3D, Estimación de incertidumbre en MD

➡️ Liu, Du, et al proponen MoleculeSDE, un nuevo marco para el preentrenamiento conjunto en datos moleculares 2D-3D. Además de la pérdida contrastiva estándar, los autores añaden dos componentes generativos: la reconstrucción de entradas 2D -> 3D y 3D -> 2D a través de la generación de difusión basada en puntuaciones. Utilizando GIN y SchNet estándar como modelos 2D y 3D, MoleculeSDE se preentrena en PCQM4M v2 y tiene un buen rendimiento en tareas de ajuste fino posteriores.

➡️ Wollschläger et al realizan un estudio exhaustivo de la Estimación de Incertidumbre en GNNs para dinámica molecular y campos de fuerza. Identificando principios clave informados por la física y enfocados en la aplicación, los autores proponen un Kernel Neural Localizado, una extensión basada en Procesos Gaussianos para cualquier GNN geométrico que funcione con cantidades invariantes y equivariantes (probado en SchNet, DimeNet y NequIP). En muchos casos, las estimaciones de LNK de un modelo son igualadas o mejores que el costoso ensamblaje en el que se necesitaría entrenar varios modelos.

Materiales y Proteínas: CLIP para proteínas, Paso de mensaje Ewald, Aumentos equivariantes

CLIP y sus descendientes se han convertido en un elemento básico estándar en modelos de texto a imagen. ¿Podemos hacer lo mismo pero para texto a proteína? ¡Sí!

➡️ Xu, Yuan, et al presentan ProtST, un marco para aprender representaciones conjuntas de descripciones de proteínas de texto (a través de PubMedBERT) y secuencias de proteínas (a través de ESM). Además de una pérdida contrastiva, ProtST tiene un objetivo de predicción de máscara multimodal, por ejemplo, cubriendo el 15% de los tokens en texto y secuencia de proteínas, y prediciéndolos conjuntamente en base a representaciones latentes, y pérdidas de predicción de máscara basadas en secuencias o solo en lenguaje. Además, los autores diseñan un nuevo conjunto de datos ProtDescribe con 550K pares alineados de secuencias y descripciones de proteínas. ProtST destaca en muchas tareas de modelado de proteínas en el benchmark PEER, incluida la anotación y localización de funciones de proteínas, pero también permite la recuperación de proteínas sin entrenamiento directamente desde la descripción textual (ver un ejemplo a continuación). Parece que ProtST tiene un futuro prometedor como base de muchos modelos generativos de proteínas 😉

En realidad, ICML presenta varios trabajos de generación de proteínas como GENIE de Lin y AlQuraishi y FrameDiff de Yim, Trippe, De Bortoli, Mathieu, et al — estos aún no están condicionados a descripciones textuales, por lo que incorporar ProtST parece ser un impulso de rendimiento obvio 📈.

⚛️ Las redes neuronales de propagación de mensajes en moléculas tienen un sesgo de localidad estricto que inhibe el modelado de interacciones a larga distancia. Kosmala et al derivan Propagación de Mensajes Ewald y aplican la idea de suma de Ewald que descompone el potencial de interacción en términos de corto y largo alcance. La interacción de corto alcance se modela mediante cualquier GNN mientras que la interacción de largo alcance es nueva y se modela con una transformada de Fourier 3D y propagación de mensajes sobre frecuencias de Fourier. Resulta que este término de largo alcance es bastante flexible y se puede aplicar a cualquier red que modele sistemas periódicos y aperiódicos (como cristales o moléculas) como SchNet, DimeNet o GemNet. El modelo se evaluó en los conjuntos de datos OC20 y OE62. Si estás interesado en más detalles, ¡echa un vistazo a la charla de 1 hora de Arthur Kosmala en el Grupo de Lectura LOG2!

Una idea similar de usar la suma de Ewald para cristales 3D se utiliza en PotNet de Lin et al, donde la conexión de largo alcance se modela con funciones de Bessel incompletas. PotNet se evaluó en el conjunto de datos del Proyecto de Materiales y JARVIS — por lo que al leer esos dos artículos puedes tener una buena comprensión de los beneficios que aporta la suma de Ewald para muchas tareas relacionadas con cristales 😉

➡️ Otra mirada a cómo dotar de equivarianza a cualquier GNN para cristales y moléculas la presentan Duval, Schmidt, et al en FAENet. Una forma estándar es incorporar ciertas simetrías y equivarianzas directamente en las arquitecturas de GNN (como en EGNN, GemNet y Propagación de Mensajes Ewald) — esta es una forma segura pero computacionalmente costosa (especialmente cuando se trata de armónicos esféricos y productos tensoriales). Otra opción que se usa a menudo en visión — mostrar muchas augmentaciones de la misma entrada y el modelo eventualmente debería aprender las mismas invariancias en las augmentaciones. Los autores optan por el segundo camino y diseñan una forma rigurosa de muestrear augmentaciones invariantes o equivariantes de datos 2D / 3D (por ejemplo, para energía o fuerzas, respectivamente) todo con pruebas elegantes ✍️. Para ello, la tubería de aumento de datos incluye proyectar las entradas 2D / 3D a una representación canónica (basada en el PCA de la matriz de covarianza de distancias) a partir de la cual podemos muestrear rotaciones de manera uniforme.

La propuesta de FAENet es un modelo simple que utiliza solo distancias pero muestra un rendimiento muy bueno con la augmentación de datos de promedio de marco estocástico, mientras que es de 6 a 20 veces más rápido. ¡Funciona también para estructuras cristalinas!

Aplicaciones interesantes: Razonamiento algorítmico, Completado inductivo de KG, GNNs para espectros de masa

Algunos artículos en esta sección no pertenecen a ninguno de los anteriores, pero aún así merecen su atención.

➡️ “Neural Algorithmic Reasoning with Causal Regularisation” de Bevilacqua et al aborda un problema común en el aprendizaje de grafos: la generalización OOD a entradas más grandes en tiempo de prueba. Estudiando la generalización OOD en problemas de razonamiento algorítmico, los autores observan que existen muchas entradas diferentes que realizan cálculos idénticos en un cierto paso. Al mismo tiempo, esto significa que algún subconjunto de entradas no afecta (no debería afectar) el resultado de la predicción. Esta suposición permite diseñar un objetivo auto-supervisado (llamado Hint-ReLIC) que prefiere un paso “significativo” a un conjunto de pasos que no afectan el resultado de la predicción. El nuevo objetivo mejora significativamente el rendimiento en muchas tareas CLRS-30, alcanzando un micro-F1 de más del 90%. Es interesante preguntarse si podríamos aprovechar el mismo principio en el paso de mensajes general y mejorar la transferencia OOD en otras tareas de aprendizaje de grafos 🤔

Si te interesa aún más el razonamiento algorítmico neural, echa un vistazo a las actas del taller de Conocimiento y Razonamiento Lógico, que cuenta con aun más trabajos sobre ese tema.

➡️ “InGram: Inductive Knowledge Graph Embedding via Relation Graphs” de Lee et al parece ser uno de los pocos artículos de grafos de conocimiento en ICML’23 (según mi búsqueda). InGram es uno de los primeros enfoques que puede generalizar de manera inductiva tanto a entidades no vistas como a relaciones no vistas en tiempo de prueba. Anteriormente, los modelos inductivos de KG necesitaban aprender al menos las incrustaciones de las relaciones de alguna forma para generalizar a nuevos nodos, y en este paradigma, las relaciones no vistas eran difíciles de modelar. InGram construye un grafo de relaciones sobre el grafo multi-relacional original, es decir, un grafo de tipos de relaciones, y aprende representaciones de relaciones basadas en este grafo ejecutando un GAT. Las representaciones de entidades se obtienen a partir de la inicialización aleatoria y un codificador GNN. Teniendo tanto las representaciones de entidades como de relaciones, se aplica un decodificador DistMult como función de puntuación. Existen buenas posibilidades de que InGram para relaciones no vistas sea tan influyente como GraIL (ICML 2020) para entidades no vistas 😉.

🌈 “Efficiently predicting high resolution mass spectra with graph neural networks” de Murphy et al es una aplicación interesante de GNNs a un problema de física real de predicción de espectros de masa. El hallazgo principal es que la mayor parte de la señal en los espectros de masa se explica por un pequeño número de componentes (fórmulas de iones de producto y pérdida neutral). Y es posible extraer un vocabulario de esas fórmulas a partir de los datos de entrenamiento. Por lo tanto, el problema se puede plantear como clasificación de grafos (o predicción de propiedades de grafos), donde, dado un grafo molecular, predecimos tokens de un vocabulario que corresponden a ciertos valores de espectros de masa. El enfoque, GRAFF-MS, construye una representación de grafo molecular a través de GIN con características de aristas, con características de Laplaciano (a través de SignNet) y se agrupa con características covariantes. En comparación con el CFM-ID de referencia, GRAFF-MS realiza la inferencia en ~19 minutos en lugar de 126 horas, logrando un rendimiento mucho mayor 👀.

La parte final del meme

¡El meme del 2022 finalmente ha convergido en Michael Bronstein!

We will continue to update Zepes; if you have any questions or suggestions, please contact us!