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

Technische Samenvatting: PeLAP-A (Adaptive Latent Pruning for Diffusion)

Probleemstelling
Latente diffusiemodellen (LDMs) bereiken hoogwaardige beeldsynthese door te opereren in een gecomprimeerde latente ruimte gegenereerd door een Variational Autoencoder (VAE). Het blijft echter een openstaande vraag of alle latente kanalen evenredig bijdragen aan het diffusieproces of dat er significante redundantie bestaat. Terwijl bestaand werk zich richt op het comprimeren van de architectuur van het diffusiemodel zelf (bijv. gewicht-pruning) of het samplingsproces, blijft de informatie-inhoud van de latente representatie zelf grotendeels onverkend. De auteurs onderzoeken of redundante latente kanalen geïdentificeerd en selectief onderdrukt kunnen worden zonder de kwaliteit van de generatie te verslechteren.

Methodologie: PeLAP-A
Het artikel stelt PeLAP-A voor (Adaptive Latent Pruning for Diffusion), een lichtgewicht framework dat de standaard latente diffusie-pipeline aanvult met een leerbare, input-adaptieve kanaal-gewijze belangrijkheidspredictor.

• Architectuur: Een twee-laagse Multi-Layer Perceptron (MLP) die slechts 292 parameters bevat, wordt tussen de VAE-encoder en de denoising UNet ingevoegd.
• Mechanisme:
  1. De latente tensor $z$ van de VAE-encoder ondergaat global average pooling om een kanaal-gewijze samenvatting te produceren.
  2. Deze samenvatting wordt verwerkt door de MLP en door een sigmoid-activatiefunctie geleid om een soft mask $m \in (0, 1)^C$ te genereren.
  3. De mask wordt ruimtelijk gebroadcast en elementgewijs toegepast op de latente tensor ($z' = z \odot m$) voordat deze de UNet ingaat.
  4. De laatste lineaire laag van de predictor wordt geïnitialiseerd met een bias van +3.0, wat ervoor zorgt dat alle kanalen bijna volledig open beginnen ($\sigma(3) \approx 0.95$) om voortijdige onderdrukking te voorkomen.
• Trainingsdoel: Het systeem wordt end-to-end getraind op CIFAR-10 met een gezamenlijke verliesfunctie:
  $$L = L_{vae} + L_{diff} + \lambda \cdot \frac{1}{C}\sum_{c=1}^{C} m_c$$
  waarbij $L_{vae}$ de reconstructie- en KL-divergentieverlies is, $L_{diff}$ de MSE van de diffusie-ruisvoorspelling is, en de derde term een sparsity-regularisator is die de gemiddelde mask-waarde naar nul stimuleert.

Belangrijkste Resultaten
Experimenten op CIFAR-10 onthulden een verrassend fenomeen dat "sparsity collapse" wordt genoemd:

• Collapse Dynamiek: Onder agressieve sparsity-regularisatie ($\lambda = 0.01$) drijft de belangrijkheidspredictor alle latente kanalen naar bijna nul-waarden binnen de eerste twee trainingsepochs.
• Denoising Prestaties: Ondanks de volledige onderdrukking van latente informatie, behaalde de denoising UNet in het PeLAP-A model lagere validatiemetrieken dan de ongepruned baseline:
  • Diffusion Loss: 0.0236 (PeLAP-A) vs. 0.0240 (Baseline).
  • VAE Reconstruction MSE: 22.59 (PeLAP-A) vs. 24.67 (Baseline).
• Generatieve Falen: Hoewel het denoising-pad verbeterde, stort de kwaliteit van de generatieve output in. De FID-score verslechterde aanzienlijk van 278.1 (Baseline) naar 362.6 (PeLAP-A). Kwalitatieve analyse toonde aan dat de decoder met bijna-nul latenten uniforme grijze outputs produceerde in plaats van gestructureerde afbeeldingen.
• Residuele Structuur: Zelfs na de collapse behielden de mask-magnitudes een klasse-gevoelige ordening (bijv. "automobile" klassen hadden iets hogere gemiddelde activaties dan "airplane" klassen), wat aangeeft dat de predictor klasse-relevante informatie codeerde, zelfs bij sub-drempel magnitudes.

Betekenis en Claims
Het artikel positioneert PeLAP-A als een verkennende studie naar sparsity-dynamiek binnen latente diffusietraining. De primaire claims en bijdragen zijn:

1. Robuustheid van UNets: De bevindingen demonstreren dat denoising UNets opmerkelijk robuust zijn tegen latente kanaal-onderdrukking. Ze kunnen effectieve onvoorwaardelijke priors leren over bijna-nul latente ruimtes, wat suggereert dat de informatievereisten voor het denoising-proces veel lager zijn dan de dimensionaliteit van de latente ruimte impliceert.
2. Beperking van Diffusion Loss: De divergentie tussen verbeterde diffusion loss en verslechterde FID benadrukt dat diffusion loss alleen een onvoldoende proxy is voor generatiekwaliteit wanneer de latente ruimte is ingestort. Perceptuele metrieken zoals FID zijn essentieel voor het evalueren van latente pruning.
3. Sparsity Collapse Fenomeen: De auteurs identificeren een scherpe fase-transitie waarbij de sparsity-druk de leerprocessen van informatieve gradiënten vroegtijdig overstemt. Ze hypothetiseren dat $\lambda = 0.01$ een kritieke drempel overschrijdt voor deze specifieke architectuur, waardoor het systeem vastloopt in een lokaal minimum waar de UNet leert te denoisen rondom bijna-nul tensoren in plaats van gestructureerde data.
4. Bescheiden Omvang: De auteurs geven expliciet aan dat de methode, zoals geïmplementeerd, het doel niet heeft bereikt om selectief informatieve kanalen te behouden terwijl anderen worden onderdrukt. Het werk dient als een fundament voor het begrijpen van de kritieke drempels van sparsity in generatieve modellen en motiveert toekomstig werk op het gebied van warm-up schema's, discrete relaxatie, of training op grotere latente ruimtes om gedeeltelijke pruning te bereiken.