XC-Cache: Atención Cruzada a Contextos en Caché para una Inferencia Eficiente en Modelos de Lenguaje Grande

XC-Cache: Cross-Attending to Cached Context for Efficient LLM Inference

April 23, 2024

Autores: João Monteiro, Étienne Marcotte, Pierre-André Noël, Valentina Zantedeschi, David Vázquez, Nicolas Chapados, Christopher Pal, Perouz Taslakian

cs.AI

Resumen

Los enfoques de aprendizaje en contexto (ICL, por sus siglas en inglés) suelen aprovechar el prompting para condicionar la generación de modelos de lenguaje de solo decodificador en función de información de referencia. El procesamiento justo a tiempo de un contexto es ineficiente debido al costo cuadrático de las operaciones de autoatención, por lo que el almacenamiento en caché es deseable. Sin embargo, almacenar en caché los estados de los transformadores puede requerir fácilmente casi tanto espacio como los parámetros del modelo. Cuando el contexto correcto no se conoce de antemano, almacenar en caché el ICL puede ser un desafío. Este trabajo aborda estas limitaciones introduciendo modelos que, inspirados en la arquitectura codificador-decodificador, utilizan atención cruzada para condicionar la generación en texto de referencia sin necesidad del prompt. Más precisamente, aprovechamos modelos preentrenados de solo decodificador y solo entrenamos un pequeño número de capas adicionales. Utilizamos la tarea de preguntas y respuestas (QA) como banco de pruebas para evaluar la capacidad de nuestros modelos para realizar generación condicional y observamos que superan al ICL, son comparables a los modelos de lenguaje grandes (LLM) ajustados mediante prompting, y reducen drásticamente la huella de espacio en comparación con el almacenamiento estándar de claves-valores (KV caching) en dos órdenes de magnitud.

English

In-context learning (ICL) approaches typically leverage prompting to condition decoder-only language model generation on reference information. Just-in-time processing of a context is inefficient due to the quadratic cost of self-attention operations, and caching is desirable. However, caching transformer states can easily require almost as much space as the model parameters. When the right context isn't known in advance, caching ICL can be challenging. This work addresses these limitations by introducing models that, inspired by the encoder-decoder architecture, use cross-attention to condition generation on reference text without the prompt. More precisely, we leverage pre-trained decoder-only models and only train a small number of added layers. We use Question-Answering (QA) as a testbed to evaluate the ability of our models to perform conditional generation and observe that they outperform ICL, are comparable to fine-tuned prompted LLMs, and drastically reduce the space footprint relative to standard KV caching by two orders of magnitude.