Cartuchos: Representaciones de contexto largo livianas y de propósito general mediante autoaprendizaje

Resumen

Los modelos de lenguaje de gran escala se utilizan frecuentemente para responder consultas basadas en grandes corpus de texto (por ejemplo, bases de código, documentos legales o historiales de chat) al incluir todo el corpus en la ventana de contexto y aprovechar el aprendizaje en contexto (ICL, por sus siglas en inglés). Aunque los modelos actuales admiten contextos de 100K a 1M de tokens, esta configuración es costosa de implementar debido a que el consumo de memoria de la caché KV escala con la longitud de la entrada. Exploramos una alternativa: entrenar una caché KV más pequeña fuera de línea para cada corpus. En el momento de la inferencia, cargamos esta caché KV entrenada, a la que llamamos Cartucho, y decodificamos una respuesta. Es crucial destacar que el costo de entrenar un Cartucho puede amortizarse en todas las consultas que hacen referencia al mismo corpus. Sin embargo, encontramos que el enfoque ingenuo de entrenar el Cartucho con la predicción del siguiente token en el corpus no es competitivo con el ICL. En su lugar, proponemos el autoestudio, una receta de entrenamiento en la que generamos conversaciones sintéticas sobre el corpus y entrenamos el Cartucho con un objetivo de destilación de contexto. Descubrimos que los Cartuchos entrenados con autoestudio replican la funcionalidad del ICL, siendo significativamente más económicos de implementar. En puntos de referencia desafiantes de contexto largo, los Cartuchos entrenados con autoestudio igualan el rendimiento del ICL mientras utilizan 38.6 veces menos memoria y permiten un rendimiento 26.4 veces mayor. El autoestudio también extiende la longitud efectiva del contexto del modelo (por ejemplo, de 128k a 484k tokens en MTOB) y, sorprendentemente, conduce a Cartuchos que pueden componerse en el momento de la inferencia sin necesidad de reentrenamiento.

English

Large language models are often used to answer queries grounded in large text corpora (e.g. codebases, legal documents, or chat histories) by placing the entire corpus in the context window and leveraging in-context learning (ICL). Although current models support contexts of 100K-1M tokens, this setup is costly to serve because the memory consumption of the KV cache scales with input length. We explore an alternative: training a smaller KV cache offline on each corpus. At inference time, we load this trained KV cache, which we call a Cartridge, and decode a response. Critically, the cost of training a Cartridge can be amortized across all the queries referencing the same corpus. However, we find that the naive approach of training the Cartridge with next-token prediction on the corpus is not competitive with ICL. Instead, we propose self-study, a training recipe in which we generate synthetic conversations about the corpus and train the Cartridge with a context-distillation objective. We find that Cartridges trained with self-study replicate the functionality of ICL, while being significantly cheaper to serve. On challenging long-context benchmarks, Cartridges trained with self-study match ICL performance while using 38.6x less memory and enabling 26.4x higher throughput. Self-study also extends the model's effective context length (e.g. from 128k to 484k tokens on MTOB) and surprisingly, leads to Cartridges that can be composed at inference time without retraining.

Cartuchos: Representaciones de contexto largo livianas y de propósito general mediante autoaprendizaje

Cartridges: Lightweight and general-purpose long context representations via self-study

Resumen

Support