PeLAP-A: Adaptive Latent Pruning for Lightweight Latent Diffusion Models

Sintesi Tecnica: PeLAP-A (Adaptive Latent Pruning for Diffusion)

Problema
I modelli di diffusione latente (LDM) ottengono una sintesi di immagini di alta qualità operando in uno spazio latente compresso generato da un Autoencoder Variazionale (VAE). Tuttavia, rimane una questione aperta se tutti i canali latenti contribuiscano equamente al processo di diffusione o se esista una significativa ridondanza. Mentre il lavoro esistente si concentra sulla compressione dell'architettura del modello di diffusione stesso (ad esempio, il pruning dei pesi) o sul processo di campionamento, il contenuto informativo della rappresentazione latente stessa rimane ampiamente inesplorato. Gli autori indagano se i canali latenti ridondanti possano essere identificati e soppressi selettivamente senza degradare la qualità della generazione.

Metodologia: PeLAP-A
Il documento propone PeLAP-A (Adaptive Latent Pruning for Diffusion), un framework leggero che aumenta la standard pipeline di diffusione latente con un predittore di importanza per canale, apprendibile e adattivo rispetto all'input.

• Architettura: Un Multi-Layer Perceptron (MLP) a due strati contenente solo 292 parametri viene inserito tra l'encoder del VAE e l'UNet di denoising.
• Meccanismo:
  1. Il tensore latente $z$ proveniente dall'encoder del VAE subisce un global average pooling per produrre un riepilogo per canale.
  2. Questo riepilogo viene elaborato dall'MLP e passato attraverso un'attivazione sigmoidale per generare una maschera morbida $m \in (0, 1)^C$.
  3. La maschera viene propagata spazialmente e applicata elemento per elemento al tensore latente ($z' = z \odot m$) prima di entrare nell'UNet.
  4. Lo strato lineare finale del predittore è inizializzato con un bias di +3.0, garantendo che tutti i canali inizino quasi completamente aperti ($\sigma(3) \approx 0.95$) per evitare una soppressione prematura.
• Obiettivo di Addestramento: Il sistema è addestrato end-to-end su CIFAR-10 utilizzando una funzione di perdita congiunta:
  $$L = L_{vae} + L_{diff} + \lambda \cdot \frac{1}{C}\sum_{c=1}^{C} m_c$$
  dove $L_{vae}$ è la perdita di ricostruzione e di divergenza KL, $L_{diff}$ è l'MSE di predizione del rumore della diffusione, e il terzo termine è un regolarizzatore di sparsità che incoraggia il valore medio della maschera verso lo zero.

Risultati Chiave
Gli esperimenti su CIFAR-10 hanno rivelato un fenomeno sorprendente denominato "collasso della sparsità" (sparsity collapse):

• Dinamiche di Collasso: Sotto una forte regolarizzazione della sparsità ($\lambda = 0.01$), il predittore di importanza spinge tutti i canali latenti verso valori vicini allo zero entro i primi due epoche di addestramento.
• Performance di Denoising: Nonostante la completa soppressione dell'informazione latente, l'UNet di denoising nel modello PeLAP-A ha ottenuto metriche di validazione inferiori rispetto alla baseline non potata:
  • Perdita di Diffusione (Diffusion Loss): 0.0236 (PeLAP-A) vs 0.0240 (Baseline).
  • MSE di Ricostruzione VAE: 22.59 (PeLAP-A) vs 24.67 (Baseline).
• Fallimento Generativo: Mentre il percorso di denoising è migliorato, la qualità dell'output generativo è crollata. Il punteggio FID è degradato significativamente da 278.1 (Baseline) a 362.6 (PeLAP-A). L'analisi qualitativa ha mostrato che, con latenti quasi nulli, il decoder produceva output grigi uniformi invece di immagini strutturate.
• Struttura Residua: Anche dopo il collasso, i valori della maschera hanno mantenuto un ordinamento sensibile alla classe (ad esempio, le classi "automobile" avevano attivazioni medie leggermente superiori rispetto alle classi "airplane"), indicando che il predittore codificava informazioni rilevanti per la classe anche a magnitudo sub-soglia.

Significato e Rivendicazioni
Il documento posiziona PeLAP-A come uno studio esplorativo sulle dinamiche di sparsità all'interno dell'addestramento della diffusione latente. Le sue rivendicazioni e i suoi contributi principali sono:

1. Robustezza delle UNet: I risultati dimostrano che le UNet di denoising sono straordinariamente robuste alla soppressione dei canali latenti. Possono apprendere prior incondizionati efficaci su spazi latenti quasi nulli, suggerendo che i requisiti informativi per il processo di denoising sono molto più bassi di quanto la dimensionalità dello spazio latente implichi.
2. Limiti della Perdita di Diffusione: La divergenza tra la perdita di diffusione migliorata e il degradato FID evidenzia che la perdita di diffusione da sola è un proxy insufficiente per la qualità della generazione quando lo spazio latente è collassato. Metriche percettive come l'FID sono essenziali per valutare il pruning latente.
3. Fenomeno del Collasso della Sparsità: Gli autori identificano una transizione di fase netta dove la pressione della sparsità sovrasta l'apprendimento di gradienti informativi nelle prime fasi dell'addestramento. Ipotizzano che $\lambda = 0.01$ superi una soglia critica per questa specifica architettura, causando il raggiungimento di un minimo locale in cui l'UNet si adatta per denoise tensori quasi nulli piuttosto che dati strutturati.
4. Ambito Modesto: Gli autori dichiarano esplicitamente che il metodo, così come implementato, non ha raggiunto l'obiettivo di mantenere selettivamente i canali informativi sopprimendo al contempo gli altri. Il lavoro serve come base per comprendere le soglie critiche della sparsità nei modelli generativi e motiva futuri lavori su programmi di warm-up, rilassamento discreto o addestramento su spazi latenti più grandi per ottenere un pruning parziale.