APE: Более быстрая и длительная генерация с увеличением контекста через адаптивное параллельное кодирование

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

February 8, 2025

Авторы: Xinyu Yang, Tianqi Chen, Beidi Chen

cs.AI

Аннотация

Техники генерации с увеличением контекста (CAG), включая RAG и ICL, требуют эффективного объединения нескольких контекстов для генерации ответов на запросы пользователей. Прямое ввод контекстов в виде последовательности вносит значительную вычислительную нагрузку, поскольку требуется повторное кодирование объединенного выбора контекстов для каждого запроса. Для решения этой проблемы мы исследуем перспективный потенциал параллельного кодирования для независимого предварительного вычисления и кэширования состояний KV каждого контекста. Этот подход позволяет непосредственно загружать кэшированные состояния во время вывода, обеспечивая повторное использование позиций среди контекстов. Однако из-за несоответствий в распределении внимания прямое применение параллельного кодирования приводит к значительному снижению производительности. Для обеспечения эффективного и эффективного CAG мы предлагаем Адаптивное Параллельное Кодирование (APE), которое внедряет общий префикс, температуру внимания и коэффициент масштабирования для согласования распределения параллельного кодирования с последовательным кодированием. Результаты на задачах RAG и ICL показывают, что APE может сохранить производительность последовательного кодирования на уровне 98% и 93%, используя те же входные данные, превзойдя при этом параллельное кодирование на 3.6% и 7.9% соответственно. Он также масштабируется для обработки многократного CAG, эффективно кодируя сотни контекстов параллельно. Оценка эффективности показывает, что APE может достичь ускорения в 4.5 раза на конечном этапе, сокращая время предварительной загрузки в 28 раз для контекста длиной 128 тыс. символов.

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.