Lightweight Chunk Selection for Mobile Retrieval-Augmented Generation

Resumo Técnico: Seleção de Fragmentos Leve para Geração Aumentada por Recuperação em Dispositivos Móveis

Declaração do Problema
A Geração Aumentada por Recuperação (RAG) aprimora Grandes Modelos de Linguagem (LLMs) ao incorporar conhecimento externo, porém a implantação de RAG em dispositivos móveis e de borda (edge) é dificultada pelos custos computacionais e de memória para processar o contexto recuperado. Uma estratégia de mitigação comum é restringir a geração a um único fragmento (chunk) recuperado. No entanto, o fragmento com a classificação mais alta retornado pela recuperação densa padrão não é necessariamente o mais evidentemente suficiente; a recuperação otimiza para similaridade semântica, não para suporte factual. Métodos existentes de redução de contexto, como compressão de prompt ou reranking por cross-encoder, frequentemente exigem modelos ou compressores adicionais de grande porte que excedem os rigorosos orçamentos de recursos de dispositivos móveis. Além disso, métricas de avaliação padrão baseadas na contenção de strings de resposta são ruidosas, pois um fragmento pode conter a string da resposta sem fornecer a evidência relacional necessária, ou vice-versa.

Metodologia
Os autores propõem uma estrutura de seleção de fragmentos leve, acoplada ao LLM, que trata o problema como uma tarefa de alinhamento de evidências. O sistema opera em paralelo com a fase de aquecimento (warmup) de apenas pergunta do LLM, reutilizando a computação para evitar custos de inferência adicionais.

1. Integração de Características: O seletor agrega três fontes de características complementares:

   • Estados Ocultos da Pergunta ($h$): O estado oculto do último token da passagem de warmup do LLM, representando a interpretação semântica da consulta pelo gerador.
   • Sinais de Roteamento MoE ($r$): Pontuações de resposta de especialistas ponderadas pelo roteamento derivadas da arquitetura Mixture-of-Experts (MoE). Estes capturam informações estruturadas sobre como o modelo distribui a computação entre os especialistas, servindo como um sinal complementar aos estados ocultos.
   • Embeddings de Fragmentos ($E$): Embeddings densos dos $K$ principais candidatos recuperados, preservando a geometria de evidência do lado do candidato.

2. Arquitetura do Seletor: Um Perceptron Multicamadas (MLP) compacto mapeia essas características combinadas para um "protótipo de evidência" previsto no mesmo espaço de embedding dos fragmentos recuperados. O fragmento final é selecionado com base na similaridade de cosseno entre este protótipo previsto e os embeddings dos candidatos. Esta abordagem evita interações explícitas de nível de token entre a consulta e cada candidato, o que é computacionalmente caro.

3. Redução de Dimensão Consciente do Orçamento: Para restrições de implantação mais rigorosas, um módulo opcional de seleção de características é introduzido. Em vez de compressão não supervisionada, o método utiliza análise de sensibilidade baseada em gradiente de um seletor de dimensão total treinado para identificar e reter as dimensões de entrada mais informativas. Isso permite um compromisso controlável entre o tamanho do seletor e a precisão.

4. Supervisão via Suficiência Semântica: Para lidar com o ruído em conjuntos de dados padrão, os autores constroem rótulos de verdade fundamental (ground-truth) baseados em "suficiência de evidência" em vez de contenção lexical de resposta. Usando um anotador automatizado, os fragmentos são rotulados como positivos apenas se contiverem informação suficiente para responder à pergunta, filtrando falsos positivos (fragmentos com strings de resposta, mas sem contexto) e falsos negativos (evidência parafraseada).

Principais Contribuições

• Formulação: O artigo formula a seleção de fragmentos para RAG móvel como um problema de alinhamento de evidências leve que reutiliza representações de consulta do lado do gerador (estados ocultos e roteamento MoE) em vez de depender apenas de sinais do lado do recuperador.
• Design do Modelo: Um modelo de seleção compacto é proposto para fundir estados ocultos de perguntas, sinais de especialistas derivados de roteamento MoE e embeddings de fragmentos para prever um protótipo de evidência.
• Construção de Dataset: Uma nova estratégia de supervisão é introduzida, criando rótulos de correção de fragmento semânticos baseados em suficiência de evidência para permitir o treinamento supervisionado para seleção de fragmentos pré-geração.
• Eficiência: O método inclui uma estratégia de seleção de características orientada à tarefa que reduz a dimensionalidade da entrada enquanto mantém um desempenho competitivo, tornando-o adequado para dispositivos de borda com recursos limitados.

Resultados
Experimentos foram conduzidos em três datasets: TriviaQA, PopQA e MS MARCO Passage Ranking.

• Desempenho: O seletor proposto alcançou a maior precisão de seleção de fragmento rank-1 em todos os três datasets entre os baselines leves avaliados. Ele melhorou sobre o baseline mais forte (KNRM) em uma média de 2,49 pontos percentuais, com ganhos específicos de 4,60 pontos em TriviaQA, 9,13 pontos em PopQA e 4,61 pontos em MS MARCO PR.
• Ablação: Os embeddings dos fragmentos recuperados forneceram o principal sinal de desempenho. No entanto, combinar estes com estados ocultos e sinais de roteamento rendeu os resultados mais robustos, particularmente no PopQA, onde múltiplos fragmentos frequentemente compartilham entidades de superfície, mas apenas um fornece suporte factual.
• Análise de Orçamento: Variantes de características reduzidas (por exemplo, usando apenas 20% das dimensões de entrada) mantiveram um desempenho competitivo, superando o baseline neural compacto mais forte (KNRM) enquanto utilizam significativamente menos parâmetros (7M vs. ~48M).

Significância
O artigo afirma que utilizar representações de consulta do lado do LLM (estados ocultos e roteamento MoE) e alinhá-las com os embeddings de candidatos do lado da recuperação é uma estratégia eficaz e eficiente em termos de parâmetros para RAG aplicável a dispositivos móveis. Ao mudar o foco do ranking de recuperação para o alinhamento de evidências e utilizar um seletor leve que evita modelos grandes adicionais, o método aborda o gargalo crítico da gestão de contexto em sistemas RAG com recursos restritos. O trabalho demonstra que uma seleção de fragmentos de alta qualidade pode ser alcançada sem o pesado overhead de inferência de cross-encoders ou compressores generativos, tornando o RAG robusto mais viável para implantação em dispositivos de borda.