APE: Generazione più veloce e prolungata con contesto ampliato tramite Codifica Parallela Adattiva

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

February 8, 2025

Autori: Xinyu Yang, Tianqi Chen, Beidi Chen

cs.AI

Abstract

Le tecniche di generazione potenziate dal contesto (CAG), inclusi RAG e ICL, richiedono l'efficiente combinazione di molteplici contesti per generare risposte alle query degli utenti. L'inserimento diretto di questi contesti come sequenza introduce un notevole onere computazionale ricodificando la selezione combinata di contesti per ogni richiesta. Per affrontare ciò, esploriamo il promettente potenziale dell'encoding parallelo per pre-calcolare e memorizzare in cache in modo indipendente gli stati KV di ciascun contesto. Questo approccio consente il caricamento diretto degli stati memorizzati durante l'inferenza consentendo il riutilizzo della posizione attraverso i contesti. Tuttavia, a causa di disallineamenti nella distribuzione dell'attenzione, l'applicazione diretta dell'encoding parallelo comporta un significativo calo delle prestazioni. Per consentire una CAG efficace ed efficiente, proponiamo l'Encoding Parallelo Adattivo (APE), che introduce prefisso condiviso, temperatura dell'attenzione e fattore di scala per allineare la distribuzione dell'encoding parallelo con l'encoding sequenziale. I risultati sulle attività RAG e ICL dimostrano che APE può preservare il 98% e il 93% delle prestazioni di encoding sequenziale utilizzando gli stessi input, superando l'encoding parallelo rispettivamente del 3.6% e del 7.9%. Si adatta anche alla generazione di CAG multipla, codificando efficacemente centinaia di contesti in parallelo. L'efficienza valutativa mostra che APE può ottenere un aumento della velocità di 4.5 volte riducendo di 28 volte il tempo di precaricamento per un contesto di lunghezza 128K dall'inizio alla fine.

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.