Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

Resumen Técnico: KGERMAR

Planteamiento del Problema
El modelado de lenguaje de contexto largo enfrenta un desafío fundamental que va más allá de simplemente extender las ventanas de contexto: mantener una comprensión coherente de los estados y las relaciones de las entidades a través de miles de tokens. Los enfoques estándar sufren dos limitaciones primarias. Primero, los modelos de atención extendida (por ejemplo, LongLoRA, YARN) sufren de la degradación de "pérdida en el medio" (lost-in-the-middle), donde la información en la parte media de un contexto largo es sistemáticamente menos atendida. Segundo, los sistemas de recuperación aumentados por memoria (por ejemplo, MemLong, ERMAR) dependen exclusivamente de la similitud semántica. Aunque son efectivos para recuperar pasajes con términos coincidentes, no logran distinguir entre pasajes que simplemente mencionan términos relacionados de forma coincidente y aquellos que describen relaciones causales o estructurales explícitas entre las mismas entidades. Esta limitación es crítica en escenarios que requieren razonamiento estructural, como el seguimiento de la progresión del estado de error de un dispositivo a lo largo de una larga sesión de resolución de problemas, donde el contexto relevante puede no compartir ninguna superposición léxica con la consulta pero sí comparte conexiones estructurales profundas.

Metodología
El artículo propone KGERMAR (Recuperación Aumentada por Memoria Mejorada por Grafos de Conocimiento), un marco que construye grafos de conocimiento dinámicos y específicos del contexto durante la inferencia para integrar incrustaciones (embeddings) estructurales de grafos con la semántica textual. El sistema opera a través de cuatro componentes integrados:

1. Construcción de Grafos de Conocimiento Contextuales:

   • Extracción: El sistema realiza un Reconocimiento de Entidades Nombradas (NER) en tiempo real utilizando un modelo BERT-large (ajustado en CoNLL-2003) y una Extracción de Relaciones utilizando un modelo BERT-base (ajustado en TACRED).
   • Construcción del Grafo: Extrae entidades y relaciones (por ejemplo, CAUSEDBY, PARTOF, USEDFOR) dentro de una ventana deslizante. Para asegurar la calidad, emplea consolidación de entidades (fusión de variantes mediante similitud de cadenas y de incrustaciones), puntuación de confianza de relación (combinando la certeza del clasificador y la frecuencia) y filtrado de grafos.
   • Naturaleza Dinámica: A diferencia de los enfoques estáticos (por ejemplo, ConceptNet), este grafo se construye a partir del propio texto de entrada, capturando entidades específicas del dominio y de la sesión (por ejemplo, un modelo de impresora específico y sus códigos de error únicos) que no existen en bases de conocimiento preconstruidas.

2. Generación de Incrustaciones Mejoradas por Grafos:

   • Codificación: Las entidades extraídas se codifican utilizando Redes Convolucionales de Grafos Relacionales (R-GCN) para capturar tanto las propiedades individuales como los roles estructurales mediante el paso de mensajes de múltiples saltos (multi-hop message passing).
   • Fusión: Estas incrustaciones estructurales de grafo se fusionan con las incrustaciones de entidades textuales (derivadas del modelo de lenguaje base) mediante atención cruzada (cross-modal attention). Esto crea una representación híbrida que captura tanto la semántica superficial como los patrones relacionales más profundos.

3. Arquitectura de Memoria de Múltiples Componentes:

   • KGERMAR mantiene tres bancos de memoria especializados:
     • Memoria Contextual: Almacena pares clave-valor a nivel de token de las capas del modelo de lenguaje (siguiendo a ERMAR).
     • Memoria Semántica: Almacena incrustaciones densas de texto de fragmentos no superpuestos (usando BGE-M3).
     • Memoria Estructural: Almacena las incrustaciones híbridas de texto-entidad de grafo.
   • Todos los bancos se actualizan incrementalmente y se gestionan mediante la expulsión LRU cuando se alcanza la capacidad.

4. Mecanismo de Recuperación Híbrido:

   • Para una consulta dada, el sistema realiza una recuperación paralela Top-K de todos los tres bancos.
   • Las señales de recuperación se combinan utilizando pesos de fusión aprendidos ($\lambda_1, \lambda_2, \lambda_3$) para producir una puntuación final.
   • Las representaciones recuperadas se inyectan en el modelo de lenguaje mediante atención causal de recuperación en capas superiores designadas.

Objetivos de Entrenamiento
El modelo emplea aprendizaje multitarea para optimizar cuatro objetivos simultáneamente:

• Pérdida de Modelado de Lenguaje ($L_{LM}$): Predicción del siguiente token utilizando el contexto recuperado.
• Pérdida de Incrustación de Grafo de Conocimiento ($L_{KG}$): Predicción de enlaces mediante entropía cruzada binaria para entrenar la R-GCN.
• Pérdida de Clasificación de Recuperación ($L_{retrieval}$): Clasificación basada en márgenes para asegurar que los contextos relevantes tengan puntuaciones más altas que los irrelevantes.
• Pérdida de Alineación Cross-Modal ($L_{alignment}$): Asegura la consistencia entre las representaciones textuales y estructurales.

Contribuciones Clave

1. Construcción de Grafos de Conocimiento Dinámicos: Un flujo de trabajo que extrae entidades y relaciones específicas del dominio del texto de entrada en tiempo real durante la inferencia, permitiendo la adaptación a contenidos especializados ausentes de bases de conocimiento estáticas.
2. Memoria de Múltiples Componentes Híbrida: Una arquitectura que integra representaciones contextuales, semánticas y estructurales. El banco de memoria estructural captura específicamente las relaciones de entidades que son complementarias a, y a menudo más discriminativas que, la pura similitud semántica.
3. Validación Empírica: Demostración de que la incorporación de conocimiento estructural produce mejoras sustanciales en el modelado de contexto largo, logrando hasta un 8.5% menos de perplejidad y una eficiencia de memoria de 2 a 2.5 veces mejor en comparación con los modelos de referencia (baselines) robustos.

Resultados Experimentales
El marco fue evaluado en SlimPajama (84.7K ejemplos para ajuste fino) y tres conjuntos de datos de referencia: WikiText-103, PG-19 y Proof-pile.

• Perplejidad: KGERMAR logró la perplejidad más baja entre todos los modelos de 3B aumentados por memoria en longitudes de contexto de 1K a 32K tokens. En WikiText-103, superó significativamente a ERMAR y MemLong (por ejemplo, 7.14 frente a 7.80 de perplejidad a 16K tokens).
• Aprendizaje en Contexto (ICL): En cinco tareas de NLU (SST-2, MR, Subj, SST-5, MPQA), KGERMAr alcanzó la precisión promedio más alta tanto en configuraciones de 4 disparos (4-shot: 78.6%) como de 20 disparos (20-shot: 81.2%), superando a ERMAR y MemLong. Las ganancias fueron particularmente notables en tareas que requieren discriminación a nivel de entidad (por ejemplo, Subj, MPQA).
• Eficiencia de Memoria: KGERMAR demostró una eficiencia de memoria superior, manteniendo un uso de memoria de 2 a 2.5 veces menor que los modelos de referencia. A 32K tokens, utilizó un 36% menos de memoria pico que ERMAR y un 41% menos que MemLong. La memoria normalizada por token mejoró a medida que la longitud del contexto aumentaba, validando su capacidad para mitigar el crecimiento cuadrático de la memoria de los mecanismos de atención estándar.
• Comportamiento de Escalamiento: Aunque KGERMAR exhibe una latencia por token ligeramente mayor debido a la construcción de grafos en tiempo real, demuestra un crecimiento de latencia casi lineal con la longitud de la secuencia, lo que lo hace más predecible para el despliegue de contexto largo que ERMAR.

Significancia
El artículo afirma que KGERMAR hace avanzar el modelado de lenguaje aumentado por memoria al pasar de la dependencia de bases de conocimiento fijas y estáticas a una representación de conocimiento adaptativa al dominio. Al construir dinámicamente grafos de conocimiento a partir del contexto de entrada, el sistema puede razonar sobre los estados de las entidades y las progresiones causales que la similitud semántica por sí sola no puede capturar. Este enfoque ofrece una solución práctica para el modelado de contexto largo en hardware comercial, proporcionando un equilibrio entre una alta calidad de modelado y una reducción significativa de la sobrecarga de memoria, lo cual es crítico para aplicaciones que involucran soporte técnico, literatura científica y análisis de documentos legales donde el razonamiento estructural es primordial.