Aplicações de Modelos de Linguagem Grandes (LLMs) que utilizam conhecimento externo, como a Geração Aumentada por Recuperação (RAG), revolucionaram a forma como interagimos com a informação, fornecendo respostas aprimoradas em conversas multi-turn e sistemas de agentes. No entanto, o uso de entradas de contexto longas — uma necessidade para aprimorar o RAG — traz um desafio fundamental: latência significativa do sistema e consumo substancial de memória para o cache Key-Value (KV), o que, por sua vez, reduz o throughput.
Em resposta a essa limitação, pesquisadores propuseram o REFRAG (REpresentation For RAG), uma estrutura de decodificação eficiente e inovadora, especificamente projetada para otimizar o desempenho em aplicações RAG.
O Desafio Único do Contexto RAG
Embora otimizar a latência de inferência em LLMs com contexto extenso seja uma área ativa de pesquisa, o REFRAG argumenta que os sistemas RAG exigem consideração especializada.
No RAG, grande parte do contexto do LLM é composta por passagens concatenadas de recuperação, mas apenas um pequeno subconjunto dessas passagens é diretamente relevante para a consulta. Além disso, devido à diversidade ou deduplicação, essas passagens frequentemente demonstram baixa similaridade semântica, resultando em padrões de atenção de bloco-diagonal (esparsos) que são incomuns em tarefas genéricas de LLM.
Tratar o Tempo para o Primeiro Token (TTFT) do RAG como um problema genérico de inferência ignora o fato de que:
- Muitos tokens no contexto são desinformativos, o que torna a alocação de memória e computação para todos esses tokens desnecessariamente dispendiosa.
- Informações pré-processadas do processo de recuperação (como vetores e re-ranqueamento) são descartadas durante a decodificação.
Como o REFRAG Acelera o Processamento
O REFRAG explora essa estrutura de esparsidade de atenção para argumentar que a maioria dos cálculos sobre o contexto RAG durante a decodificação pode ser eliminada com impacto mínimo no desempenho.
O REFRAG opera através de um processo de compressão, detecção (sensing) e expansão. Sua arquitetura principal consiste em um modelo base decoder-only (como o LLaMA) e um modelo encoder leve (como o RoBERTa).
Em vez de usar tokens integrais das passagens recuperadas, o REFRAG aproveita embeddings de chunks (pedaços) pré-calculados e compactados, que servem como representações aproximadas, alimentando-os diretamente ao decodificador.
Este método de compressão oferece três vantagens principais:
- Encurta o comprimento da entrada do decodificador, melhorando a eficiência da alocação de tokens.
- Permite o reuso de embeddings de chunk pré-calculados da recuperação, eliminando a computação redundante.
- Reduz a complexidade do cálculo de atenção, que passa a escalar quadraticamente com o número de chunks, em vez do número de tokens no contexto.
Compressão Seletiva Orientada por RL
Uma característica inovadora do REFRAG é sua capacidade de suportar a compressão seletiva de tokens. Uma política leve de Reinforcement Learning (RL) é treinada para determinar se é necessário fornecer a entrada completa do chunk (expandir) ou se um embedding de chunk compactado de baixo custo é suficiente (comprimir).
Essa política é crucial, pois segmentos de contexto diferentes contribuem de forma desigual para a previsão da resposta. Ao treinar a política de RL, que usa a perplexity de próxima frase como uma recompensa negativa, o REFRAG pode expandir seletivamente chunks importantes, preservando a propriedade autorregressiva do decodificador e permitindo a colocação flexível da compressão em qualquer posição da entrada.
Experimentos mostram que a política de compressão seletiva baseada em RL consistentemente alcança desempenho superior em comparação com heurísticas baseadas em perplexity ou seleção aleatória.
Resultados Surpreendentes de Velocidade e Performance
Os resultados empíricos demonstram a eficácia do REFRAG em diversas tarefas de contexto longo, incluindo RAG, conversas multi-turn e sumarização de documentos longos.
| Métrica | Aceleração (vs. LLaMA) | Aceleração (vs. CEPE – SOTA anterior) |
|---|---|---|
| TTFT (Tempo para o Primeiro Token) | Até 30.85× (com k=32) | 3.75× (com k=32) |
| Aceleração TTFT (k=16, com cache) | 16.53× | 2.01× |
| Extensão de Contexto | 16× | N/A |
Além da aceleração, a estrutura otimizada do REFRAG permite estender o tamanho do contexto dos LLMs em 16×. Mesmo mantendo um desempenho comparável ou superior ao LLaMA, o REFRAG consegue incorporar mais informação dentro do mesmo orçamento de latência.
RAG e Conversação Multi-Turn
- Melhor Desempenho com Latência Equivalente: Em cenários de RAG, sob restrições de latência equivalentes (por exemplo, REFRAG8 usando 8 passagens versus LLaMAFT usando 1 passagem), o REFRAG supera o desempenho do LLaMA.
- Retriever Fraco: Em um cenário de retriever fraco (que se assemelha mais a sistemas do mundo real), o REFRAG supera significativamente o LLaMA na mesma latência, pois o contexto expandido permite extrair mais informações úteis mesmo quando as passagens são menos relevantes.
- Conversas Longas: Em tarefas de conversação multi-turn, onde o LLaMAFT (limitado a 4K tokens) pode precisar truncar o histórico, o REFRAG mantém um desempenho robusto e superior, demonstrando os benefícios de sua abordagem de compressão.
Conclusão
REFRAG representa um avanço significativo para a inferência eficiente de LLMs em cenários de uso intensivo de conhecimento. Ao adaptar a decodificação para a estrutura esparsa inerente aos contextos RAG, o REFRAG fornece uma solução prática e escalável para implantar LLMs em aplicações sensíveis à latência, atingindo acelerações substanciais sem comprometer a precisão.
Autor: André Cardia
