Cartucce: Rappresentazioni di contesto lungo leggere e generiche attraverso l'auto-apprendimento

Abstract

I modelli linguistici di grandi dimensioni sono spesso utilizzati per rispondere a query basate su ampi corpora testuali (ad esempio codebase, documenti legali o cronologie di chat) inserendo l'intero corpus nella finestra contestuale e sfruttando l'apprendimento in contesto (ICL). Sebbene i modelli attuali supportino contesti di 100K-1M token, questa configurazione è costosa da gestire perché il consumo di memoria della cache KV scala con la lunghezza dell'input. Esploriamo un'alternativa: addestrare offline una cache KV più piccola su ciascun corpus. Al momento dell'inferenza, carichiamo questa cache KV addestrata, che chiamiamo Cartridge, e decodifichiamo una risposta. In modo cruciale, il costo di addestramento di un Cartridge può essere ammortizzato su tutte le query che fanno riferimento allo stesso corpus. Tuttavia, scopriamo che l'approccio ingenuo di addestrare il Cartridge con la previsione del token successivo sul corpus non è competitivo rispetto all'ICL. Proponiamo invece il self-study, una ricetta di addestramento in cui generiamo conversazioni sintetiche sul corpus e addestriamo il Cartridge con un obiettivo di distillazione contestuale. Scopriamo che i Cartridge addestrati con il self-study replicano la funzionalità dell'ICL, pur essendo significativamente più economici da gestire. Su benchmark impegnativi con contesti lunghi, i Cartridge addestrati con il self-study eguagliano le prestazioni dell'ICL utilizzando 38,6 volte meno memoria e consentendo un throughput 26,4 volte maggiore. Il self-study estende inoltre la lunghezza contestuale effettiva del modello (ad esempio da 128k a 484k token su MTOB) e, sorprendentemente, porta a Cartridge che possono essere composti al momento dell'inferenza senza bisogno di riaddestramento.

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.

Cartucce: Rappresentazioni di contesto lungo leggere e generiche attraverso l'auto-apprendimento

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

Abstract

Support