APE: Schnellere und längere kontextangereicherte Generierung durch adaptives paralleles Codieren.

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

February 8, 2025

Autoren: Xinyu Yang, Tianqi Chen, Beidi Chen

cs.AI

Zusammenfassung

Techniken zur kontextgestützten Generierung (Context-augmented generation, CAG), einschließlich RAG und ICL, erfordern die effiziente Kombination mehrerer Kontexte, um Antworten auf Benutzeranfragen zu generieren. Das direkte Eingeben dieser Kontexte als Sequenz führt zu einer erheblichen Rechenlast, da die kombinierte Auswahl von Kontexten für jede Anfrage erneut codiert werden muss. Um dies zu lösen, untersuchen wir das vielversprechende Potenzial der parallelen Codierung, um unabhängig jeden KV-Zustand jedes Kontexts vorzuberechnen und zu zwischenspeichern. Dieser Ansatz ermöglicht das direkte Laden zwischengespeicherter Zustände während der Inferenz und ermöglicht eine bessere Handhabung von mehreren Kontexten durch die Wiederverwendung von Positionen über verschiedene Kontexte hinweg. Allerdings führt die direkte Anwendung der parallelen Codierung aufgrund von Missabstimmungen in der Aufmerksamkeitsverteilung zu einem signifikanten Leistungsabfall. Um eine effektive und effiziente CAG zu ermöglichen, schlagen wir Adaptive Parallel Encoding (APE) vor, das gemeinsame Präfixe, Aufmerksamkeitstemperatur und Skalierungsfaktoren verwendet, um die Verteilung der parallelen Codierung mit der sequenziellen Codierung abzustimmen. Ergebnisse bei RAG- und ICL-Aufgaben zeigen, dass APE 98% bzw. 93% der Leistung der sequenziellen Codierung beibehalten kann, während es die parallele Codierung um 3,6% bzw. 7,9% übertrifft. Es ist auch für die Verarbeitung von vielen Kontexten in einem Durchgang geeignet. Die Effizienzbewertung zeigt, dass APE eine End-to-End-Beschleunigung um das 4,5-fache erreichen kann, indem die Vorabfüllzeit für einen Kontext der Länge 128K um das 28-fache reduziert wird.

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.