Размышления в скрытом пространстве с помощью дифференцируемого расширения кэша

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

Авторы: Luyang Liu, Jonas Pfeiffer, Jiaxing Wu, Jun Xie, Arthur Szlam

cs.AI

Аннотация

Техники, позволяющие крупным языковым моделям (LLM) "думать больше", генерируя и обращая внимание на промежуточные шаги рассуждений, показали перспективы в решении сложных проблем. Однако стандартные подходы генерируют последовательности дискретных токенов непосредственно перед ответом, что может привести к значительным задержкам и затруднить оптимизацию. В данной работе мы демонстрируем, что замороженную LLM можно дополнить оффлайн-копроцессором, который работает с кэшем ключ-значение (kv) модели. Этот копроцессор дополняет кэш набором латентных вложений, разработанных для улучшения достоверности последующего декодирования. Мы обучаем этот копроцессор, используя потери языкового моделирования от декодера на стандартных данных предварительного обучения, при этом декодер сам остается замороженным. Этот подход позволяет модели учиться, в энд-ту-энд дифференцируемом стиле, как дистиллировать дополнительные вычисления в свой kv-кэш. Поскольку декодер остается неизменным, копроцессор может работать оффлайн и асинхронно, и языковая модель может функционировать нормально, если копроцессор недоступен или если данный кэш считается не требующим дополнительных вычислений. Мы экспериментально показываем, что при дополнении кэша декодер достигает более низкой перплексии на многочисленных последующих токенах. Более того, даже без какого-либо задаче-специфического обучения, наши эксперименты демонстрируют, что дополнение кэша последовательно снижает перплексию и улучшает производительность на широком спектре задач, требующих рассуждений.

English

Techniques enabling large language models (LLMs) to "think more" by generating and attending to intermediate reasoning steps have shown promise in solving complex problems. However, the standard approaches generate sequences of discrete tokens immediately before responding, and so they can incur significant latency costs and be challenging to optimize. In this work, we demonstrate that a frozen LLM can be augmented with an offline coprocessor that operates on the model's key-value (kv) cache. This coprocessor augments the cache with a set of latent embeddings designed to improve the fidelity of subsequent decoding. We train this coprocessor using the language modeling loss from the decoder on standard pretraining data, while keeping the decoder itself frozen. This approach enables the model to learn, in an end-to-end differentiable fashion, how to distill additional computation into its kv-cache. Because the decoder remains unchanged, the coprocessor can operate offline and asynchronously, and the language model can function normally if the coprocessor is unavailable or if a given cache is deemed not to require extra computation. We show experimentally that when a cache is augmented, the decoder achieves lower perplexity on numerous subsequent tokens. Furthermore, even without any task-specific training, our experiments demonstrate that cache augmentation consistently reduces perplexity and improves performance across a range of reasoning-intensive tasks.