Latent Space Communication via K-V Cache Alignment

TL;DR

This paper proposes learning a shared representation space that aligns the k-v caches of multiple models, creating a high-bandwidth channel for collaboration without altering the underlying pre-trained parameters.

Abstract

Solving increasingly complex problems with large language models (LLMs) necessitates a move beyond individual models and towards multi-model systems that can effectively collaborate. While text has traditionally served as the medium for inter-model communication, a richer and more efficient exchange is possible if models can access each other's internal states directly. In this paper, we propose learning a shared representation space that aligns the k-v caches of multiple models, creating a high-bandwidth channel for collaboration without altering the underlying pre-trained parameters. We do so by augmenting each model with adapters to translate its state into and out of this shared space. Via a suite of experiments with Gemma-2 models, we demonstrate that this approach not only enables seamless inter-model communication but also improves individual model performance. We also show that the shared space allows for the direct transfer of learned skills, such as soft prompts, between different models. Our work represents a significant step towards a future where models can fluidly share knowledge and capabilities.

Figures (6)

• Figure 1: Overview of an example instantiation of our framework. [Upper] Each model (here Russian and Spanish LMs) is augmented with adapters for translating a k-v cache block into and out of the shared latent space. E.g Es-Out translates a k-v cache out of the Es model's latent space and into the shared space. [Lower] We can perform mixed language modelling (for context switching settings), by translating generated prefix caches from a source model and continuing to decode from the target model.
• Figure 2: Cross attention architecture for mapping into and out of the shared latent space. The first layer k-or-v cache serves as seed input to the model. Each layer in the module cross-attends to the corresponding layer in the input k-or-v (not both at the same time) block, meaning that the number of layers in the translator is equal to the number of layers in the corresponding model. We use the superscript $^{*}$ to represent vectors in $\mathbf{\Sigma}$.
• Figure 3: We can learn a shared latent space between different models with different degrees of overlap in their training trajectories. Evaluation loss is negative log likelihood on held-out text.
• Figure 4: We extend a pool of 3 models with different seeds-{1, 2, 3} with a fourth model. Solid lines mark translation paths that were not trained on. Without explicitly training on a these translation paths, we are still able to zero-shot translate the between the k-v caches of the two models -- seeds-{1, 4} -- with only mild performance degradation compared to fully trained paths.
• Figure 5: We adapt the shared global space between two models trained on different seeds to allow for zero-shot module portability. Shaded areas correspond to standard deviation over the 200 meta-eval task. We compare using pre-trained latent space adapters versus randomly initialized ones (cyan).