Deliberación en Espacio Latente a través de la Ampliación Diferenciable de la Caché

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

Autores: Luyang Liu, Jonas Pfeiffer, Jiaxing Wu, Jun Xie, Arthur Szlam

cs.AI

Resumen

Las técnicas que permiten a los modelos de lenguaje grandes (LLMs) "pensar más" mediante la generación y atención a pasos de razonamiento intermedios han mostrado promesa en la resolución de problemas complejos. Sin embargo, los enfoques estándar generan secuencias de tokens discretos inmediatamente antes de responder, lo que puede acarrear costos significativos de latencia y resultar desafiante de optimizar. En este trabajo, demostramos que un LLM congelado puede ser ampliado con un coprocesador sin conexión que opera en la caché clave-valor (kv) del modelo. Este coprocesador aumenta la caché con un conjunto de incrustaciones latentes diseñadas para mejorar la fidelidad de la decodificación subsiguiente. Entrenamos este coprocesador utilizando la pérdida de modelado de lenguaje del decodificador en datos estándar de preentrenamiento, manteniendo el decodificador en sí congelado. Este enfoque permite que el modelo aprenda, de manera diferenciable de extremo a extremo, cómo destilar cálculos adicionales en su kv-cache. Debido a que el decodificador permanece sin cambios, el coprocesador puede operar sin conexión y de forma asíncrona, y el modelo de lenguaje puede funcionar normalmente si el coprocesador no está disponible o si se considera que una caché determinada no requiere cálculos adicionales. Mostramos experimentalmente que cuando se amplía una caché, el decodificador logra una menor perplejidad en numerosos tokens subsiguientes. Además, incluso sin ningún entrenamiento específico de tarea, nuestros experimentos demuestran que la ampliación de la caché reduce consistentemente la perplejidad y mejora el rendimiento en una variedad de tareas intensivas en razonamiento.

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.