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

Resumo Técnico: PeLAP-A (Adaptive Latent Pruning for Diffusion)

Definição do Problema
Modelos de difusão latente (LDMs) alcançam síntese de imagens de alta qualidade ao operar em um espaço latente comprimido gerado por um Autoencoder Variacional (VAE). No entanto, permanece uma questão em aberto se todos os canais latentes contribuem igualmente para o processo de difusão ou se existe redundância significativa. Enquanto o trabalho existente foca na compressão da própria arquitetura do modelo de difusão (ex: poda de pesos) ou no processo de amostragem, o conteúdo informacional da própria representação latente permanece amplamente inexplorado. Os autores investigam se canais latentes redundantes podem ser identificados e suprimidos seletivamente sem degradar a qualidade da geração.

Metodologia: PeLAP-A
O artigo propõe o PeLAP-A (Adaptive Latent Pruning for Diffusion), um framework leve que aumenta o pipeline padrão de difusão latente com um preditor de importância por canal, aprendível e adaptativo à entrada.

• Arquitetura: Um Perceptron de Camadas Multicamadas (MLP) de duas camadas contendo apenas 292 parâmetros é inserido entre o codificador VAE e o UNet de denoising.
• Mecanismo:
  1. O tensor latente $z$ do codificador VAE passa por um pooling de média global para produzir um resumo por canal.
  2. Este resumo é processado pelo MLP e passa por uma ativação sigmoide para gerar uma máscara suave $m \in (0, 1)^C$.
  3. A máscara é propagada espacialmente e aplicada elemento a elemento ao tensor latente ($z' = z \odot m$) antes de entrar no UNet.
  4. A camada linear final do preditor é inicializada com um viés (bias) de +3.0, garantindo que todos os canais comecem quase totalmente abertos ($\sigma(3) \approx 0.95$) para evitar a supressão prematura.
• Objetivo de Treinamento: O sistema é treinado de ponta a ponta no CIFAR-10 usando uma função de perda conjunta:
  $$L = L_{vae} + L_{diff} + \lambda \cdot \frac{1}{C}\sum_{c=1}^{C} m_c$$
  onde $L_{vae}$ é a perda de reconstrução e divergência KL, $L_{diff}$ é o MSE de predição de ruído da difusão, e o terceiro termo é um regularizador de esparsidade que incentiva o valor médio da máscara em direção a zero.

Resultos Principais
Experimentos no CIFAR-10 revelaram um fenômeno surpreendente denominado "colapso de esparsidade" (sparsity collapse):

• Dinâmica de Colapso: Sob uma regularização de esparsidade agressiva ($\lambda = 0.01$), o preditor de importância conduz todos os canais latentes a valores próximos de zero dentro das duas primeiras épocas de treinamento.
• Desempenho de Denoising: Apesar da supressão completa da informação latente, o UNet de denoising no modelo PeLAP-A alcançou métricas de validação menores do que o baseline não podado:
  • Perda de Difusão: 0.0236 (PeLAP-A) vs. 0.0240 (Baseline).
  • MSE de Reconstrução VAE: 22.59 (PeLAP-A) vs. 24.67 (Baseline).
• Falha Generativa: Embora o caminho de denoising tenha melhorado, a qualidade da saída generativa colapsou. O score FID degradou significativamente de 278.1 (Baseline) para 362.6 (PeLAP-A). A análise qualitativa mostrou que, com latentes próximos de zero, o decodificador produziu saídas cinzas uniformes em vez de imagens estruturadas.
• Estrutura Residual: Mesmo após o colapso, as magnitudes da máscara mantiveram uma ordenação sensível à classe (ex: classes "automóvel" tinham ativações médias ligeiramente superiores às classes "avião"), indicando que o preditor codificou informações relevantes à classe mesmo em magnitudes sub-limiares.

Significância e Alegações
O artigo posiciona o PeLAP-A como um estudo exploratório sobre a dinâmica de esparsidade no treinamento de difusão latente. Suas principais alegações e contribuições são:

1. Robustez dos UNets: As descobertas demonstram que os UNets de denoising são notavelmente robustos à supressão de canais latentes. Eles podem aprender priors incondicionais eficazes sobre espaços latentes próximos de zero, sugerindo que os requisitos de informação para o processo de denoising são muito menores do que a dimensionalidade do espaço latente implica.
2. Limitação da Perda de Difusão: A divergência entre a melhora na perda de difusão e a degradação no FID destaca que a perda de difusão sozinha é um proxy insuficiente para a qualidade de geração quando o espaço latente colapsou. Métricas perceptuais como o FID são essenciais para avaliar a poda latente.
3. Fenômeno de Colapso de Esparsidade: Os autores identificam uma transição de fase abrupta onde a pressão de esparsidade sobrecarrega o aprendizado de gradientes informativos precocemente no treinamento. Eles hipotetizam que $\lambda = 0.01$ excede um limiar crítico para esta arquitetura específica, fazendo com que o sistema se estabilize em um mínimo local onde o UNet se adapta para realizar o denoising de tensores próximos de zero em vez de dados estruturados.
4. Escopo Modesto: Os autores declaram explicitamente que o método, conforme implementado, não atingiu o objetivo de selecionar a retenção de canais informativos enquanto suprime outros. O trabalho serve como uma fundação para compreender os limiares críticos de esparsidade em modelos generativos e motiva trabalhos futuros sobre cronogramas de aquecimento (warm-up), relaxação discreta ou treinamento em espaços latentes maiores para alcançar uma poda parcial.