Quantizzazione delle attivazioni per gruppi di token compatibile con l’hardware per la super-risoluzione basata su Mamba a pochi bit

Abstract degli autori

I modelli di super-risoluzione basati su Mamba sono alternative efficienti alle architetture di base per il ripristino basate su Transformer, ma la loro quantizzazione post-addestramento a pochi bit resta difficile: la scansione selettiva, l’operatore a spazio di stato di Mamba che dipende dall’input, produce intervalli di attivazione diversi tra i token spaziali. Un unico intervallo di clipping per livello si adatta quindi male a tutti i token. Proponiamo Token-Group Quantization, un sistema di quantizzazione post-addestramento delle attivazioni compatibile con l’hardware che approssima staticamente un oracolo di raggruppamento successivo alla scansione poco adatto all’hardware: i token che richiedono intervalli di clipping simili vengono assegnati allo stesso gruppo. Per evitare di raggruppare i token dopo la scansione durante l’inferenza statica a pochi bit, un predittore leggero deduce i gruppi dalle caratteristiche disponibili prima del percorso di scansione quantizzato. Durante la calibrazione, creiamo pseudoetichette raggruppando i descrittori delle attivazioni successive alla scansione che presentano comportamenti simili di arrotondamento e clipping. Nell’inferenza, il predittore, i cui parametri sono fissati, associa le caratteristiche precedenti alla scansione ai gruppi e seleziona le voci dalle tabelle dei limiti di clipping per gruppo calibrate offline, senza raccogliere statistiche durante l’esecuzione, aggiornare dinamicamente i limiti di clipping o assegnare i gruppi dopo la scansione. Per contrastare il degrado del ripristino a pochi bit, introduciamo inoltre una funzione obiettivo Frequency-Preserving Refinement (FPR) per calibrare i limiti di clipping. FPR allinea l’ampiezza e la fase di Fourier degli output del modello insegnante a piena precisione e del modello studente quantizzato per preservare le strutture ad alta frequenza. In diversi protocolli di PTQ, numeri di bit, architetture di base e compiti, il nostro metodo supera sistematicamente i precedenti metodi di riferimento per la PTQ di Mamba e della super-risoluzione, si trasferisce alla super-risoluzione basata su Transformer e alla classificazione con Vision Mamba, si combina con la quantizzazione basata su rotazioni e ottiene un’accelerazione di 2,23 volte su Jetson Orin Nano rispetto a FP16. Il codice è disponibile qui.