cs.LGDec 28, 2025
SaveReDiF: Resource-Efficient Few-Step Diffusion Distillation via Reinforcement Learning
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
Figure 1: ReDiF training on the Stable Diffusion v1.5 branch ( 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 | R | D | C | FID | Pub. FID |
| (a) EDM teacher, ImageNet-64 ( ) | ||||||||||
| 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 ) | 32 | 15.39 | ||||||||
| Consistency Distill. ( Song et al., 2023 ) | 32 | 6.20 | ||||||||
| DMD ( Yin et al., 2024b ) | 128 | 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 sample sd, last checkpoint; per branch, best in bold, second best bold and dash-underlined, teacher rows unranked.
| Method | Mode | Aux. net | PFLOPs | P | R | D | C | FID |
|---|---|---|---|---|---|---|---|---|
| Progressive Distill. ( Salimans and Ho, 2022 ) | – | interm. teachers | ||||||
| + RL (ours) | structural | interm. teachers | ||||||
| Consistency Distill. ( Song et al., 2023 ) | – | EMA target | ||||||
| + RL (ours) | structural | EMA target | ||||||
| DMD2 ( Yin et al., 2024a ) | – | score net + disc. | ||||||
| + RL (ours) | loss | score net + disc. |
Table 2: ReDiF’s PPO objective combined with existing methods on the EDM teacher, ImageNet-64 ( ), 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 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 . 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
| Method | Step Distill. | One-Step | Data-Free | RL | Few Pr. | Non-Diff. | Low-Comp. |
|---|---|---|---|---|---|---|---|
| Trajectory matching | |||||||
| 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 ) | ✓ | ✓ | ✗ | ✗ | ✗ | ✗ | ✗ |
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 ( , EDM teacher). Markers denote evaluated checkpoints; the dashed line is the -step teacher ( ) 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 ).
Figure 4: The SD v1.5 student of Figure 1 trained for epochs, well past the -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.
Figure 5: Teacher and student samples, SD v1.5, first three held-out prompts.
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 ( ) | UNet attention ( ) |
| Learning rate | ||
| PPO clip range | ||
| Inner epochs per rollout | 4 | 4 |
Table 4: Hyperparameters used in all reported ReDiF runs. Both branches use AdamW with , , weight decay , and gradient-norm clipping at . The preference runs of Appendix K use the COCO column with the preference weight added.
| Reward | P | R | D | C | CLIP-T | FID |
|---|---|---|---|---|---|---|
| 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 |
| Single-encoder semantic rewards | ||||||
| CLIP (R 1 ) | ||||||
| DINOv3 | ||||||
| PE | ||||||
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, , , generated samples against real validation images; the R 1 and R 2 rows are the ReDiF rows of Table 1 . Mean 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.
| Reward | P | R | D | C | FID | CLIPScore |
|---|---|---|---|---|---|---|
| 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 |
| Single-encoder semantic rewards | ||||||
| CLIP (R 1 ) | ||||||
| DINOv3 | ||||||
| PE | ||||||
Table 6: Reward ablation, Stable Diffusion v1.5 on COCO ( ). Same design as Table 5 : identical optimization, only the reward differs. Last checkpoint, , COCO validation prompts; the R 1 and R 2 rows are the ReDiF rows of Table 1 . Mean sample sd over seeds. Best entry bold in bold, second best bold and dash-underlined; teacher rows are not ranked.
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 ( steps); the corresponding curves on ImageNet-64 differ, and the text below says how. These are -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 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.
Figure 8: The COCO student trained with the distillation reward plus one non-differentiable preference, for two fixed prompts and seeds: compressibility equation 2 above, sharpness equation 4 below. Init is the student before any update. Only the policy’s weights differ between columns.
| Student | LR | P | R | D | C | FID | |
|---|---|---|---|---|---|---|---|
| LoRA ( ) | |||||||
| LoRA ( ) | |||||||
| Full | |||||||
| Full |
Table 7: Student parameterization crossed with the PPO clip range, EDM teacher, ImageNet-64 ( ), R 2 reward, epochs. Mean 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.
Figure 9: Effect of the student step count 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 -step teacher and the dotted line the operating point . is the number of runs at each and the band is s.d. over runs.
| Method | Family | Steps | Aux. net | Trainable | Batch | Devices / training cost (reported) | FID |
|---|---|---|---|---|---|---|---|
| (a) ImageNet-64 | |||||||
| EDM teacher ( Karras et al., 2022 ) | – | 256 | – | – | – | – | 1.36 |
| Progressive Distill. ( Salimans and Ho, 2022 ) | trajectory | interm. teachers | full | 2048 | 64 GPUs; 5,533 A100-h † | 15.39 ‡ | |
| Consistency Distill. ( Song et al., 2023 ) | trajectory | target copy | full | 2048 | 64 GPUs; 7,867 A100-h † | 6.20 ‡ | |
| CTM ( Kim et al., 2024 ) | trajectory | target copy + disc. | full | 2048 | 64 GPUs; 902 A100-h † | 1.92 | |
| DMD ( Yin et al., 2024b ) | distribution | score net | full | n/r | n/r | 2.62 | |
Table 8: Published results at convergence. FID as reported by the original authors, except where marked: ImageNet-64 with k balanced samples, COCO-2014 zero-shot with k samples. Batch is the effective training batch and Devices the reported hardware, both as disclosed; marks a third-party A100-hour estimate from Zhou et al. (2024b) ; 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; 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 | – | fixed teacher target | |
| Consistency Distillation | – | fixed teacher target | ||
| DMD | moving target (aux. network) | |||
| DMD2 | moving target (aux. network) | |||
| SiD | moving target (aux. network) | |||
| SiDA | moving target (aux. network) |
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.