Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

Resumo Técnico: KGERMAR

Definição do Problema
A modelagem de linguagem de contexto longo enfrenta um desafio fundamental que vai além da simples extensão das janelas de contexto: manter uma compreensão coerente dos estados e relacionamentos de entidades ao longo de milhares de tokens. As abordagens padrão sofrem de duas limitações primárias. Primeiro, os modelos de atenção estendida (ex: LongLoRA, YARN) sofrem de degradação de "perda no meio" (lost-in-the-middle), onde informações no meio de um contexto longo são sistematicamente menos atendidas. Segundo, sistemas de recuperação aumentados por memória (ex: MemLong, ERMAR) dependem exclusivamente de similaridade semântica. Embora eficazes para recuperar passagens com termos correspondentes, eles falham em distinguir entre passagens que meramente mencionam termos relacionados de forma coincidente e aquelas que descrevem relações causais ou estruturais explícitas entre as mesmas entidades. Essa limitação é crítica em cenários que exigem raciocínio estrutural, como o rastreamento da progressão do estado de erro de um dispositivo durante uma longa sessão de solução de problemas, onde o contexto relevante pode não compartilhar sobreposição lexical com a consulta, mas compartilha conexões estruturais profundas.

Metodologia
O artigo propõe o KGERMAR (Knowledge Graph Enhanced Memory-Augmented Retrieval), um framework que constrói grafos de conhecimento dinâmicos e específicos ao contexto durante a inferência para integrar embeddings estruturais de grafos com semântica textual. O sistema opera através de quatro componentes integrados:

1. Construção de Grafo de Conhecimento Contextual:

   • Extração: O sistema realiza Reconhecimento de Entidades Nomeadas (NER) em tempo real usando um modelo BERT-large (ajustado no CoNLL-2003) e Extração de Relações usando um modelo BERT-base (ajustado no TACRED).
   • Construção do Grafo: Ele extrai entidades e relações (ex: CAUSEDBY, PARTOF, USEDFOR) dentro de uma janela deslizante. Para garantir a qualidade, emprega consolidação de entidades (fusão de variantes via similaridade de string e embedding), pontuação de confiança de relação (combinando certeza do classificador e frequência) e filtragem de grafo.
   • Natureza Dinâmica: Ao contrário de abordagens estáticas (ex: ConceptNet), este grafo é construído a partir do próprio texto de entrada, capturando entidades específicas do domínio e da sessão (ex: um modelo específico de impressora e seus códigos de erro únicos) que não existem em bases de conhecimento pré-construídas.

2. Geração de Embedding Aumentada por Grafo:

   • Codificação: Entidades extraídas são codificadas usando Redes Convolucionais de Grafo Relacional (R-GCN) para capturar tanto propriedades individuais quanto papéis estruturais via passagem de mensagens de múltiplos saltos (multi-hop message passing).
   • Fusão: Esses embeddings estruturais de grafo são fundidos com embeddings de entidades textuais (derivados do modelo de linguagem base) via atenção cross-modal. Isso cria uma representação híbrida que captura tanto a semântica de superfície quanto padrões relacionais mais profundos.

3. Arquitetura de Memória de Múltiplos Componentes:

   • O KGERMAR mantém três bancos de memória especializados:
     • Memória Contextual: Armazena pares chave-valor ao nível de token das camadas do modelo de linguagem (seguindo o ERMAR).
     • Memória Semântica: Armazena embeddings densos de texto de fragmentos não sobrepostos (usando BGE-M3).
     • Memória Estrutural: Armazena os embeddings fundidos de texto-entidade do grafo.
   • Todos os bancos são atualizados incrementalmente e gerenciados via evasão LRU (Least Recently Used) quando a capacidade é atingida.

4. Mecanismo de Recuperação Híbrido:

   • Para uma determinada consulta, o sistema realiza a recuperação paralela top-K de todos os três bancos.
   • Os sinais de recuperação são combinados usando pesos de fusão aprendidos ($\lambda_1, \lambda_2, \lambda_3$) para produzir uma pontuação final.
   • As representações recuperadas são injetadas no modelo de linguagem via atenção causal de recuperação em camadas superiores designadas.

Objetivos de Treinamento
O modelo emprega aprendizagem multitarefa para otimizar quatro objetivos simultaneamente:

• Perda de Modelagem de Linguagem ($L_{LM}$): Predição do próximo token usando o contexto recuperado.
• Perda de Embedding de Grafo de Conhecimento ($L_{KG}$): Predição de ligação via entropia cruzada binária para treinar o R-GCN.
• Perda de Ranking de Recuperação ($L_{retrieval}$): Ranking baseado em margem para garantir que contextos relevantes recebam pontuações mais altas do que os irrelevantes.
• Perda de Alinhamento Cross-Modal ($L_{alignment}$): Garante a consistência entre representações textuais e estruturais.

Principais Contribuições

1. Construção Dinâmica de Grafo de Conhecimento: Um pipeline que extrai entidades e relações específicas do domínio do texto de entrada em tempo real durante a inferência, permitindo a adaptação a conteúdos especializados ausentes de bases de conhecimento estáticas.
2. Memória de Múltiplos Componentes Híbrida: Uma arquitetura que integra representações contextuais, semânticas e estruturais. O banco de memória estrutural captura especificamente relações de entidades que são complementares a, e frequentemente mais discriminativas do que, a pura similaridade semântica.
3. Validação Empírica: Demonstração de que a incorporação de conhecimento estrutural produz melhorias substanciais na modelagem de contexto longo, alcançando até 8,5% de menor perplexidade e 2–2,5x melhor eficiência de memória em comparação com fortes baselines.

Resultados Experimentais
O framework foi avaliado no SlimPajama (84,7K exemplos para fine-tuning) e em três conjuntos de dados de referência: WikiText-103, PG-19 e Proof-pile.

• Perplexidade: O KGERMAR alcançou a menor perplexidade entre todos os modelos de 3B aumentados por memória em comprimentos de contexto de 1K a 32K tokens. No WikiText-103, superou significativamente o ERMAR e o MemLong (ex: 7,14 vs. 7,80 de perplexidade em 16K tokens).
• Aprendizagem em Contexto (ICL): Em cinco tarefas de NLU (SST-2, MR, Subj, SST-5, MPQA), o KGERMAR alcançou a maior precisão média tanto em configurações de 4-shot (78,6%) quanto de 20-shot (81,2%), superando o ERMAR e o MemLong. Os ganhos foram particularmente notáveis em tarefas que exigem discriminação ao nível de entidade (ex: Subj, MPQA).
• Eficiência de Memória: O KGERMAR demonstrou eficiência de memória superior, mantendo 2–2,5x menos uso de memória do que os baselines. Em 32K tokens, utilizou 36% menos memória de pico que o ERMAR e 41% menos que o MemLong. A memória normalizada por token melhorou conforme o comprimento do contexto aumentava, validando sua capacidade de mitigar o crescimento quadrático de memória dos mecanismos de atenção padrão.
• Comportamento de Escalonamento: Embora o KGERMAR exiba uma latência por token ligeiramente maior devido à construção de grafo em tempo real, ele demonstra um crescimento de latência quase linear com o comprimento da sequência, tornando-o mais previsível para implantação de contexto longo do que o ERMAR.

Significância
O artigo afirma que o KGERMAR avança a modelagem de linguagem aumentada por memória ao mudar da dependência de bases de conhecimento fixas e estáticas para uma representação de conhecimento adaptável ao domínio. Ao construir dinamicamente grafos de conhecimento a partir do contexto de entrada, o sistema pode raciocinar sobre estados de entidades e progressões causais que a similaridade semântica sozinha não consegue capturar. Esta abordagem oferece uma solução prática para a modelagem de contexto longo em hardware comum, proporcionando um equilíbrio entre alta qualidade de modelagem e redução significativa do overhead de memória, o que é crítico para aplicações envolvendo suporte técnico, literatura científica e análise de documentos legais, onde o raciocínio estrutural é primordial.