Cartouches : Représentations de contexte long léger et polyvalent via l'auto-apprentissage

Résumé

Les grands modèles de langage sont souvent utilisés pour répondre à des requêtes basées sur de vastes corpus de texte (par exemple, des bases de code, des documents juridiques ou des historiques de chat) en plaçant l'intégralité du corpus dans la fenêtre de contexte et en exploitant l'apprentissage en contexte (ICL). Bien que les modèles actuels prennent en charge des contextes de 100 000 à 1 million de tokens, cette configuration est coûteuse à déployer car la consommation mémoire du cache KV augmente avec la longueur de l'entrée. Nous explorons une alternative : entraîner un cache KV plus petit hors ligne pour chaque corpus. Au moment de l'inférence, nous chargeons ce cache KV entraîné, que nous appelons une Cartouche, et décodons une réponse. De manière cruciale, le coût de l'entraînement d'une Cartouche peut être amorti sur toutes les requêtes faisant référence au même corpus. Cependant, nous constatons que l'approche naïve consistant à entraîner la Cartouche avec une prédiction de token suivant sur le corpus n'est pas compétitive par rapport à l'ICL. À la place, nous proposons l'auto-apprentissage, une méthode d'entraînement dans laquelle nous générons des conversations synthétiques sur le corpus et entraînons la Cartouche avec un objectif de distillation de contexte. Nous constatons que les Cartouches entraînées avec l'auto-apprentissage reproduisent la fonctionnalité de l'ICL, tout en étant nettement moins coûteuses à déployer. Sur des benchmarks exigeants en contexte long, les Cartouches entraînées avec l'auto-apprentissage égalent les performances de l'ICL tout en utilisant 38,6 fois moins de mémoire et en permettant un débit 26,4 fois plus élevé. L'auto-apprentissage étend également la longueur de contexte effective du modèle (par exemple, de 128 000 à 484 000 tokens sur MTOB) et, de manière surprenante, conduit à des Cartouches qui peuvent être composées au moment de l'inférence sans nécessiter de réentraînement.

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.

Cartouches : Représentations de contexte long léger et polyvalent via l'auto-apprentissage

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

Résumé

Support