How do we transfer context between models?
When a transformer reads a context, it produces a set of key and value tensors at each layer, namely the kv-cache. This kv-cache is attended over when the LLM generates tokens via decoding.
The crux of long-context inference is building this cache. Prefilling is quadratic, while decoding is linear in length. Our research aims to find a cheap function that can translate between caches, such that the target model can attend over a cache that a completely different model built.
Transferring caches is difficult because each model represents the prompt differently due to differences in architecture. So, learned features don’t naturally align with another model.
The core question then becomes: how do we let a smaller/cheaper model (A) read the prompt, translate its KV cache and let a larger target model (B) answer from that translated memory?
Learned translators
Layer Correspondence
To map caches, we must find which source layers are most relevant for recreating the target model’s cache at layer t. We precompute this using cosine similarities and by minimizing KL divergence of the output logit tokens. The base layers that corresponded the best to target layer generally followed this form (with a few learned exceptions):
We remove the source model’s rotary positional embedding (RoPE) from the keys, translate it, and apply the target model’s RoPE using the initial absolute position.
Training the Translator
Stage 1: Initialization step
We solve the uncentered ridge system (X^TX+0.01|)W_0 = X^TY for both keys and values separately. This gives us a (stable) map that should approximate the target cache before gradient descent.
Stage 2: Gradient Descent
Our initial thoughts on the approach was training to minimize how “different” cache tensors were. After many weeks experimenting with many quantitative definitions of this “difference”, we learned to replicate what the target model does with cache since downstream accuracy of tokens holds importance in practice.
To do so, we split a document into a long prefix and a shorter continuation. We let the source model prefill the prefix, translate the cache, and then give it to the frozen target model.
Distillation:
A teacher model reads the full prefix and the continuation, but the student model sees the continuation on top of the translated prefix cache. We minimize KL divergence on the continuation over cache MSE, cosine similarity, or other methods.
Interesting optimizations:
Each model pair has interesting quirks. In our translation, we compute the first context token and last k tokens with the larger model (sink1). Increasing k improves quality, but at the obvious tradeoff of added latency (k vs time/quality is an interesting graph, join us to see :)
During inference, we simultaneously prefill on the smaller model and translate the caches, which adds some complexity to the streaming algorithm but it saves us a few seconds.
Post Mapping:
We add back the positional encoding according to the target model and then stream the response normally from a combination of the mapped and native cache (for specific tokens).
Relationship between time, context size, and quality
Because prefill is O(n^2), our cache transfer tech becomes exponentially more useful as context size increases.
As context size increases, the quality of our mapper remains stable
What transfers well?
In short, semantic meaning, summaries, and broad factual information are excellent candidates for transfer. Therefore, tasks like classification, retrieval, and document-level decisions are very successful. Because the big model is likely pretrained on a larger dataset, it has more innate knowledge to assess its context during decode.
What can get lost/what needs to be read by the large model:
Token at position 0 is unusually sensitive because the attention-sink key can propagate error (a common ML finding). Exact strings and numbers are fragile, and better off read by the larger model natively. We are working on building a harness to decide this with low latency.
Inference engineering:
With our KV cache transfer technique in inference, it becomes important to keep the source model and the mapper on the same (or nearby) high-bandwidth GPU, interconnected if possible.
Post-training the mapper:
A generic mapper is a strong baseline, but domain-specific or fine-tuned models can/should use a custom mapper, trained on a representative set of data from their workloads. Training takes tens of minutes to a few hours, can be automated, and requires O(100) documents/traces. In order to post-train the mapper for specialized purposes, we keep the source and target models frozen and only optimize the rank-16 residual matrices, using the target model’s output distribution as the teacher. Companies can periodically repeat this process as their models/their tasks change. We ran a controlled experiment, transferring caches from Qwen 3.5-4b to a fine-tuned Qwen 3.8-27b model, and post-training the mapper improved the held-out accuracy from 65.6% to 78.1%. Post-training also helps with reasoning-intensive tasks: on the Zebralogic benchmark, we improved the generic mapper from 0.185 to 0.79 accuracy, matching the 8b (target model) on 465 task documents. The good thing is that we don’t need labeled data for fine-tuning because the target model already provides this. We just need example documents and workloads.
FAQ:
What is more influential? Keys or values?
Recall accuracy for the larger model’s exact keys and mapped values was ≥99.6% while mapped keys with exact values dropped to 73-89%. Most of our work centers around mapping keys accurately.
What quantization are you working in?
All matmuls/maps are done in fp32. We store the resulting target cache in bf16 to match the runtime path.
What are your bottlenecks?
Would love access to more compute to further speed up our research :)
What were your experiment settings?
Our standard setup uses eight documents per optimizer step, AdamW, clips the global gradient norm at 1.0, and uses a short linear warmup followed by a constant learning rate. We also use Muon.