Картриджи: Легковесные и универсальные представления длинных контекстов через самообучение

Аннотация

Крупные языковые модели часто используются для ответов на запросы, основанные на больших текстовых корпусах (например, базах кода, юридических документах или истории чатов), путем размещения всего корпуса в контекстном окне и использования обучения в контексте (ICL). Хотя современные модели поддерживают контексты размером 100K-1M токенов, такое решение является дорогостоящим для обслуживания, поскольку потребление памяти кэша ключей и значений (KV) масштабируется с длиной входных данных. Мы исследуем альтернативный подход: обучение меньшего кэша KV оффлайн для каждого корпуса. Во время вывода мы загружаем этот обученный кэш KV, который мы называем "Картридж", и декодируем ответ. Важно, что стоимость обучения Картриджа может быть распределена на все запросы, ссылающиеся на один и тот же корпус. Однако мы обнаружили, что наивный подход обучения Картриджа с предсказанием следующего токена на корпусе не конкурирует с ICL. Вместо этого мы предлагаем метод "самообучения" — рецепт обучения, в котором мы генерируем синтетические диалоги о корпусе и обучаем Картридж с целью контекстной дистилляции. Мы обнаружили, что Картриджи, обученные с помощью самообучения, воспроизводят функциональность ICL, при этом будучи значительно дешевле в обслуживании. На сложных бенчмарках для длинных контекстов Картриджи, обученные с помощью самообучения, соответствуют производительности ICL, используя в 38,6 раз меньше памяти и обеспечивая в 26,4 раза большую пропускную способность. Самообучение также расширяет эффективную длину контекста модели (например, с 128k до 484k токенов на MTOB) и, что удивительно, позволяет комбинировать Картриджи во время вывода без необходимости их переобучения.

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.

Картриджи: Легковесные и универсальные представления длинных контекстов через самообучение

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

Аннотация

Support