Flexiffusion: Training-Free Segment-Wise Neural Architecture Search for Efficient Diffusion Models

TL;DR

Flexiffusion introduces a training-free, segment-wise neural architecture search method that significantly accelerates diffusion models by optimizing generation schedules and architectures without retraining, using a lightweight evaluation metric.

Contribution

It proposes a novel segment-wise NAS framework that reduces search complexity and evaluation time, enabling efficient diffusion model acceleration without retraining or quality loss.

Findings

Achieves at least 2x speedup on various diffusion models with less than 5% FID degradation.

Reduces evaluation time by over 90% using the relative FID metric.

Attains 5.1x speedup on Stable Diffusion with near-identical CLIP scores.

Abstract

Diffusion models (DMs) are powerful generative models capable of producing high-fidelity images but are constrained by high computational costs due to iterative multi-step inference. While Neural Architecture Search (NAS) can optimize DMs, existing methods are hindered by retraining requirements, exponential search complexity from step-wise optimization, and slow evaluation relying on massive image generation. To address these challenges, we propose Flexiffusion, a training-free NAS framework that jointly optimizes generation schedules and model architectures without modifying pre-trained parameters. Our key insight is to decompose the generation process into flexible segments of equal length, where each segment dynamically combines three step types: full (complete computation), partial (cache-reused computation), and null (skipped computation). This segment-wise search space reduces the candidate pool exponentially compared to step-wise NAS while preserving architectural diversity. Further, we introduce relative FID (rFID), a lightweight evaluation metric for NAS that measures divergence from a teacher model's outputs instead of ground truth, slashing evaluation time by over $90\%$. In practice, Flexiffusion achieves at least $2\times$ acceleration across LDMs, Stable Diffusion, and DDPMs on ImageNet and MS-COCO, with FID degradation under $5\%$, outperforming prior NAS and caching methods. Notably, it attains $5.1\times$ speedup on Stable Diffusion with near-identical CLIP scores. Our work pioneers a resource-efficient paradigm for searching high-speed DMs without sacrificing quality.