Riepilogo Tecnico: Compressione dei Pesi con Compensazione del Movimento (MCWC)

Enunciato del Problema

Il deployment delle reti neurali è sempre più limitato non dalla potenza di calcolo, ma dallo storage, dalla larghezza di banda e dall'overhead di avvio a freddo necessari per caricare i checkpoint negli acceleratori. Le pipeline di compressione esistenti, dominate dalla quantizzazione e dal pruning post-allenamento, trattano tipicamente ogni strato come un oggetto indipendente. Questo approccio trascura la sostanziale ridondanza inter-strato indotta da simmetrie che preservano la funzione (ad esempio, la capacità di re-indicizzare unità nascoste o testine di attenzione senza modificare la funzione input-output del modello). Di conseguenza, questi metodi non riescono a sfruttare la forte similarità e prevedibilità inter-strato intrinseca nelle architetture moderne come i Transformer.

Metodologia

MCWC propone un codec basato solo sui pesi che tratta la profondità del modello come un segnale sequenziale, analogo ai fotogrammi video. La premessa fondamentale è che, rimuovendo il "movimento" indotto dalle simmetrie di permutazione, l'evoluzione residua dei pesi attraverso gli strati diventa altamente prevedibile, consentendo una codifica efficiente dei residui. La pipeline consta di quattro fasi principali:

1. Compensazione del Movimento tramite Allineamento Funzionale

L'encoder identifica permutazioni che preservano la funzione per blocchi di pesi (ad esempio, unità nascoste FFN, testine di attenzione) che massimizzano la corrispondenza inter-strato.

• Meccanismo: Per ogni strato $\ell$, il metodo seleziona una permutazione $\Pi_\ell$ che allinea i blocchi con i loro corrispettivi nello strato precedente $\ell-1$. Ciò viene ottenuto risolvendo un problema di assegnazione lineare per massimizzare la similarità (un ibrido tra similarità coseno basata sui pesi e similarità basata sulle statistiche di attivazione) tra i blocchi candidati e il riferimento precedentemente allineato.
• Risultato: Questo trasforma la sequenza dei pesi in un "sistema di coordinate allineato" dove gli strati adiacenti sono geometricamente confrontabili, riducendo significativamente l'energia delle differenze tra di essi.

2. Codifica Predittiva Sequenziale per Strati

Una volta allineati, i pesi sono trattati come una sequenza temporale.

• Predizione: Un predittore leggero e sequenziale per strati ($g_\theta$) stima i blocchi allineati dello strato corrente basandosi sul contesto precedentemente decodificato (blocchi allineati dallo strato $\ell-1$), condizionato da embedding dello strato e del tipo di blocco.
• Residui: Il codec codifica solo i residui di predizione quantizzati ($R = \text{TargetAllineato} - \text{Predizione}$).
• Fotogrammi Chiave: Per prevenire l'accumulo di errori (deriva) su sequenze lunghe, periodici "fotogrammi chiave" sono codificati come valori assoluti anziché come residui. Questi agiscono come punti di sincronizzazione, permettendo al decoder di resettare il proprio contesto.

3. Quantizzazione Appresa e Modellazione dell'Entropia

• Quantizzazione: I residui sono mappati su indici di codice discreti utilizzando un quantizzatore scalare appreso con passi e medie specifici per canale.
• Codifica dell'Entropia: Un modello di entropia appreso assegna probabilità ai simboli quantizzati, abilitando la codifica aritmetica. Il modello è condizionato dall'indice dello strato, dal tipo di blocco e dalle statistiche locali derivate dall'output del predittore.
• Addestramento: L'intero sistema (predittore, quantizzatore, modello di entropia) è addestrato end-to-end sotto un obiettivo di distorsione-tasso ($L = L_{task} + \lambda E[-\log p_\psi(c)]$), dove $L_{task}$ è una perdita di distillazione che garantisce che i pesi ricostruiti mantengano le prestazioni del compito.

4. Informazioni Laterali e Decodifica

Il flusso di bit include i residui/fotogrammi chiave codificati in entropia, le informazioni laterali sulle permutazioni (codificate tramite codici di Lehmer codificati in delta per ridurre l'overhead) e i parametri del quantizzatore. Il decoder inverte il processo: decodifica in entropia, dequantizza, ricostruisce i residui tramite predizione e applica le permutazioni inverse per ripristinare la struttura originale del tensore per il deployment.

Contributi Chiave

1. Compensazione del Movimento nello Spazio dei Pesi: L'introduzione della selezione di permutazioni a blocchi per aumentare la corrispondenza inter-strato mantenendo la funzione rappresentata, analogo alla stima del movimento nella compressione video.
2. Codec Predittivo per la Profondità: Un'architettura innovativa che tratta la profondità come un asse temporale, utilizzando fotogrammi chiave espliciti e codifica dei residui per trasformare gli strati allineati in "fotogrammi" prevedibili.
3. Addestramento End-to-End di Distorsione-Tasso: L'addestramento congiunto di modelli di quantizzazione ed entropia per garantire che il flusso di residui discreti sia sia accurato dopo la decodifica sia a bassa entropia sotto il modello di probabilità appreso.
4. Validazione Empirica: Dimostrazione di fronti di Pareto migliorati tra tasso e accuratezza sia su modelli linguistici Transformer (Pythia) che su classificatori visivi (ViT), con studi di ablazione che confermano che allineamento, predizione e modellazione dell'entropia sono tutti necessari per ottenere i guadagni completi.

Risultati

MCWC è stato valutato su Pythia-1.4B (modellazione linguistica) e ViT-B/16 (classificazione visiva) confrontandolo con basi solide tra cui quantizzazione uniforme INT8/INT4, SmoothQuant, GPTQ, AWQ e quantizzazione scalare appresa.

• Modellazione Linguistica: A tassi di compressione aggressivi (ad esempio, 2.8 bit/parametro), MCWC raggiunge una perplessità di 16.40, superando Prune+Quant (18.10) e la quantizzazione scalare appresa (17.30). Nel regime di tasso medio (4.2 bit/parametro), raggiunge una perplessità di 15.95, superando l'INT4 uniforme (16.85) e rimanendo competitivo con metodi consapevoli dell'Hessiano come GPTQ (16.40) e AWQ (16.25).
• Classificazione Visiva: A 2.8 bit/parametro, MCWC mantiene un'accuratezza top-1 del 80.15% su ImageNet-1k, superando Prune+Quant (77.60%) e la quantizzazione scalare appresa (78.70%). A 4.2 bit/parametro, raggiunge un'accuratezza del 80.92%, superando l'INT4 uniforme (79.35%).
• Efficienza di Decodifica: I tempi di decodifica rimangono competitivi rispetto alle basi di quantizzazione standard (ad esempio, ~2.5s per Pythia-1.4B su CPU), poiché la ricostruzione predittiva e la permutazione inversa sono operazioni lineari nel tempo, con il costo dominante che è la decodifica in entropia e la materializzazione del tensore.
• Diagnostica: Studi di ablazione confermano che l'allineamento funzionale aumenta la similarità coseno tra strati adiacenti (da ~0.28 a ~0.67 per Pythia-1.4B) e riduce l'energia normalizzata del residuo, validando che l'allineamento espone un sistema di coordinate più prevedibile.

Significato e Affermazioni

Il documento posiziona MCWC non come un sostituto per i kernel di inferenza a bassa precisione (ad esempio, moltiplicazione di matrici INT8/INT4), ma come un codec specializzato per lo storage e il trasporto dei checkpoint. Il suo valore principale risiede nella riduzione dell'impronta di storage, della larghezza di banda di rete e della latenza di avvio a freddo associate alla distribuzione e alla materializzazione ripetuta di pesi di modelli di grandi dimensioni.

Gli autori affermano che, sfruttando la ridondanza inter-strato attraverso la compensazione del movimento, MCWC affronta un collo di bottiglia ortogonale alla quantizzazione tradizionale. Il metodo è più efficace nei regimi di deployment in cui i checkpoint vengono spostati, archiviati o replicati molte volte (ad esempio, avvii a freddo serverless, registri di modelli, distribuzione edge), permettendo al costo di codifica offline una tantum di essere ammortizzato sui risparmi di deployment ripetuti. Il lavoro suggerisce che trattare i pesi neurali come traiettorie strutturate attraverso la profondità, piuttosto che come tensori isolati, offre una nuova via per la compressione efficiente dei modelli.