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

TL;DR

D-OPSD introduces an on-policy self-distillation training method for step-distilled diffusion models, allowing continuous fine-tuning without losing their efficient inference capabilities.

Contribution

The paper proposes a novel on-policy self-distillation paradigm for fine-tuning step-distilled diffusion models while preserving their few-step inference efficiency.

Findings

Enables models to learn new concepts without sacrificing inference speed.

Formulates training as an on-policy self-distillation process.

Leverages encoder's in-context capabilities for improved training.

Abstract

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.