Riepilogo Tecnico: Trasferimento della KV Cache tra Modelli della stessa Famiglia di LLM

1. Formulazione del Problema

Il serving di Large Language Models (LLM) in produzione si affida sempre più all'orchestrazione multi-modello per il cascading costo-qualità, lo switching a metà conversazione e il routing. Questi workflow scambiano frequentemente membri di dimensioni diverse della stessa famiglia di modelli (ad esempio, da Qwen3 14B a 32B). Sebbene questi modelli condividano scelte architettoniche fondamentali e dati di pretraining, differiscono per scala, numero di layer e ricette di training.

Il collo di bottiglia principale in tali workflow è il costo di prefill. Ogni volta che un contesto di conversazione viene passato da un modello sorgente a un modello target, il target deve rielaborare l'intero contesto accumulato per popolare la propria KV cache. Questo passaggio forward scala sia con la dimensione del modello che con la lunghezza del prompt, annullando i vantaggi di efficienza dello swapping dei modelli. Le soluzioni esistenti come il prefix caching funzionano solo all'interno di un singolo modello, e i precedenti approcci cross-modello (ad esempio, neural fusers, latent adapters) richiedono tipicamente training basati su gradienti o vincoli architettonici stretti (ad esempio, conteggio dei layer identico).

Questo articolo propone il Trasferimento della KV Cache Cross-Model, un metodo per riutilizzare la KV cache pre-compilata del modello sorgente direttamente nel modello target, saltando interamente il prefill del target. L'obiettivo è trovare una mappatura $f: C_S \to \hat{C}_T$ tale che il decoding del modello target partendo dalla cache mappata $\hat{C}_T$ produca output equivalenti al decoding dalla propria cache ground-truth $C_T$.

2. Metodologia

Gli autori ipotizzano che le relazioni KV tra modelli diversi possiedano una sostanziale struttura lineare, consentendo una mappatura a forma chiusa e priva di gradienti. Il framework proposto opera su coppie "matched-KV", dove il modello sorgente e il target condividono lo stesso numero di teste KV e la stessa dimensione per testa, anche se il loro numero totale di layer o le dimensioni dei parametri differiscono.

Il mapper è progettato con tre componenti core:

A. Regressione Ridge Per-Testa

Invece di una rete neurale globale, il metodo adatta una mappatura lineare indipendente per ogni layer e testa del target.

• Input: Caratteristiche KV concatenate da un sottoinsieme selezionato di layer sorgente.
• Output: I vettori Key ($K$) e Value ($V$) attesi del target.
• Ottimizzazione: I pesi sono risolti tramite regressione ridge (soluzione a forma chiusa) su un piccolo set di calibrazione (500 sequenze di 1.024 token da FineWeb-Edu). Viene utilizzata la regolarizzazione ridge ($\lambda=0.01$) per gestire la natura quasi singolare della matrice delle caratteristiche quando i layer sorgente sono correlati.

B. Selezione Cross-Layer della Sorgente

Poiché i modelli sorgente e target hanno profondità diverse ($L_S \neq L_T$), il metodo deve determinare quali layer sorgente informano ciascun layer del target.

• Strategia di Selezione: Per ogni layer del target, i top-$k$ layer sorgente vengono selezionati in base al loro potere predittivo (misurato tramite $R^2$ mediato sulle teste su un set di calibrazione).
• Meccanismo: Le caratteristiche KV di questi top-$k$ layer sorgente vengono concatenate per formare l'input per la regressione ridge. L'analisi empirica mostra che l'informazione complementare è distribuita su più layer sorgente; un singolo layer è insufficiente, ma le prestazioni saturano rapidamente (ad esempio, $k=6$ cattura la maggior parte della varianza di $k=tutti$).

C. Mappatura dello Spazio del Contenuto (Fattorizzazione RoPE)

Gli LLM standard utilizzano l'Embedding di Posizione Rotatoria (RoPE), che applica rotazioni dipendenti dalla posizione alle chiavi.

• Problema: Adattare un mapper direttamente su chiavi ruotate lega i pesi alla specifica distribuzione di posizione vista durante la calibrazione, limitando la trasferibilità a diverse lunghezze di contesto.
• Soluzione: Il metodo rimuove il RoPE dalle chiavi sorgente prima della mappatura e riapplica il RoPE del target dopo la mappatura.
  • $K_{stripped} = K_{source} \cdot R^{-1}_{\Theta_s}$
  • $\hat{K}_{target} = (K_{stripped} \cdot W + b) \cdot R_{\Theta_t}$
• Questo decoupling separa la rotazione posizionale dal contenuto semantico, rendendo il fit lineare indipendente dalla posizione e riutilizzabile su diverse lunghezze di contesto (fino a 32k token).

D. Estensione Non Lineare (MLP)

Per le coppie in cui il mapper lineare ridge non riesce a mantenere un'accuratezza sufficiente, gli autori propongono un sostituto "drop-in" utilizzando un piccolo Multi-Layer Perceptron (MLP). Questo MLP ridistribuisce l'errore residuo lontano dai sottospazi sensibili all'attenzione, recuperando le prestazioni sulle coppie difficili.

3. Contributi Chiave

1. Framework di Mappatura a Forma Chiusa: Un framework di regressione ridge per-testa, privo di gradienti, che consente il trasferimento della KV cache cross-modello senza training. Si basa sulla selezione di top-$k$ layer sorgente e sulla mappatura del contenuto con rimozione del RoPE.
2. Validazione Multi-Famiglia: Validazione su tre famiglie di modelli (Qwen3, Llama 3.1, Ministral 3) e sei coppie matched-KV.
3. Recupero Non Lineare: Dimostrazione che un leggero MLP può recuperare le prestazioni sulle coppie "difficili" dove la mappatura lineare fallisce, specificamente ridistribuendo l'errore lontano dalle dimensioni critiche per l'attenzione.
4. Metrica Diagnostica: Identificazione del fatto che la similarità del coseno dell'output dell'attenzione è un predittore migliore della ritenzione del trasferimento rispetto alle metriche di ricostruzione standard ($R^2$). Gli autori mostrano una correlazione di Pearson $r=+0.57$ tra la similarità del coseno dell'output dell'attenzione e la ritenzione di HellaSwag, rispetto a $r=-0.20$ per $R^2$.

4. Risultati Sperimentali

Il framework è stato valutato su sei coppie matched-KV attraverso tre famiglie utilizzando cinque benchmark di accuratezza (ARC-C, HellaSwag, WinoGrande, MMLU, GSM8K) e la perplessità di WikiText-2.

• Ritenzione dell'Accuratezza:
  • Tier 1 (Successo): Quattro coppie (Qwen3 14B→32B, 8B→32B; Llama 3.1 8B→70B; Ministral 3B→8B) hanno mantenuto il 73–98% dell'accuratezza standalone del modello target.
  • Tier 2 (Fallimento): Due coppie (Ministral 3B→14B, 8B→14B) hanno subito un degrado significativo (ritenzione del 42–44%). Tuttavia, sostituendo il mapper lineare con un MLP, la ritenzione di HellaSwag è stata recuperata di +24.3 e +36.8 punti percentuali, portando queste coppie sopra il 90%.
• Latenza:
  • Il mapper gira da 2.7 a 25 volte più velocemente rispetto al re-prefill del modello target.
  • Gli accelerazioni sono più pronunciate per i trasferimenti da piccoli a grandi modelli e per lunghezze di contesto maggiori (ad esempio, 25x di velocità per 32k token per Qwen3 14B→32B).
• Stabilità Multi-Turn:
  • Negli scenari di handoff multi-turno (CoQA), il drift rimane contenuto. Il drift da piccolo a grande aumenta di soli 1.7 pp su 10 turni; il drift da grande a piccolo cresce linearmente a 0.33 pp/turno, suggerendo stabilità per le tipiche sessioni.
• Efficienza di Calibrazione:
  • Il mapper richiede solo 500 sequenze per la calibrazione.
  • Il fitting richiede circa 47–87 minuti per coppia su un singolo nodo 8×H100, senza training basato su gradienti.

5. Significato e Rivendicazioni

L'articolo sostiene che il trasferimento della KV cache cross-modello all'interno di una famiglia di LLM ammette un fit a forma chiusa e senza training. Gli autori sostengono che la relazione KV cross-modello sia largamente lineare, sfidando l'assunto che siano necessari complessi adapter neurali per tali trasferimenti.

La significatività di questo lavoro risiede nella sua:

1. Praticità: Permette il cascading costo-qualità e lo switching a metà conversazione senza il costo proibitivo del re-prefill, rendendo fattibile l'orchestrazione dinamica dei modelli in produzione.
2. Spostamento della Valutazione: I risultati suggeriscono che i criteri di valutazione per i mapper cross-modello dovrebbero spostarsi dalle metriche di ricostruzione grezza ($R^2$) a diagnostici consapevoli del sottospazio (come la similarità del coseno dell'output dell'attenzione), poiché questi ultimi predicono meglio le prestazioni sui task a valle.
3. Ridistribuzione dell'Errore: Il lavoro evidenzia che il fallimento del trasferimento è spesso dovuto a dove l'errore finisce (nei sottospazi sensibili all'attenzione) piuttosto che all'entità dell'errore, motivando futuri obiettivi allineati all'attenzione.

Gli autori rimangono modesti riguardo ai limiti, notando che l'approccio è attualmente limitato ai modelli a densità completa con attenzione completa e configurazioni KV compatibili all'interno della stessa famiglia, e che la calibrazione è eseguita su un singolo dominio (FineWeb-Edu). Lasciano il trasferimento cross-famiglia e le condizioni con KV mismatch come questioni aperte per ricerche future.