Deliberação no Espaço Latente via Augmentação de Cache Diferenciável

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

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

cs.AI

Resumo

Técnicas que permitem que grandes modelos de linguagem (LLMs) "pensem mais" gerando e atendendo a etapas intermediárias de raciocínio têm mostrado promessa na resolução de problemas complexos. No entanto, as abordagens padrão geram sequências de tokens discretos imediatamente antes de responder, o que pode acarretar custos significativos de latência e ser desafiador de otimizar. Neste trabalho, demonstramos que um LLM congelado pode ser aumentado com um coprocessador offline que opera no cache chave-valor do modelo. Esse coprocessador aumenta o cache com um conjunto de embeddings latentes projetados para melhorar a fidelidade da decodificação subsequente. Treinamos esse coprocessador usando a perda de modelagem de linguagem do decodificador em dados padrão de pré-treinamento, mantendo o próprio decodificador congelado. Essa abordagem permite que o modelo aprenda, de forma diferenciável de ponta a ponta, como destilar computação adicional em seu cache chave-valor. Como o decodificador permanece inalterado, o coprocessador pode operar offline e de forma assíncrona, e o modelo de linguagem pode funcionar normalmente se o coprocessador não estiver disponível ou se um determinado cache não precisar de computação extra. Mostramos experimentalmente que, quando um cache é aumentado, o decodificador alcança menor perplexidade em numerosos tokens subsequentes. Além disso, mesmo sem nenhum treinamento específico da tarefa, nossos experimentos demonstram que a ampliação do cache consistentemente reduz a perplexidade e melhora o desempenho em uma variedade de tarefas intensivas em raciocínio.

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.