TACT-KV: trasformata coseno triassiale per comprimere le cache KV volumetriche nei VLM medici

Abstract degli autori

I modelli vision-language (VLM) applicati alle scansioni mediche 3D si scontrano con un limite strutturale legato alla memoria. Un singolo volume di TC o RM genera decine di migliaia di token visivi e la cache key-value (KV) risultante cresce linearmente sia con il batch sia con il contesto. I metodi di compressione delle cache KV esistenti operano sulle statistiche delle caratteristiche o dei token della cache, oppure applicano trasformate 1D e 2D, senza sfruttare la struttura spaziale triassiale della cache. Proponiamo TACT-KV (trasformata coseno triassiale), un metodo di compressione delle cache KV che non richiede calibrazione e sfrutta le forti correlazioni spaziali degli stati KV visivi lungo tutti e tre gli assi anatomici delle scansioni 3D. Questa correlazione concentra l’energia della cache KV dopo una trasformazione nel dominio delle frequenze. TACT-KV applica alla cache KV una trasformata coseno discreta 3D separabile, distribuisce un budget di bit per ciascun livello tra i bin di frequenza tramite reverse water-filling (assegnando più bit ai componenti ad alta energia) e quantizza ciascun bin con un codebook scalare fisso. Deriviamo una relazione in forma chiusa: la riduzione della distorsione ottenuta con l’allocazione adattiva dei bit è determinata soprattutto dalla concentrazione spettrale della cache trasformata. In una valutazione su sei VLM medici e due benchmark di TC 3D, nel confronto principale TACT-KV ottiene il tasso più basso di inversione delle predizioni a 1 bit in undici delle dodici coppie modello-dataset. Lo stesso vantaggio si ripresenta su quattro VLM per uso generale valutati su attività video, il che indica che i benefici dello sfruttamento della struttura triassiale si estendono oltre l’imaging medico. A parità di memoria ad ampia banda (HBM), TACT-KV aumenta fino a $8\times$ la dimensione del batch di decodifica su una singola GPU.