APE : Génération plus rapide et prolongée de contexte augmenté via Encodage Parallèle Adaptatif

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

February 8, 2025

Auteurs: Xinyu Yang, Tianqi Chen, Beidi Chen

cs.AI

Résumé

Les techniques de génération augmentée par contexte (CAG), telles que RAG et ICL, nécessitent une combinaison efficace de plusieurs contextes pour générer des réponses aux requêtes des utilisateurs. Entrer directement ces contextes sous forme de séquence introduit une charge computationnelle considérable en réencodant la sélection combinée de contextes pour chaque demande. Pour remédier à cela, nous explorons le potentiel prometteur du codage parallèle pour précalculer et mettre en cache de manière indépendante les états KV de chaque contexte. Cette approche permet le chargement direct des états mis en cache pendant l'inférence tout en permettant d'accueillir davantage de contextes grâce à la réutilisation de positions à travers les contextes. Cependant, en raison de désalignements dans la distribution de l'attention, l'application directe du codage parallèle entraîne une baisse significative des performances. Pour permettre une CAG efficace et efficiente, nous proposons le Codage Parallèle Adaptatif (APE), qui introduit un préfixe partagé, une température d'attention et un facteur d'échelle pour aligner la distribution du codage parallèle avec le codage séquentiel. Les résultats sur les tâches RAG et ICL démontrent que l'APE peut préserver 98 % et 93 % des performances de codage séquentiel en utilisant les mêmes entrées tout en surpassant le codage parallèle de 3,6 % et 7,9 % respectivement. Il est également adapté à la génération de CAG à plusieurs prises, permettant d'encoder efficacement des centaines de contextes en parallèle. L'évaluation de l'efficacité montre que l'APE peut obtenir un gain de vitesse de 4,5 fois en réduisant de 28 fois le temps de préremplissage pour un contexte de 128 000 caractères.

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.