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

Abstract: Latent diffusion models achieve strong generative performance by operating in a compressed latent space produced by a variational autoencoder (VAE). However, it remains unclear whether all latent channels contribute equally to the diffusion process, or whether significant redundancy exists. We introduce PeLAP-A (Adaptive Latent Pruning for Diffusion), a lightweight framework that augments a standard latent diffusion pipeline with a learnable channel-wise importance predictor. A two-layer MLP operating on globally pooled latent features produces a soft mask that suppresses unimportant latent channels before they enter the denoising UNet. The entire system is trained jointly on CIFAR-10 under a combined diffusion, reconstruction, and sparsity loss. Experiments reveal a striking result: under aggressive sparsity regularization (lambda = 0.01), the importance predictor drives all latent channels to near-zero yet the denoising UNet achieves lower diffusion loss (0.0236 vs. 0.0240) and lower VAE reconstruction MSE (22.59 vs. 24.67) compared to the unpruned baseline. We term this the sparsity collapse phenomenon and provide an analysis of why it occurs and what it reveals about the information requirements of latent diffusion models. These findings constitute an exploratory study of sparsity dynamics in latent diffusion training, and demonstrate that denoising UNets can remain remarkably robust to latent channel suppression even under aggressive regularization. Code is available at: https://github.com/kissasium/PeLAP-A.git.

• The paper demonstrates that adaptive latent pruning significantly lowers denoising and reconstruction errors in latent diffusion models.
• It integrates a compact, learnable MLP-based importance predictor that dynamically applies channel-wise sparsity during training.
• Aggressive sparsity leads to a collapse of latent information, yielding improved loss metrics but degraded generative fidelity as shown by FID scores.

PeLAP-A: Adaptive Latent Pruning in Latent Diffusion Models

Introduction and Motivation

The paper "PeLAP-A: Adaptive Latent Pruning for Lightweight Latent Diffusion Models" (2606.23086) presents a study into the efficiency and redundancy of latent representations within latent diffusion models (LDMs), which have become central in generative modeling due to their strong performance and computational benefits. Latent spaces, typically constructed via a VAE, may include channels that are non-essential for the denoising process. PeLAP-A introduces an adaptive, learnable, input-conditioned module for channel-wise latent pruning with the objective of understanding and possibly reducing this redundancy in LDMs.

Methodology

PeLAP-A augments a standard latent diffusion pipeline by interposing a compact, 292-parameter, two-layer MLP-based importance predictor between the VAE encoder and the denoising UNet. This module performs global average pooling across latent features, producing a per-channel soft mask via sigmoid activation. During training, a sparsity regularizer—in conjunction with VAE reconstruction and diffusion losses—drives the mask toward suppressing unneeded channels. Importantly, the masking is image-dependent, enabling dynamic selection of channels rather than enforcing a static sparsity pattern across the dataset.

The full training objective combines the VAE pixel-wise reconstruction loss, KL-divergence regularization, noise-prediction loss for DDPM, and a channel-wise mean sparsity penalty weighted by parameter $\lambda$. Notably, the predictor is initialized so that all pruned channel masks start near 1, minimizing accidental early suppression. Architectural details, including UNet and VAE configuration, remain identical between the baseline and PeLAP-A for direct comparability.

Empirical Analysis and Results

On CIFAR-10, PeLAP-A reveals a salient phenomenon when applying strong sparsity regularization ($\lambda=0.01$): rapid, complete suppression of all latent channels to near-zero (termed "sparsity collapse") within two epochs, followed by persistence in this collapsed regime throughout extended training. Despite aggressive pruning, PeLAP-A achieves:

• Lower diffusion noise-prediction MSE (0.0236 vs. 0.0240)
• Lower VAE reconstruction MSE (22.59 vs. 24.67)
• Lower overall validation loss (0.0330 vs. 0.0390)

compared to the unpruned baseline. These quantitative improvements are achieved with the UNet operating on zeroed latent inputs. However, generated samples post-collapse are nearly uniform gray images, and FID degrades substantially (362.6 for PeLAP-A vs. 278.1 for baseline).

Class-wise analysis of the soft mask values—even post-collapse—exposes non-trivial structure: mask activations preserve a consistent, class-dependent order at values several orders of magnitude below threshold, suggesting that sub-threshold mask variations still encode class-specific signal.

Theoretical and Practical Implications

The study demonstrates that the denoising UNet can adapt to, and be robust under, a regime where all latent inputs collapse to near zero. The denoising loss thus becomes decoupled from generative quality under full suppression: the network optimizes for denoising pure Gaussian noise, which is simpler than denoising structured high-entropy latents. This uncovers a baseline solution that optimizes the objective but fails at meaningful generative modeling, as evidenced by FID collapse.

This observation underscores that diffusion loss is not a reliable stand-alone metric for sample quality in the extreme pruning regime. It mandates that perceptual metrics, such as FID, be used in conjunction for proper assessment. The observed "sparsity collapse" suggests the existence of a critical regularization threshold, above which informative latents are entirely suppressed before the denoising network can leverage them. The paper hypothesizes on the gradient dynamics driving this collapse: the sparsity penalty dominates early epochs, with the denoising objective unable to counteract its influence due to slower learning dynamics.

Limitations and Future Directions

Several practical limitations are identified:

• Sparsity collapse precludes selective (partial) channel pruning, failing the objective of retaining only key informative latents while suppressing others.
• The degradation of generative quality, as measured by FID, confirms that pruned latents at high regularization compromise downstream image fidelity.
• Empirical findings (including the collapse threshold) are currently only substantiated for a minimal latent dimension (4 channels on CIFAR-10), with preliminary ablations implying that even lower regularization triggers collapse.

The study proposes future remedies, including regularization warm-up, top-$k$ masking, frozen VAE encoders, evaluation on higher-dimensional latent spaces, and curriculum-based learning schemes to balance sparsity and generative fidelity. Broader generalization across architectures, datasets, and latent sizes is highlighted as essential for further theoretical clarification.

Conclusion

PeLAP-A advances the study of latent redundancy and adaptive channel pruning within LDMs by introducing a lightweight, input-adaptive importance predictor. The empirical discovery of sparsity collapse—a regime where all information in the latent is suppressed yet the denoising path remains robust—raises critical questions about the actual information requirements of generative diffusion models. While strong numerical improvements on denoising losses are observed under collapse, generative fidelity may be fundamentally compromised if pruning is too aggressive. This work motivates deeper analysis of the interaction between latent sparsity dynamics and generative modeling, including the search for practical mechanisms that enable true selectivity in adaptive pruning for scalable LDMs.