ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation

TL;DR

ALTER introduces a unified approach combining layer pruning and temporal expert routing with a trainable hypernetwork, significantly reducing computational costs of diffusion models while maintaining high image quality.

Contribution

It presents a novel single-stage optimization framework that jointly prunes layers and manages expert routing during diffusion model fine-tuning, addressing prior sub-optimal sequential strategies.

Findings

Achieves 3.64x speedup with 35% sparsity.

Maintains original model fidelity with only 20 inference steps.

Reduces MACs to 25.9% of the original while preserving quality.

Abstract

Diffusion models have demonstrated exceptional capabilities in generating high-fidelity images. However, their iterative denoising process results in significant computational overhead during inference, limiting their practical deployment in resource-constrained environments. Existing acceleration methods often adopt uniform strategies that fail to capture the temporal variations during diffusion generation, while the commonly adopted sequential pruning-then-fine-tuning strategy suffers from sub-optimality due to the misalignment between pruning decisions made on pretrained weights and the model's final parameters. To address these limitations, we introduce ALTER: All-in-One Layer Pruning and Temporal Expert Routing, a unified framework that transforms diffusion models into a mixture of efficient temporal experts. ALTER achieves a single-stage optimization that unifies layer pruning, expert routing, and model fine-tuning by employing a trainable hypernetwork, which dynamically generates layer pruning decisions and manages timestep routing to specialized, pruned expert sub-networks throughout the ongoing fine-tuning of the UNet. This unified co-optimization strategy enables significant efficiency gains while preserving high generative quality. Specifically, ALTER achieves same-level visual fidelity to the original 50-step Stable Diffusion v2.1 model while utilizing only 25.9% of its total MACs with just 20 inference steps and delivering a 3.64x speedup through 35% sparsity.