Technische Samenvatting: HeatKV - Head-tuned KV-cache Compressie voor Visuele Autoregressieve Modellering

Probleemstelling

Visuele Autoregressieve (VAR) modellen zijn een concurrerend alternatief geworden voor diffusiemodellen bij het genereren van afbeeldingen. Ze bieden snelle inferentiesnelheden en gunstige schaalbaarheidseigenschappen door afbeeldingen te genereren als een reeks multi-schaal tokenkaarten. Deze modellen kampen echter met ernstige geheugenbeperkingen vanwege de aard van hun KV-cache. In tegenstelling tot Large Language Models (LLM's), die doorgaans één token per stap genereren, genereren VAR-modellen volledige tokenkaarten op toenemende resoluties. Bijgevolg groeit het aantal tokens in de KV-cache superlineair met de schaal van de generatie.

De meest geavanceerde VAR-modellen, zoals Infinity-2B en Infinity-8B, kunnen tot 10 GB aan KV-cache-geheugen vereisen om één enkele afbeelding van 1024 × 1024 te genereren. Dit hoge geheugengebruik beperkt de maximale batchgrootte en het aantal gelijktijdige gebruikers dat een dienstverlener kan ondersteunen op gespecialiseerde hardware-versnellers. Bestaande compressiemethoden voor VAR's, zoals ScaleKV en HACK, vertrouwen vaak op grofkorrelige strategieën, zoals het classificeren van attention-heads of lagen in binaire categorieën (bijvoorbeeld "hoog-budget" versus "laag-budget"). De auteurs betogen dat deze binaire classificaties te beperkend zijn, aangezien attentionpatronen binnen dezelfde laag vaak aanzienlijk variëren tussen verschillende heads, wat leidt tot verspilling van geheugen of suboptimale pruning.

Methodologie: HeatKV

De auteurs stellen HeatKV (Head-tuned KV-cache Compressie) voor, een nieuwe strategie die cache-toewijzing aanpast op de granulariteit van individuele attention-heads en specifieke afbeeldingsschalen. De methode bestaat uit drie hoofdbestanddelen:

1. Head-Schaal Belang (Offline Kalibratie)

In tegenstelling tot eerdere methoden die gebruikmaken van variantie-gevoelige metrieken, rangschikt HeatKV attention-heads op basis van hun Schaal-afhankelijke Cumulatieve Attention Score (S-CAS).

• Mechanisme: Met behulp van een kleine offline kalibratiedataset aggregeert de methode attention-kansen van tokens op toekomstige schalen terug naar tokens op eerdere schalen.
• Metriek: Voor een gegeven laag $\ell$, head $h$, en bronschaal $k$, meet de S-CAS de gemiddelde attention-massa van alle toekomstige schalen ($\tau > k$) naar de tokens op schaal $k$.
• Resultaat: Dit genereert een specifieke rangschikking van "head-schaal"-paren voor elke bronschaal, wat aangeeft hoe kritiek het cache van een specifieke head op een specifieke eerdere schaal is voor toekomstige tokengeneratie. Dit stelt het systeem in staat vast te stellen dat een specifieke head Schaal 1 moet behouden maar Schaal 2 kan verwijderen, een niveau van granulariteit dat binaire methoden niet kunnen bereiken.

2. Budgettoewijzing

HeatKV berekent het aantal heads ($N_k$) dat moet worden verwijderd na het genereren van elke schaal $k$ om ervoor te zorgen dat de totale cachegrootte binnen een door de gebruiker gedefinieerd geheugenbudget $b$ blijft.

• Beperking: De methode houdt zich aan het "Streaming-LLM"-principe van het behouden van de eerste $s$ schalen als een "attention sink" (essentieel voor stabiliteit) en cacheert nooit de huidige schaal die wordt gegenereerd.
• Berekening: Het algoritme bepaalt het minimale aantal heads dat moet worden verwijderd om aan de budgetbeperking te voldoen, waardoor een dynamische verkleining van de cache mogelijk is naarmate de generatie vordert.

3. Optimalisatie van Verwijderingsvolgorde (Gierige Vroege Pruning)

Een kritieke uitdaging bij VAR's is dat de cachegrootte superlineair groeit binnen één enkele generatiestap (over transformerlagen heen). Standaardmethoden die alleen tussen schalen pruning toepassen, kunnen leiden tot tijdelijke geheugenoverschrijdingen die de hardwarelimieten overschrijden.

• Oplossing: HeatKV introduceert een Gierige Vroege Pruning-algoritme (Algoritme 1). Het dwingt de budgetbeperking af na elke transformerlaag, niet alleen tussen schalen.
• Strategie: Als de cachegrootte na een specifieke laag het budget overschrijdt, selecteert het algoritme gierig extra head-schalen om te verwijderen uit die laag of eerdere lagen. Het prioriteert het verwijderen van heads met lagere S-CAS-rangschikkingen en, in geval van gelijkstand, prioriteert het het verwijderen uit latere lagen of latere bronschalen om geheugenontlasting te maximaliseren.
• Optimaliteit: De auteurs bewijzen (Propositie 1) dat deze gierige aanpak een verzameling met minimale kardinaliteit van vroeg verwijderde heads oplevert, waardoor wordt gegarandeerd dat het budget nooit wordt geschonden terwijl het beschikbare geheugen optimaal wordt benut.

Belangrijkste Bijdragen

1. Gedetailleerde Observatie: De auteurs tonen aan dat attentionpatronen in VAR-modellen aanzienlijk variëren tussen verschillende schalen binnen dezelfde attention-head, maar relatief stabiel blijven over verschillende input-prompten. Dit motiveert een verschuiving van binaire head-classificatie naar een continue, schaal-afhankelijke budgettoewijzing.
2. S-CAS Metriek: De introductie van de Schaal-afhankelijke Cumulatieve Attention Score maakt het mogelijk om een statisch pruning-schema te construeren dat is toegespitst op een specifiek geheugenbudget, waardoor elke head niet-essentiële schalen kan verwijderen terwijl kritieke schalen behouden blijven.
3. Geoptimaliseerd Pruning-algoritme: De ontwikkeling van een gierig evictie-algoritme dat budgetbeperkingen af dwingt op laagniveau (in plaats van alleen op schaalniveau) voorkomt tijdelijke geheugenoverlopen, wat cruciaal is voor het superlineaire groeikarakteristiek van VAR's.
4. State-of-the-Art Prestaties: De methode bereikt een nieuwe state-of-the-art voor KV-cache-compressie van VAR-modellen en toont aan dat gedetailleerde toewijzing aanzienlijk beter presteert dan grofkorrelige benaderingen.

Experimentele Resultaten

De auteurs evalueerden HeatKV op de Infinity-2B en Infinity-8B modellen met behulp van de MS-COCO 2017 validatieset. Ze vergeleken HeatKV met baselines zoals StreamingLLM, SnapKV, HACK en ScaleKV over geheugenbudgetten van 4%, 10% en 20%.

• Fideliteit: HeatKV presteerde consequent beter dan alle baselines. Opmerkelijk is dat op het Infinity-2B-model HeatKV met een 10% budget betere FID-, LPIPS- en PSNR-scores behaalde dan ScaleKV met een 20% budget. Bij een budget van 4% behield HeatKV een FID van 2,70 (Infinity-2B), aanzienlijk lager dan de volgende beste methode (ScaleKV met 4,92).
• Prompt-uitlijning: Met behulp van GenEval behield HeatKV een sterke prompt-uitlijning, zelfs bij lage budgetten, met scores die vergelijkbaar waren met de volledige-cache-baseline.
• Perceptuele Kwaliteit: Resultaten van de Human Preference Score (HPSv2.1) gaven aan dat HeatKV hoge perceptuele kwaliteit behield en andere methoden versloeg bij een budget van 4% op Infinity-2B.
• Efficiëntie: HeatKV bereikte een 2× hogere compressieverhouding in geheugentoewijzing in vergelijking met bestaande methoden, terwijl het vergelijkbare of betere afbeeldingsfideliteit behield. De methode maakte ook grotere batchgroottes mogelijk, wat effectief de generatielatentie verlaagde in vergelijking met inferentie met volledige cache.

Betekenis en Claims

Het artikel beweert dat HeatKV een aanzienlijke vooruitgang betekent in het mogelijk maken van VAR-modellen voor implementatie op hardware met beperkte geheugencapaciteit, zoals consumentengpu's, of voor het bedienen van een groter aantal gelijktijdige gebruikers op bestaande infrastructuur.

De auteurs benadrukken dat hun methode geen wijzigingen vereist in de modelarchitectuur of het trainingsproces; het vertrouwt op een kleine offline kalibratiedataset (zoveel als 1–10 prompten) om een statisch pruning-schema te genereren. De resultaten suggereren dat gedetailleerde, head-specifieke cache-toewijzing essentieel is voor het behoud van de semantische structuur en de generatiekwaliteit van VAR-modellen onder agressieve geheugenbeperkingen, en dat het methoden verslaat die vertrouwen op grove binaire classificaties van attention-heads of lagen.

Het artikel concludeert met de opmerking dat, hoewel HeatKV het geheugen optimaliseert, toekomstig werk zou kunnen onderzoeken om dit te combineren met andere optimalisatietechnieken, zoals het benutten van attention-sparseheid of modelquantisatie, om de efficiëntie verder te verbeteren.