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

Resumen Técnico: PeLAP-A (Poda Adaptativa de Latentes para Difusión)

Planteamiento del Problema
Los modelos de difusión latente (LDM) logran una síntesis de imágenes de alta calidad operando en un espacio latente comprimido generado por un Autoencoder Variacional (VAE). Sin embargo, sigue siendo una cuestión abierta si todos los canales latentes contribuyen por igual al proceso de difusión o si existe una redundancia significativa. Mientras que el trabajo existente se centra en comprimir la arquitectura del modelo de difusión en sí misma (por ejemplo, la poda de pesos) o el proceso de muestreo, el contenido de información del propio espacio de representación latente permanece mayormente inexplorado. Los autores investigan si los canales latentes redundantes pueden identificarse y suprimirse selectivamente sin degradar la calidad de la generación.

Metodología: PeLAP-A
El artículo propone PeLAP-A (Poda Adaptativa de Latentes para Difusión), un marco ligero que aumenta el flujo estándar de difusión latente con un predictor de importancia por canal, aprendible y adaptativo a la entrada.

• Arquitectura: Se inserta un Perceptrón Multicapa (MLP) de dos capas que contiene solo 292 parámetros entre el codificador del VAE y el UNet de eliminación de ruido (denoising).
• Mecanismo:
  1. El tensor latente $z$ del codificador del VAE se somete a un agrupamiento promedio global (global average pooling) para producir un resumen por canal.
  2. Este resumen es procesado por el MLP y pasado a través de una activación sigmoide para generar una máscara suave $m \in (0, 1)^C$.
  3. La máscara se transmite espacialmente y se aplica elemento a elemento al tensor latente ($z' = z \odot m$) antes de entrar al UNet.
  4. La capa lineal final del predictor se inicializa con un sesgo (bias) de +3.0, asegurando que todos los canales comiencen casi totalmente abiertos ($\sigma(3) \approx 0.95$) para evitar una supresión prematura.
• Objetivo de Entrenamiento: El sistema se entrena de extremo a extremo en CIFAR-10 utilizando una función de pérdida conjunta:
  $$L = L_{vae} + L_{diff} + \lambda \cdot \frac{1}{C}\sum_{c=1}^{C} m_c$$
  donde $L_{vae}$ es la pérdida de reconstrucción y divergencia KL, $L_{diff}$ es el error cuadrático medio (MSE) de la predicción de ruido de la difusión, y el tercer término es un regularizador de dispersión (sparsity) que incentiva el valor medio de la máscara hacia cero.

Resultados Clave
Los experimentos en CIFAR-10 revelaron un fenómeno sorprendente denominado "colapso de dispersión" (sparsity collapse):

• Dinámica de Colapso: Bajo una regulación de dispersión agresiva ($\lambda = 0.01$), el predictor de importancia impulsa todos los canales latentes hacia valores cercanos a cero dentro de las dos primeras épocas de entrenamiento.
• Rendimiento de Denoising: A pesar de la supresión completa de la información latente, el UNet de eliminación de ruido en el modelo PeLAP-A logró métricas de validación menores que la línea base sin podar:
  • Pérdida de Difusión: 0.0236 (PeLAP-A) vs. 0.0240 (Línea Base).
  • MSE de Reconstrucción del VAE: 22.59 (PeLAP-A) vs. 24.67 (Línea Base).
• Fallo Generativo: Aunque la vía de eliminación de ruido mejoró, la calidad de la salida generativa colapsó. La puntuación FID se degradó significativamente de 278.1 (Línea Base) a 362.6 (PeLAP-A). El análisis cualitativo mostró que, con latentes cercanos a cero, el decodificador producía salidas grises uniformes en lugar de imágenes estructuradas.
• Estructura Residual: Incluso después del colapso, las magnitudes de la máscara retuvieron un orden sensible a la clase (por ejemplo, las clases "automóvil" tenían activaciones medias ligeramente superiores a las de "avión"), lo que indica que el predictor codificó información relevante de la clase incluso en magnitudes sub-umbral.

Significancia y Reivindicaciones
El artículo posiciona a PeLAP-A como un estudio exploratorio sobre la dinámica de la dispersión dentro del entrenamiento de la difusión latente. Sus principales reivindicaciones y contribuciones son:

1. Robustez de los UNets: Los hallazgos demuestran que los UNets de eliminación de ruido son notablemente robustos a la supresión de canales latentes. Pueden aprender prioris incondicionales efectivos sobre espacios latentes cercanos a cero, lo que sugiere que los requisitos de información para el proceso de eliminación de ruido son mucho menores de lo que la dimensionalidad del espacio latente implica.
2. Limitación de la Pérdida de Difusión: La divergencia entre la mejora en la pérdida de difusión y la degradación del FID resalta que la pérdida de difusión por sí sola es un proxy insuficiente para la calidad de la generación cuando el espacio latente ha colapsado. Las métricas perceptuales como el FID son esenciales para evaluar la poda de latentes.
3. Fenómeno de Colapso de Dispersión: Los autores identifican una transición de fase brusca donde la presión de dispersión supera el aprendizaje de gradientes informativos al inicio del entrenamiento. Hipotetizan que $\lambda = 0.01$ excede un umbral crítico para esta arquitectura específica, causando que el sistema se asiente en un mínimo local donde el UNet se adapta para eliminar el ruido de tensores cercanos a cero en lugar de datos estructurados.
4. Alcance Modesto: Los autores declaran explícitamente que el método, tal como se implementó, no logró el objetivo de retener selectivamente canales informativos mientras se suprimían otros. El trabajo sirve como fundamento para comprender los umbrales críticos de la dispersión en modelos generativos y motiva trabajos futuros sobre esquemas de calentamiento (warm-up), relajación discreta o entrenamiento en espacios latentes más grandes para lograr una poda parcial.