Source-linked AI summary
SwiftBrush v2: Make Your One-step Diffusion Model Better Than Its Teacher
Trung Dao, Thuan Hoang Nguyen, Thanh Le, Duc Vu, Khoi Nguyen, Cuong Pham, Anh Tran
TL;DR
One-step text-to-image models have lagged multi-step Stable Diffusion on FID, despite their speed advantage. SwiftBrush v2 changes SwiftBrush training and combines full-training and LoRA model weights. The merged model reaches FID 8.14 after extra image regularization, surpassing its multi-step teacher and GAN-based approaches.
Problem
One-step text-to-image models still trail multi-step Stable Diffusion on FID, while multi-step sampling is slow and computationally expensive for real-time or on-device use.
Source: §1 · pp. 1–2Method
The method initializes SwiftBrush with SD Turbo, adds a clamped CLIP loss, uses full-student and LoRA training schemes, and merges their weights.
Source: §1 · p. 3Results
FID 8.14 after extra image regularization lets the merged one-step model surpass its multi-step teacher and GAN-based text-to-image approaches.
Source: §1 · p. 3Takeaways & Limitations
The authors present the merged model as a one-step student that sets a new standard for efficient, high-quality text-to-image generation.
Source: §1 · p. 3Naively applying CLIP loss can degrade image quality, with blurriness, increased color saturation, and textual artifacts observed.
Source: §4.3 · p. 8
Abstract
from arXiv · showhide
In this paper, we aim to enhance the performance of SwiftBrush, a prominent one-step text-to-image diffusion model, to be competitive with its multi-step Stable Diffusion counterpart. Initially, we explore the quality-diversity trade-off between SwiftBrush and SD Turbo: the former excels in image diversity, while the latter excels in image quality. This observation motivates our proposed modifications in the training methodology, including better weight initialization and efficient LoRA training. Moreover, our introduction of a novel clamped CLIP loss enhances image-text alignment and results in improved image quality. Remarkably, by combining the weights of models trained with efficient LoRA and full training, we achieve a new state-of-the-art one-step diffusion model, achieving an FID of 8.14 and surpassing all GAN-based and multi-step Stable Diffusion models. The project page is available at https://swiftbrushv2.github.io.
1 Introduction
SwiftBrush v2 targets the quality gap between one-step and multi-step text-to-image models by combining SwiftBrush’s image-free training with SD Turbo initialization and a clamped CLIP loss. The merged model reaches FID 8.14 after extra image regularization, surpassing its multi-step teacher and GAN-based approaches.
1 Introduction
- SDv2.1 reaches FID-30K 9.64 on COCO 2014, compared with 11.49 for the best reported one-step model of equivalent parameter scale.The paper frames this gap as a challenge to the belief that a one-step model cannot exceed its multi-step teacher.Source: §1 · p. 2
- SwiftBrush’s image-free, dynamic training yields more diverse outputs, while SD Turbo produces higher-quality but mode-collapse-like outputs.This quality-diversity comparison motivates initializing SwiftBrush training with SD Turbo.Source: §1 · p. 3
- The proposed changes include SD Turbo weight initialization, an extra clamped CLIP loss, and training on a larger text-prompt dataset.The paper presents these changes as ways to improve one-step student performance and teacher-student knowledge transfer.Source: §1 · p. 3
- 8.77 FID-30K for the merged student surpasses its SDv2.1 teacher, and extra image regularization further lowers the FID score to 8.14.The authors report the merged model also exceeds GAN-based text-to-image approaches.Source: §1 · p. 3
- A100 40GB and affordable-hour constraints motivate full-student training without the extra loss and LoRA training with the clamped CLIP loss.Both schemes produce high-quality models that surpass prior one-step diffusion approaches on most metrics.Source: §1 · p. 3
2 Related Work
Text-to-image generation spans several model families, while diffusion models are popular for image quality but commonly require many sampling steps. Work on accelerated diffusion includes one-step distillation, with SwiftBrush distinguished by prompt-only, image-free training.
2.1 Text-to-Image Generation
- Text-to-image approaches include autoregressive, mask-based transformer, GAN-based, and diffusion models, alongside vision-language models and large datasets.The related-work overview traces the task’s shift from constrained domains such as CUB and COCO to general domains such as LAION-5B.Source: §2.1 · pp. 3–4
- Diffusion models are popular for high-quality image generation, but their many-step sampling limits real-time and on-device applications.Source: §2.1 · p. 4
2.2 Accelerating Text-to-Image Diffusion Models
- One-step diffusion methods distill a pretrained multi-step model using techniques such as Rectified Flows, reconstruction and distribution matching losses, or adversarial objectives.The cited methods often still underperform multi-step models and require large-scale text-image pairs.Source: §2.2 · p. 4
- SwiftBrush trains from prompt inputs alone and transfers teacher knowledge through an intermediate LoRA teacher, allowing its image-free training data to scale.The authors note that SwiftBrush still falls short of teacher quality, motivating further development.Source: §2.2 · p. 4
3 Background
The background describes diffusion as a forward noise-adding process followed by a learned reverse process. It then introduces classifier-free guidance and SwiftBrush’s two-teacher, VSD-based distillation setup.
3 Background
- The forward process adds noise to a clean image over T steps, while the reverse process learns to reconstruct the original data from noise.The forward-process expression is given as x_t = α_t x_0 + σ_t ε.Source: §3 · p. 4
- Classifier-free guidance combines text-conditioned and unconditional model predictions, with guidance scale γ controlling the conditioning contribution.The unconditional prediction uses null text.Source: §3 · p. 5
- SwiftBrush uses frozen and LoRA teachers to guide a one-step student with Variational Score Distillation, alternating student and LoRA-teacher training until convergence.The LoRA teacher is trained with diffusion loss and bridges the frozen teacher and student.Source: §3 · p. 5
4 Proposed Methods
The section examines the quality-diversity trade-off among one-step models and develops ways to combine SwiftBrush’s diversity with SD Turbo’s image quality. It proposes initialization, training, and weight-merging strategies for improving SwiftBrush while addressing resource and image-text alignment challenges.
4.1 Quality-Diversity Trade-off in Existing Models
- SwiftBrush produces more diverse but lower-quality outputs than SD Turbo, whose realistic outputs have low recall; both one-step models have FID scores around 15–16.For SDv2.1, cfg 2 gives FID 9.64 and recall 0.53, while cfg 7.5 gives CLIP score 0.33, recall 0.36, and FID 15.93; cfg 4.5 gives the highest precision.Source: §4.1 · p. 6, Table 1 · p. 6
4.2 SwiftBrush and SD Turbo Integration
- Initializing SwiftBrush with SD Turbo weights, without using its adversarial training, improves FID and recall while image-free training gradually enhances diversity.The authors avoid combining adversarial and score-distillation training because of computational demands, resource requirements, and the risk of constraining diversity.Source: §4.2 · p. 7
4.3 In-training Improvements
- Adding 2M LAION prompts to JourneyDB’s 1.5M prompts improves FID and precision but slightly reduces recall, while text alignment still has room to improve.The authors describe image-free training as enabling scalable prompt datasets and auxiliary losses.Source: §4.3 · p. 7, Table 2 · p. 8
4.5 Post-training improvements
- SwiftBrush v2 merges a fully finetuned VSD-only student with a LoRA-finetuned student trained using VSD and CLIP losses.The two training schemes produce models with different behaviors, which the authors use as ingredients for the final model.Source: Figure 2 · p. 8, §4.5 · p. 9
4.3 In-training Improvements
- Clamped CLIP loss uses a ReLU threshold and a schedule that reduces its weight to zero, aiming to improve text alignment while preserving image quality.The threshold limits pressure to increase image-text cosine similarity at the expense of image quality.Source: §4.3 · p. 8
4.4 Resource-efficient training schemes
- LoRA limits fine-tuning to small-rank parameters, and TinyVAE reduces decoder costs while preserving structure and object identity, though some fine detail is lost.The approach aims to keep training efficiency close to full fine-tuning while accommodating CLIP loss under resource constraints.Source: §4.4 · p. 9
4.5 Post-training improvements
- Linear weight interpolation between SD Turbo and SwiftBrush improves all evaluated metrics except CLIP score at some optimal point, without increasing model size or inference costs.Figure 3 reports FID, CLIP score, precision, and recall on zero-shot MS COCO-2014; its endpoints are SD Turbo at 0.0 and original SwiftBrush at 1.0.Source: §4.5 · p. 9, Figure 3 · p. 10
5 Experiments
The experiments evaluate SwiftBrush v2 on zero-shot MS COCO-2014 and compare its quality, diversity, and text alignment with prior models. The ablations examine initialization, training data, and clamped CLIP loss, with the final model reaching an FID of 8.14 after image regularization.
5.1 Experimental Setup
- Evaluation uses zero-shot MS COCO-2014, with FID for image quality, precision, recall, CLIP score, and HPSv2 for complementary measures.Training prompts comprise 1.3M JourneyDB prompts, with an expanded version adding 2M LAION prompts; image regularization is optional.Source: §5.1 · p. 10
5.2 Comparison with Prior Approaches
- SD Turbo has the best image quality among the counterparts but the worst diversity, while the method produces varied views, colors, and environments.The comparison describes SD Turbo samples as repeating similar details and colors across samples for a prompt.Source: §5.2 · p. 12, Figure 6 · p. 13
- In a survey of 250 participants across 20 prompts, the method matched or exceeded SDv2.1 in image quality and diversity.Participants compared images generated by the method and its multi-step teacher.Source: §5.2 · p. 12, Figure 4 · p. 12
5.3 Ablation Studies
- Initializing from SD Turbo significantly improves FID in both full and efficient training; adding LAION prompts yields FID scores of 11.27 and 11.02, respectively.Combining the two training schemes reduces FID to 8.77, and image regularization further lowers it to 8.14.Source: §5.3 · pp. 12–13, Table 5 · p. 14
5.2 Comparison with Prior Approaches
- FID 8.14 after image regularization; the paper reports that its approach outperforms competitors and its distilled students surpass SDv2.1 in FID.The comparison covers GAN-based, multi-step diffusion, and distillation methods on zero-shot MS COCO-2014.Source: §5.2 · p. 11, Table 3 · p. 11, §5.3 · p. 13
5.3 Ablation Studies
- Naive CLIP loss worsens FID by 5 points, whereas clamped CLIP loss reduces it to 13.19 and a weight scheduler further improves it to 11.70.The CLIP-loss and resource-efficient training ablation is summarized in Table 6.Source: §5.3 · pp. 13–14, Table 6 · p. 14
6 Discussion and Conclusion
The paper presents SwiftBrush v2 as a one-step text-to-image model that combines SD Turbo initialization, efficient training, clamped CLIP loss, and optional image regularization. Its reported FID is 8.14, while compositional problems inherited from the teacher remain.
Conclusion
- The fused model reaches an FID of 8.14 and is reported to surpass existing GAN-based and one-step diffusion approaches while maintaining near real-time inference.The method uses weight merging and optional image regularization.Source: §6 · p. 14
- The approach enhances SwiftBrush with SD Turbo initialization, efficient in-training techniques, margin CLIP loss, and large-scale dataset training.The final model is produced through weight merging, with image regularization as an option.Source: §6 · p. 14
Limitations
- The distilled model still inherits limitations from its teacher, including compositional problems; the paper points to cross-attention auxiliary losses as future work.Source: §6 · p. 14
7 Additional details
Additional analyses examine how model interpolation, CLIP-loss design, initialization, and training cost affect SwiftBrush v2. They report a useful interpolation weight, improved image quality with the proposed CLIP-loss treatment, and near real-time inference speed.
7.1 Weight Interpolation
- At interpolation weight λ = 0.5, the merged models optimize FID without trading off too much with other metrics; both FID and recall improve when the models are fused.CLIP and precision scores change monotonically between the two models, and the weight is selected to optimize FID while balancing other metrics.Source: §7.1 · p. 19, Figure 7 · p. 19
7.2 CLIP loss
- The proposed CLIP-loss integration yields better image quality and diversity than naive CLIP loss, whose outputs show over-saturation, over-smoothing, and textual artifacts.The comparison is shown qualitatively in Fig. 9.Source: §7.2 · pp. 20–21, Figure 9 · p. 21
7.3 Dependency on the existing one-step diffusion
- With SwiftBrush initialization, the merged model reaches FID 11.69 without image regularization and 11.17 with it, though its gain is less significant than when an existing one-step model is available.The paper reports that this setting still reaches state-of-the-art one-step performance and remains more efficient than other one-step distillation methods, even counting SwiftBrush training time.Source: §7.3 · p. 21, Table 7 · p. 22
7.4 Compare to SD Turbo training with SB initialization.
- SD Turbo training requires a large text-image dataset, lacks publicly available training code, and is prone to mode collapse; SwiftBrush training is image-free and easier to reimplement.The authors report mode-collapse issues in early epochs when reimplementing SD Turbo training.Source: §7.4 · p. 22, Figure 10 · p. 22
7.5 Training Cost and Inference Speed
- The method retains SwiftBrush’s 0.13-second inference time while improving FID, with a slightly longer training time than SwiftBrush.Inference timing was rerun under a standardized procedure; reported training times for most other methods came from their papers, which did not specify GPU memory size.Source: §7.5 · pp. 22–23, Table 8 · p. 23
8 Analysis and further applications
Further analyses test auxiliary HPSv2 guidance, composition, latent and prompt interpolation, and arbitrary image sizes. They show both capabilities and boundaries, including persistent compositional difficulty and an inherited architecture constraint on image dimensions.
8.1 About the robustness of the training scheme
- Adding HPSv2 loss improves zero-shot HPSv2 performance while maintaining other metrics, though the authors note slight FID trade-offs.The loss is integrated into the resource-efficient distillation process that already includes CLIP loss.Source: §8.1 · pp. 24–25
8.2 Composition
- The model has low compositional ability and text-image misalignment for prompts requiring multiple objects with attributes; iterative Divide-and-Bind latent updates improve prompt-faithful generation.The authors use 100 latent updates with the original method's scale size of 20, while noting that the optimization is slower than one-step inference but faster than the multi-step teacher.Source: §8.2 · p. 25
8.3 Latent Interpolation
- Spherical interpolation between random input noises produces gradual visual transformations while preserving adherence to the prompt.Source: §8.3 · p. 25
8.4 Prompt Interpolation
- Interpolating prompt embeddings that differ by one word yields intermediate images blending characteristics of both prompts.Source: §8.4 · pp. 25–26
8.5 Arbitrary size and aspect ratio generation with ScaleCrafter
- The SDv2.1-based architecture limits image sizes and aspect ratios; applying ScaleCrafter's U-Net re-dilation enables generation at varied ratios and resolutions.The paper presents this as another application of multi-step diffusion techniques to its one-step model.Source: §8.5 · p. 26
9 More qualitative results
This section adds qualitative comparisons with SD Turbo, SDv2.1, and SwiftBrush, and illustrates image diversity and additional uncurated samples.
9 More qualitative results
- The section adds comparisons with SD Turbo, SDv2.1, and SwiftBrush, alongside diversity examples and uncurated samples.Fig. 16 presents the comparisons, Fig. 17 illustrates diversity across object categories, styles, and scenes, and Fig. 18 shows additional samples.Source: §9 · p. 26, Figure 16 · p. 29, Figure 17 · p. 30