Riepilogo Tecnico: AIRE-Prune

Problema

I modelli Deep State Space (SSM) hanno dimostrato efficacia nella modellazione di dati sequenziali comprimendo la storia degli input in stati interni. Tuttavia, la dimensione dello stato $n$ rimane un driver primario dei costi computazionali e di memoria. Le architetture moderne spesso utilizzano parametrizzazioni diagonali o diagonalizzabili (ad esempio, stack multi-SISO o MIMO come S5) per garantire stabilità ed efficienza, ma questi modelli soffrono frequentemente di sovra-parametrizzazione. I meccanismi post-training esistenti per ottimizzare $n$ sono limitati. Gli approcci precedenti, come il pruning dei pesi non strutturato, non riducono direttamente la dimensione dello stato. Altri metodi strutturati, come l' recente LAST (Gwak et al., 2025), si basano sul guadagno nel dominio della frequenza nel caso peggiore ($H_\infty$) per valutare gli stati. Sebbene efficaci, le misure del caso peggiore possono essere eccessivamente conservative per carichi di lavoro tipici, poiché enfatizzano picchi di amplificazione che i task raramente eccitano, portando potenzialmente a rapporti di compressione subottimali.

Metodologia: AIRE-Prune

Gli autori introducono AIRE-Prune (Asymptotic Impulse-Response Energy for State PRUN(E)ing), un metodo di pruning strutturato, post-training e adattivo per livello, progettato per SSM diagonali (o diagonalizzabili). A differenza dei metodi precedenti che si concentrano sul guadagno del caso peggiore, AIRE-Prune allinea il pruning all'energia asintotica della risposta all'impulso, un criterio del caso tipico che rappresenta l'energia totale in uscita che uno stato contribuisce su un orizzonte temporale infinito.

Meccanismo Core

1. Classificazione basata sull'Energia:
   Per un layer SSM diagonale discreto stabile definito da $(\Lambda, B, C)$, dove $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_n)$ e $|\lambda_i| < 1$, il metodo calcola l'energia totale quadratica in uscita generata da un input a impulso unitario per ogni stato (modo) $i$.
   Il punteggio in forma chiusa per il modo $i$ è derivato come:
   $$\text{EnergyScore}_{\text{local}}(x_i) = E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1 - |\lambda_i|^2}$$
   Questo punteggio integra tre fattori:

   • Controllabilità: Rappresentata da $\|B_{i,:}\|_2^2$ (facilità di eccitazione).
   • Osservabilità: Rappresentata da $\|C_{:,i}\|_2^2$ (facilità di misurazione).
   • Smorzamento/Memoria: Rappresentata dal polo $\lambda_i$; gli stati con poli più vicini al cerchio unitario (decadimento più lento) contribuiscono maggiormente all'energia a lungo termine.

2. Normalizzazione Cross-Layer:
   Poiché le scale di energia possono variare significativamente tra i diversi layer (ad esempio, a causa dei guadagni encoder/decoder), gli autori propongono una strategia di punteggio con normalizzazione del prefisso per consentire il confronto globale:

   • All'interno di ogni layer $\ell$, i modi sono ordinati in base alla loro energia locale $E^{(\ell)}_i$ in ordine decrescente.
   • Viene calcolata una somma prefissa $S^{(\ell)}_{(i)}$ per le energie ordinate.
   • Il punteggio globale è definito come il rapporto "hazard-rate":
     $$\text{AIREPrune}(x^{(\ell)}_{(i)}) = \frac{E^{(\ell)}_{(i)}}{S^{(\ell)}_{(i)} + \epsilon}$$
     Questa normalizzazione assicura che i punteggi siano monotonicamente non crescenti all'interno di un layer, permettendo l'applicazione di una singola soglia globale $\tau$ in tutto il modello.

3. Esecuzione del Pruning:
   Dato un rapporto di pruning o una soglia target, l'algoritmo mantiene il prefisso più lungo di stati in ogni layer in cui il punteggio normalizzato supera $\tau$. Ciò consente un pruning adattivo per layer, dove alcuni layer possono essere pesantemente potati o persino rimossi interamente se i loro stati cadono sotto la soglia globale, mentre altri rimangono ampiamente intatti.

Contributi Chiave

• Punteggio di Energia Asintotica in Forma Chiusa: Il paper fornisce un punteggio in forma chiusa, derivato matematicamente, per ogni modo appreso basato sull'energia in uscita all'orizzonte infinito. Questo evita la necessità di costose simulazioni numeriche o ri-addestramento.
• Estensione della Troncatura Modale: Il lavoro estende la classica troncatura modale (tipicamente applicata a singoli sistemi lineari) a stack profondi di SSM non lineari, preservando la struttura diagonale essenziale per l'efficienza.
• Caso Tipico vs. Caso Peggiore: Spostando la metrica di importanza dal guadagno $H_\infty$ del caso peggiore all'energia totale della risposta all'impulso, il metodo mira al comportamento aggregato del sistema piuttosto che a rare amplificazioni di picco.
• Agnosticismo dell'Architettura: Il metodo è applicabile sia a costruzioni multi-SISO che a stack MIMO (come S5-style), così come ad altre famiglie di SSM come S4D e Mamba.

Risultati Sperimentali

Gli autori hanno valutato AIRE-Prune sul benchmark Long Range Arena (LRA) e sul dataset Speech Commands utilizzando backbone S5, S4D e Mamba.

• Efficienza di Compressione: Su diversi benchmark, AIRE-Prune ha raggiunto un pruning medio degli stati del 60,8% con una perdita di accuratezza media di solo lo 0,29% senza alcun ri-addestramento.
• Confronto con i Baseline:
  • Sulla suite LRA, AIRE-Prune ha superato significativamente i baseline inclusi Random, Uniform Magnitude, Global Magnitude, LAMP, Uniform $H_\infty$, Global $H_\infty$ e LAST.
  • Sotto un vincolo stretto di $\le 1$ punto percentuale (pp) di perdita di accuratezza, AIRE-Prune ha permesso rapporti di pruning sostanzialmente più elevati rispetto al baseline più forte (LAST). Ad esempio, sul task Path-X, AIRE-Prune ha rimosso in sicurezza il 70% degli stati rispetto al 30% di LAST.
  • Sul task Speech Commands, il 45% degli stati è stato rimosso con una perdita di accuratezza trascurabile (96,40% vs. 96,43% per il modello completo).
• Efficienza di Inferenza: Il pruning si è tradotto direttamente in guadagni a livello di sistema. Su una GPU NVIDIA H100, il metodo ha prodotto accelerazioni di inferenza compresi tra 1,23$\times$ e 2,86$\times$ e riduzioni dei parametri dal 19% al 65%, a seconda del task e del rapporto di pruning.
• Generalizzazione: Il metodo ha dimostrato una robusta performance sulle architetture S4D e Mamba, raggiungendo fino al 90% di pruning sul task Text per S4D e l'80% sul task Image per Mamba con una degradazione minima dell'accuratezza.

Significato e Claims

Il paper sostiene che AIRE-Prune rivela una sostanziale ridondanza negli spazi di stato degli SSM addestrati, che può essere rimossa senza ri-addestramento. Allineando il pruning all'energia della risposta asintotica invece che al guadagno del caso peggiore, il metodo ottiene una compressione aggressiva con un costo trascurabile per le prestazioni del task.

Gli autori evidenziano come il profilo di caduta dell'accuratezza "a gradino" osservato con AIRE-Prune (dove l'accuratezza rimane stabile finché non viene raggiunta un'alta soglia di pruning, per poi scendere bruscamente) indichi una forte separazione tra stati importanti e non importanti. Questo contrasta con i baseline che mostrano una degradazione fluida, suggerendo che i loro metodi di punteggio mescolano stati importanti e non importanti.

In termini pratici, il metodo consente:

1. Budgeting Globale: Impostare un singolo budget target senza necessità di tuning per ogni layer.
2. Riduzione a Livello di Blocco: La capacità di rimuovere interi layer quando i loro stati scendono sotto la soglia, convertendo la granularità fine in riduzioni strutturate a blocchi, più compatibili con l'hardware.
3. Ottimizzazione Senza Training: Un meccanismo post-training che non richiede ulteriore fine-tuning, rendendolo immediatamente applicabile ai modelli distribuiti.

Il lavoro conclude che AIRE-Prune offre un approccio fondato sull'energia per la riduzione dell'ordine dei modelli nei deep SSM, superando le attuali strategie di pruning in termini di rapporto di compressione e preservazione dell'accuratezza.