YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference

1. Il Problema

L'inferenza dei Large Language Models (LLM) basati su Transformer è spesso limitata dalla memoria necessaria per la KV Cache (Key-Value Cache). Per gestire sequenze lunghe, la cache KV cresce linearmente, diventando un collo di bottiglia.
Le tecniche di compressione KV cross-layer (come YOCO) sono state proposte per ridurre questo consumo di memoria condividendo le chiavi e i valori (KVs) di un sottoinsieme di layer con gli altri. Tuttavia, questi metodi introducono un degrado delle prestazioni non trascurabile. La causa principale è che la condivisione diretta dei KVs da un layer all'altro indebolisce la capacità del modello, poiché i layer superiori perdono la loro rappresentazione specifica a favore di quella condivisa.

2. Metodologia: YOCO++

Gli autori propongono YOCO++, un'evoluzione del metodo YOCO ("You Only Cache Once") che mira a mantenere l'efficienza della compressione KV del 50% senza sacrificare le prestazioni.

• Meccanismo di Base (YOCO): In YOCO, i primi $L/2$ layer (bottom-half) sono "self-decoder" che generano la cache KV, mentre i successivi $L/2$ layer (top-half) sono "cross-decoder" che riutilizzano la cache del layer mediano senza generarne di nuove. Questo riduce la memoria e la latenza del 50%.
• Innovazione di YOCO++: Per mitigare il degrado delle prestazioni, YOCO++ introduce connessioni residue pesate tra i KVs.
  • Invece di condividere direttamente i KVs del layer mediano, YOCO++ combina i KVs di ogni layer della metà inferiore ($i$) con i KVs del layer di base (primo layer, indice 1).
  • La combinazione è definita come:
    $$K'_i = \alpha_{i,1}^K K_1 + \alpha_{i,2}^K K_i$$
    $$V'_i = \alpha_{i,1}^V V_1 + \alpha_{i,2}^V V_i$$
    dove $\alpha$ sono pesi residui apprendibili.
  • I layer superiori riutilizzano quindi le KVs combinate e arricchite del layer mediano ($K'_{L/2}, V'_{L/2}$), che contengono più informazioni rispetto alle KVs originali.
• Ottimizzazione dell'Efficienza:
  • Caching: Per evitare overhead di I/O durante il decoding, il modello memorizza in cache le KVs combinate ($K', V'$) invece di quelle originali.
  • Fattore di Scaling ($\lambda$): Per facilitare l'apprendimento dei pesi residui (che tendono a rimanere vicini ai valori iniziali), viene introdotto un fattore di scaling $\lambda > 1$ (sperimentalmente $\lambda=35$). Questo permette ai pesi di divergere più facilmente dai valori iniziali durante l'addestramento.
  • Inizializzazione: I pesi sono inizializzati in modo che il modello iniziale sia identico a YOCO ($\alpha_{i,1}=0, \alpha_{i,2}=1/\lambda$).

3. Contributi Chiave

1. Miglioramento delle Prestazioni senza Costo Computazionale: YOCO++ aumenta la capacità del modello integrando informazioni dal layer di base in tutti i layer inferiori, superando le prestazioni dello standard Transformer, pur mantenendo la stessa efficienza di inferenza di YOCO.
2. Architettura Efficiente: A differenza di metodi precedenti come FusedKV che introducono overhead di I/O aggiuntivo, YOCO++ cachea le KVs combinate, garantendo che l'overhead di I/O durante il decoding rimanga invariato rispetto a YOCO.
3. Scoperta sull'Asimmetria K/V: L'analisi dei pesi appresi mostra che i pesi per i valori ($V$) tendono a essere più significativi rispetto a quelli per le chiavi ($K$), confermando l'asimmetria tra Key e Value osservata in lavori precedenti.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su un modello da 1.1B parametri (TinyLlama) con 22 layer, addestrato su un subset di 100B token di SlimPajama.

• Efficienza di Inferenza:
  • Prefilling Latency: Ridotta del 50% rispetto al Transformer standard (come YOCO), poiché i layer superiori vengono saltati.
  • Throughput di Decoding: YOCO++ mantiene lo stesso throughput elevato di YOCO, superando il Transformer standard grazie alla minore occupazione di memoria che permette batch size più grandi. FusedKV, al contrario, vede un calo di throughput a causa dell'overhead I/O.
• Prestazioni del Modello:
  • Su task di ragionamento commonsense (HellaSwag, OBQA, WinoGrande, ARC, PIQA), YOCO++ ottiene la migliore performance tra tutti i metodi di compressione KV cross-layer (superando YOCO, FusedKV e FusedKV-Lite).
  • Sorprendentemente, YOCO++ supera anche il Transformer standard a una velocità di compressione KV del 50%, ottenendo un punteggio medio di accuratezza zero-shot di 48.99 contro 48.37 del Transformer.
• Studi Ablativi:
  • Senza il fattore di scaling $\lambda$, le prestazioni crollano a livelli simili a YOCO, dimostrando che i pesi non vengono appresi efficacemente.
  • Rimuovere le connessioni residue per le chiavi ($K$) degrada le prestazioni, confermando che sia $K$ che $V$ devono essere combinati.

5. Significato

YOCO++ rappresenta un avanzamento significativo nel campo dell'inferenza efficiente degli LLM. Dimostra che è possibile ridurre drasticamente l'uso della memoria KV (50%) senza compromettere, e anzi migliorando, la qualità del modello.
La proposta di utilizzare connessioni residue pesate tra i KVs, unita a un'attenta gestione dell'overhead di I/O, offre una soluzione pratica per l'implementazione di LLM su hardware con risorse limitate, rendendo l'inferenza più veloce e accessibile senza sacrificare l'intelligenza del modello. Il codice è stato reso disponibile pubblicamente, favorendo la riproducibilità e l'adozione di queste tecniche.