Délibération dans l'espace latent via l'augmentation de cache différenciable.

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

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

cs.AI

Résumé

Les techniques permettant aux grands modèles de langage (LLMs) de "penser davantage" en générant et en se concentrant sur des étapes de raisonnement intermédiaires ont montré des promesses dans la résolution de problèmes complexes. Cependant, les approches standard génèrent des séquences de jetons discrets immédiatement avant de répondre, ce qui peut entraîner des coûts de latence significatifs et être difficile à optimiser. Dans ce travail, nous démontrons qu'un LLM figé peut être augmenté avec un coprocesseur hors ligne qui opère sur le cache clé-valeur du modèle. Ce coprocesseur enrichit le cache avec un ensemble d'incrustations latentes conçues pour améliorer la fidélité du décodage ultérieur. Nous formons ce coprocesseur en utilisant la perte de modélisation de langage du décodeur sur des données de pré-entraînement standard, tout en maintenant le décodeur lui-même figé. Cette approche permet au modèle d'apprendre, de manière différenciable de bout en bout, comment distiller des calculs supplémentaires dans son cache clé-valeur. Étant donné que le décodeur reste inchangé, le coprocesseur peut fonctionner hors ligne et de manière asynchrone, et le modèle de langage peut fonctionner normalement si le coprocesseur est indisponible ou si un cache donné ne nécessite pas de calcul supplémentaire. Nous montrons expérimentalement que lorsque qu'un cache est augmenté, le décodeur atteint une perplexité plus faible sur de nombreux jetons ultérieurs. De plus, même sans aucun entraînement spécifique à la tâche, nos expériences démontrent que l'augmentation du cache réduit de manière constante la perplexité et améliore les performances sur une gamme de tâches intensives en raisonnement.

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.