Überlegung im latenten Raum mittels differenzierbarer Cache-Erweiterung.

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

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

cs.AI

Zusammenfassung

Techniken, die es großen Sprachmodellen (LLMs) ermöglichen, durch Generierung und Beachtung von Zwischenschritten im Denkprozess "besser zu denken", haben das Potenzial gezeigt, komplexe Probleme zu lösen. Die Standardansätze generieren jedoch unmittelbar vor der Antwort Sequenzen diskreter Tokens, was zu erheblichen Latenzkosten führen kann und die Optimierung herausfordernd macht. In dieser Arbeit zeigen wir, dass ein eingefrorenes LLM mit einem Offline-Coprozessor ergänzt werden kann, der auf dem Schlüssel-Wert (kv)-Cache des Modells arbeitet. Dieser Coprozessor erweitert den Cache um eine Reihe latenter Einbettungen, die darauf abzielen, die Genauigkeit des nachfolgenden Decodings zu verbessern. Wir trainieren diesen Coprozessor unter Verwendung des Sprachmodellierungsverlusts des Decoders auf standardmäßigen Pretraining-Daten, während der Decoder selbst eingefroren bleibt. Dieser Ansatz ermöglicht es dem Modell, auf differenzierbare Weise zu lernen, wie zusätzliche Berechnungen in seinen kv-Cache destilliert werden können. Da der Decoder unverändert bleibt, kann der Coprozessor offline und asynchron arbeiten, und das Sprachmodell kann normal funktionieren, wenn der Coprozessor nicht verfügbar ist oder festgestellt wird, dass ein bestimmter Cache keine zusätzlichen Berechnungen erfordert. Experimentell zeigen wir, dass der Decoder bei einer Cache-Erweiterung eine geringere Perplexität bei zahlreichen nachfolgenden Tokens erreicht. Darüber hinaus zeigen unsere Experimente, dass die Cache-Erweiterung konsistent die Perplexität reduziert und die Leistung bei einer Vielzahl von auf Reasoning basierenden Aufgaben verbessert, selbst ohne task-spezifisches Training.

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.