Lightweight Chunk Selection for Mobile Retrieval-Augmented Generation

Sintesi Tecnica: Selezione Leggera di Chunk per la Generazione Aumentata da Recupero su Dispositivi Mobile

Problematica
La Generazione Aumentata da Recupero (RAG) potenzia i Large Language Models (LLM) incorporando conoscenze esterne, tuttavia l'implementazione della RAG su dispositivi mobili ed edge è ostacolata dai costi computazionali e di memoria necessari per elaborare il contesto recuperato. Una strategia di mitigazione comune consiste nel limitare la generazione a un singolo chunk recuperato. Tuttavia, il chunk con il ranking più alto restituito dal recupero denso standard non è necessariamente il più sufficiente dal punto di vista evidenziale; il recupero ottimizza per la similarità semantica, non per il supporto fattuale. I metodi esistenti di riduzione del contesto, come la compressione dei prompt o il reranking tramite cross-encoder, richiedono spesso modelli aggiuntivi o compressori che superano i rigorosi budget di risorse dei dispositivi mobili. Inoltre, le metriche di valutazione standard basate sulla presenza di stringhe di risposta sono rumorose, poiché un chunk può contenere la stringa della risposta senza fornire l'evidenza relazionale necessaria, o viceversa.

Metodologia
Gli autori propongono un framework leggero di selezione dei chunk accoppiato all'LLM che tratta il problema come un compito di allineamento delle evidenze. Il sistema opera in parallelo con la fase di warmup della sola domanda dell'LLM, riutilizzando il calcolo per evitare costi di inferenza aggiuntivi.

1. Integrazione delle Caratteristiche: Il selettore aggrega tre fonti di caratteristiche complementari:

   • Stati Nascosti della Domanda ($h$): Lo stato nascosto dell'ultimo token dalla passata di warmup dell'LLM, che rappresenta l'interpretazione semantica della query da parte del generatore.
   • Segnali di Routing MoE ($r$): Punteggi di risposta degli esperti pesati dal routing derivati dall'architettura Mixture-of-Experts (MoE). Questi catturano informazioni strutturate su come il modello distribuisce il calcolo tra gli esperti, fungendo da segnale complementare agli stati nascosti.
   • Embedding dei Chunk ($E$): Embedding densi dei candidati recuperati top-$K$, preservando la geometria delle evidenze lato candidato.

2. Architettura del Selettore: Un MLP (Multilayer Perceptron) compatto mappa queste caratteristiche combinate verso un "prototipo di evidenza" predetto nello stesso spazio di embedding dei chunk recuperati. Il chunk finale viene selezionato in base alla similarità del coseno tra questo prototipo predetto e gli embedding dei candidati. Questo approccio evita interazioni esplicite a livello di token tra la query e ogni candidato, che sarebbero computazionalmente costose.

3. Riduzione della Dimensionalità Consapevole del Budget: Per vincoli di implementazione più stringenti, viene introdotto un modulo opzionale di selezione delle caratteristiche. Invece della compressione non supervisionata, il metodo utilizza un'analisi di sensibilità basata sul gradiente su un selettore a dimensionalità completa già addestrato per identificare e mantenere le dimensioni di input più informative. Ciò consente un compromesso controllabile tra dimensione del selettore e accuratezza.

4. Supervisione tramite Sufficienza Semantica: Per affrontare il rumore nei dataset standard, gli autori costruiscono etichette di ground-truth basate sulla "sufficienza evidenziale" piuttosto che sulla mera presenza lessicale della risposta. Utilizzando un annotatore automatizzato, i chunk sono etichettati come positivi solo se contengono informazioni sufficienti per rispondere alla domanda, filtrando i falsi positivi (chunk con stringhe di risposta ma senza contesto) e i falsi negativi (evidenze parafrasate).

Contributi Chiave

• Formulazione: Il documento formula la selezione dei chunk per la RAG mobile come un problema leggero di allineamento delle evidenze che riutilizza le rappresentazioni della query lato generatore (stati nascosti e routing MoE) invece di fare affidamento esclusivamente sui segnali lato recuperatore.
• Design del Modello: Viene proposto un modello di selezione compatto che fonde gli stati nascosti della domanda, i segnali degli esperti derivati dal routing MoE e gli embedding dei chunk per predire un prototipo di evidenza.
• Costruzione del Dataset: Viene introdotta una nuova strategia di supervisione, creando etichette di correttezza semantica dei chunk basate sulla sufficienza evidenziale per abilitare l'addestramento supervisionato per la selezione dei chunk pre-generazione.
• Efficienza: Il metodo include una strategia di selezione delle caratteristiche consapevole del compito che riduce la dimensionalità dell'input mantenendo prestazioni competitive, rendendolo adatto a dispositivi edge con risorse limitate.

Risultati
Gli esperimenti sono stati condotti su tre dataset: TriviaQA, PopQA e MS MARCO Passage Ranking.

• Prestazioni: Il selettore proposto ha raggiunto la più alta accuratezza di selezione del chunk rank-1 su tutti e tre i dataset tra le baseline leggere valutate. Ha migliorato la baseline più forte (KNRM) di una media di 2,49 punti percentuali, con guadagni specifici di 4,60 punti su TriviaQA, 9,13 punti su PopQA e 4,61 punti su MS MARCO PR.
• Ablazione: Gli embedding dei chunk recuperati hanno fornito il segnale di prestazione primario. Tuttavia, la combinazione con gli stati nascosti e i segnali di routing ha prodotto i risultati più robusti, particolarmente su PopQA dove più chunk spesso condividono entità superficiali ma solo uno fornisce supporto fattuale.
• Analisi del Budget: Le varianti con caratteristiche ridotte (ad esempio, utilizzando solo il 20% delle dimensioni di input) hanno mantenuto prestazioni competitive, superando la baseline neurale compatta più forte (KNRM) pur utilizzando significativamente meno parametri (7M vs. ~48M).

Significatività
Il documento sostiene che l'utilizzo delle rappresentazioni della query lato LLM (stati nascosti e routing MoE) e l'allineamento con gli embedding dei candidati lato recupero sia una strategia efficace e con un numero contenuto di parametri per la RAG applicabile ai dispositivi mobili. Spostando l'attenzione dal rango di recupero all'allineamento delle evidenze e utilizzando un selettore leggero che evita ulteriori grandi modelli, il metodo affronta il collo di bottiglia critico della gestione del contesto nei sistemi RAG con risorse limitate. Il lavoro dimosta che una selezione dei chunk di alta qualità può essere ottenuta senza l'elevato overhead di inferenza di cross-encoder o compressori generativi, rendendo la RAG robusta più fattibile per l'implementazione su edge.