Technisch overzicht: Cross-Model KV Cache Transfer in LLM-families

1. Probleemformulering

De productie van Large Language Model (LLM) serving leunt steeds meer op multi-model orchestratie voor kosten-kwaliteit cascading, wisselingen midden in een gesprek en routing. Deze workflows schakelen frequent tussen verschillende groottes binnen dezelfde modelfamilie (bijv. Qwen3 14B naar 32B). Hoewel deze modellen dezelfde kernarchitecturale keuzes en pretraining-data delen, verschillen ze in schaal, aantal lagen en trainingsrecepten.

De primaire bottleneck in dergelijke workflows is de prefill-kosten. Elke keer dat de context van een gesprek wordt overgedragen van een bronmodel naar een doelmodel, moet het doelmodel de volledige geaccumuleerde context opnieuw verwerken om zijn eigen Key-Value (KV) cache te vullen. Deze forward pass schaalt met zowel de modelgrootte als de promptlengte, wat de efficiëntiewinst van het wisselen van modellen tenietdoet. Bestaande oplossingen zoals prefix caching werken alleen binnen één enkel model, en eerdere cross-model benaderingen (bijv. neurale fusers, latent adapters) vereisen doorgaans gradiënt-gebaseerde training of strikte architecturale beperkingen (bijv. identiek aantal lagen).

Dit paper stelt Cross-Model KV Cache Transfer voor, een methode om de vooraf gevulde KV-cache van het bronmodel direct te hergebruiken in het doelmodel, waardoor de prefill van het doelmodel volledig wordt overgeslagen. Het doel is om een mapping $f: C_S \to \hat{C}_T$ te vinden zodanig dat de output van het doelmodel bij het decoderen vanaf de gemapte cache $\hat{C}_T$ equivalent is aan het decoderen vanaf zijn eigen grondwaarheid-cache $C_T$.

2. Methodologie

De auteurs hypothetiseren dat cross-model KV-relaties een substantiële lineaire structuur bezitten, wat een closed-form, gradiënt-vrije mapping mogelijk maakt. Het voorgestelde framework opereert op "matched-KV" paren, waarbij het bron- en doelmodel hetzelfde aantal KV-heads en dezelfde per-head dimensie delen, zelfs als hun totale aantal lagen of parameteromvang verschilt.

De mapper is ontworpen met drie kernelementen:

A. Per-Head Ridge Regression

In plaats van een globale neurale netwerk, past de methode een onafhankelijke lineaire mapping toe voor elke doel-laag en head.

• Input: Geconcateneerde KV-features van een geselecteerde subset van bron-lagen.
• Output: De verwachte Key ($K$) en Value ($V$) vectoren van het doelmodel.
• Optimalisatie: De gewichten worden opgelost via ridge regressie (closed-form oplossing) op een kleine kalibratieset (500 sequenties van 1.024 tokens uit FineWeb-Edu). Ridge regularisatie ($\lambda=0.01$) wordt gebruikt om de bijna-singuliere aard van de feature matrix te hanteren wanneer bron-lagen gecorreleerd zijn.

B. Cross-Layer Bronselectie

Aangezien bron- en doelmodellen verschillende dieptes hebben ($L_S \neq L_T$), moet de methode bepalen welke bron-lagen elke doel-laag informeren.

• Selectiestrategie: Voor elke doel-laag worden de top-$k$ bron-lagen geselecteerd op basis van hun voorspellende kracht (gemeten door head-gemiddelde $R^2$ op een kalibratieset).
• Mechanisme: De KV-features van deze top-$k$ bron-lagen worden geconcateneerd om de input te vormen voor de ridge regressie. Empirische analyse toont aan dat complementaire informatie verspreid is over meerdere bron-lagen; één enkele laag is onvoldoende, maar de prestaties verzadigen snel (bijv. $k=6$ vangt het grootste deel van de variantie van $k=alle$).

C. Content-Space Mapping (RoPE Factoring)

Standaard LLM's gebruiken Rotary Position Embeddings (RoPE), die positie-afhankelijke rotaties toepassen op keys.

• Probleem: Het fitten van een mapper direct op geroteerde keys bindt de gewichten aan de specifiek tijdens de kalibratie geziene positieverdeling, wat de transferabiliteit naar verschillende contextlengtes beperkt.
• Oplossing: De methode verwijdert RoPE uit de bron-keys voordat er gemapt wordt, en past de doel-RoPE pas toe na de mapping.
  • $K_{stripped} = K_{source} \cdot R^{-1}_{\Theta_s}$
  • $\hat{K}_{target} = (K_{stripped} \cdot W + b) \cdot R_{\Theta_t}$
• Dit ontkoppelt de positionele rotatie van de semantische inhoud, waardoor de lineaire fit positie-onafhankelijk en herbruikbaar is over verschillende contextlengtes (tot 32k tokens).

D. Niet-lineaire Extensie (MLP)

Voor paren waarbij de lineaire ridge mapper onvoldoende nauwkeurigheid behoudt, stellen de auteurs een drop-in vervanging voor met een kleine Multi-Layer Perceptron (MLP). Deze MLP herverdeelt de residuele fout weg van aandacht-gevoelige subspaces, waardoor de prestaties op moeilijke paren worden hersteld.

3. Belangrijkste Bijdragen

1. Closed-Form Mapping Framework: Een gradiënt-vrij, per-head ridge regressie framework dat cross-model KV transfer mogelijk maakt zonder training. Het steunt op top-$k$ bron-laag selectie en RoPE-gestripte content mapping.
2. Validatie over Meerdere Families: Validatie over drie modelfamilies (Qwen3, Llama 3.1, Ministral 3) en zes matched-KV paren.
3. Niet-lineaire Herstel: Demonstratie dat een lichtgewicht MLP de prestaties kan herstellen op "moeilijke" paren waar lineaire mapping faalt, specifiek door de fout te herverdelen weg van aandacht-kritieke dimensies.
4. Diagnostische Metriek: Identificatie dat attention-output cosine similarity een betere voorspeller is voor downstream transfer retentie dan standaard reconstructie metrieken ($R^2$). De auteurs tonen een Pearson correlatie van $r=+0.57$ aan tussen attention-output cosine en HellaSwag retentie, vergeleken met $r=-0.20$ voor $R^2$.

4. Experimentele Resultaten

Het framework werd geëvalueerd op zes matched-KV paren over drie families met behulp van vijf accuratiewaarden (ARC-C, HellaSwag, WinoGrande, MMLU, GSM8K) en WikiText-2 perplexiteit.

• Accuratie Retentie:
  • Tier 1 (Succes): Vier paren (Qwen3 14B→32B, 8B→32B; Llama 3.1 8B→70B; Ministral 3B→8B) behielden 73–98% van de standalone accuratie van het doelmodel.
  • Tier 2 (Falen): Twee paren (Ministral 3B→14B, 8B→14B) vertoonden significante degradatie (42–44% retentie). Echter, het vervangen van de lineaire mapper door een MLP herstelde de HellaSwag retentie met +24.3 tot +36.8 procentpunten, waardoor deze paren boven de 90% uitkwamen.
• Latentie:
  • De mapper draait 2.7–25× sneller dan het opnieuw prefillen van het doelmodel.
  • De versnellingen zijn het meest uitgesproken bij small-to-large transfers en langere contextlengtes (bijv. 25× versnelling bij 32k tokens voor Qwen3 14B→32B).
• Multi-Turn Stabiliteit:
  • In multi-turn handoff scenario's (CoQA) blijft de drift klein. Small-to-large drift neemt slechts 1.7 pp toe over 10 turns; large-to-small drift groeit lineair met 0.33 pp/turn, wat duidt op stabiliteit voor typische sessielengtes.
• Kalibratie Efficiëntie:
  • De mapper vereist slechts 500 sequenties voor kalibratie.
  • Het fitten duurt ~47–87 minuten per paar op een enkele 8×H100 node, zonder gradiënt-gebaseerde training.

5. Betekenis en Claims

Het paper claimt dat cross-model KV cache transfer binnen een LLM-familie een closed-form, training-vrije fit toelaat. De auteurs argumenteren dat de cross-model KV relatie grotendeels lineair is, wat de assumptie uitdaagt dat complexe neurale adapters nodig zijn voor dergelijke transfers.

De betekenis van dit werk ligt in:

1. Praktische bruikbaarheid: Het maakt kosten-kwaliteit cascading en mid-conversation switching mogelijk zonder de prohibitieve kosten van re-prefilling, wat dynamische model orchestratie haalbaar maakt in productie.
2. Verschuiving in Evaluatie: De bevindingen suggereren dat evaluatiecriteria voor cross-model mappers moeten verschuiven van ruwe reconstructie metrieken ($R^2$) naar subspace-bewuste diagnostiek (zoals attention-output cosine), aangezien deze laatste de downstream taakprestaties beter voorspelt.
3. Foutherverdeling: Het werk benadrukt dat transferfalen vaak komt door waar de fout landt (in aandacht-gevoelige subspaces) in plaats van de magnitude van de fout, wat motiveert voor toekomstige aandacht-gealigneerde doelstellingen.

De auteurs blijven bescheiden over beperkingen, waarbij zij opmerken dat de aanpak momenteel beperkt is tot binnen-familie, dense full-attention modellen met gematchte KV configuraties, en dat kalibratie wordt uitgevoerd op een enkele domein (FineWeb-Edu). Zij laten cross-familie transfer en mismatched-KV condities als openstaande vragen voor toekomstig werk.