PaperlayerView exact revision on arXiv

Source-linked AI summary

D-OPSD: On-Policy Self-Distillation for Continuously Tuning Step-Distilled Diffusion Models

Dengyang Jiang, Xin Jin, Dongyang Liu, Zanyi Wang, Mingzhe Zheng, Ruoyi Du, Xiangpeng Yang, Qilong Wu, Zhen Li, Peng Gao, Harry Yang, Steven Hoi

arXiv:2605.05204v3cs.CV

TL;DR

Continually fine-tuning step-distilled diffusion models while preserving their few-step capability remains unclear because conventional supervision targets states outside the model’s own sampler trajectories. D-OPSD uses on-policy self-distillation to adapt these models, and experiments show it learns new concepts and styles while preserving few-step generation quality and prior knowledge.

  • Problem

    Continually fine-tuning step-distilled diffusion models while preserving their original few-step inference capability remains unclear because standard supervision uses externally induced states rather than the model’s own sampler states.

  • Method

    D-OPSD distills multimodal teacher predictions into a text-conditioned student along the student’s own few-step roll-outs without external rewards or auxiliary modules.

  • Results

    D-OPSD learns new concepts, styles, and domain preferences while preserving few-step generation quality and prior knowledge across LoRA and full fine-tuning settings.

  • Takeaways & Limitations

    D-OPSD supports continual adaptation of step-distilled diffusion models while retaining their original few-step inference capability and generalizing acquired knowledge across unseen prompts.

  • Takeaways & Limitations

    D-OPSD’s success depends on the base model’s in-context abilities to provide meaningful multimodal supervision, without which training can fail.

Abstract

from arXiv · show

The landscape of high-performance image generation models is currently shifting from the inefficient multi-step ones to the efficient few-step counterparts (e.g, Z-Image-Turbo and FLUX.2-klein). However, these models present significant challenges for direct continuous supervised fine-tuning. For example, applying the commonly used fine-tuning technique would compromise their inherent few-step inference capability. To address this, we propose D-OPSD, a novel training paradigm for step-distilled diffusion models that enables on-policy learning during supervised fine-tuning. We first find that the modern diffusion models, where the LLM/VLM serves as the encoder, can inherit its encoder's in-context capabilities. This enables us to formulate the training as an on-policy self-distillation process. Specifically, during training, we make the model act as both the teacher and the student with different contexts, where the student is conditioned only on the text feature, while the teacher is conditioned on the multimodal feature of both the text prompt and the target image. Training minimizes the two predicted distributions over the student's own roll-outs. By optimizing on the model's own trajectory and under its own supervision, D-OPSD enables the model to learn new concepts, styles, etc., without sacrificing the original few-step capacity.

1 Introduction

Step-distilled diffusion models reduce sampling cost but are difficult to continuously fine-tune without disrupting their few-step dynamics. D-OPSD addresses this by combining on-policy self-distillation with paired image-text supervision on the model’s own roll-outs, enabling new knowledge while preserving few-step inference.

  • Motivation: Step-distillation reduces the latency and computational cost of text-to-image generation, but continuous fine-tuning of these models remains unclear.Conventional text-to-image sampling typically requires numerous iterative denoising steps, motivating step-distillation techniques.
  • Motivation: Standard supervised fine-tuning uses externally induced noised target-image states, whose mismatch with distilled few-step dynamics can degrade inference quality.The passage states that standard SFT often compromises the model’s original distilled few-step ability to generate high-quality images.
  • Method: D-OPSD combines on-policy updates on the model’s own roll-outs with supervision from paired image-text data on those visited states.This formulation follows on-policy self-distillation, retaining on-policy learning while avoiding explicit reward design.
  • Method: Modern diffusion models with LLM/VLM encoders can inherit encoder in-context capability, allowing multimodal target-image and text features to guide the teacher role.The paper reports this capability empirically in Figure 1 using Z-Image Turbo with 8 steps.
  • Results: D-OPSD supports LoRA training on small customized datasets and full fine-tuning on larger datasets while acquiring concepts or styles and preserving few-step inference.The acquired knowledge also generalizes across unseen prompts rather than merely overfitting to the training pair.

2 Method

D-OPSD adapts on-policy self-distillation to continuously tune step-distilled diffusion models while preserving few-step inference. It uses text-only student roll-outs, multimodal teacher supervision, and trajectory-consistent optimization to reduce train–test mismatch.

  • Motivation: D-OPSD trains the student on states from its own few-step denoising trajectories rather than noised target images or externally supplied target velocities.This directly addresses the mismatch introduced by vanilla supervised fine-tuning.
  • On-policy formulation: The method inherits on-policy learning because both optimization states and supervision are computed along trajectories produced by the current student policy.This avoids borrowing supervision from an external offline distribution and reduces the mismatch between training and inference.
  • Condition construction: The student receives a text-only feature, while the teacher receives a multimodal feature encoding the prompt and target image to provide stronger concept, style, or preference supervision.The student pathway therefore matches the original text-to-image inference process, whereas the teacher has access to additional target information.
  • Training objective: For each training pair, D-OPSD samples a few-step student trajectory, evaluates teacher and student velocities on the same states, and updates the student by minimizing their discrepancy.The teacher uses the richer multimodal context, while stop-gradient prevents its predictions from receiving the student update.
  • Inference: After training, the teacher branch is discarded, and inference uses the unchanged few-step text-to-image pathway.The framework defines the inference trajectory with the same few-step solver used at test time, such as 4 or 8 steps.

3 Experiment

Experiments evaluate D-OPSD with step-distilled Z-Image-Turbo and FLUX.2-klein against representative fine-tuning baselines using fidelity, generalization, quality, and aesthetic metrics. Across LoRA and full-finetuning settings, D-OPSD learns target concepts or domains while preserving few-step generation and prior knowledge.

  • Experimental setup: Experiments use Z-Image-Turbo 6B and FLUX.2-klein 4B as baseline models, with original step-distilled inference settings retained across methods.Evaluation includes DINO-D, LPIPS-D, FID, VLM-J, CLIP-S, Quality-S, and Aesthetic-S.
  • Experimental setup: Compared methods include Vanilla SFT, SFT with LoRA on the distilled model, DreamBooth, and PSO.Two-stage strategies are not shown.
  • D-OPSD for LoRA training on small customized dataset: In small-dataset LoRA training, D-OPSD learns new concepts from only a few image–text pairs while generalizing beyond the training set.Experiments use the DreamBooth dataset together with a small amount of stylized data; D-OPSD substantially outperforms SFT-style training.
  • D-OPSD for LoRA training on small customized dataset: Vanilla SFT and DreamBooth lose few-step image quality, whereas PSO captures target concepts but overfits and degrades novel-instruction following.These effects are reflected in Quality-S, Aesthetic-S, and CLIP-S, alongside blurry outputs for SFT and DreamBooth.
  • D-OPSD for full finetuning on larger scale dataset: In full finetuning on anime data, D-OPSD substantially outperforms the base model and other methods on FID, DINO-D, and LPIPS-D.It adapts to the new distribution while retaining original knowledge and few-step inference ability.
  • D-OPSD for full finetuning on larger scale dataset: D-OPSD is preferred over PSO, Vanilla SFT, and the base model for image quality, aesthetics, and prompt following in the user study.Its preference margin is particularly large against SFT, and it also improves prompt following over the base model.
  • Effect of on-policy self-distillation: Ablations show that vanilla SFT gradually impairs few-step image quality, while on-policy distillation corresponds to the proposed method.The study compares target-image SFT, teacher-sample SFT, off-policy distillation, and on-policy distillation.
  • Construction of the teacher model: Using the frozen base model as teacher provides stable, effective training, whereas directly using a student copy causes training collapse.EMA-based teachers require a large momentum coefficient to stabilize training.

4 Discussion on Limitations and Future Works

D-OPSD incurs substantially higher training cost than vanilla SFT and depends on the base model producing meaningful teacher supervision under multimodal conditioning. Future work centers on richer teacher contexts, additional training targets or constraints, and multi-expert OPD.

  • Computation cost: D-OPSD requires student roll-outs and teacher inference during training, costing roughly 4× FLOPs and 2× training time per iteration versus vanilla SFT.The authors consider this cost acceptable because vanilla SFT degrades few-step generation capability.
  • Requirements for teacher capability: D-OPSD’s success depends on the base model’s in-context abilities to generate meaningful supervision from the target image and text prompt.If the diffusion teacher cannot provide meaningful supervision under multimodal conditioning, training fails.
  • Future Works: Future work could construct richer teacher-side contexts by incorporating stronger conditional signals from image-editing or video-generation models.The paper identifies image editing and video generation models as possible sources of stronger teacher conditioning.
  • Future Works: Future work could explore other training targets, including combinations of D-OPSD with additional training constraints.The paper presents this as a separate direction from improving teacher-side conditioning.
  • Future Works: Future work could investigate introducing multi-expert OPD into the post-training stage.This is listed as a third direction for extending the framework.

5 Related Work

Prior work accelerates diffusion inference through timestep distillation, typically by distilling trajectories or distributions. In language models, on-policy self-distillation addresses off-policy training mismatch by letting the model serve as its own teacher through in-context capabilities.

  • Step-distilled diffusion models: Timestep distillation compresses original diffusion models into generators capable of few-step sampling to accelerate inference.
  • Step-distilled diffusion models: Existing timestep-distillation methods typically distill either the diffusion trajectory or the model distribution.
  • On-policy self-distillation: On-policy distillation in large language models mitigates train-test mismatch caused by off-policy supervised fine-tuning or knowledge distillation.
  • On-policy self-distillation: On-policy self-distillation removes the need for a stronger external teacher by enabling the model to act as teacher using its own in-context capabilities.

6 Conclusion · Appendix

D-OPSD is presented as an on-policy self-distillation framework for continually tuning step-distilled diffusion models. It uses emergent in-context capability to make one model serve as text-conditioned student and multimodally conditioned teacher during the student’s own few-step roll-outs.

  • 6 Conclusion: D-OPSD is an on-policy self-distillation framework for continually tuning step-distilled diffusion models.The framework is introduced as the paper’s central contribution.
  • 6 Conclusion: The framework builds on an emergent in-context capability in modern diffusion models with LLM/VLM encoders.This capability underpins the method’s teacher–student formulation.
  • 6 Conclusion: The same model serves as a student under text-only conditioning.The student receives text-only conditioning in the proposed formulation.
  • 6 Conclusion: The same model serves as a teacher under stronger multimodal conditioning.The teacher uses multimodal conditioning in the proposed formulation.
  • 6 Conclusion: D-OPSD distills the teacher’s predictions along the student’s own few-step roll-outs.Distillation is performed on trajectories generated by the student.
  • 6 Conclusion: The method formulates training around the student’s own few-step trajectories and the teacher’s multimodally conditioned predictions.This combines on-policy roll-outs with teacher supervision within the same model.

A Investigation of FLUX.2-klein

The FLUX.2-klein investigation shows that multimodal conditioning produces image variations while preserving the target image’s underlying concept or stylistic identity, mirroring the behavior observed in Z-Image-Turbo.

  • A Investigation of FLUX.2-klein: FLUX.2-klein exhibits similar inherited in-context behavior, suggesting applicability across diffusion models using LLM/VLM encoders.The analysis parallels the investigation presented in Figure 1.
  • A Investigation of FLUX.2-klein: Using multimodal features instead of text-only features enables FLUX.2-klein-4B to produce image variations while maintaining the target image’s underlying concept or stylistic identity.This behavior is observed with 4-step generation.

B Discussion and Comparison of Different Training Paradigms

The discussion compares training paradigms for continually tuning step-distilled diffusion models by supervision source, on-policy status, auxiliary reward requirements, and training–inference consistency. D-OPSD is presented as an on-policy approach whose supervision and optimization states are tied to the student’s few-step roll-outs, addressing the stricter train–test alignment needs of step-distilled models.

  • Comparison framework: Table 3 compares representative paradigms by supervision source, on-policy learning, auxiliary reward requirements, and training–inference state consistency.These four dimensions define the design-space comparison.
  • Vanilla SFT: Vanilla SFT uses the target image for supervision and, in flow-matching models, trains on ground-truth-image noised states toward the corresponding ground-truth velocity.The passage characterizes this objective as standard for training diffusion models from scratch but mismatched to continual tuning of step-distilled models.
  • Offline RL-style methods: Offline RL-style methods such as Diffusion-DPO and PSO use preference-style or pairwise supervision derived from fixed datasets rather than the model’s current sampling process.Their inputs use target-image-related states and ground-truth-velocity-based pairwise supervision.
  • Online RL-style methods: Online RL methods such as ReFL and flow-GRPO optimize sampled on-policy trajectories whose resulting images are scored by a reward function or reward model.The passage describes these methods as conceptually more suitable for preserving step-distilled behavior because optimization states follow the model’s own samples.
  • Our method: D-OPSD performs optimization on the student’s own few-step roll-outs, while using the target image only to enrich the teacher’s multimodal in-context condition.Thus, both the optimization states and supervision signal are defined on states the student visits during inference, without introducing the target image as an external denoising target.
  • Why this distinction matters for step-distilled models: Step-distilled models are less forgiving of train–test mismatch than multi-step models because few denoising steps leave less room for error correction.The passage states that even small deviations in learned dynamics can directly harm image quality.

C Comparison with Two-stage Training Pipeline

D-OPSD is compared with direct few-step SFT and a two-stage pipeline that fine-tunes a multi-step model before DMD distillation. The results indicate that D-OPSD preserves few-step behavior and performance more effectively, while adding training overhead relative to vanilla SFT.

  • Experimental setup: The comparison uses open-source Z-Image as the multi-step model and Z-Image-Turbo as its few-step counterpart, with DMD for the distillation stage.The two-stage pipeline applies multi-step SFT followed by step-distillation, using LoRA or full-model updates depending on the setting.
  • Direct few-step SFT: Direct SFT on Z-Image-Turbo substantially weakens few-step generation, especially with LoRA, because it optimizes target-image-induced states rather than the model sampler’s visited states.The passage specifically notes a drop in Quality-S for the LoRA setting, though the supplied text does not provide the table value.
  • Two-stage training pipeline: The two-stage pipeline reduces inference to 8 NFEs but incurs substantial extra training cost and yields lower Quality-S and GenEval scores than D-OPSD.DMD distillation after multi-step SFT does not recover the overall performance of the tuned model and degrades few-step quality and retained prior knowledge.
  • D-OPSD: D-OPSD performs continual tuning in one stage while optimizing and supervising on student roll-outs, preserving the original few-step sampling behavior.Its on-policy design learns target-domain information without switching optimization to target-image-induced states.
  • Training cost: 4× FLOPs and about 2× wall-clock training time per iteration make D-OPSD more expensive than vanilla SFT, a cost deemed acceptable for continual tuning.The overhead comes from on-policy student roll-outs and teacher inference during training.

D Implementation Details.

D-OPSD uses distinct text-only and multimodal encoder settings for student and teacher models, with Qwen3-based encoders for both Z-Image-Turbo and FLUX.2-klein. Training supports LoRA on small datasets and full backbone finetuning on larger datasets, with specified batch sizes, learning rates, EMA decay, iterations, and hardware.

  • Encoder settings of student and teacher: The student uses the original text encoder, while the teacher uses Qwen3-VL-4B to incorporate target-image information for both diffusion models.For Z-Image-Turbo and FLUX.2-klein, prompts use Qwen3-4B in the student; replacing it with Qwen3-VL-4B for the teacher can create inference-training feature-space mismatch.
  • LoRA training on small customized dataset: LoRA finetuning uses rank 64, alpha 128, batch size 4, EMA decay 0.9999, and 1K iterations on one H800 GPU for both models.The learning rate is 4e-55 for Z-Image-Turbo and 1e-5 for FLUX.2-klein; LoRA-only EMA saves memory by operating on LoRA weights.
  • Full finetuning on larger scale dataset: For full finetuning, Z-Image-Turbo unlocks all diffusion-transformer-backbone parameters, using batch size 256, learning rate 3e-5, and 10k iterations on 32 H800 GPUs.EMA decay is set to 0.9999 by default.

E More Details of Evaluation

The evaluation uses separate protocols for small-scale LoRA customization, large-scale full fine-tuning, and human assessment. It measures concept or style fidelity, prompt alignment, perceptual similarity, image quality, and distributional quality using automated metrics and user studies.

  • LoRA training on a small customized dataset: Small-scale evaluation averages scores across 30 concept classes and 10 style classes using DreamBooth-style text–image pairs and LoRA fine-tuning.The setup follows common secondary fine-tuning with fewer than 10 pairs per concept or style.
  • LoRA training on a small customized dataset: DINO-D measures cosine distance between DINOv3-ViT-S-plus features of generated images and corresponding target images after caption paraphrasing.Paraphrased captions preserve core semantics while changing the training prompts used for generation.
  • LoRA training on a small customized dataset: VLM-J evaluates subject or style consistency on four groups of novel prompts that retain the learned concept while varying scenes, compositions, or image contents.The prompts are generated by an LLM and differ from the training prompts.
  • LoRA training on a small customized dataset: CLIP-S measures image–text alignment for the same generated images, replacing the DreamBooth class token [V] with the original class name.DFN-CLIP-H is used for this metric.
  • Full finetuning on larger scale dataset: Large-scale full fine-tuning uses an in-house dataset of 25K high-quality anime images, while FID compares generated and training-image feature distributions from 2K sampled examples.FID features are extracted with Inception-v3 under a multivariate-Gaussian assumption.
  • User study: The user study evaluates 100 prompts total: 50 from the small-scale LoRA setting and 50 from the large-scale full-finetuning setting.Small-scale evaluations additionally provide reference images or style exemplars to annotators.

F More Related Works

The related-work landscape spans diffusion-based image generation, knowledge distillation—including self-distillation—and fine-tuning pretrained diffusion models for new concepts, styles, and domains.

  • Knowledge distillation for diffusion model: Knowledge distillation in diffusion models includes step-distillation for faster sampling and guidance from more powerful pretrained models during training.The passage distinguishes step-distillation from other knowledge-distillation approaches used in diffusion models.
  • Knowledge distillation for diffusion model: Self-distillation frameworks improve diffusion models without external stronger models or other external components.The cited example, SRA, aligns earlier-layer high-noise representations with later-layer low-noise representations to enhance representation learning and accelerate convergence.
  • Diffusion model fine-tuning: Diffusion-model fine-tuning adapts pretrained models to new concepts, styles, or downstream domains through full-parameter supervision, subject-driven customization, or parameter-efficient methods.Representative methods include DreamBooth, textual inversion, and LoRA; full-parameter supervised fine-tuning remains common with sufficient paired data.
Loading 2605.05204v3…