A larger model can hand prompt work to a smaller one
BridgeKV: Large-Model Prefill and Small-Model Decoding for LLM Inference via KV Cache Translation · arXiv:2611.11761
In plain words
A large language model can be good at understanding a prompt, while a smaller one can be quicker at writing the answer. BridgeKV is a method that lets the large model read the prompt, then translates the information it saved so the small model can continue from it.
In tests of six model pairs across five datasets, BridgeKV scored higher than small-model versions trained for the same tasks with the same number of trainable parameters. The timing results cover one GPU request and separate simulated workloads.
The research
A larger model may handle a prompt well, but using it to generate every answer token can take longer than using a smaller model. BridgeKV tests whether the smaller model can take over after the larger one reads the prompt, while retaining some of the larger model’s task performance.
BridgeKV translates a larger language model’s prompt cache into a smaller model’s format, letting the smaller model generate the answer after the larger model reads the prompt.
The evaluation compared BridgeKV with small-model baselines trained for the same tasks using matched numbers of trainable parameters.
Across six model pairs and five datasets, BridgeKV scored higher in all 30 test settings, using answer-word overlap for question answering and shared word sequences for summarization.
This supports the translated handoff as a way to improve the smaller model’s task scores in the evaluated settings. The size of the gain varied by pair and dataset.
The task-score comparison covers six model pairs and five datasets. The timing results come from one request on one GPU and simulations of specified workloads, and the method requires a shared tokenizer and a separately trained connector for each model pair and dataset.
What the evidence shows
- Translate the prompt cache
BridgeKV is an inference method for handing generation from a larger language model to a smaller one. The larger model processes the prompt; a learned connector translates the information it saved into the cache layout the smaller model expects.
The smaller model then generates the answer. This addresses a concrete obstacle: caches differ across models in their layer counts, head layouts and internal representations.
- Teach the smaller model
The connector is trained to do more than imitate the smaller model’s usual cache. With both base models fixed, training aligns the translated cache with the smaller model’s cache and teaches the smaller model to match the larger model’s next-token predictions.
Validation can select a second stage that also trains adapters on the smaller model for the task. That staged design gives the handoff a way to carry information from the larger model into generation.
- Compare task scores
On test splits, BridgeKV scored higher than small-model baselines trained for the same tasks with matched numbers of trainable parameters in all 30 evaluated model-pair-by-dataset settings. The comparison spans six pairs from two instruction-tuned model families and five datasets.
This supports a quality improvement across those tested settings, while the amount of improvement varied by task and pair.
- Test the layer choices
A separate layer-mapping comparison helps explain one design choice: on one model pair and question-answering task, learned weighted combinations of nearby layers scored above the tested fixed layer assignments. This points to the value of choosing layer contributions during training in that setting. It does not show that the same mapping wins across all pairs or tasks.
- Check request timing
Timing results examine whether the handoff’s quality gains come with a time cost. In one measured request, the 8B-to-1B handoff finished close to the 1B model alone and faster than the 8B model alone.
In a separate simulated workload, it also reduced response time relative to the 8B model alone. These are distinct results from a single-request GPU measurement and a workload simulation.
Deep
- Original paper authors
- Payman Behnam, Debopam Sanyal, Pranav Ponnusamy, Kusum Meda Ravi, Pranav Sastry
- Submitted to arXiv
- License
- CC BY 4.0