Resumo Técnico: GraViti – Autoencoders Variacionais de Nível de Grafo com Invariância de Permutação Relaxada

Declaração do Problema

A auto-codificação de nível de grafo visa mapear grafos inteiros em vetores latentes compactos para facilitar tarefas como interpolação, busca guiada por propriedades e modelagem generativa. Embora os Autoencoders Variacionais (VAEs) de nível de nó sejam bem estabelecidos, eles sofrem de uma limitação fundamental: o espaço latente é definido no nível do nó e depende de uma ordenação arbitrária dos nós. Consequentemente, dois grafos isomorfos podem mapear para conjuntos de vetores latentes totalmente diferentes, impedindo a definição de distâncias significativas entre grafos ou a captura de variedades de nível de grafo.

As abordagens existentes de nível de grafo enfrentam um trade-off. Métodos como PIGVAE e GRALE tentam impor invariância de permutação estrita, frequentemente dependendo de procedimentos de correspondência computacionalmente caros (por exemplo, correspondência de grafos, algoritmos de Sinkhorn) para alinhar grafos de entrada e reconstruídos. Essas etapas de alinhamento introduzem uma sobrecarga computacional significativa (tipicamente complexidade cúbica) e podem limitar a escalabilidade em grandes conjuntos de dados. Além disso, muitos decodificadores de nível de grafo dependem de estratégias de preenchimento que impõem um tamanho máximo fixo de grafo ou introduzem classes de "nó fictício", o que restringe a generalização para grafos maiores e viésa tarefas de classificação.

Metodologia: A Arquitetura GraViti

O GraViti é um VAE de nível de grafo baseado em transformers projetado para codificar grafos inteiros em vetores latentes únicos, relaxando o requisito estrito de reconstrução invariante a permutações quando existe uma ordenação canônica consistente de nós.

1. Representação do Grafo

O modelo opera em grafos com classes de nós e arestas discretas (por exemplo, tipos de átomos, tipos de ligações).

• Entrada: Uma matriz de classes de nós $X$, um tensor de classes de arestas $E$ e um vetor de características globais $y$ codificando estatísticas de nível de grafo (contagens de tipos de nós/arestas e tamanho do grafo).
• Codificações Posicionais: Para injetar informações estruturais, o modelo utiliza autovetores do Laplaciano e características de passeio aleatório como codificações posicionais ($P$).
• Conhecimento de Domínio: Para domínios específicos como moléculas, características adicionais específicas de nós ($K$), como massa atômica ou eletronegatividade, são incorporadas.
• Agrupamento (Batching): Os grafos são preenchidos até o tamanho do maior grafo em um lote, mas o modelo é projetado para ignorar completamente esses elementos preenchidos, evitando a necessidade de classes de nós fictícios.

2. Codificador

O codificador mapeia a triplet de entrada $(X, E, y)$ para os parâmetros da distribuição latente $(z_\mu, z_\sigma)$.

• Embedding: As classes de nós, classes de arestas e características globais são projetadas em embeddings.
• Pilha de Transformers de Grafo: Os embeddings são processados através de uma pilha de camadas de transformers de grafo (baseadas nas arquiteturas DiGress e CatFlow). Isso permite atualizações simultâneas de características de nível de nó, aresta e grafo, garantindo equivariância de permutação durante a fase de codificação.
• Leitura (Readout): Diferentemente do agrupamento por soma/média tradicional, o GraViti emprega uma leitura baseada em transformer de sequência. Dois decodificadores de sequência transformers paralelos processam as representações de nós e arestas, respectivamente, usando o vetor de características globais $y_{rep}$ como alvo para injetar contexto global. Isso garante que a saída seja independente da ordenação dos nós na etapa de leitura.

3. Previsão do Tamanho do Grafo

Uma inovação crucial é o desentrelaçamento explícito do tamanho do grafo do restante da representação latente.

• O modelo força as primeiras $d_s$ dimensões do espaço latente a codificar o log-tamanho do grafo.
• Um MLP prevê o log-tamanho $\hat{s}$ a partir dessas dimensões.
• Durante a decodificação, o log-tamanho real $s$ é re-codificado e reinjetado na representação latente para guiar a geração do número correto de nós, garantindo que o decodificador conheça o tamanho alvo do grafo sem depender de mascaramento.

4. Decodificador

O decodificador reconstrói o grafo a partir do vetor latente $z$.

• Transmissão (Broadcasting): O vetor latente é duplicado $N$ vezes (onde $N$ é o tamanho previsto ou real) e aumentado com codificações posicionais sinusoidais para quebrar a simetria.
• Transformer de Sequência: Um encoder transformer de sequência processa essa sequência, dividindo-a em sequências orientadas a nós e orientadas a arestas.
• Refinamento: Uma segunda pilha de transformers de grafo refina as representações de nós e arestas.
• Cabeças de Previsão: MLPs finais classificam os tipos de nós e tipos de arestas. O modelo prevê o tamanho do grafo diretamente, eliminando a necessidade de uma classe fictícia "sem aresta".

5. Objetivo de Treinamento

O modelo é treinado usando um objetivo VAE padrão:

• Perda de Reconstrução: Uma combinação de Perda Focal para classificação de nós e arestas (para lidar com desequilíbrio de classes) e uma perda L1 suave para previsão do tamanho do grafo.
• Perda Variacional: Um termo de divergência KL regulariza o espaço latente em direção a uma prior normal padrão.
• Invariância de Permutação Relaxada: Crucialmente, a perda de reconstrução é aplicada com base na ordem dos nós de entrada. Os autores argumentam que, em domínios com uma ordenação canônica confiável (por exemplo, moléculas, redes bayesianas), impor invariância de permutação estrita via correspondência é prejudicial. Ao confiar em uma ordenação consistente, o modelo alcança complexidade quadrática e maior fidelidade de reconstrução.

Contribuições Principais

1. Arquitetura Escalável: O GraViti alcança precisão de reconstrução state-of-the-art em grandes conjuntos de dados de grafos (PubChem16, PubChem32) ao evitar procedimentos caros de correspondência invariante a permutações.
2. Espaço Latente Estruturado: O espaço latente aprendido mostra-se bem estruturado, suportando interpolação suave, edição controlada e otimização orientada a propriedades. Em benchmarks moleculares, o modelo recupera princípios químicos estabelecidos (por exemplo, a relação entre o comprimento da cadeia alquila e a hidrofobicidade) diretamente do espaço latente.
3. Invariância de Permutação Relaxada: O artigo demonstra empiricamente que a invariância de permutação estrita nem sempre é necessária. Quando uma ordenação canônica consistente de nós está disponível, relaxar essa restrição melhora significativamente a qualidade da reconstrução e a escalabilidade em comparação com métodos que impõem invariância através de correspondência complexa.

Resultados Experimentais

Os autores avaliaram o GraViti em vários benchmarks:

• Reconstrução: Em conjuntos de dados moleculares (QM9, PubChem16, PubChem32), o GraViti supera as linhas de base (GRALE, PIGVAE) em até 25% na precisão de reconstrução, particularmente em grafos maiores onde métodos baseados em correspondência têm dificuldades.
• Estudos de Ablação: Remover o encoder transformer ou a leitura baseada em transformer degrada significativamente o desempenho, confirmando a importância desses componentes.
• Dessensibilização (Denoising): O modelo restaura efetivamente estruturas quimicamente válidas a partir de entradas ruidosas, com a variante VAE mostrando robustez superior em comparação à variante AE.
• Otimização de Propriedades: Usando CMA-ES no espaço latente, o modelo maximiza com sucesso o logP (lipofilicidade), gerando moléculas com cadeias de carbono progressivamente mais longas ou grupos hidrofóbicos específicos, alinhando-se com tendências químicas conhecidas.
• Interpolação: A interpolação linear no espaço latente produz transições suaves entre estruturas moleculares.
• Tarefas de Descendente: Vetores latentes treinados no PubChem preveem com alta precisão propriedades moleculares (momento de dipolo, gap HOMO-LUMO) e escores do Critério de Informação Bayesiano (BIC) para redes bayesianas.
• Geração: Embora a validade da geração seja menor no PubChem (aproximadamente 50-70%, dependendo da representação aromática) em comparação com o QM9, o modelo produz moléculas diversas e novas. Os autores atribuem a validade menor em conjuntos de dados maiores à compactação global do espaço latente, e não a falhas estruturais locais.

Significado e Alegações

O artigo alega que o GraViti fornece uma base prática para aprendizado de representação de nível de grafo. Seu principal significado reside em demonstrar que a invariância de permutação estrita não é um pré-requisito para a auto-codificação eficaz de nível de grafo quando existe uma ordenação de nós estável. Ao relaxar essa restrição, o GraViti evita os gargalos computacionais da correspondência de grafos, permitindo treinamento escalável em conjuntos de dados massivos enquanto mantém alta qualidade de reconstrução.

Os autores posicionam o GraViti como uma alternativa leve a pipelines de geração complexos, oferecendo um mecanismo de decodificação de etapa única que equilibra desempenho generativo com qualidade prática de amostras. O trabalho sugere que, para domínios como grafos moleculares e redes bayesianas, onde ordenações canônicas são inerentes ou podem ser construídas, modelos de nível de grafo podem alcançar desempenho superior sem a sobrecarga de perdas baseadas em alinhamento. Propõe-se trabalho futuro para integrar restrições conscientes do domínio para melhorar ainda mais a validade da geração e explorar a transferência para outros domínios de grafos, como circuitos e proteínas.