Invariant-Based Weight Sharing for Message Passing

Resumo Técnico: Compartilhamento de Pesos Baseado em Invariantes para Passagem de Mensagens

Declaração do Problema
Redes neurais de passagem de mensagens (MPNNs) são um framework dominante para aprender representações de dados estruturados em grafos. No entanto, MPNNs padrão e transformadores de grafos geralmente dependem de compartilhamento de pesos baseado exclusivamente em características de nós ou arestas, ou aprendem coeficientes de atenção específicos da instância. Essa abordagem limita a capacidade do modelo de capturar padrões estruturais inerentes, pois os pesos não estão explicitamente ligados à topologia do grafo. Além disso, MPNNs padrão são frequentemente restritos à agregação de vizinhança local, tornando-os suscetíveis ao excesso de suavização e ao excesso de compressão em arquiteturas profundas, enquanto transformadores de grafos frequentemente carecem de reutilização sistemática de parâmetros entre subgrafos estruturalmente equivalentes.

Metodologia: ShareGNNs
Os autores introduzem ShareGNNs, uma arquitetura inovadora fundamentada em compartilhamento de pesos baseado em invariantes. O princípio central é indexar parâmetros aprendíveis diretamente por invariantes de grafos — funções preservadas sob permutações de nós — em vez de dimensões de características fixas ou posições específicas da instância.

1. Compartilhamento de Pesos Baseado em Invariantes:

   • Codificador (Passagem de Mensagens): Em vez de matrizes de pesos fixas, o modelo constrói dinamicamente pesos de passagem de mensagens para cada par de nós $(v, w)$ com base em uma tripla invariante: $(l(v), l(w), d(v, w))$, onde $l$ é uma função de rotulagem (por exemplo, número atômico, rótulo de Weisfeiler-Leman, contagem de padrões) e $d$ é a distância do caminho mais curto. Todos os pares de nós que compartilham a mesma tripla invariante compartilham um único peso escalar aprendível $\omega_\tau$. Isso permite que a informação flua entre nós arbitrários em uma única camada, não apenas entre vizinhos imediatos.
   • Decodificador (Agrupamento): Da mesma forma, termos de viés e pesos de agrupamento para agregação em nível de grafo são indexados pelos rótulos invariantes de nós individuais.
   • Construção Dinâmica: Para qualquer grafo de entrada, as matrizes de pesos relevantes são montadas sob demanda a partir de um pool de parâmetros compartilhados. Isso garante que o modelo seja equivariante à permutação no codificador e invariante à permutação no decodificador por construção.

2. Funções de Rotulagem:
   A expressividade do modelo é controlada pela escolha das funções de rotulagem $l$. O artigo explora:

   • Rotulagem Molecular: Números atômicos.
   • Rotulagem Weisfeiler-Leman (WL): Refinamento iterativo de rótulos de nós com base em multiconjuntos de vizinhança.
   • Rotulagem por Padrões: Contagens de padrões específicos de subgrafos (por exemplo, triângulos, ciclos) envolvendo um nó.
     Estes podem ser combinados para capturar informações estruturais de múltiplos níveis.

3. Arquitetura:
   ShareGNNs utilizam uma estrutura simples de codificador-decodificador. O codificador propaga informações através de distâncias arbitrárias em uma única camada, e o decodificador agrega incorporações de nós em uma representação de grafo usando agrupamento indexado por invariantes. A arquitetura suporta extensões de múltiplas cabeças onde diferentes cabeças utilizam diferentes conjuntos de invariantes.

Principais Contribuições

1. Princípio de Compartilhamento de Pesos Baseado em Invariantes: Um novo mecanismo de compartilhamento de pesos que incorpora explicitamente a estrutura do grafo, permitindo reutilização sistemática de parâmetros entre pares de nós estruturalmente equivalentes dentro e entre diferentes grafos.
2. Arquitetura ShareGNN: Uma instanciação prática desse princípio que combina passagem de mensagens com conectividade semelhante à de transformadores (interações todos-para-todos), mas com parâmetros ligados a invariantes estruturais em vez de pontuações de atenção aprendidas.
3. Validação Teórica e Empírica:
   • Expressividade: Os autores provam que ShareGNNs são pelo menos tão expressivos quanto a função de rotulagem escolhida. Se dois grafos são distinguíveis pela função de rotulagem (por exemplo, rótulos WL), o ShareGNN pode distingui-los.
   • Controle: O modelo oferece controle explícito sobre complexidade e expressividade através da seleção de invariantes.
   • Desempenho: Experimentos demonstram que ShareGNNs superam MPNNs padrão (por exemplo, GIN, GAT) e são competitivos com transformadores de grafos em várias benchmarks, incluindo classificação molecular, análise de redes sociais e contagem de subestruturas.

Resultados

• Classificação de Grafos: Em conjuntos de dados diversos (NCI1, NCI109, MUTAG, IMDB-B/M), ShareGNNs alcançam precisão state-of-the-art ou altamente competitiva. Notavelmente, alcançam precisão quase perfeita (100%) em benchmarks sintéticos (RingTransfer, Snowflakes) projetados para desafiar MPNNs padrão, demonstrando sua capacidade de capturar dependências de longo alcance e distinguir grafos indistinguíveis pelo teste 1-WL.
• Contagem de Subestruturas: Na Benchmark de Contagem de Subestruturas, ShareGNNs alcançam resultados state-of-the-art em configurações de tarefa única e permanecem competitivos em configurações de múltiplas tarefas, validando que invariantes expressivos levam a modelos expressivos mesmo com arquiteturas rasas (camada única de codificador).
• Escalabilidade: O modelo escala para grandes conjuntos de dados como ZINC (12k e 250k grafos) enquanto mantém uma arquitetura rasa, evitando problemas de excesso de suavização comuns em MPNNs profundos.
• Eficiência: Estudos de ablação mostram que aumentar a distância máxima de passagem de mensagens $D$ tem um impacto mais significativo no desempenho do que adicionar mais camadas de codificador. Além disso, podar triplas invariantes raras reduz significativamente a contagem de parâmetros sem sacrificar a precisão.

Significado e Alegações
O artigo alega que o compartilhamento de pesos baseado em invariantes fornece uma alternativa principial ao compartilhamento de pesos baseado em características e à atenção específica da instância. Ao vincular pesos a invariantes estruturais, ShareGNNs oferecem:

• Interpretabilidade Aprimorada: Pesos aprendidos podem ser interpretados diretamente em termos de padrões estruturais (por exemplo, distâncias específicas entre tipos específicos de átomos).
• Transferibilidade: Parâmetros são naturalmente transferíveis entre grafos com padrões estruturais semelhantes, sugerindo potencial para modelos fundamentais de grafos.
• Controle de Expressividade: O modelo permite que pesquisadores controlem explicitamente o poder discriminativo da rede selecionando invariantes apropriados, fechando a lacuna entre expressividade teórica (por exemplo, testes k-WL) e implementação prática.

Os autores reconhecem limitações, incluindo o custo de pré-processamento de computação de invariantes e distâncias pares, e a complexidade de memória quadrática no pior caso, embora observem que isso é comparável a transformadores de grafos e pode ser mitigado por poda. Eles também notam que a abordagem atual requer invariantes definidos pelo usuário, deixando espaço para trabalhos futuros em seleção automatizada de invariantes.