Invariant-Based Weight Sharing for Message Passing

Resumen Técnico: Compartición de Pesos Basada en Invariantes para el Paso de Mensajes

Enunciado del Problema
Las redes neuronales de paso de mensajes (MPNNs) son un marco dominante para aprender representaciones de datos con estructura de grafo. Sin embargo, las MPNNs estándar y los transformadores de grafos suelen depender de una compartición de pesos basada únicamente en características de nodos o aristas, o aprenden coeficientes de atención específicos de la instancia. Este enfoque limita la capacidad del modelo para capturar patrones estructurales inherentes, ya que los pesos no están explícitamente vinculados a la topología del grafo. Además, las MPNNs estándar a menudo se restringen a la agregación de vecindarios locales, haciéndolas susceptibles al sobre-alisado y al sobre-aplastamiento en arquitecturas profundas, mientras que los transformadores de grafos a menudo carecen de una reutilización sistemática de parámetros entre subgrafos estructuralmente equivalentes.

Metodología: ShareGNNs
Los autores introducen ShareGNNs, una arquitectura novedosa fundamentada en la compartición de pesos basada en invariantes. El principio central consiste en indexar los parámetros aprendibles directamente mediante invariantes de grafos —funciones preservadas bajo permutaciones de nodos— en lugar de hacerlo mediante dimensiones de características fijas o posiciones específicas de la instancia.

1. Compartición de Pesos Basada en Invariantes:

   • Codificador (Paso de Mensajes): En lugar de matrices de pesos fijas, el modelo construye dinámicamente los pesos de paso de mensajes para cada par de nodos $(v, w)$ basándose en una tripleta invariante: $(l(v), l(w), d(v, w))$, donde $l$ es una función de etiquetado (por ejemplo, número atómico, etiqueta de Weisfeiler-Leman, conteo de patrones) y $d$ es la distancia del camino más corto. Todos los pares de nodos que comparten la misma tripleta invariante comparten un único peso escalar aprendible $\omega_\tau$. Esto permite que la información fluya entre nodos arbitrarios en una sola capa, no solo entre vecinos inmediatos.
   • Decodificador (Agrupación): De manera similar, los términos de sesgo y los pesos de agrupación para la agregación a nivel de grafo se indexan mediante las etiquetas invariantes de los nodos individuales.
   • Construcción Dinámica: Para cualquier grafo de entrada, las matrices de pesos relevantes se ensamblan sobre la marcha a partir de un pool de parámetros compartido. Esto asegura que el modelo sea equivariante a la permutación en el codificador e invariante a la permutación en el decodificador por construcción.

2. Funciones de Etiquetado:
   La expresividad del modelo se controla mediante la elección de las funciones de etiquetado $l$. El artículo explora:

   • Etiquetado Molecular: Números atómicos.
   • Etiquetado de Weisfeiler-Leman (WL): Refinamiento iterativo de las etiquetas de los nodos basado en multiconjuntos de vecinos.
   • Etiquetado de Patrones: Conteos de patrones de subgrafos específicos (por ejemplo, triángulos, ciclos) que involucran a un nodo.
     Estos pueden combinarse para capturar información estructural de múltiples niveles.

3. Arquitectura:
   ShareGNNs utilizan una estructura simple de codificador-decodificador. El codificador propaga información a través de distancias arbitrarias en una sola capa, y el decodificador agrega los embebidos de nodos en una representación de grafo utilizando agrupación indexada por invariantes. La arquitectura soporta extensiones de múltiples cabezas donde diferentes cabezas utilizan diferentes conjuntos de invariantes.

Contribuciones Clave

1. Principio de Compartición de Pesos Basada en Invariantes: Un nuevo mecanismo de compartición de pesos que incorpora explícitamente la estructura del grafo, permitiendo una reutilización sistemática de parámetros entre pares de nodos estructuralmente equivalentes dentro y entre diferentes grafos.
2. Arquitectura ShareGNN: Una instanciación práctica de este principio que combina el paso de mensajes con conectividad tipo transformador (interacciones de todos a todos), pero con parámetros vinculados a invariantes estructurales en lugar de puntuaciones de atención aprendidas.
3. Validación Teórica y Empírica:
   • Expresividad: Los autores demuestran que ShareGNNs son al menos tan expresivos como la función de etiquetado elegida. Si dos grafos son distinguibles por la función de etiquetado (por ejemplo, etiquetas WL), ShareGNN puede distinguirlos.
   • Control: El modelo ofrece un control explícito sobre la complejidad y la expresividad mediante la selección de invariantes.
   • Rendimiento: Los experimentos demuestran que ShareGNNs superan a las MPNNs estándar (por ejemplo, GIN, GAT) y son competitivos con los transformadores de grafos en diversas pruebas, incluyendo clasificación molecular, análisis de redes sociales y conteo de subestructuras.

Resultados

• Clasificación de Grafos: En conjuntos de datos diversos (NCI1, NCI109, MUTAG, IMDB-B/M), ShareGNNs logran una precisión de vanguardia o altamente competitiva. Notablemente, alcanzan una precisión casi perfecta (100%) en pruebas sintéticas (RingTransfer, Snowflakes) diseñadas para desafiar a las MPNNs estándar, demostrando su capacidad para capturar dependencias de largo alcance y distinguir grafos indistinguibles por la prueba 1-WL.
• Conteo de Subestructuras: En la Prueba de Conteo de Subestructuras, ShareGNNs logran resultados de vanguardia en configuraciones de tarea única y permanecen competitivos en configuraciones de múltiples tareas, validando que invariantes expresivos conducen a modelos expresivos incluso con arquitecturas poco profundas (una sola capa de codificador).
• Escalabilidad: El modelo escala a conjuntos de datos grandes como ZINC (12k y 250k grafos) manteniendo una arquitectura poco profunda, evitando los problemas de sobre-alisado comunes en las MPNNs profundas.
• Eficiencia: Los estudios de ablación muestran que aumentar la distancia máxima de paso de mensajes $D$ tiene un impacto más significativo en el rendimiento que agregar más capas de codificador. Además, la poda de tripletas invariantes raras reduce significativamente la cantidad de parámetros sin sacrificar la precisión.

Significado y Afirmaciones
El artículo afirma que la compartición de pesos basada en invariantes proporciona una alternativa fundamentada a la compartición de pesos basada en características y a la atención específica de la instancia. Al vincular los pesos a invariantes estructurales, ShareGNNs ofrecen:

• Mejor Interpretabilidad: Los pesos aprendidos pueden interpretarse directamente en términos de patrones estructurales (por ejemplo, distancias específicas entre tipos de átomos específicos).
• Transferibilidad: Los parámetros son naturalmente transferibles entre grafos con patrones estructurales similares, sugiriendo un potencial para modelos fundacionales de grafos.
• Control de Expresividad: El modelo permite a los investigadores controlar explícitamente el poder discriminatorio de la red mediante la selección de invariantes apropiados, cerrando la brecha entre la expresividad teórica (por ejemplo, pruebas k-WL) y la implementación práctica.

Los autores reconocen limitaciones, incluido el costo de preprocesamiento para calcular invariantes y distancias por pares, y la complejidad de memoria cuadrática en el peor de los casos, aunque señalan que esto es comparable a los transformadores de grafos y puede mitigarse mediante la poda. También notan que el enfoque actual requiere invariantes definidos por el usuario, dejando espacio para trabajos futuros en la selección automatizada de invariantes.