Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

Riepilogo Tecnico: KGERMAR

Problematica
La modellazione del linguaggio a lungo contesto affronta una sfida fondamentale che va oltre il semplice estendere le finestre di contesto: mantenere una comprensione coerente degli stati delle entità e delle loro relazioni attraverso migliaia di token. Gli approoli standard soffrono di due limitazioni primarie. In primo luogo, i modelli a attenzione estesa (ad es., LongLoRA, YARN) soffrono di un degrado del tipo "lost-in-the-middle", dove le informazioni nel mezzo di un contesto lungo vengono sistematicamente meno attentate. In secondo luogo, i sistemi di recupero con memoria aumentata (ad es., MemLong, ERMAR) si affidano esclusivamente alla somiglianza semantica. Sebbene siano efficaci nel recuperare passaggi con termini corrispondenti, essi non riescono a distinguere tra passaggi che menzionano semplicemente termini correlati per coincidenza e quelli che descrivono espliciti relazioni causali o strutturali tra le stesse entità. Questa limitazione è critica in scenari che richiedono ragionamento strutturale, come il monitoraggio della progressione dello stato di errore di un dispositivo durante una lunga sessione di risoluzione dei problemi, dove il contesto rilevante può non condividere alcuna sovrapposenza lessicale con la query, ma condivide profonde connessioni strutturali.

Metodologia
Il documento propone KGERMAR (Knowledge Graph Enhanced Memory-Augmented Retrieval), un framework che costruisce grafi di conoscenza dinamici e specifici per il contesto durante l'inferenza per integrare embedding strutturali-grafici con la semantica testuale. Il sistema opera attraverso quattro componenti integrate:

1. Costruzione del Grafo di Conoscenza Contestuale:

   • Estrazione: Il sistema esegue il riconoscimento delle entità nominate (NER) in tempo reale utilizzando un modello BERT-large (fine-tuned su CoNLL-2003) e l'estrazione delle relazioni utilizzando un modello BERT-base (fine-tuned su TACRED).
   • Costruzione del Grafo: Estrae entità e relazioni (ad es., CAUSEDBY, PARTOF, USEDFOR) all'interno di una finestra scorrevole. Per garantire la qualità, impiega la consolidazione delle entità (fusione delle varianti tramite somiglianza di stringa e di embedding), il punteggio di confidenza delle relazioni (combinando la certezza del classificatore e la frequenza) e il filtraggio del grafo.
   • Natura Dinamica: A differenza degli approcci statici (ad es., ConceptNet), questo grafo è costruito dal testo in input stesso, catturando entità specifiche del dominio e della sessione (ad es., un modello di stampante specifico e i suoi codici di errore unici) che non esistono nelle basi di conoscenza pre-costruite.

2. Generazione di Embedding Potenziati dal Grafo:

   • Codifica: Le entità estratte sono codificate utilizzando Reti Convoluzionali Grafiche Relazionali (R-GCN) per catturare sia le proprietà individuali che i ruoli strutturali tramite il passaggio di messaggi multi-hop.
   • Fusione: Questi embedding strutturali-grafici sono fusi con gli embedding delle entità testuali (derivati dal modello linguistico di base) tramite attenzione cross-modale. Ciò crea una rappresentazione ibrida che cattura sia la semantica superficiale che i pattern relazionali più profondi.

3. Architettura di Memoria Multi-Componente:

   • KGERMAR mantiene tre banche di memoria specializzate:
     • Memoria Contestuale: Memorizza coppie chiave-valore a livello di token dai livelli del modello linguistico (seguendo ERMAR).
     • Memoria Semantica: Memorizza embedding testuali densi di chunk non sovrapposti (utilizzando BGE-M3).
     • Memoria Strutturale: Memorizza gli embedding fusi testo-grafo delle entità.
   • Tutte le banche sono aggiornate incrementalmente e gestite tramite l'evizione LRU quando la capacità viene raggiunta.

4. Meccanismo di Recupero Ibrido:

   • Per una data query, il sistema esegue il recupero top-K in parallelo da tutte e tre le banche.
   • I segnali di recupero sono combinati utilizzando pesi di fusione appresi ($\lambda_1, \lambda_2, \lambda_3$) per produrre un punteggio finale.
   • Le rappresentazioni recuperate vengono iniettate nel modello linguistico tramite attenzione causale di recupero in livelli superiori designati.

Obiettivi di Addestramento
Il modello impiega l'apprendimento multi-task per ottimizzare simultaneamente quattro obiettivi:

• Perdita di Modellazione del Linguaggio ($L_{LM}$): Predizione del token successivo utilizzando il contesto recuperato.
• Perdita di Embedding del Grafo di Conoscenza ($L_{KG}$): Predizione dei collegamenti tramite cross-entropy binaria per addestrare l'R-GCN.
• Perdita di Classificazione del Recupero ($L_{retrieval}$): Ranking basato sul margine per garantire che i contesti rilevanti ricevano punteggi più alti rispetto a quelli irrilevanti.
• Perdita di Allineamento Cross-Modale ($L_{alignment}$): Garantisce la coerenza tra le rappresentazioni testuali e strutturali.

Contributi Chiave

1. Costruzione Dinamica del Grafo di Conoscenza: Una pipeline che estrae entità e relazioni specifiche del dominio dal testo in input in tempo reale durante l'inferenza, consentendo l'adattamento a contenuti specializzati assenti dalle basi di conoscenza statiche.
2. Memoria Multi-Componente Ibrida: Un'architettura che integra rappresentazioni contestuali, semantiche e strutturali. La banca di memoria strutturale cattura specificamente le relazioni tra entità che sono complementari a, e spesso più discriminanti di, la pura somiglianza semantica.
3. Validazione Empirica: Dimostrazione che l'incorporazione della conoscenza strutturale produce miglioramenti sostanziali nella modellazione del lungo contesto, ottenendo una perplexity fino all'8,5% inferiore e un'efficienza di memoria da 2 a 2,5 volte migliore rispetto ai forti baseline.

Risultati Sperimentali
Il framework è stato valutato su SlimPajama (84,7K esempi per il fine-tuning) e tre dataset di benchmark: WikiText-103, PG-19 e Proof-pile.

• Perplexity: KGERMAR ha ottenuto la perplexity più bassa tra tutti i modelli a 3B con memoria aumentata su lunghezze di contesto da 1K a 32K token. Su WikiText-103, ha superato significativamente ERMAR e MemLong (ad es., 7,14 vs 7,80 di perplexity a 16K token).
• In-Context Learning (ICL): Su cinque compiti NLU (SST-2, MR, Subj, SST-5, MPQA), KGERMAR ha raggiunto la massima accuratezza media sia nelle impostazioni 4-shot (78,6%) che 20-shot (81,2%), superando ERMAR e MemLong. I guadagni sono stati particolarmente notevoli in compiti che richiedono la discriminazione a livello di entità (ad es., Subj, MPQA).
• Efficienza della Memoria: KGERMAR ha dimostrato una superiore efficienza della memoria, mantenendo un uso della memoria da 2 a 2,5 volte inferiore rispetto ai baseline. A 32K token, ha utilizzato il 36% in meno di memoria di picco rispetto a ERMAR e il 41% in meno rispetto a MemLong. La memoria normalizzata per token è migliorata all'aumentare della lunghezza del contesto, validando la sua capacità di mitigare la crescita quadratica della memoria dei meccanismi di attenzione standard.
• Comportamento di Scalabilità: Sebbene KGERMAR presenti una latenza per-token leggermente superiore a causa della costruzione del grafo in tempo reale, dimostra una crescita della latenza quasi lineare con la lunghezza della sequenza, rendendolo più prevedibile per il deployment a lungo contesto rispetto a ERMAR.

Significatività
Il documento afferma che KGERMAR fa progredire la modellazione del linguaggio con memoria aumentata passando dal affidarsi a basi di conoscenza fisse e statiche a una rappresentazione della conoscenza adattiva al dominio. Costruendo dinamicamente grafi di conoscenza dal contesto in input, il sistema può ragionare sugli stati delle entità e sulle progressioni causali che la sola somiglianza semantica non può catturare. Questo approccio offre una soluzione pratica per la modellazione a lungo contesto su hardware comune, fornendo un equilibrio tra alta qualità di modellazione e un overhead di memoria significativamente ridotto, aspetto critico per applicazioni che coinvolgono il supporto tecnico, la letteratura scientifica e l'analisi di documenti legali dove il ragionamento strutturale è fondamentale.