ReDiF: Resource-Efficient Few-Step Diffusion Distillation via Reinforcement Learning

Amirhossein Tighkhorshid; Zahra Dehghanian; Hamid R. Rabiee

cs.LGDec 28, 2025

ReDiF: Resource-Efficient Few-Step Diffusion Distillation via Reinforcement Learning

Authors: Amirhossein Tighkhorshid, Zahra Dehghanian, Hamid R. Rabiee

Organizations: Computer Engineering Sharif University of Technology

Abstract

Step distillation accelerates diffusion sampling by training a few-step student to imitate a many-step teacher, but distillation itself remains expensive. Typically, this requires thousands of GPU-hours and a large pre-generated trajectory dataset. We introduce ReDiF, which casts step distillation as terminal-reward policy optimization rather than step-wise regression. The student is optimized against a reward computed on the terminal sample, measuring alignment with the teacher's output, instead of matching the teacher's intermediate trajectory under a reconstruction or consistency loss. Because the reward need not be differentiable or trajectory-aligned, ReDiF admits non-differentiable objectives, multi-objective combinations, and preferences the teacher does not express, while exploration lets the student find sampling paths matched to its own step schedule. ReDiF converges in 400 policy updates with 3200 rollouts on a single A100 GPU with 1,000 noise-class pairs and no paired dataset: about 4 GPU-hours, against roughly 336 A100-hours reported for DMD2 on the same EDM teacher. At 8 steps on ImageNet-64, it achieves an FID 3.64 points better than the strongest retrained distillation baseline in the low-training regime under the same single-GPU budget. The formulation is also orthogonal to existing distillation objectives: added to the DMD2 loss, it further improves DMD2's FID by 4.4.

Figures & tables

Panel of Figure 1 — Figure 1: ReDiF training on the Stable Diffusion v1.5 branch ( $50\rightarrow 5$ steps). Within a few epochs a student sampling at a tenth of the teacher’s step count reaches a quality close to the teacher’s reference (rightmost column), although the only supervision is a scalar computed on the final image. Columns are successive epochs of one run at fixed conditioning. Figure 4 continues this run past the budget.

Figure 2: Overview of ReDiF. The same noise and conditioning are provided to both the frozen teacher and the trainable student. Their final images are embedded using frozen, pretrained encoders, and the similarity between the embeddings serves as the terminal reward, which PPO maximizes by updating only the student. Everything to the right of the two generators is frozen and operates in a forward-only manner, while the teacher reference is generated by the teacher’s own sampler within the training loop. Aesthetic and Text-Image Alignment are two auxiliary preference rewards.

Method	Steps	Batch	GPU-h	PFLOPs	P $\uparrow$	R $\uparrow$	D $\uparrow$	C $\uparrow$	FID $\downarrow$	Pub. FID
(a) EDM teacher, ImageNet-64 ( $256\rightarrow 8$ )
Teacher (256 steps, reference)	256	–	–	–	0.8310	0.7103	0.8930	0.7831	13.74	1.36
Teacher (8 steps)	8	–	–	–	0.5308	0.4780	0.3910	0.5326	35.04	–
Progressive Distill. ( Salimans and Ho, 2022 )	$256\!\rightarrow\!8$	32	$3.39$	$1137.6$	$0.5995\pm 0.0017$	$0.5591\pm 0.0073$	$0.6171\pm 0.0010$	$0.6770\pm 0.0001$	$38.02\pm 0.19$	15.39
Consistency Distill. ( Song et al., 2023 )	$256\!\rightarrow\!8$	32	$3.30$	$1138.5$	$0.7110\pm 0.0028$	$0.5387\pm 0.0190$	$0.4004\pm 0.0043$	$0.6219\pm 0.0051$	$22.24\pm 0.68$	6.20
DMD ( Yin et al., 2024b )	$256\!\rightarrow\!8$	128	$3.25$	$1191.4$	$0.7327\pm 0.0043$	$0.4596\pm 0.0072$	$0.7751\pm 0.0031$	$0.7602\pm 0.0043$	$32.83\pm 0.66$	2.62

Table 1: The low-training regime. All rows are trained by us on one A100, from the same teacher initialization and for the same number of epochs, under an otherwise identical protocol (Appendix P ). PFLOPs is the analytical cost of reproducing a row from scratch, including the teacher-sample pool that is synthesized once and read by every row and that dominates it; Appendix C defines it. GPU-h is wall-clock on the same device; cost columns are not ranked. Pub. FID is the authors’ own number under their protocol and step count and is not comparable to the FID column. Mean

\pm

sample sd, last checkpoint; per branch, best in bold, second best bold and dash-underlined, teacher rows unranked.

Method	Mode	Aux. net	PFLOPs	P $\uparrow$	R $\uparrow$	D $\uparrow$	C $\uparrow$	FID $\downarrow$
Progressive Distill. ( Salimans and Ho, 2022 )	–	interm. teachers	$1137.6$	$0.5995\pm 0.0017$	$\mathbf{0.5591\pm 0.0073}$	$0.6171\pm 0.0010$	$0.6770\pm 0.0001$	$38.02\pm 0.19$
+ RL (ours)	structural	interm. teachers	$1342.4$	$\mathbf{0.7321\pm 0.0424}$	$0.5250\pm 0.1136$	$\mathbf{0.6472\pm 0.0601}$	$\mathbf{0.7950\pm 0.0451}$	$\mathbf{34.08\pm 2.71}$
Consistency Distill. ( Song et al., 2023 )	–	EMA target	$1138.5$	$0.7110\pm 0.0028$	$0.5387\pm 0.0190$	$0.4004\pm 0.0043$	$0.6219\pm 0.0051$	$\mathbf{22.24\pm 0.68}$
+ RL (ours)	structural	EMA target	$1213.5$	$\mathbf{0.7620\pm 0.0221}$	$\mathbf{0.5932\pm 0.0520}$	$\mathbf{0.6490\pm 0.0680}$	$\mathbf{0.8160\pm 0.0300}$	$22.68\pm 6.80$
DMD2 ( Yin et al., 2024a )	–	score net + disc.	$1365.9$	$0.7791\pm 0.0039$	$0.6162\pm 0.0058$	$0.5349\pm 0.0016$	$0.7881\pm 0.0034$	$25.67\pm 0.45$
+ RL (ours)	loss	score net + disc.	$1420.3$	$\mathbf{0.7982\pm 0.0126}$	$\mathbf{0.6915\pm 0.0109}$	$\mathbf{0.7672\pm 0.0171}$	$\mathbf{0.8742\pm 0.0051}$	$\mathbf{21.32\pm 0.86}$

Table 2: ReDiF’s PPO objective combined with existing methods on the EDM teacher, ImageNet-64 (

256\rightarrow 8

), under the same epoch budget as Section 4.2 . Each “+RL” row keeps its baseline’s architectural modification. PFLOPs is total analytical training compute, teacher forwards included. Mean

\pm

sample sd over seeds, last checkpoint. The DMD2 rows keep the auxiliary score network and are not cost-comparable to the rest; ReDiF alone reaches

18.60

. Within each baseline/“+RL” pair the better entry is bold; rows are compared within a pair, not across the table.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Trajectory matching
Method	Step Distill.	One-Step	Data-Free	RL	Few Pr.	Non-Diff.	Low-Comp.
Progressive Distill. ( Salimans and Ho, 2022 )	✓	✓	✗	✗	✗	✗	✗
Guided Distill. ( Meng et al., 2023 )	✓	✓	✗	✗	✗	✗	✗
Consistency ( Song et al., 2023 )	✓	✓	✗	✗	✗	✗	✗
LCM / LCM-LoRA ( Luo et al., 2023a ; Luo et al., 2023b )	✓	✓	✗	✗	✗	✗	✗
CTM ( Kim et al., 2024 )	✓	✓	✗	✗	✗	✗	✗

Appendix

Table 3: Comparison of diffusion acceleration and distillation methods. Data-Free indicates that no real images are used for training: SiD and Guided SiD qualify but retain an auxiliary score network; SiDA and ADD do not, because their discriminators are trained against real data; TFD does not, because its drifting field is estimated against a real-data reference set where available; RG-LCD does not, because its consistency loss is computed on real image–text data. The reward-based alignment methods use prompts only. ReDiF uses real images only for evaluation and to fix the training length once (Appendix H ). One-Step indicates that the source paper demonstrates a single-step student; – marks methods for which it is not reported or not applicable, and ReDiF is marked ✗ because it targets the few-step rather than the one-step regime (Appendix M ). Low-Comp. indicates methods that avoid heavy supervised reconstruction and large-scale synthetic datasets.

Figure 3: FID versus training epoch on ImageNet-64 ( $256\rightarrow 8$ , EDM teacher). Markers denote evaluated checkpoints; the dashed line is the $256$ -step teacher ( $\mathrm{FID}=13.74$ ) under the same protocol. The horizontal axis counts epochs; because the rest of the protocol is held fixed, an epoch is the same number of training samples for every curve (Appendix C ).

Panel of Figure 4 — Figure 4: The SD v1.5 student of Figure 1 trained for $50$ epochs, well past the $20$ -epoch budget of the reported COCO runs, for the same conditioning and seed; the rightmost column is the teacher’s reference. Quality stays stable for a few epochs and then degrades as the student exploits directions the reward does not penalize.

Panel of Figure 5 — Figure 5: Teacher and student samples, SD v1.5, first three held-out prompts.

Panel of Figure 6 — Figure 6: Teacher and student samples, SD v1.5, remaining three held-out prompts.

Hyperparameter	EDM / ImageNet-64	SD v1.5 / COCO
Optimizer	AdamW	AdamW
Trainable parameters	LoRA, rank / alpha 4 / 4	LoRA, rank / alpha 4 / 4
LoRA target modules	UNet attention ( $q,k,v,\text{out}$ )	UNet attention ( $q,k,v,\text{out}$ )
Learning rate	$3\times 10^{-4}$	$3\times 10^{-4}$
PPO clip range $\epsilon$	$10^{-4}$	$10^{-4}$
Inner epochs per rollout	4	4

Appendix

Table 4: Hyperparameters used in all reported ReDiF runs. Both branches use AdamW with

\beta_{1}{=}0.9

\beta_{2}{=}0.999

, weight decay

10^{-4}

\epsilon{=}10^{-8}

and gradient-norm clipping at

1.0

. The preference runs of Appendix K use the COCO column with the preference weight added.

Single-encoder semantic rewards
Reward	P $\uparrow$	R $\uparrow$	D $\uparrow$	C $\uparrow$	CLIP-T $\uparrow$	FID $\downarrow$
Teacher (256 steps, reference)	0.8310	0.7103	0.8930	0.7831	–	13.74
Teacher (8 steps)	0.5308	0.4780	0.3910	0.5326	–	35.04
CLIP (R 1 )	$0.7529\pm 0.0180$	$0.6265\pm 0.0213$	$0.6318\pm 0.0291$	$0.8223\pm 0.0169$	$0.8836\pm 0.0028$	$18.84\pm 2.01$
DINOv3	$0.7198\pm 0.0161$	$0.5654\pm 0.0367$	$0.5972\pm 0.0418$	$0.7828\pm 0.0191$	$0.8675\pm 0.0058$	$33.11\pm 9.95$
PE	$0.6992\pm 0.0152$	$0.5685\pm 0.0182$	$0.5323\pm 0.0247$	$0.7452\pm 0.0226$	$0.8522\pm 0.0054$	$25.67\pm 3.62$

Appendix

Table 5: Reward ablation, EDM teacher on ImageNet-64. Identical PPO settings, initialization, conditioning pool, budget and seeds across rows; only the reward differs. Last checkpoint,

M=8

\eta=0

10{,}000

generated samples against

10{,}000

real validation images; the R 1 and R 2 rows are the ReDiF rows of Table 1 . Mean

\pm

sample sd over seeds. CLIP-T is the CLIP image-to-image similarity to the teacher sample; for rewards containing a CLIP term it is part of the training objective and not an independent measure. Best entry in bold, second best bold and dash-underlined; teacher rows are not ranked.

Single-encoder semantic rewards
Reward	P $\uparrow$	R $\uparrow$	D $\uparrow$	C $\uparrow$	FID $\downarrow$	CLIPScore $\uparrow$
Teacher (50 steps, reference)	0.6103	0.2029	0.8878	0.6832	24.59	0.3137
Teacher (5 steps)	0.4420	0.1762	0.6291	0.5813	31.45	0.3041
CLIP (R 1 )	$0.4890\pm 0.0062$	$0.1864\pm 0.0047$	$0.7411\pm 0.0087$	$0.6682\pm 0.0065$	$27.42\pm 0.84$	$0.3038\pm 0.0007$
DINOv3	$0.4310\pm 0.0084$	$0.1809\pm 0.0061$	$0.6406\pm 0.0112$	$0.5903\pm 0.0075$	$27.40\pm 0.91$	$0.3046\pm 0.0004$
PE	$0.4498\pm 0.0072$	$0.1770\pm 0.0057$	$0.6733\pm 0.0104$	$0.5927\pm 0.0069$	$27.97\pm 0.76$	$0.3044\pm 0.0001$

Appendix

Table 6: Reward ablation, Stable Diffusion v1.5 on COCO (

50\rightarrow 5

). Same design as Table 5 : identical optimization, only the reward differs. Last checkpoint,

M=5

5{,}000

COCO validation prompts; the R 1 and R 2 rows are the ReDiF rows of Table 1 . Mean

\pm

sample sd over seeds. Best entry bold in bold, second best bold and dash-underlined; teacher rows are not ranked.

Panel of Figure 7 — Figure 7: Training reward curves for ten of the fourteen reward configurations of Table 5 , under identical PPO optimization, conditioning pool and seeds. All panels are from the Stable Diffusion v1.5 branch ( $50\rightarrow 5$ steps); the corresponding curves on ImageNet-64 differ, and the text below says how. These are $20$ -epoch runs, shorter than the runs reported in Table 5 , and are shown for the shape of the curves rather than for their end points. Solid lines show the mean reward per epoch, shaded regions $\pm 1$ standard deviation across conditioning signals, and the dashed vertical line the reward-maximizing epoch. Panels (a)–(d) are single-encoder rewards, (e)–(h) combinations, and (i)–(j) the pixel and divergence controls.

Panel of Figure 8 — Figure 8: The COCO student trained with the distillation reward plus one non-differentiable preference, for two fixed prompts and seeds: compressibility $R_{\mathrm{jpeg}}$ equation 2 above, sharpness $R_{\mathrm{sharp}}$ equation 4 below. Init is the student before any update. Only the policy’s weights differ between columns.

Student	$\bm{\epsilon}$	LR	P $\uparrow$	R $\uparrow$	D $\uparrow$	C $\uparrow$	FID $\downarrow$
LoRA ( $r=4$ )	$10^{-3}$	$3\times 10^{-4}$	$0.7257\pm 0.0013$	$0.6003\pm 0.0075$	$0.5751\pm 0.0102$	$0.7857\pm 0.0010$	$20.01\pm 0.46$
LoRA ( $r=4$ )	$10^{-4}$	$3\times 10^{-4}$	$0.7391\pm 0.0016$	$0.6145\pm 0.0171$	$0.7829\pm 0.0084$	$0.8043\pm 0.0001$	$18.96\pm 1.66$
Full	$10^{-3}$	$1\times 10^{-5}$	$0.7574\pm 0.0535$	$0.4652\pm 0.1089$	$0.6131\pm 0.1521$	$0.6229\pm 0.2581$	$29.78\pm 0.85$
Full	$10^{-4}$	$1\times 10^{-5}$	$0.7572\pm 0.0147$	$0.5677\pm 0.0250$	$0.6624\pm 0.0137$	$0.8129\pm 0.0199$	$22.22\pm 4.81$

Appendix

Table 7: Student parameterization crossed with the PPO clip range, EDM teacher, ImageNet-64 (

256\rightarrow 8

), R 2 reward,

100

epochs. Mean

\pm

sample sd over two seeds, last checkpoint. The full-parameter student uses a learning rate thirty times smaller than the LoRA student. Best in bold, second best bold and dash-underlined.

Panel of Figure 9 — Figure 9: Effect of the student step count $M$ on the EDM/ImageNet-64 branch under a fixed training budget. Top: FID (left) and Recall (right). Bottom: Coverage (left) and Density (right). Dashed lines mark the $256$ -step teacher and the dotted line the operating point $M=8$ . $n$ is the number of runs at each $M$ and the band is $\pm 1$ s.d. over runs.

(a) ImageNet-64
Method	Family	Steps	Aux. net	Trainable	Batch	Devices / training cost (reported)	FID $\downarrow$
EDM teacher ( Karras et al., 2022 )	–	256	–	–	–	–	1.36
Progressive Distill. ( Salimans and Ho, 2022 )	trajectory	$N\!\rightarrow\!1$	interm. teachers	full	2048	$\sim$ 64 GPUs; $\sim$ 5,533 A100-h †	15.39 ‡
Consistency Distill. ( Song et al., 2023 )	trajectory	$N\!\rightarrow\!1$	target copy	full	2048	$\sim$ 64 GPUs; $\sim$ 7,867 A100-h †	6.20 ‡
CTM ( Kim et al., 2024 )	trajectory	$N\!\rightarrow\!1$	target copy + disc.	full	2048	$\sim$ 64 GPUs; $\sim$ 902 A100-h †	1.92
DMD ( Yin et al., 2024b )	distribution	$N\!\rightarrow\!1$	score net	full	n/r	n/r	2.62

Appendix

Table 8: Published results at convergence. FID as reported by the original authors, except where marked: ImageNet-64 with

50

k balanced samples, COCO-2014 zero-shot with

30

k samples. Batch is the effective training batch and Devices the reported hardware, both as disclosed;

\dagger

marks a third-party A100-hour estimate from Zhou et al. (2024b) ;

\ddagger

marks an FID we quote from Song et al. (2023) , who evaluate progressive and consistency distillation on this teacher, rather than from the method’s own paper;

\ast

marks the batch of the default ablation setting reported by Sauer et al. (2024) , which the paper does not state to be the setting of its final run; n/r not reported; – no number under this protocol. ReDiF has not been evaluated under these protocols, so its FID is omitted; its numbers under the single-device protocol are in Table 1 . Results published only under other COCO protocols are excluded rather than converted. Best published FID in bold, second best bold and dash-underlined; teacher rows are not ranked.

Branch	Method	LR	Aux. LR	Target type
ImageNet-64	Progressive Distillation	$3\times 10^{-4}$	–	fixed teacher target
	Consistency Distillation	$3\times 10^{-4}$	–	fixed teacher target
	DMD	$1\times 10^{-4}$	$3\times 10^{-4}$	moving target (aux. network)
	DMD2	$1\times 10^{-4}$	$3\times 10^{-4}$	moving target (aux. network)
	SiD	$1\times 10^{-4}$	$3\times 10^{-4}$	moving target (aux. network)
	SiDA	$1\times 10^{-4}$	$3\times 10^{-4}$	moving target (aux. network)

Appendix

Table 9: Learning rate used for each retrained baseline in Table 1 . Values are taken from each method’s own paper where it reports a setting for this teacher, and otherwise from the work reporting its strongest published result; none was selected by us. Objectives that regress onto a fixed teacher target take a larger rate than those whose target moves with an auxiliary network. Direct reward backpropagation on COCO uses the ReDiF rate because it optimizes the same reward. “Aux.” is the rate of the fake-score network where one exists. All students are LoRA adapters trained with AdamW.

Explore similar work

Apr 21, 2026cs.LG

Guiding Distribution Matching Distillation with Gradient-Based Reinforcement Learning

Diffusion distillation, exemplified by Distribution Matching Distillation (DMD), has shown great promise in few-step generation but often sacrifices quality for sampling speed. While integrating Reinforcement Learning (RL) into distillation offers potential, a naive fusion of these two objectives relies on suboptimal raw sample evaluation. This sample-based scoring creates inherent conflicts with the distillation trajectory and produces unreliable rewards due to the noisy nature of early-stage generation. To overcome these limitations, we propose GDMD, a novel framework that redefines the reward mechanism by prioritizing distillation gradients over raw pixel outputs as the primary signal for optimization. By reinterpreting the DMD gradients as implicit target tensors, our framework enables existing reward models to directly evaluate the quality of distillation updates. This gradient-level guidance functions as an adaptive weighting that synchronizes the RL policy with the distillation objective, effectively neutralizing optimization divergence. Empirical results show that GDMD sets a new SOTA for few-step generation. Specifically, our 4-step models outperform the quality of their multi-step teacher and substantially exceed previous DMDR results in GenEval and human-preference metrics, exhibiting strong scalability potential.

Linwei Dong, Ruoyu Guo, Ge Bai +3

1Zhejiang University · 2Bytedance Inc. · 3Zhejiang Lab

Few-Step Distillation Distribution Matching Distillation Symbolic Distillation Few-Step Generation

Ask PDF

Jul 7, 2026cs.CV

Bridging Diffusion Pruning and Step Distillation with Teacher-Aligned Repair

Diffusion models generate high-quality images, but their inference cost comes from two sources: large denoising networks and repeated denoising steps. Existing compression pipelines usually attack these costs separately. Pruning reduces the network, but most pruning methods still rely on a long post-pruning retraining stage to recover a many-step sampler. Step distillation reduces the number of denoising steps, but it usually assumes a student that can already follow the teacher well enough to receive useful distillation gradients. This paper asks whether post-pruning retraining can be replaced by step distillation. We find that the direct replacement fails: after pruning an EDM2-XS teacher, starting SiDA from the pruned checkpoint produces unusable samples. We introduce a short teacher-alignment repair stage as a bridge between pruning and step distillation. The bridge matches the pruned generator to the teacher on noisy real-image latents, then hands the repaired checkpoint to one-step distillation. On ImageNet-512, the original EDM2-XS baseline uses 124.713M parameters and 63 network evaluations, reaching an FID of 3.53. With a suitable distillation objective, our 20% pruned one-step generator uses 98.826M parameters and one network evaluation, reaching an FID of 3.12. With 30% pruning, the model uses 88.029M parameters and one network evaluation, with an FID of 4.26.

Jincheng Ying, Li Wenlin, Minghui Xu +1

School of Big Data and Artificial Intelligence Guangdong University of Finance and Economics Guangzhou, China · School of Computer Science Shandong University Qingdao, China · GDUFE-UCM Joint Institute Guangdong University of Finance and Economics Guangzhou, China

Diffusion Models Teacher Inference Cost

Ask PDF

Apr 29, 2026cs.CV

AdvDMD: Adversarial Reward Meets DMD For High-Quality Few-Step Generation

Diffusion models offer superior generation quality at the expense of extensive sampling steps. Distillation methods, with Distribution Matching Distillation (DMD) as a popular example, can mitigate this issue, but performance degradation remains pronounced when sampling steps are limited. Reinforcement learning (RL) has been leveraged to improve the few-step generation quality during distillation, with the potential to even surpass the performance of the teacher model. However, existing approaches are combinatorial in nature, merely integrating an RL process with the distillation process, which introduces unnecessary complexities. To address this gap, we propose AdvDMD, a method that seamlessly unifies DMD distillation and RL. Specifically, AdvDMD employs the adversarially trained discriminator from DMD2 as the reward model, which assigns low scores to generated images and high scores to real ones. It is trained on both intermediate and final states of the denoising process and updated online with the distilled model, enabling a holistic supervision of the sampling trajectories and mitigating reward hacking. We adopt a unified SDE backward simulation and a different training schedule for DMD and RL to enable a more stable and efficient training. Experimental results demonstrate that the 4-step AdvDMD outperforms the original 40-step model for SD3.5 on DPG-Bench, while achieving significant performance gains for SD3 on the GenEval. On Qwen-Image, our 2-step AdvDMD achieves superior performance over TwinFlow.

Xu Wang, Zexian Li, Litong Gong +2

Shanghai Jiao Tong University · Alimama Tech

Few-Step Distillation Few-Step Generation Distribution Matching Distillation Residual Distillation Adversarial Robustness

Ask PDF