Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

Resumen Técnico: Escalamiento de la Expresividad del Análisis de Weisfeiler–Leman a Grafos Masivos con GPUs

Planteamiento del Problema
La prueba de Weisfeiler–Leman (1-WL) y su coloración estable asociada son fundamentales para comprender el poder de expresión de las Redes Neuronales de Grafos (GNN). La coloración estable 1-WL proporciona un límite superior en la capacidad de las GNN de paso de mensajes para distinguir entre estructuras de grafos; los nodos que reciben el mismo color estable son indistinguibles para tales redes. Sin embargo, computar esta coloración estable para grafos del mundo real y masivos enfrenta dos cuellos de botella críticos:

1. Ejecución Secuencial: Los algoritmos clásicos basados en el refinamiento de particiones (por ejemplo, el enfoque de divisor basado en $O(m \log n)$) son inherentemente secuenciales y no pueden explotar el paralelismo masivo de las arquitecturas de GPU modernas.
2. Requerimientos de Memoria Global: Estos algoritmos requieren que el grafo completo resida en memoria simultáneamente. Esta naturaleza "global" impide el análisis de grafos a escala web (por ejemplo, aquellos con decenas de miles de millones de aristas) que exceden la capacidad de un solo dispositivo o incluso de grandes clústeres de CPU.

Los intentos existentes por paralelizar o linealizar el problema a menudo dependen de la aritmética de punto flotante, lo que introduce inestabilidad numérica a gran escala, o no abordan la restricción de memoria de grafos más grandes que la memoria del dispositivo.

Metodología
Los autores proponen un enfoque de dos vertientes que desacopla la escalabilidad de cómputo de la escalabilidad de memoria, utilizando una interpretación algebraico-lineal de la coloración estable 1-WL.

1. Refinamiento Algebraico-Lineal Aleatorizado (Escalabilidad de Cómputo)
En lugar del refinamiento de particiones secuencial, los autores reformulan el problema utilizando la propiedad de invarianza de las particiones equitativas: una partición es estable si y solo si la imagen de cualquier vector constante por bloques bajo la matriz de adyacencia permanece constante por bloques.

• Algoritmo: El método itera muestreando vectores aleatorios simétricos sobre la partición actual, multiplicándolos por la matriz de adjacencia ($z = Aw$), y refinando la partición basándose en la simetría del vector resultante.
• Dominio Aritmético: Para evitar la inestabilidad numérica de la iteración de potencia de punto flotante (que puede causar fusiones incorrectas en grafos grandes), el algoritmo opera enteramente dentro del anillo modular $\mathbb{Z}/2^{64}\mathbb{Z}$ utilizando aritmética de enteros sin signo de 64 bits.
• Garantía Teórica: Aprovechando la restricción estructural de que las sumas de filas de la matriz de adyacencia están acotadas por $2^{64}/2$, los autores derivan una garantía de corrección probabilística. Un sondeo aleatorio detecta cualquier refinamiento faltante con una probabilidad de error de al menos $10^{-9}$, evitando la necesidad del clásico lema de Schwartz–Zippel que requiere un cuerpo.
• Mapeo a GPU: La operación central se reduce a una única multiplicación de matriz dispersa por vector (SpMV) seguida de un ordenamiento lexicográfico, procesos que son altamente eficientes en el hardware de GPU.

2. Procesamiento por Lotes Preservador de la Corrección (Escalabilidad de Memoria)
Para manejar grafos que no caben en la memoria del dispositivo, los autores introducen un esquema de procesamiento por lotes (batching) que descompone el conjunto de aristas en subconjuntos disjuntos preservando la corrección de la coloración estable global.

• Nodos Internos vs. Nodos Fronterizos: Para un lote dado de aristas, los nodos se clasifican como "internos" (todas las aristas salientes residen dentro del lote) o "fronterizos" (las aristas salientes cruzan los límites del lote).
• Refinamiento Local: El algoritmo procesa cada lote de forma independiente. Crucialmente, la partición inicial para cada lote coloca a todos los nodos internos en un solo bloque y a todos los nodos fronterizos en bloques de un solo elemento (singletons). Esto asegura que el refinamiento solo fusione nodos cuya información completa de vecindad esté disponible dentro de ese lote.
• Contracción Iterativa: Los refinamientos locales se fusionan en una partición global. El algoritmo construye entonces un grafo cociente (donde los nodos representan los bloques de la partición actual) y repite el proceso.
• Terminación: El proceso itera hasta que el grafo reducido quepa en un solo lote (obteniendo la coloración estable más gruesa demostrable) o hasta alcanzar un punto fijo. Si se alcanza un punto fijo antes de que el grafo quepa en memoria, el resultado es una coloración estable válida que es un refinamiento de la más gruesa.

Contribuciones Clave

1. Algoritmo de Refinamiento Aleatorizado: Un marco de Monte Carlo para la coloración estable 1-WL que opera en aritmética de enteros con garantías probabilísticas estrictas (error $<10^{-9}$), evitando la pérdida de precisión de los métodos de iteración de potencia de punto flotante.
2. Esquema de Procesamiento por Lotes: Una estrategia de descomposición que permite la computación 1-WL en grafos que exceden la memoria del dispositivo mediante el procesamiento de subconjuntos de aristas independientes y la fusión de resultados vía grafos cociente, manteniendo demostrablemente la propiedad de coloración estable.
3. Implementación en GPU: Una implementación en CUDA que mapea estos algoritmos directamente a primitivas de GPU (SpMV, ordenamiento, reducción), logrando aceleraciones masivas sobre las bases de comparación de CPU.

Resultados Experimentales
Los autores evaluaron su enfoque en once grafos de escala media y tres grafos de escala web (33–48 mil millones de aristas).

• Corrección: En grafos de escala media donde el grafo completo cabe en memoria ($N_b=1$), el algoritmo aleatorizado recuperó la coloración estable más gruesa exacta, coincidiendo con las bases deterministas.
• Estabilidad Numérica: El método propuesto de aritmética de enteros distinguió con éxito los nodos que fueron fusionados incorrectamente por un método previo de iteración de potencia de punto flotante en grafos grandes (por ejemplo, el conjunto de datos it-2004 mostró una discrepancia de más de 800,000 bloques).
• Calidad del Procesamiento por Lotes: Al usar múltiples lotes ($N_b=2, 4$) para simular restricciones de memoria, las particiones resultantes estuvieron empíricamente dentro del 5% de la coloración estable más gruesa. Para varios conjuntos de datos, la brecha fue cero.
• Rendimiento:
  • Velocidad: La implementación de una sola GPU logró aceleraciones de hasta 138× respecto al refinamiento de particiones clásico de un solo hilo en CPU y hasta 449× respecto a las bases de comparación multinúcleo de CPU.
  • Escalabilidad: En grafos de escala web con más de 30 mil millones de aristas, todas las bases de comparación de CPU (incluyendo la reducción simbólica multinúcleo) agotaron el tiempo de ejecución o se quedaron sin memoria. El método propuesto computó con éxito las coloraciones estables para estas instancias, logrando reducciones de aproximadamente ~39–46% en el número de bloques.

Significancia y Reivindicaciones
El artículo afirma que este trabajo es el primero en computar con éxito coloraciones estables 1-WL en grafos de escala web con más de 30 mil millones de aristas, una tarea previamente inviable debido a las limitaciones de memoria y cómputo. La significancia radica en:

• Habilitación del Análisis de Expresividad: Permite el análisis empírico de la expresividad de las GNN en conjuntos de datos reales y masivos, en lugar de limitarse a grafos pequeños o medianos.
• Ajuste Arquitectónico: Demuestra que el problema 1-WL, tradicionalmente visto como secuencial y global, puede ser efectivamente paralelizado y distribuido utilizando formulaciones algebraico-lineales y estrategias de procesamiento por lotes adaptadas a las arquitecturas de GPU.
• Robustez: Provee un método que es tanto numéricamente estable (vía aritmética de enteros) como escalable en memoria (vía procesamiento por lotes), superando los modos de falla específicos de enfoques anteriores.

Los autores concluyen que su método permite el análisis de la expresividad 1-WL a escalas previamente inalcanzables, señalando que la correspondencia entre la coloración estable y las particiones equitativas ofrece vías potenciales para aplicar estas técnicas a tareas de aprendizaje de grafos descendentes y reducción de modelos para sistemas dinámicos no lineales.