Overweging in Latente Ruimte via Differentieerbare Cache-augmentatie

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

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

cs.AI

Samenvatting

Technieken die grote taalmodellen (LLM's) in staat stellen om "meer te denken" door het genereren en aandacht te schenken aan tussenliggende redeneerstappen, hebben veelbelovende resultaten laten zien bij het oplossen van complexe problemen. De standaard benaderingen genereren echter sequenties van discrete tokens direct voor het reageren, wat aanzienlijke latentiekosten met zich mee kan brengen en moeilijk te optimaliseren kan zijn. In dit werk tonen we aan dat een bevroren LLM kan worden uitgebreid met een offline coprocessor die werkt op de sleutel-waarde (kv) cache van het model. Deze coprocessor breidt de cache uit met een reeks latente embeddings die zijn ontworpen om de geloofwaardigheid van daaropvolgende decodering te verbeteren. We trainen deze coprocessor met behulp van het taalmodelleringsverlies van de decoder op standaard voorafgaande trainingsgegevens, terwijl de decoder zelf bevroren blijft. Deze benadering stelt het model in staat om op een end-to-end differentieerbare manier te leren hoe extra berekeningen te destilleren in zijn kv-cache. Omdat de decoder ongewijzigd blijft, kan de coprocessor offline en asynchroon werken, en kan het taalmodel normaal functioneren als de coprocessor niet beschikbaar is of als een bepaalde cache niet extra berekening vereist. We tonen experimenteel aan dat wanneer een cache wordt uitgebreid, de decoder een lagere perplexiteit bereikt op tal van daaropvolgende tokens. Bovendien laten onze experimenten zelfs zonder enige taakspecifieke training zien dat cache-uitbreiding consequent de perplexiteit verlaagt en de prestaties verbetert bij een reeks redeneerintensieve taken.

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.