One-Way Ticket:Time-Independent Unified Encoder for Distilling Text-to-Image Diffusion Models

TL;DR

This paper tackles slow inference in text-to-image diffusion by distilling into a fast one-pass model. It introduces Time-independent Unified Encoder (TiUE), which uses a single encoder pass while a multi-step decoder runs in parallel across $K$ steps, enabling loop-free sampling. Training combines variational score distillation with a KL-divergence regularizer on UNet outputs and leverages a three-network setup (teacher, SD-LoRA, and TiUE student). Empirically, TiUE outperforms strong one-step baselines in diversity and realism while maintaining competitive inference speed, illustrating the practical benefits of encoder feature sharing and loop-free diffusion distillation.

Abstract

Text-to-Image (T2I) diffusion models have made remarkable advancements in generative modeling; however, they face a trade-off between inference speed and image quality, posing challenges for efficient deployment. Existing distilled T2I models can generate high-fidelity images with fewer sampling steps, but often struggle with diversity and quality, especially in one-step models. From our analysis, we observe redundant computations in the UNet encoders. Our findings suggest that, for T2I diffusion models, decoders are more adept at capturing richer and more explicit semantic information, while encoders can be effectively shared across decoders from diverse time steps. Based on these observations, we introduce the first Time-independent Unified Encoder TiUE for the student model UNet architecture, which is a loop-free image generation approach for distilling T2I diffusion models. Using a one-pass scheme, TiUE shares encoder features across multiple decoder time steps, enabling parallel sampling and significantly reducing inference time complexity. In addition, we incorporate a KL divergence term to regularize noise prediction, which enhances the perceptual realism and diversity of the generated images. Experimental results demonstrate that TiUE outperforms state-of-the-art methods, including LCM, SD-Turbo, and SwiftBrushv2, producing more diverse and realistic results while maintaining the computational efficiency.

Figures (17)

• Figure 1: The correlation between image generation quality (Fig.a-b) and encoder feature similarity (Fig.c). Above a certain threshold of steps, such as 15 steps in SD2.1, the model maintains image generation quality (Fig.a-b) while the features show high similarity (Fig.c). Below this threshold, feature similarity deteriorates along with worse generation quality, accompanied by a degradation in image generation quality. Furthermore, the encoder features consistently exhibit higher similarity than the decoder across all sampling steps (Fig.c).
• Figure 2: The generated images and the corresponding UNet decoder features using different time steps. The images generated in one step mainly focus on structural information, with a minimal amount of semantic information present in the feature layers (Left). By comparison, the decoder across multiple time-steps contain richer and more explicit semantic information in the feature level (Right), leading to better image generation.
• Figure 3: Generated high-fidelity $512^2$ images by the one-step model distilled with our method TiUE. Compared with the baselines, our method produces higher image quality and more diverse results. Here the diversity means the degree of variation of different output images with same semantic information when given the same prompt and different seed values. SwiftBrushv2 dao2024swiftbrushv2, which is initialized with SD-Turbo sauer2023adversarial, achieves T2I generation with limited diversity as the SD-Turbo model.
• Figure 4: Previous methods use the time-dependent UNet encoder-decoder design and differ from the distillation losses. (a) LCM luo2023latent updates with consistency distillation loss $\mathcal{L}_{cd}$. (b) SD-Turbo sauer2023adversarial is trained with adversarial loss $\mathcal{L}_{adv}$ and score distillation loss $\mathcal{L}_{sds}$. (c) SwiftBrushv1/v2 nguyen2023swiftbrushdao2024swiftbrushv2 distills with variational score distillation $\mathcal{L}_{vsd}$ for one-step model. (d) Our method TiUE first proposes the Time-independent Unified Encoder (TiUE) design to achieve the one-pass loop-free generation. TiUE distills the T2I model with variational score distillation $\mathcal{L}_{vsd}$ and KL divergence $\mathcal{L}_{kl}$. The loop-free nature of the architecture allows the decoder blocks run in parallel across various time steps, thus saving much inference time.
• Figure 5: Loop-free distillation based on our Time-independent Unified Encoder (TiUE) architecture. Our framework is composed of three networks: a student generator $\epsilon^{SG}_{\theta}$, a pretrained teacher SD generator $\epsilon^{TG}_{\psi}$ and a SD-LoRA generator $\epsilon^{LG}_{\psi}$. The encoder of $\epsilon^{SG}_{\theta}$ is only used once while passing through the decoder for K time-steps. The student network is updated with both variational score distillation loss $\mathcal{L}_{vsd}$ and KL divergence loss $\mathcal{L}_{kl}$. We also optimize the SD-LoRA $\epsilon^{LG}_{\psi}$ by diffusion loss. The SD-LoRA teacher and the student model are updated alternately, following the previous methods wang2023prolificdreamernguyen2023swiftbrush.