Сквозное сжатие контекста в масштабе

Аннотация

Вывод (инференс) языковых моделей с длинным контекстом ограничен объемом памяти, так как KV-кэш растет с увеличением длины контекста. Современные методы сжатия KV-кэша не оправдывают ожиданий: они либо существенно ухудшают качество модели, либо требуют значительного времени и вычислительных ресурсов для сжатия одного длинного промпта. Кроме того, многие методы требуют, чтобы входные данные помещались в окно контекста целевой модели, и обычно несовместимы с современными производственными движками инференса. Компрессоры типа «кодировщик-декодировщик», которые преобразуют длинную последовательность токенов в более короткую последовательность скрытых эмбеддингов, потребляемых декодировщиком, в принципе являются привлекательной альтернативой. Однако существующие подходы неконкурентоспособны по сравнению со сжатием KV-кэша на границе точности и эффективности. В данной работе мы пересматриваем сжатие на основе архитектуры «кодировщик-декодировщик» и устраняем этот разрыв. Сначала мы проводим поиск архитектуры, предварительно обучая множество вариантов с нуля, чтобы определить наилучший способ проектирования и обучения компрессоров данного типа. Руководствуясь полученными результатами, мы последовательно предварительно обучаем семейство моделей с кодировщиком на 0,6 миллиарда параметров и декодировщиком на 4 миллиарда параметров на более чем 350 миллиардах токенов каждая, с коэффициентами сжатия 1:4, 1:8 и 1:16. Мы представляем Latent Context Language Models (LCLMs) — семейство компрессоров, улучшающих границу Парето по производительности на общих задачах, скорости сжатия и пиковому использованию памяти. Мы демонстрируем, что LCLM служат эффективными магистралями для агентов с длительным горизонтом планирования, позволяя агенту бегло просматривать сжатый длинный контекст и адаптивно расширять релевантные сегменты по запросу.

English

Long-context language model inference is bottlenecked by memory, as the KV cache grows with context length. Recent techniques to compress the KV cache fall short: they either degrade model quality substantially or require considerable time and compute to compress a single long prompt. Furthermore, many methods require the input to fit within the target model's context window, and are generally incompatible with modern production inference engines. Encoder-decoder compressors, which map a long token sequence to a shorter sequence of latent embeddings consumed by a decoder, are an appealing alternative in principle. However, existing approaches are not competitive with KV cache compression on the accuracy-efficiency frontier. In this work, we revisit encoder-decoder compression and close this gap. We first perform an architecture search, pre-training many variants from scratch to determine how best to design and train encoder-decoder compressors. Guided by our findings, we continually pre-train a family of 0.6B-encoder, 4B-decoder models on over 350B tokens each, at compression ratios of 1:4, 1:8, and 1:16. We introduce Latent Context Language Models (LCLMs), a family of compressors that improve the Pareto frontier across general-task performance, compression speed, and peak memory usage. We demonstrate that LCLMs serve as efficient backbones for long-horizon agents, letting the agent skim through a compressed long context and adaptively expand relevant segments on demand.