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

Probleemstelling

De inferentie van grote taalmodellen (LLM's) wordt vaak beperkt door het geheugengebruik van de Key-Value (KV) cache, een noodzakelijk onderdeel van de Transformer-architectuur. Bestaande methoden voor cross-layer KV-compressie (zoals YOCO) lossen dit op door de KV-cache van een subset van lagen te hergebruiken in andere lagen, wat het geheugengebruik met 50% verlaagt.

Echter, deze methoden introduceren een niet-verwaarloosbare prestatiedaling. De directe deling van KV-waarden tussen lagen verzwakt de modelcapaciteit omdat informatie uit specifieke lagen verloren gaat of niet optimaal wordt verwerkt. Bestaande oplossingen zoals FusedKV proberen dit op te lossen door gewogen combinaties te gebruiken, maar dit introduceert extra I/O-overhead tijdens het decoderen, wat de inferentie-efficiëntie negatief beïnvloedt.

Methodologie: YOCO++

De auteurs stellen YOCO++ voor, een verbeterde versie van YOCO die de prestaties verbetert zonder de efficiëntie te verliezen. De kern van de methode is het introduceren van gewichtige residuale verbindingen tussen de KV-cache van lagen.

1. Architectuur:

   • Het model bestaat uit een onderste helft (self-decoder) en een bovenste helft (cross-decoder).
   • In de onderste helft (lagen $i \in \{2, ..., L/2\}$) worden de KV-cache's niet alleen lokaal berekend, maar gecombineerd met de KV-cache van de bodemlaag (laag 1).
   • De gecombineerde KV's worden als volgt berekend:
     $$K'_i = \lambda(\alpha^K_{i,1}K_1 + \alpha^K_{i,2}K_i)$$
     $$V'_i = \lambda(\alpha^V_{i,1}V_1 + \alpha^V_{i,2}V_i)$$
   • Hierbij zijn $\alpha$ leerbare gewichten en $\lambda$ een schalingsfactor.

2. Schalingsfactor ($\lambda$):

   • Om het leren van de residuale gewichten te vergemakkelijken, wordt een vaste hyperparameter $\lambda > 1$ gebruikt (in de experiments ingesteld op 35).
   • De gewichten worden geïnitieerd zodat het model aanvankelijk identiek is aan het originele YOCO ($\alpha_{i,1}=0, \alpha_{i,2}=1/\lambda$), maar tijdens het training proces kunnen ze aanpassen.

3. Caching-strategie:

   • In plaats van de originele KV's te cacheën, cacheert YOCO++ de gecombineerde KV's ($K'_i, V'_i$).
   • Dit voorkomt extra I/O-overhead tijdens het decoderen, omdat de lagen in de bovenste helft direct de reeds gecombineerde informatie kunnen ophalen, in plaats van meerdere bronlagen te hoeven lezen.

Belangrijkste Bijdragen

• Residuale KV-verbindingen: Het introduceren van een mechanisme waarbij de KV-cache van de bodemlaag als een residu wordt toegevoegd aan de KV-cache van elke laag in de onderste helft van het model.
• Efficiëntiebehoud: In tegenstelling tot eerdere methoden die gewogen combinaties toepassen (zoals FusedKV), cacheert YOCO++ de gecombineerde resultaten. Dit behoudt dezelfde I/O-efficiëntie als YOCO, terwijl het meer modelcapaciteit biedt.
• Schalingsmechanisme: De introductie van de schalingsfactor $\lambda$ om het trainen van de residuale gewichten te stabiliseren en te verbeteren.
• State-of-the-art Prestaties: Het bereiken van superieure prestaties ten opzichte van standaard Transformers en andere compressiemethoden bij een compressiegraad van 50%.

Resultaten

De auteurs hebben experimenten uitgevoerd met een 1.1B parameter model (TinyLlama configuratie) op de SlimPajama dataset.

• Inferentie-efficiëntie:
  • Prefilling-latentie: YOCO++ is ongeveer 2x sneller dan een standaard Transformer (door vroege exit van de bovenste lagen), vergelijkbaar met YOCO.
  • Decoderingsthruput: YOCO++ behoudt een hoge doorvoer, vergelijkbaar met YOCO en FusedKV-Lite, en significant hoger dan de standaard Transformer. FusedKV (zonder Lite) zakt in door extra I/O-overhead, wat YOCO++ vermijdt.
• Modelprestaties (Zero-shot):
  • Op diverse redeneertaken (HellaSwag, OBQA, WinoGrande, ARC, PIQA) behaalt YOCO++ de hoogste gemiddelde nauwkeurigheid (48.99%).
  • Het presteert beter dan de standaard Transformer (48.37%) en alle andere compressiemethoden (YOCO, FusedKV, FusedKV-Lite).
• Ablatiestudies:
  • Zonder de schalingsfactor $\lambda$ presteert het model net zo slecht als YOCO, wat aantoont dat de schalingsfactor essentieel is voor het leren van de gewichten.
  • Het verwijderen van residuale verbindingen voor de 'keys' leidt tot prestatieverlies, wat aangeeft dat zowel keys als values gecombineerd moeten worden.

Betekenis

YOCO++ biedt een belangrijke doorbraak in de efficiënte inferentie van LLM's. Het lost het fundamentele compromis op tussen geheugenefficiëntie en modelkwaliteit dat bij cross-layer compressie vaak aanwezig is. Door slimme residuale verbindingen te gebruiken zonder extra I/O-kosten, stelt YOCO++ onderzoekers en ingenieurs in staat om LLM's met 50% minder KV-cache-geheugen in te zetten, terwijl ze zelfs beter presteren dan de oorspronkelijke, niet-gecompressieerde Transformer. Dit maakt het een zeer aantrekkelijke oplossing voor productiesystemen met beperkte hardwarebronnen.