Cartuchos: Representações de contexto longo leves e de propósito geral via autoestudo

Resumo

Modelos de linguagem de grande escala são frequentemente usados para responder a consultas baseadas em grandes corpora de texto (por exemplo, bases de código, documentos jurídicos ou históricos de chat) ao colocar todo o corpus na janela de contexto e aproveitar o aprendizado em contexto (ICL, in-context learning). Embora os modelos atuais suportem contextos de 100K a 1M de tokens, essa configuração é custosa para servir, pois o consumo de memória do cache KV escala com o comprimento da entrada. Exploramos uma alternativa: treinar um cache KV menor offline para cada corpus. No momento da inferência, carregamos esse cache KV treinado, que chamamos de Cartucho, e decodificamos uma resposta. Criticamente, o custo de treinar um Cartucho pode ser amortizado em todas as consultas que referenciam o mesmo corpus. No entanto, descobrimos que a abordagem ingênua de treinar o Cartucho com a previsão do próximo token no corpus não é competitiva com o ICL. Em vez disso, propomos o autoestudo, uma receita de treinamento na qual geramos conversas sintéticas sobre o corpus e treinamos o Cartucho com um objetivo de destilação de contexto. Descobrimos que os Cartuchos treinados com autoestudo replicam a funcionalidade do ICL, sendo significativamente mais baratos de servir. Em benchmarks desafiadores de contexto longo, os Cartuchos treinados com autoestudo igualam o desempenho do ICL, enquanto usam 38,6 vezes menos memória e permitem um throughput 26,4 vezes maior. O autoestudo também estende o comprimento efetivo do contexto do modelo (por exemplo, de 128k para 484k tokens no MTOB) e, surpreendentemente, leva a Cartuchos que podem ser compostos no momento da inferência sem retreinamento.

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: Representações de contexto longo leves e de propósito geral via autoestudo

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

Resumo

Support