Resumen Técnico: CoRe-GNN: Paso de Mensajes Multinivel en Grafos Coartados

Declaración del Problema

El entrenamiento de Redes Neuronales de Grafos (GNN) en grafos de gran escala está fundamentalmente limitado por el costo de memoria requerido para almacenar las representaciones intermedias de los nodos a través de las capas. En las arquitecturas estándar de paso de mensajes, la representación de un nodo en la capa $l$ depende de su vecindad de $l$ saltos. Para grafos grandes, esta vecindad crece exponencialmente con la profundidad, lo que genera cuellos de botella de memoria que saturan la memoria de la GPU tanto durante el entrenamiento como durante la inferencia.

Los enfoques escalables existentes generalmente se dividen en dos categorías, ambas de las cuales sufren limitaciones complementarias:

1. Coartación de Grafos (Graph Coarsening): Métodos como Coarsen-GNN mapean los nodos originales a "supernodos" para crear un grafo más pequeño. Si bien esto permite una propagación eficiente y ofrece garantías espectrales, obliga a que todos los nodos dentro de un supernodo compartan la misma representación. Esta "uniformidad" limita la discriminabilidad por nodo, particularmente en entornos heterofílicos donde los nodos adyacentes pueden tener etiquetas diferentes.
2. Restricción Basada en Clústeres (Cluster-Based Restriction): Métodos como Cluster-GCN restringen el paso de mensajes únicamente a las conexiones intra-clúster, descartando las conexiones inter-clúster. Esto permite un procesamiento por lotes (batching) eficiente al tratar los clústeres como subgrafos independientes. Sin embargo, el corte de las conexiones inter-clúster elimina las dependencias de largo alcance, degradando el rendimiento en tareas que requieren la agregación de información a través de muchos saltos.

Ambos enfoques dependen de una descomposición del grafo en grupos de nodos, pero no logran equilibrar la compensación entre preservar la estructura de largo alcance y mantener la discriminabilidad local de los nodos.

Metodología: CoRe-GNN

Los autores proponen CoRe-GNN (Coarsen and Restore), una arquitectura diseñada para abordar simultáneamente el techo de precisión de la coartación de grafos y la ceguera de largo alcance de Cluster-GCN.

Formulación Matricial Unificada

El artículo establece primero una perspectiva unificada sobre las GNN de grafos comprimidos viendo estas como modificaciones estructuradas de la matriz de propagación estándar $S$.

• Coarsen-GNN reemplaza $S$ con $\Pi S \Pi$, donde $\Pi = QP$ es un operador de constante por bloques que mapea los nodos a supernodos y viceversa.
• Cluster-GCN reemplaza $S$ con $S_{intra}$, una matriz de bloques diagonales que retiene solo las conexiones intra-clúster.

La Arquitectura CoRe-GNN

CoRe-GNN realiza dos propagaciones paralelas en cada capa, acoplando el espacio coartado y el espacio de nodos originales:

1. Propagación Inter-Clúster (Coartada): Un término $S_{inter}^c$ opera sobre el grafo coartado (supernodos). Esto captura la información estructural de largo alcance de manera eficiente.
2. Propagación Intra-Clúster (Local): Un término $S_{intra}$ opera en el espacio de nodos originales, preservando la discriminabilidad por nodo dentro de los clústeres.

La regla de actualización para la capa $l$ es:
$$H^{(l)} = Q \sigma(S_{inter}^c P H^{(l-1)} \theta_c^{(l)}) + \sigma(S_{intra} H^{(l-1)} \theta^{(l)})$$
Aquí, $P$ y $Q$ son matrices de reducción y elevación derivadas de la partición del grafo. La arquitectura mantiene dos estados paralelos: $H_c^{(l)}$ en el espacio coartado y $H^{(l)}$ en el espacio de nodos originales. Los pesos $\theta^{(l)}$ y $\theta_c^{(l)}$ son independientes, lo que permite a la red aprender transformaciones distintas para las interacciones locales y de largo alcance.

Esquema de Procesamiento por Lotes (Batching)

Para escalar a grafos con millones de nodos, CoRe-GNN introduce una estrategia de batching natural:

• Término inter-clúster: Se calcula sobre todos los $n$ supernodos (donde $n \ll N$) en cada paso, lo cual es computacionalmente económico.
• Término intra-clúster: Se calcula solo para los clústeres en el lote actual $B$. Para los supernodos no incluidos en el lote, el término intra-clúster se aproxima mediante un proxy de bucle de auto-conexión (self-loop proxy) en el espacio coartado durante el entrenamiento.
  Este enfoque reduce la complejidad de memoria a $O(K(n+b)d)$, independiente del número total de nodos $N$, permitiendo el entrenamiento en grafos que de otro modo excederían la memoria de la GPU.

Inferencia Exacta: Aunque el esquema de batching de entrenamiento utiliza una aproximación, los autores derivan un procedimiento de inferencia exacta no trivial que evita este cuello de botella. En este modo, el término inter-clúster se calcula sobre todos los supernodos en la GPU, mientras que el término intra-clúster se calcula clúster por clúster, cargando bloques de la CPU a la GPU bajo demanda y descartándolos inmediatamente. Esto permite la inferencia exacta en grafos muy grandes sin retener el grafo computacional completo a través de las capas.

Garantías Teóricas

Los autores demuestran que CoRe-GNN hereda garantías de aproximación similares a la coartación de grafos. Específicamente, acotan el error entre el paso de mensajes en el grafo original y CoRe-GNN utilizando la constante de Aproximación Espectral Restringida (RSA). El límite teórico para CoRe-GNN depende de las propiedades del subgrafo inter-clúster ($S_{inter}$) en lugar del grafo completo, lo que sugiere que a medida que aumenta la razón de coartación, los límites de error pueden mejorar en comparación con la coartación estándar.

Contribuciones Clave

1. Perspectiva Unificada: El artículo proporciona una formulación matricial que muestra que la coartación de grafos y Cluster-GCN son modificaciones estructuradas complementarias de la matriz de propagación, aclarando sus respectivas limitaciones.
2. Arquitectura CoRe-GNN: Una arquitectura novedosa que acopla las propagaciones inter-clúster (coartada) e intra-clúster (local) en paralelo, mitigando la pérdida de discriminabilidad en la coartación y la pérdida de información de largo alcance en la agrupación (clustering).
3. Batching Escalable: Un esquema de batching natural que permite a CoRe-GNN entrenar en grafos con millones de nodos, junto con un procedimiento de inferencia exacta que evita la aproximación.
4. Análisis Teórico: Derivación de las garantías de propagación para CoRe-GNN, demostrando que preserva señales suaves con constantes dependientes de la estructura inter-clúster.
5. Validación Empírica: Evaluación exhaustiva en benchmarks homofílicos, heterofílicos, de gran escala y de largo alcance.

Resultados Experimentales

Los autores evaluaron CoRe-GNN en un conjunto diverso de conjuntos de datos, incluyendo:

• Homofílicos: Cora, CiteSeer, PubMed.
• Heterofílicos: Chameleon, Squirrel, Amazon Ratings.
• Gran Escala: ogbn-arxiv, Reddit, ogbn-products.
• Largo Alcance: Redes viales de París y Londres.

Hallazgos Clave:

• Rendimiento: CoRe-GNN supera consistentemente a las líneas base de Coarsen-GNN y Cluster-GCN en todas las categorías.
• Tareas de Largo Alcance: En los conjuntos de datos de París y Londres, CoRe-GNN alcanza la mayor precisión, demostrando su capacidad para capturar dependencias de largo alcance que Cluster-GCN pierde y la discriminabilidad que carece la Coarsen-GNN pura.
• Heterofilia: En conjuntos de datos heterofílicos (ej. Squirrel), CoRe-GNN supera significativamente a Coarsen-GNN, que tiene dificultades cuando los nodos en un mismo supernodo tienen etiquetas diferentes.
• Escalabilidad: En conjuntos de datos de gran escala como Reddit y ogbn-products, el GCN estándar y la versión no batcheada de CoRe-GNN agotan la memoria (OOM). La variante con batching de CoRe-GNN logra entrenar con éxito y alcanzar una precisión competitiva o superior a la de Cluster-GCN.
• Ablación: Los experimentos confirman que el uso de pesos independientes para las dos ramas y el uso de la matriz de propagación específica para inter-clúster ($S_{inter}^c$) en lugar de la matriz coartada completa produce un mejor rendimiento.

Significado y Reivindicaciones

El artículo sostiene que CoRe-GNN logra unificar los beneficios de la coartación de grafos y la restricción basada en clústeres. Al realizar propagaciones paralelas, supera los "fallos complementarios" de los métodos escalables existentes: la pérdida de detalle a nivel de nodo en la coartación y el corte de las dependencias de largo alcance en la agrupación.

Los autores enfatizan que CoRe-GNN proporciona un marco que:

1. Mantiene garantías espectrales para la propagación de señales suaves.
2. Preserva la discriminabilidad por nodo en entornos heterofílicos.
3. Captura dependencias de largo alcance esenciales para tareas de grafos complejas.
4. Escala a millones de nodos mediante un esquema de batching eficiente en memoria.

El trabajo sugiere que la descomposición del grafo en grupos no tiene por qué ser un intercambio entre precisión y escalabilidad; más bien, un enfoque multinivel puede lograr ambos. Los autores señalan que el marco está actualmente limitado a GNN con matrices de propagación fijas (excluyendo mecanismos de atención) y que el trabajo futuro podría explorar extensiones hacia conexiones residuales y escenarios de aprendizaje federado asíncrono.