APE: Generación más rápida y con contextos más largos mediante Codificación 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

Resumen

Las técnicas de generación aumentada por contexto (CAG), incluyendo RAG e ICL, requieren la combinación eficiente de múltiples contextos para generar respuestas a las consultas de los usuarios. Introducir directamente estos contextos como una secuencia impone una carga computacional considerable al re-codificar la selección combinada de contextos para cada solicitud. Para abordar esto, exploramos el potencial prometedor del codificado paralelo para pre-calcular y almacenar en caché de forma independiente los estados KV de cada contexto. Este enfoque permite la carga directa de estados almacenados en caché durante la inferencia, al mismo tiempo que permite acomodar más contextos a través de la reutilización de posiciones entre contextos. Sin embargo, debido a desalineaciones en la distribución de atención, la aplicación directa de codificación paralela resulta en una caída significativa del rendimiento. Para habilitar una CAG efectiva y eficiente, proponemos la Codificación Paralela Adaptativa (APE), que incorpora prefijo compartido, temperatura de atención y factor de escala para alinear la distribución de codificación paralela con la codificación secuencial. Los resultados en tareas de RAG e ICL demuestran que APE puede preservar el rendimiento de la codificación secuencial en un 98% y 93% utilizando las mismas entradas, superando la codificación paralela en un 3.6% y 7.9%, respectivamente. Además, se adapta a la generación de CAG de múltiples disparos, codificando de manera efectiva cientos de contextos en paralelo. La evaluación de eficiencia muestra que APE puede lograr una aceleración de extremo a extremo de 4.5 veces al reducir el tiempo de precarga en 28 veces para un contexto de longitud de 128K.

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.