Reproduce “Deliberation in Latent Space via Differentiable Cache Augmentation” on OpenResearch

About this paper

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

This research introduces differentiable cache augmentation, a method to enhance the reasoning abilities of frozen large language models by enabling deliberation in their latent space. A trainable coprocessor generates context-dependent latent embeddings from the LLM's key-value cache, which are then used to augment the LLM's internal state, leading to consistent reductions in perplexity and improved performance on complex reasoning benchmarks without modifying the base model.

Key insights

• Deliberation can be effectively performed in the continuous latent space of LLMs by augmenting the KV-cache, rather than relying on discrete token generation.
• Training a separate coprocessor to generate these latent embeddings, while keeping the base LLM frozen, enables end-to-end differentiable optimization using standard language modeling loss.
• Initializing the coprocessor with pretrained LLM weights and using an 'ahead token prediction' strategy during pretraining helps it learn useful representations for future text generation.

Summary by alphaXiv.