Latent Space Communication via K-V Cache Alignment

1. Introduction

• Growing pools of generalist and domain-specific models create opportunities for multi-model collaboration and for sharing skills such as soft prompts or prefix k-v caches.The introduction frames collaboration as a way to address problems beyond the scope of any single model.Source: §1 · p. 1
• Because a k-v cache represents a model’s internal state, accessing another model’s cache could offer higher-bandwidth communication than token-level text exchange.The introduction notes that this could support extracting knowledge, following soft reasoning chains, or reusing partial computations and skills.Source: §1 · p. 1
• Models trained under different conditions cannot be expected to have interchangeable latent spaces, because cache discrepancies compound with depth.Tokens create dependencies on preceding layers’ k-v spaces, contributing to this compounding mismatch.Source: §1 · p. 1
• In the context-switching example, a source model’s generated prefix cache is translated before the target model continues decoding.Figure 1 illustrates mixed-language modelling with Russian and Spanish models.Source: Figure 1 · p. 2
• The proposed global shared k-v cache space uses two adapters per model to translate cache blocks into and out of that space.The architecture is designed as an implicit shared space whose parameter overhead grows linearly with the number of models.Source: §1 · p. 2
• The experiments use Gemma-2 models trained under varied conditions and learn the shared space without modifying their pretrained parameters.The paper describes a scalable architecture and efficient learning algorithm.Source: §1 · p. 2
• The shared space enables inter-model cache communication, improves individual language-model performance with sufficient data, and supports transferring learned soft prompts between models.The paper reports these findings from experiments with Gemma-2 models.Source: §1 · p. 2

2.1. Setup

• Each model maps its cache into a fixed-size shared space and reconstructs model-specific layer information on the way out; adding a model requires only its own pair of maps.The method assumes a shared vocabulary for the simple setup, but does not require it.Source: §2.1 · pp. 2–3

2.2. Translator Architecture

• Layer-wise cross-attention translates cache blocks, using the previous layer’s output for queries and the corresponding input-cache layer for keys and values.The translator has as many layers as the corresponding model, and key and value inputs share cross-attention parameters.Source: §2.2 · pp. 3–4, Figure 2 · p. 3

2.3.1. Reconstruction Loss

• The learning algorithm uses two losses; reconstruction translates one model’s cache through the shared space into another model’s cache, while longer translation chains are computationally expensive to train through.Reducing the intermediate layer count can lower adapter latency at the cost of expressiveness.Source: §2.3 · p. 4, §2.3.1 · p. 4

2.3.2. Suffix Language Modelling

• Suffix language modelling trains on a target model’s remaining text after conditioning on a prefix cache translated from another model.Because the models operate on disjoint text sections, this loss can be used with different vocabularies; reconstruction requires a vocabulary mapping in that case.Source: §2.3.2 · p. 4

3. Experimental Setting

• Experiments use frozen Gemma-2 models sized 100M–400M parameters and multilingual C4 sequences up to 512 tokens, focusing on English, Russian, and Spanish unless otherwise specified.The parameter count excludes token embeddings.Source: §3 · pp. 4–5
• Adapters use 32 heads of dimension 64; the default embed_dim_factor is 2 and translation_dim_factor is 1, yielding translators about one-quarter the base-model size.Source: §3 · p. 5
• Translator training runs for 50k steps, and unless otherwise specified uses the first 128 tokens as prefix and the remaining tokens as the target-model suffix.The learning rate warms up for 2.5k steps and then follows cosine decay.Source: §3 · p. 5

4. Models can be taught to communicate via a global k-v latent space

• The experiments use suffix language-modelling loss to train communication across models with shared origins or disjoint training trajectories.The models in the experiments share a token vocabulary.Source: §4 · p. 5

4.1. Same trajectory, different checkpoints

• Prefix caches translated through the shared space improve on the base model for different checkpoints, and the stronger-to-weaker translation reaches that level with less training data.The reported gain also appears when the prefix cache is cyclically translated into the shared space and back into its local space.Source: §4.1 · p. 5, Figure 3 · p. 6

4.2. Same origin model, different fine-tuning distributions

• Even with greater latent-space divergence between Russian- and Spanish-finetuned models, cyclic translation can outperform decoding with the untouched cache when trained on enough data.The setup uses parallel text and skips the target-language prefix-length tokens before decoding.Source: §4.2 · p. 6, Figure 3 · p. 6

4.3. Same data distribution but different random initialization

• Models with different random seeds but the same data distribution can be taught to use one another’s k-v caches.The experiment uses a 64-token prefix and 128-token sequence length.Source: Figure 3 · p. 6, §4.3 · p. 7

4.4. Different Model Sizes

• Adapters translate between the 400M model’s 16-layer cache and the 100M model’s 4-layer cache, improving the weaker model with translated prefix caches.The models also have disjoint training trajectories because their random initializations and amounts of training data differ.Source: §4.4 · p. 7

5. The global k-v latent space is extensible

• Adding a model requires learning its adapters into and out of the shared space, but does not necessarily require retraining existing adapters.The experiments demonstrate extension using a fourth model with a different random seed.Source: §5 · p. 7
• Seed-1 and Seed-4 can zero-shot translate between their caches with performance comparable to explicitly trained paths, although their paths were held out during training.The Seed-4 adapters were trained using paths involving Seeds 2 and 3.Source: §5 · p. 7, Figure 4 · p. 7
• When the pool grows, not all existing models may be needed to train an incoming model’s translator, reducing extension memory and compute overhead.Source: §5 · p. 7

6. A global k-v latent space enables module portability

• A task-specific module expressed as a k-v cache can be translated for use by other models, amortizing skill-learning compute across the pool and limiting direct exposure of task data to one model.Differentially private training of the original module is described as a way to achieve stronger privacy guarantees.Source: §6 · p. 8

6.1. Setup

• Prompt recovery learns a soft prompt from example completions of a prompted model and evaluates whether it generalizes to held-out completions.The setup uses 5.8K meta-training tasks and 200 meta-test tasks.Source: §6.1 · p. 8

6.2. Results

• Zero-shot ported modules outperform random untrained target soft prompts and approach the upper bound of learning a soft prompt directly on each target model.The experiments use a few thousand meta-training tasks; the authors suggest more data could match or exceed the direct-learning baseline.Source: §6.2 · p. 8, Figure 5 · p. 9
• Pre-training the global latent space before meta-training helps, with the best reported performance from suffix-language-modelling pre-training without reconstruction loss.Source: §6.2 · p. 8

7.1. Performance improves with shared space adapter model size

• Identity mapping performs poorly even when models differ only by random initialization; a simple linear map matches or exceeds base-model performance.Identity mapping also requires matching cache dimensions exactly.Source: §7.1 · p. 9
• Linear maps expand cache dimensions by 8× for 269M models and 24× for 806M models.These are the upward projection factors used to create the shared latent space.Source: §7.1 · p. 9
• The cross-attention translator outperforms the simple linear map with fewer parameters and improves consistently as latent-space adapter size increases.Unlike the linear map, its performance scales consistently with adapter size.Source: §7.1 · p. 9

7.2. We do not need to see all all paths in a single batch during training

• For some translation directions, 2 paths per step outperform 4, so using more paths is not uniformly better.The figure caption reports this direction-dependent comparison.Source: Figure 6 · p. 10
• 3.53 NLL across all 9 possible paths approaches the 3.460 NLL average base performance, even with 1 translation path per training step.Seed-1 and Seed-3 perform best with 2 paths per batch, while Seed-2 performs best with 4.Source: §7.2 · p. 10

7.3. Reconstruction loss is unnecessary in the presence of the suffix language modelling loss

• For the primary results, the authors exclude reconstruction loss because suffix language-modeling loss makes it unnecessary.At fixed data, Table 2 reports that reconstruction loss did not improve learning efficiency.Source: §7.3 · p. 10, Table 2 · p. 10

Related communication approaches

• LLM agents commonly communicate through natural language or structured formats, while multi-agent reinforcement learning and emergent communication often use low-dimensional vectors [Wang et al., 2024, 2023; Gronauer and Diepold, 2022; Lazaridou and Baroni, 2020].The paper describes its framework as a higher-bandwidth channel based on direct access to k-v cache latent spaces.Source: §8 · p. 10

Mixture-of-experts applications

• For mixture-of-experts systems that periodically switch models, translating cache spaces could reduce the cost of recomputing prefix caches after a switch [Douillard et al., 2024; Filippova et al., 2024].The cited architectures are described as incurring high inference latency from that recomputation.Source: §8 · pp. 10–11

Latent-space alignment

• Prior latent-space alignment work includes linear transformations and relative representations, often for statically stitching vision models or simpler architectures [Lähner and Moeller (2024); Moschella et al. (2022)].This paper instead motivates alignment for dynamic, stateful collaboration.Source: §8 · p. 11

Conclusion

• The paper presents a scalable method for direct latent-space communication between large language models using a shared k-v cache representation.The conclusion says the approach bridges models trained under different conditions.Source: §9 · p. 11

Future work

• The current work focuses on language modeling as a proof of concept; future work should test more heterogeneous architectures and downstream tasks such as agentic research and software engineering.Source: §9 · p. 11

A. Model Details

• Table 3 provides an overview of model parameters and design choices.Source: Table 3 · p. 14

B. Prompt Examples

• The prompt-completion examples include requests for stories, math instruction, and historical descriptions.Source: B. Prompt Examples · p. 14
• One example asks for a pterodactyl trying to convince a friend to buy a car.Source: B. Prompt Examples · p. 15

• Table 4 presents zero-shot module portability performance before meta-training; its caption identifies meta-training as key to adapting the framework.Source: Table 4 · p. 15