Resumo Técnico: Difusão Latente Principiada para Grafos via Autoencoders Laplacianos

Declaração do Problema

Modelos de difusão em grafos alcançaram desempenho de última geração em tarefas generativas em domínios como design molecular, otimização combinatória e busca de arquitetura neural. No entanto, as abordagens existentes tipicamente operam diretamente no espaço de grafos, incorrendo em complexidade computacional quadrática ($O(n^2)$) em relação ao número de nós. Esse gargalo surge porque esses modelos devem aprender tanto representações de nós quanto de arestas, frequentemente desperdiçando capacidade ao modelar a ausência de arestas em grafos esparsos. Além disso, tentativas de adaptar modelos autoregressivos escaláveis para grafos frequentemente quebram a equivariância a permutações ao linearizar estruturas de grafos.

Uma alternativa promissora é a difusão latente, que comprime dados em um espaço de baixa dimensionalidade antes de aplicar o processo generativo (como visto na geração de imagens com VAEs). No entanto, aplicar isso a grafos apresenta um desafio único: ao contrário de imagens, onde erros menores de decodificação são visualmente negligenciáveis, a geração de grafos requer reconstrução quase sem perdas. Um único erro na reconstrução de uma matriz de adjacência pode invalidar toda a amostra (por exemplo, quebrando ligações químicas ou violando restrições estruturais). Modelos latentes de difusão em grafos existentes frequentemente falham em satisfazer esse requisito de fidelidade ou sacrificam eficiência ao usar arquitetões ad hoc que ainda dependem de tokens de arestas quadráticos.

Metodologia: LG-Flow

Os autores propõem o LG-Flow, um framework de difusão latente em grafos projetado para superar esses obstáculos enquanto mantém alta fidelidade de reconstrução. O framework consiste em dois componentes principais: um autoencoder principiado e um modelo generativo de correspondência de fluxo latente.

1. Autoencoder Variacional de Grafo Laplaciano (LG-VAE)

A inovação central é um autoencoder equivariante a permutações que mapeia grafos para um espaço latente onde a dimensionalidade escala linearmente com o número de nós ($n \times d$), em vez de quadraticamente.

• Base Teórica: A arquitetura é fundamentada em codificações posicionais identificadoras de adjacência. Os autores exploram a propriedade de que, se os embeddings de nós forem suficientemente "identificadores de adjacência", a matriz de adjacência pode ser reconstruída a partir desses embeddings.
• Codificador: O codificador utiliza Codificações Posicionais Laplacianas (LPE) derivadas da decomposição espectral do Laplaciano do grafo ($L = U\Lambda U^\top$). Ele incorpora características de nós e arestas em uma distribuição posterior Gaussiana. Para lidar com Grafos Acíclicos Direcionados (DAGs), o framework estende-se ao Laplaciano Magnético ($L_q$), uma matriz Hermitiana que preserva a direcionalidade das arestas, resultando em uma LPE Magnética (mLPE).
• Decodificador: O decodificador é projetado para reconstrução de alta fidelidade.
  • Rótulos de Nós: Decodificados diretamente via uma camada linear e softmax.
  • Matriz de Adjacência: Reconstruída via uma camada bilinear aplicada aos embeddings latentes ($Z W_Q (Z W_K)^\top$). A detecção de arestas é tratada como um problema de classificação binária usando uma arquitetura DeepSet para garantir equivariância a permutações e lidar com tamanhos de grafos variáveis.
  • Rótulos de Arestas: Decodificados usando uma camada bilinear multi-cabeça.
• Garantia de Reconstrução: O design assegura que as representações latentes contenham informações estruturais suficientes para recuperar a matriz de adjacência, abordando a necessidade crítica de reconstrução quase sem perdas na geração de grafos.

2. Correspondência de Fluxo Latente (LG-Flow)

Uma vez que o LG-VAE é treinado, seus parâmetros são congelados. Um modelo generativo é então treinado no espaço latente $Z \in \mathbb{R}^{n \times d}$.

• Arquitetura: Os autores empregam um Transformador de Difusão (DiT) com objetivos de Correspondência de Fluxo (FM).
• Processo: O modelo aprende a prever o campo de velocidade que transporta ruído para a distribuição latente.
• Eficiência: Ao operar no espaço latente, o DiT evita a complexidade quadrática de prever a matriz de adjacência completa diretamente. A única etapa quadrática é a decodificação final, que é tratada por um decodificador leve. Isso permite o treinamento de backbones generativos substancialmente maiores.

Contribuições Principais

1. LG-VAE: Um autoencoder principiado e equivariante a permutações que alcança reconstrução quase sem perdas tanto para grafos não direcionados quanto para DAGs. É o primeiro a fornecer garantias teóricas (via propriedades identificadoras de adjacência) de que a matriz de adjacência é recuperável a partir do espaço latente.
2. Escalabilidade Linear: O framework mapeia nós para embeddings de dimensão fixa, garantindo que o tamanho da representação latente escale linearmente com o número de nós, removendo efetivamente o gargalo quadrático da difusão no espaço de grafos.
3. Framework Unificado: A abordagem unifica a geração de grafos não direcionados e DAGs dentro de um único framework modular, permitindo o uso de arquiteturas de remoção de ruído genéricas (DiT) e técnicas da geração de imagens.
4. Desempenho: O método alcança desempenho generativo competitivo enquanto entrega acelerações significativas (até 1000$\times$) em comparação com modelos de difusão em grafos de última geração.

Resultados Experimentais

Os autores avaliaram o LG-Flow em seis conjuntos de dados: três sintéticos (Planar Estendido, Árvore Estendida, Ego), dois moleculares (MOSES, GuacaMol) e dois conjuntos de dados de DAGs (TPU Tile, Price).

• Fidelidade de Reconstrução (Q1): Em benchmarks estruturais (Planar Estendido e Árvore Estendida), o LG-VAE alcançou reconstrução quase sem perdas (Precisão de Amostra de 0,9961 e 0,9844, respectivamente), superando significativamente baselines como DGVAE e GLAD, que falharam em reconstruir amostras exatamente.
• Desempenho Generativo (Q2):
  • Grafos Sintéticos: O LG-Flow alcançou resultados competitivos em métricas de distribuição (Grau, Órbita, Agrupamento, Espectro) e em Validade, Unicidade e Novidade (V.U.N.).
  • Moléculas: No GuacaMol e MOSES, o método alcançou desempenho de última geração ou competitivo em validade e correspondência de distribuição química (FCD), frequentemente superando modelos de difusão discreta como DiGress e DisCo.
  • DAGs: No conjunto de dados TPU Tile, o LG-Flow superou o autoregressivo LayerDAG e o baseado em difusão Directo em métricas de distribuição, mantendo alta validade.
• Eficiência (Q3):
  • O LG-Flow demonstrou acelerações massivas, variando de 10$\times$ a 1000$\times$ mais rápido na inferência em comparação com baselines como DeFoG e Cometh.
  • O uso de memória escalou linearmente com o tamanho do lote, permitindo que todo o lote de teste de 151 amostras coubesse em uma única GPU, enquanto o DeFoG esgotou a memória em tamanhos de lote tão baixos quanto 32.

Significado e Alegações

O artigo alega que o LG-Flow demonstra que mover a difusão em grafos para um espaço latente de nós de alta fidelidade pode melhorar substancialmente a eficiência da inferência, particularmente para grafos pequenos a médios. O framework desloca o foco do design de modelos de difusão complexos e específicos para grafos para o desenvolvimento de autoencoders em grafos fortes e mais baratos de treinar.

Ao unificar a geração de grafos não direcionados e direcionados e permitir a transferência de arquiteturas poderosas como o Transformador de Difusão da geração de imagens, os autores sugerem um caminho para escalar modelos generativos para grafos muito maiores que anteriormente eram intratáveis devido à complexidade quadrática. O trabalho destaca que, com um autoencoder principiado garantindo fidelidade de reconstrução, os benefícios da difusão latente (eficiência, modularidade e capacidade) podem ser realizados com sucesso no domínio de grafos.