Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

Resumo Técnico: Escalando a Análise de Expressividade de Weisfeiler–Leman para Grafos Massivos com GPUs

Declaração do Problema
O teste de Weisfeiler–Leman (1-WL) e sua coloração estável associada são fundamentais para compreender o poder expressivo das Redes Neurais de Grafos (GNNs). A coloração estável 1-WL fornece um limite superior na capacidade de distinguir estruturas de grafos por redes neurais de passagem de mensagens; nós que recebem a mesma cor estável são indistinguíveis para tais redes. No entanto, computar essa coloração estável para grafos reais e massivos enfrenta dois gargalos críticos:

1. Execução Sequencial: Algoritmos clássicos baseados em refinamento de partições (por exemplo, a abordagem de splitter de $O(m \log n)$) são inerentemente sequenciais e não podem explorar o paralelismo massivo das arquiteturas modernas de GPU.
2. Requisitos de Memória Global: Esses algoritmos exigem que o grafo inteiro resida na memória simultaneamente. Essa natureza "global" impede a análise de grafos de escala web (por exemplo, aqueles com dezenas de bilhões de arestas) que excedem a capacidade de dispositivos únicos ou mesmo de grandes clusters de CPU.

Tentativas existentes de paralelizar ou linearizar o problema frequentemente dependem de aritmética de ponto flutuante, o que introduz instabilidade numérica em escala, ou falham em abordar a restrição de memória de grafos maiores que a memória do dispositivo.

Metodologia
Os autores propõem uma abordagem de duas frentes que desacopla a escalabilidade de computação da escalabilidade de memória, utilizando uma interpretação algébrica linear da coloração estável 1-WL.

1. Refinamento Algébrico Linear Aleatório (Escalabilidade de Computação)
Em vez de um refinamento de partição sequencial, os autores reformulam o problema usando a propriedade de invariância de partições equitativas: uma partição é estável se, e somente se, a imagem de qualquer vetor constante por bloco sob a matriz de adjacência permanecer constante por bloco.

• Algoritmo: O método itera amostrando vetores aleatórios simétricos na partição atual, multiplicando-os pela matriz de adjacência ($z = Aw$) e refinando a partição com base na simetria do vetor resultante.
• Domínio Aritmético: Para evitar a instabilidade numérica da potenciação de ponto flutuante (que pode causar fusões incorretas em grafos grandes), o algoritmo opera inteiramente dentro do anel modular $\mathbb{Z}/2^{64}\mathbb{Z}$ usando aritmética de inteiros sem sinal de 64 bits.
• Garantia Teórica: Ao aproveitar a restrição estrutural de que as somas das linhas da matriz de adjacência são limitadas por $2^{64}/2$, os autores derivam uma garantia de correção probabilística. Uma sonda aleatória detecta qualquer refinamento ausente com uma probabilidade de falha de no máximo $10^{-9}$, contornando a necessidade do clássico lema de Schwartz–Zippel, que requer um corpo (field).
• Mapeamento para GPU: A operação central reduz-se a uma única multiplicação de matriz esparsa por vetor (SpMV) seguida de uma ordenação lexicográfica, ambas altamente eficientes em hardware de GPU.

2. Processamento em Lotes Preservando a Correção (Escalabilidade de Memória)
Para lidar com grafos que não cabem na memória do dispositivo, os autores introduzem um esquema de lote (batching) que decompõe o conjunto de arestas em subconjuntos disjuntos, preservando a correção da coloração estável global.

• Nós Internos vs. Nós de Fronteira: Para um determinado lote de arestas, os nós são classificados como "internos" (todas as arestas de saída residem dentro do lote) ou de "fronteira" (as arestas de saída cruzam os limites do lote).
• Refinamento Local: O algoritmo processa cada lote de forma independente. Crucialmente, a partição inicial para cada lote coloca todos os nós internos em um único bloco e todos os nós de fronteira em blocos de um único elemento (singleton). Isso garante que o refinamento apenas funda nós cuja informação completa de vizinhança esteja disponível dentro daquele lote.
• Contração Iterativa: Os refinamentos locais são fundidos em uma partição global. O algoritmo então constrói um grafo quociente (onde os nós representam os blocos da partição atual) e repete o processo.
• Terminação: O processo itera até que o grafo reduzido caiba em um único lote (resultando na coloração estável mais grosseira) ou atinja um ponto fixo. Se um ponto fixo for alcançado antes do grafo caber na memória, o resultado é uma coloração estável válida que é um refinamento da mais grosseira.

Principais Contribuições

1. Algoritmo de Refinamento Aleatório: Um framework de Monte Carlo para coloração estável 1-WL que opera em aritmética de inteiros com garantias probabilísticas rigorosas (erro $< 10^{-9}$), evitando a perda de precisão dos métodos de potenciação de ponto flutuante.
2. Esquema de Lote (Batching): Uma estratégia de decomposição que permite o cálculo 1-WL em grafos que excedem a memória do dispositivo, processando subconjuntos independentes de arestas e fundindo resultados via grafos quocientes, mantendo comprovadamente a propriedade de coloração estável.
3. Implementação em GPU: Uma implementação em CUDA que mapeia esses algoritmos diretamente para primitivas de GPU (SpMV, ordenação, redução), alcançando acelerações massivas sobre as bases de comparação (baselines) de CPU.

Resultados Experimentais
Os autores avaliaram sua abordagem em onze grafos de média escala e três grafos de escala web (33–48 bilhões de arestas).

• Correção: Em grafos de média escala onde o grafo completo cabe na memória ($N_b=1$), o algoritmo aleatório recuperou a coloração estável mais grosseira exata, correspondendo às bases determinísticas.
• Estabilidade Numérica: O método de aritmética de inteiros proposto distinguiu com sucesso nós que foram incorretamente fundidos por um método anterior de potenciação de ponto flutuante em grafos grandes (por exemplo, o conjunto de dados it-2004 mostrou uma discrepância de mais de 800.000 blocos).
• Qualidade do Lote: Ao usar múltiplos lotes ($N_b=2, 4$) para simular restrições de memória, as partições resultantes ficaram empiricamente dentro de 5% da coloração estável mais grosseira. Para vários conjuntos de dados, a diferença foi zero.
• Desempenho:
  • Velocidade: A implementação de GPU única alcançou acelerações de até 138× sobre o refinamento de partição clássico de thread única em CPU e até 449× sobre as bases de comparação multicore de CPU.
  • Escalabilidade: Em grafos de escala web com mais de 30 bilhões de arestas, todas as bases de comparação de CPU (incluindo redução simbólica multicore) sofreram timeout ou falta de memória. O método proposto calculou com sucesso as colorações estáveis para essas instâncias, alcançando reduções de aproximadamente 39–46% no número de blocos.

Significância e Alegações
O artigo afirma que este trabalho é o primeiro a calcular com sucesso colorações estáveis 1-WL em grafos de escala web com mais de 30 bilhões de arestas, uma tarefa anteriormente inviável devido às limitações de memória e computação. A significância reside em:

• Habilitar a Análise de Expressividade: Permitir a análise empírica da expressividade de GNNs em conjuntos de dados reais e massivos, em vez de ser limitado a grafos pequenos ou médios.
• Ajuste Arquitetural: Demonstrar que o problema 1-WL, tradicionalmente visto como sequencial e global, pode ser efetivamente paralelizado e distribuído usando formulações algébricas lineares e estratégias de lote adaptadas para arquiteturas de GPU.
• Robustez: Fornecer um método que é tanto numericamente estável (via aritmética de inteiros) quanto escalável em memória (via lote), superando os modos de falha específicos de abordagens anteriores.

Os autores concluem que seu método permite a análise de expressividade 1-WL em escalas anteriormente inalcançáveis, observando que a correspondência entre coloração estável e partições equitativas oferece caminhos futuros potenciais para aplicar essas técnicas em tarefas de aprendizado de máquina em grafos e redução de modelos para sistemas dinâmicos não lineares.