Using offline coprocessor enhances LLM's KV-cache by adding extra "latent embeddings"

Google DeepMind proposed a method that enhances LLMs with an offline coprocessor that works with the models' internal memory (kv-cache).

What's the coprocessor's role?

It enhances the model's KV-cache by adding extra "latent embeddings" (compressed representations) for more accurate outputs.

? What is good about it?

- The coprocessor operates independently, and the base LLM remains frozen.

- It operates offline and asynchronously, meaning it can improve the model’s memory in the background.

- If the coprocessor isn’t available or extra computation isn’t needed, the model still functions as usual.

- The model achieves lower perplexity.

- This method works across various tasks without additional fine-tuning.

Here are the details:

The interaction between the LLM and the coprocessor happens in 3 main steps:

1. KV-cache generation: The frozen LLM processes the input and creates a kv-cache representing its internal state for that input. The LLM itself remains unchanged.
2. Augmentation (the main process): The kv-cache is passed to the coprocessor, which shares the same architecture as the LLM but is trainable. The coprocessor also receives soft tokens (abstract, trainable prompts) that don’t represent actual words but guide the processing. Then the coprocessor uses the kv-cache and soft tokens to generate latent embeddings for added reasoning or context.
3. LLM generation with augmented context: The latent embeddings are added back to the original kv-cache. The LLM processes this augmented kv-cache along with the input to generate the final enhanced output.

Results of using coprocessor:

? Testing on reasoning-heavy tasks showed:

- 10.05% improvement on math reasoning (GSM8K).

- 4.70% improvement on multitask language understanding (MMLU) (multitask language understanding).

? These gains were achieved without any fine-tuning for specific tasks, highlighting the versatility of the method.

? This method also showed significantly lower perplexity compared to baseline models.

Paper: https://arxiv.org/pdf/2412.17747