Deliberazione nello spazio latente tramite l'incremento differenziabile della cache.

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

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

cs.AI

Abstract

Le tecniche che consentono ai grandi modelli linguistici (LLM) di "pensare di più" generando e prestando attenzione a passaggi intermedi di ragionamento hanno mostrato promesse nella risoluzione di problemi complessi. Tuttavia, gli approcci standard generano sequenze di token discreti immediatamente prima di rispondere, e quindi possono comportare costi di latenza significativi e essere difficili da ottimizzare. In questo lavoro, dimostriamo che un LLM congelato può essere potenziato con un coprocessore offline che opera sulla cache chiave-valore (kv) del modello. Questo coprocessore potenzia la cache con un insieme di embedding latenti progettati per migliorare la fedeltà della decodifica successiva. Alleniamo questo coprocessore utilizzando la perdita di modellazione del linguaggio dal decodificatore su dati standard di preaddestramento, mantenendo il decodificatore stesso congelato. Questo approccio consente al modello di imparare, in modo differenziabile end-to-end, come distillare ulteriori calcoli nella sua kv-cache. Poiché il decodificatore rimane invariato, il coprocessore può operare offline e in modo asincrono, e il modello linguistico può funzionare normalmente se il coprocessore non è disponibile o se una determinata cache non richiede calcoli aggiuntivi. Mostriamo sperimentalmente che quando una cache è potenziata, il decodificatore raggiunge una minore perplessità su numerosi token successivi. Inoltre, anche senza alcun addestramento specifico per compiti, i nostri esperimenti dimostrano che l'aumento della cache riduce costantemente la perplessità e migliora le prestazioni in una serie di compiti intensivi di ragionamento.

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.