APE: Snellere en Langere Context-Verrijkte Generatie via Adaptieve Parallelle Codering

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

February 8, 2025

Auteurs: Xinyu Yang, Tianqi Chen, Beidi Chen

cs.AI

Samenvatting

Context-verrijkte generatie (CAG) technieken, waaronder RAG en ICL, vereisen de efficiënte combinatie van meerdere contexten om antwoorden op gebruikersvragen te genereren. Het direct invoeren van deze contexten als een sequentie brengt aanzienlijke rekenlast met zich mee door de gecombineerde selectie van contexten voor elke aanvraag opnieuw te coderen. Om dit aan te pakken, verkennen we het veelbelovende potentieel van parallelle codering om onafhankelijk de KV-staten van elke context voor te berekenen en te cachen. Deze benadering maakt het mogelijk om tijdens inferentie de gecachte staten direct te laden, terwijl meer contexten worden ondergebracht door hergebruik van posities over contexten. Echter, vanwege misalignments in aandachtsverdeling, leidt directe toepassing van parallelle codering tot een aanzienlijke prestatiedaling. Om effectieve en efficiënte CAG mogelijk te maken, stellen we Adaptive Parallel Encoding (APE) voor, die gedeelde prefix, aandachtstemperatuur en schaalfactor introduceert om de verdeling van parallelle codering af te stemmen op sequentiële codering. Resultaten op RAG en ICL taken tonen aan dat APE 98% en 93% van de prestaties van sequentiële codering kan behouden met dezelfde invoer, terwijl het respectievelijk parallelle codering met 3.6% en 7.9% overtreft. Het schaalt ook naar many-shot CAG, waarbij effectief honderden contexten parallel worden gecodeerd. Efficiëntie-evaluatie toont aan dat APE een end-to-end versnelling van 4.5 keer kan bereiken door de voorbereidingstijd met 28 keer te verminderen voor een context van 128K lengte.

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.