APE: Geração mais rápida e com contexto ampliado por meio de Codificação Paralela Adaptativa

APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding

February 8, 2025

Autores: Xinyu Yang, Tianqi Chen, Beidi Chen

cs.AI

Resumo

As técnicas de geração aumentada por contexto (CAG), incluindo RAG e ICL, requerem a combinação eficiente de múltiplos contextos para gerar respostas às consultas dos usuários. Inserir diretamente esses contextos como uma sequência introduz uma considerável carga computacional ao recodificar a seleção combinada de contextos para cada solicitação. Para lidar com isso, exploramos o potencial promissor da codificação paralela para pré-calcular e armazenar em cache independentemente os estados KV de cada contexto. Esse método possibilita o carregamento direto dos estados em cache durante a inferência, ao mesmo tempo que acomoda mais contextos por meio da reutilização de posições entre contextos. No entanto, devido a desalinhamentos na distribuição de atenção, a aplicação direta da codificação paralela resulta em uma queda significativa de desempenho. Para viabilizar a CAG eficaz e eficiente, propomos a Codificação Paralela Adaptativa (APE), que incorpora prefixo compartilhado, temperatura de atenção e fator de escala para alinhar a distribuição da codificação paralela com a codificação sequencial. Resultados em tarefas de RAG e ICL demonstram que o APE pode preservar 98% e 93% do desempenho de codificação sequencial usando as mesmas entradas, superando a codificação paralela em 3,6% e 7,9%, respectivamente. Além disso, ele se adapta à geração de CAG com muitos contextos, codificando efetivamente centenas de contextos em paralelo. A avaliação de eficiência mostra que o APE pode alcançar uma aceleração de 4,5 vezes no processo completo, reduzindo em 28 vezes o tempo de pré-preenchimento para um contexto de 128 mil caracteres.

English

Context-augmented generation (CAG) techniques, including RAG and ICL, require the efficient combination of multiple contexts to generate responses to user queries. Directly inputting these contexts as a sequence introduces a considerable computational burden by re-encoding the combined selection of contexts for every request. To address this, we explore the promising potential of parallel encoding to independently pre-compute and cache each context's KV states. This approach enables the direct loading of cached states during inference while accommodating more contexts through position reuse across contexts. However, due to misalignments in attention distribution, directly applying parallel encoding results in a significant performance drop. To enable effective and efficient CAG, we propose Adaptive Parallel Encoding (APE), which brings shared prefix, attention temperature, and scaling factor to align the distribution of parallel encoding with sequential encoding. Results on RAG and ICL tasks demonstrate that APE can preserve 98% and 93% sequential encoding performance using the same inputs while outperforming parallel encoding by 3.6% and 7.9%, respectively. It also scales to many-shot CAG, effectively encoding hundreds of contexts in parallel. Efficiency evaluation shows that APE can achieve an end-to-end 4.5times speedup by reducing 28times prefilling time for a 128K-length context.