An LLM-Guided Query-Aware Inference System for GNN Models on Large Knowledge Graphs

O artigo apresenta o KG-WISE, um sistema de inferência orientado a consultas que utiliza Grandes Modelos de Linguagem (LLMs) para decompor e carregar dinamicamente componentes específicos de modelos de Redes Neurais em Grafos (GNNs) em grandes Grafos de Conhecimento, alcançando ganhos significativos de velocidade e redução de uso de memória sem comprometer a precisão.

O essencial

Imagine que você tem uma biblioteca gigante (o "Conhecimento" ou Knowledge Graph) com milhões de livros, artigos e conexões entre eles. Agora, imagine que você tem um super-inteligente (a Rede Neural ou GNN) que pode responder perguntas sobre essa biblioteca, como "Quem escreveu este livro?" ou "Quais são as conexões entre dois cientistas?".

O problema é que, para responder a uma pergunta simples, o sistema tradicional precisa:

1. Carregar todos os milhões de livros na mesa de trabalho.
2. Ler toda a enciclopédia inteira, mesmo que a pergunta seja apenas sobre um único autor.
3. Usar uma quantidade enorme de energia e espaço, deixando o computador lento e caro.

É como tentar achar uma agulha em um palheiro, mas em vez de procurar só no palheiro, você carrega o palheiro inteiro, o celeiro inteiro e o fazendeiro inteiro para dentro da sua sala apenas para olhar uma única agulha.

A Solução: KG-WISE (O "Detetive Inteligente")

Os autores deste paper criaram um sistema chamado KG-WISE. Pense nele como um detetive super-eficiente que usa uma Inteligência Artificial (LLM) como assistente.

Aqui está como funciona, passo a passo, com analogias simples:

1. O Plano de Busca (O Template LLM)

Antes de começar a trabalhar, o KG-WISE usa a IA (o "Detetive") para ler a pergunta e planejar exatamente o que precisa.

• Antes: O sistema carregava tudo.
• Agora: O Detetive diz: "Para responder a essa pergunta sobre 'Física Quântica', eu só preciso dos livros da seção de Física, dos últimos 10 anos, e ignorar tudo sobre culinária ou esportes".
• Ele cria um mapa de busca (um template de consulta) que diz exatamente quais partes da biblioteca são relevantes.

2. A Biblioteca Desmontável (Decomposição do Modelo)

O sistema de inteligência artificial (o modelo treinado) é como uma caixa de ferramentas gigante.

• Antes: A caixa era um bloco de concreto. Você tinha que levantar o bloco inteiro (carregar 20GB de dados) para usar apenas um martelo.
• Agora: O KG-Wise desmonta a caixa. Ele separa os martelos, as chaves de fenda e as serras em caixas menores e organizadas.
• Quando você precisa apenas de um martelo (responder a uma pergunta específica), o sistema pega apenas a caixa do martelo. Ele não carrega as serras ou as chaves que você não vai usar.

3. A Entrega Sob Demanda (Inferência Consciente da Consulta)

Quando chega uma nova pergunta:

1. O sistema olha o mapa de busca criado pelo Detetive.
2. Ele vai à biblioteca e traz apenas os livros relevantes (o subgrafo semântico).
3. Ele monta a caixa de ferramentas mínima necessária apenas para aquela tarefa.
4. Ele responde a pergunta instantaneamente.

Por que isso é incrível? (Os Resultados)

Os autores testaram isso em bibliotecas gigantescas (com até 42 milhões de "livros" ou nós). Os resultados foram impressionantes:

• Velocidade: O sistema ficou até 28 vezes mais rápido. É como ir de carro em vez de andar a pé.
• Memória: O uso de memória caiu 98%. Em vez de precisar de um caminhão de mudanças (memória gigante), você precisa apenas de uma mochila pequena.
• Precisão: Mesmo usando menos coisas, a resposta continua tão boa ou até melhor do que antes, porque o sistema não se distrai com informações inúteis (como ler sobre culinária para responder uma pergunta de física).
• Energia: Como trabalha menos, ele gasta muito menos energia e polui menos o planeta (menos emissão de CO2).

Resumo em uma frase

O KG-WISE é como ter um assistente que, em vez de te entregar a enciclopédia inteira para você procurar uma palavra, abre o livro exato na página correta, entrega apenas a frase que você precisa e joga o resto fora, economizando tempo, espaço e energia.

Isso torna possível usar inteligência artificial avançada em grandes bancos de dados do mundo real sem precisar de computadores superpotentes e caros.

Resumo técnico

1. O Problema

A inferência eficiente de Redes Neurais em Grafos (GNNs) em grandes Grafos de Conhecimento (KGs) é um desafio crítico para aplicações do mundo real, como recomendação, descoberta de fármacos e detecção de fraudes. O problema central identificado pelos autores é a ineficiência de recursos nos métodos atuais de aceleração de inferência:

• Carga de Dados Excessiva: Os sistemas tradicionais carregam o modelo completo (incluindo todos os parâmetros e embeddings de nós) e o grafo inteiro na memória, mesmo quando a consulta de inferência envolve apenas um pequeno subconjunto de nós-alvo. Em KGs heterogêneos, os embeddings pré-computados de nós não-alvo podem representar mais de 95% do tamanho do modelo, gerando um gargalo de memória e I/O.
• Computação Redundante: Métodos existentes (como poda, quantização e destilação de conhecimento) reduzem o tamanho dos pesos do modelo, mas não adaptam a estrutura do modelo ou o carregamento de dados à semântica e estrutura específica de cada consulta. Eles tratam todas as consultas de forma uniforme, processando vizinhanças fixas que podem conter nós semanticamente irrelevantes.
• Escalabilidade Pobre: À medida que os KGs crescem (milhões de nós e bilhões de arestas), a sobrecarga de carregar o modelo monolítico e realizar agregação densa de tensores torna a inferência lenta e inviável em termos de custo computacional e energético.

2. Metodologia: KG-WISE

O artigo propõe o KG-WISE, um paradigma de inferência orientado a tarefas que alinha o treinamento e a inferência através da extração de subgrafos semanticamente relevantes. O sistema opera em três etapas principais:

A. Fase de Treinamento Guiada por LLM

• Extração de Subgrafos Semânticos: Em vez de treinar no grafo completo, o KG-WISE utiliza um Grande Modelo de Linguagem (LLM) para analisar a descrição da tarefa e o esquema do KG.
• Geração de Templates SPARQL: O LLM gera um template de consulta SPARQL reutilizável que identifica entidades e predicados semanticamente relevantes para a tarefa específica. Isso resulta na extração de um subgrafo denso e alinhado à tarefa, significativamente menor que o KG original.
• Decomposição do Modelo: Após o treinamento no subgrafo extraído, o modelo GNN é decomposto em componentes finos:
  • Parâmetros do Modelo: Pesos e vieses das camadas convolucionais.
  • Embeddings de Nós: Vetores de representação dos nós.
• Armazenamento Modular: Os componentes são armazenados separadamente. Os embeddings são divididos em "chunks" (blocos) por tipo de nó e armazenados em uma loja de chave-valor (KV Store, especificamente usando Zarr), enquanto os parâmetros são mantidos em um arquivo separado. Isso permite o carregamento parcial.

B. Fase de Inferência Consciente de Consulta

• Recuperação do Subgrafo: Ao receber uma consulta de inferência (com um conjunto de nós-alvo), o sistema reutiliza o template SPARQL gerado anteriormente (sem necessidade de chamar o LLM novamente) para extrair o subgrafo semântico relevante.
• Instanciação de Modelo Compacto: O sistema instancia um modelo específico para a consulta (f_M) carregando apenas os embeddings e pesos necessários para os nós presentes no subgrafo extraído.
• Agregação Esparsa: Dependendo da densidade do subgrafo recuperado, o sistema seleciona dinamicamente entre agregação de tensores esparsos ou densos, otimizando o cálculo.
• Vantagem: Isso evita o carregamento de embeddings irrelevantes e reduz drasticamente a memória necessária para a inferência.

3. Principais Contribuições

1. Primeiro Sistema End-to-End: O primeiro sistema projetado para armazenamento escalável e inferência adaptativa de GNNs em grandes KGs heterogêneos.
2. Decomposição e Armazenamento Fino: Um mecanismo que decompõe modelos GNN em componentes armazenados em lojas de chave-valor, permitindo o carregamento parcial de embeddings com sobrecarga de pré-processamento mínima.
3. Inferência Guiada por LLM: Uso de LLMs para gerar templates de consulta que extraem subgrafos semanticamente relevantes, permitindo a instanciação de modelos compactos com custo de recuperação quase nulo.
4. Avaliação Abrangente: Validação em seis grandes KGs (até 42 milhões de nós e 166 milhões de arestas), demonstrando que o desempenho não depende de um LLM específico (funciona tanto com modelos proprietários como GPT-4/Gemini quanto com modelos de peso aberto como Qwen/DeepSeek).

4. Resultados Experimentais

Os autores avaliaram o KG-WISE comparando-o com aceleradores de inferência de última geração (GCNP, Degree-Quant, GKD) e métodos de treinamento (GraphSAINT, IBMB, KG-TOSA).

• Desempenho de Velocidade: O KG-WISE alcançou até 28x mais rápido na inferência em comparação com os sistemas mais avançados (SOTA).
• Eficiência de Memória: Redução de até 98% no uso de memória. Enquanto métodos baselines carregam modelos de 6GB a 20GB, o KG-WISE conseguiu reduzir o footprint para 17MB - 121MB dependendo do tamanho da consulta.
• Precisão: O sistema manteve ou melhorou a precisão (acurácia em classificação de nós e Hits@10 em previsão de links) em comparação com os métodos baselines. A esparsidade do subgrafo atua como uma forma de regularização, excluindo ruído.
• Escalabilidade:
  • Weak Scaling: Melhor eficiência de memória em cenários limitados por memória.
  • Strong Scaling: Melhor velocidade computacional em cenários limitados por CPU.
  • O sistema escala sub-linearmente com o aumento do número de nós-alvo, ao contrário dos métodos baselines que são lineares ou constantes (devido ao carregamento total).
• Impacto Ambiental: Redução de 60% nas emissões de CO2 e 62% no consumo de energia por inferência, devido à eliminação de computação e transferência de dados desnecessários.

5. Significado e Conclusão

O KG-WISE representa uma mudança de paradigma na inferência de GNNs em grandes grafos. Ao demonstrar que a maioria do custo computacional e de memória vem do carregamento de embeddings de nós irrelevantes, o trabalho propõe uma solução que desacopla o modelo do grafo completo.

A integração de LLMs não apenas para gerar dados, mas para orquestrar a infraestrutura de inferência (definindo quais partes do modelo e do grafo são necessárias), abre caminho para a implantação sustentável e escalável de GNNs em ambientes de produção com grandes volumes de dados. A capacidade de funcionar com diferentes LLMs e manter alta eficiência em hardware limitado (CPU e GPU) torna a abordagem particularmente relevante para aplicações comerciais e científicas que exigem baixa latência e alto throughput.