PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimizing Few-Step Generation with Adaptive Matching Distillation

Lichen Bai, Zikai Zhou, Shitong Shao, Wenliang Zhong, Shuo Yang, Shuo Chen, Bojun Chen, Zeke Xie

arXiv:2602.07345v2cs.CVcs.LG

TL;DR

Few-step DMD can become unstable when distorted samples enter Forbidden Zones, where real-teacher guidance is unreliable and fake-teacher repulsion is insufficient. AMD addresses this with reward-aware diagnostics and adaptive corrective dynamics, and experiments report improved fidelity and robustness across image and video generation benchmarks.

  • Problem

    DMD stability is compromised when student samples enter Forbidden Zones, where real-teacher score estimates become unreliable and fake-teacher repulsion is insufficient.

  • Method

    AMD uses reward models as diagnostic proxies and dynamically adapts distillation signals to prioritize escape from Forbidden Zones.

  • Results

    AMD demonstrates robust effectiveness across image and video benchmarks and achieves the highest ImageReward (88.37) and HPSv2 (31.25) in the reported image-generation evaluation.

  • Takeaways & Limitations

    Explicitly addressing Forbidden Zones can improve training robustness and sample fidelity in few-step generative models.

  • Takeaways & Limitations

    AMD relies on pretrained teacher and reward models, so the distilled student’s safety and ethical alignment depend on the quality of these foundational signals.

Abstract

from arXiv · show

Distribution Matching Distillation (DMD) is a powerful acceleration paradigm, yet its stability is often compromised in Forbidden Zone, regions where the real teacher provides unreliable guidance while the fake teacher exerts insufficient repulsive force. In this work, we propose a unified optimization framework that reinterprets prior art as implicit strategies to avoid these corrupted regions. Based on this insight, we introduce Adaptive Matching Distillation (AMD), a self-correcting mechanism that utilizes reward proxies to explicitly detect and escape Forbidden Zones. AMD dynamically prioritizes corrective gradients via structural signal decomposition and introduces Repulsive Landscape Sharpening to enforce steep energy barriers against failure mode collapse. Extensive experiments across image and video generation tasks (e.g., SDXL, Wan2.1) and rigorous benchmarks (e.g., VBench, GenEval) demonstrate that AMD significantly enhances sample fidelity and training robustness. For instance, AMD improves the HPSv2 score on SDXL from 30.64 to 31.25, outperforming state-of-the-art baselines. These findings validate that explicitly rectifying optimization trajectories within Forbidden Zones is essential for pushing the performance ceiling of few-step generative models.

1. Introduction

Diffusion models offer high-quality generation but iterative denoising creates substantial inference latency, motivating few-step distillation. The paper identifies Forbidden Zones as a DMD failure regime and proposes AMD, a reward-aware self-correcting framework that improves stability and fidelity across image and video generation.

  • Iterative denoising can require dozens or hundreds of NFEs during inference, creating substantial latency.
  • DMD compresses multi-step generation into few steps but depends on reliable real-teacher guidance and sufficient fake-teacher repulsion.
  • Forbidden Zones arise when distorted student samples receive unreliable real-teacher scores and insufficient fake-teacher repulsion, causing degradation and hindering convergence.
  • Prior DMD methods implicitly avoid corrupted regions, but they do not explicitly detect Forbidden Zones or adapt dynamics after entry.
  • AMD uses reward-guided diagnostics, targeted repulsive guidance, and dynamic signal prioritization to push students out of Forbidden Zones.
  • AMD mitigates training collapse and improves sample fidelity across large-scale image and video tasks, including SDXL and Wan2.1.

2. Preliminaries

The preliminaries describe DMD as compressing a pretrained diffusion model into a few-step generator by matching target and student-induced distributions. They also introduce its real/fake teacher pull-push dynamics and a fixed black-box reward model.

  • DMD compresses a pretrained multi-step diffusion model into a few-step student by minimizing KL divergence between target and student-induced distributions.
  • The student sample is generated as x = Gθ(z), with z sampled from a standard normal distribution and diffusion time t sampled uniformly.
  • The real teacher attracts student samples toward the target distribution, while the fake teacher repels them from the student’s current distribution.
  • DMD uses a pretrained reward model R that assigns scalar scores to generated samples and remains fixed as a black-box scoring function during training.

3. Method

The method reframes DMD as latent-space optimization, identifies Forbidden Zones where teacher forces fail, and introduces AMD to adapt gradients and strengthen repulsion in those regions.

  • An Optimization Perspective on DMD: DMD can be reformulated as effective gradient descent on generated latent samples, clarifying distillation as push-pull navigation between teacher targets.The update moves toward the real teacher’s target through dreal and away from the fake teacher’s estimated region through dfake.
  • An Optimization Perspective on DMD: Forbidden Zones are high-energy regions beyond the real teacher’s empirical support, where distorted student samples receive unreliable attraction and weak repulsion.The real teacher’s energy becomes fractured or ill-posed, while the fake teacher’s energy is flat at the student distribution’s extreme tails.
  • An Optimization Perspective on DMD: Existing DMD variants can be interpreted as configurations of a generalized displacement operator that implicitly mitigate Forbidden Zones through adaptation, noise, or auxiliary forces.The framework includes external recovery forces and categorizes prior methods as strategies for avoiding corrupted regions.
  • Dynamic Score Adaptation via AMD: AMD dynamically adapts the combination of attractive and repulsive signals to correct distorted samples and prioritize escape from Forbidden Zones.Dynamic Score Adaptation uses fine-grained intervention because failure involves structural gradient collapse and destructive interference, not only insufficient magnitude.
  • Repulsive Landscape Sharpening: Repulsive Landscape Sharpening focuses fake-teacher training on failure cases, creating steeper likelihood slopes and stronger repulsive gradients for escape.Advantage-aware weighting penalizes low-advantage samples more heavily, transforming the fake teacher into an active failure detector.

4. Experiments

Experiments evaluate AMD across image and video generation settings, benchmarks, and ablations. Results show improved fidelity, stability, reward alignment, and motion quality, with reward-aware mechanisms guiding students away from low-quality regions.

  • Experimental Setup: AMD is evaluated across SiT, SDXL, Wan2.1, and LongLive using image- and video-generation benchmarks including MS-COCO, ImageNet, GenEval, and VBench.The experiments span diverse diffusion backbones, modalities, and architectural designs.
  • Image Generation: AMD achieves the highest ImageReward (88.37) and HPSv2 (31.25) on the 10k COCO2014-val text-to-image evaluation.It also ranks first among distilled models on GenEval with an overall score of 0.57.
  • Image Generation: AMD achieves superior FID (3.4690) and sFID (5.7464) to DMD on 50K-ImageNet, while avoiding DMDR’s severe mode collapse.The reported comparison emphasizes balanced image quality and diversity rather than reward hacking.
  • Video Generation: AMD surpasses LongLive with a Total Score of 82.21 vs. 81.42 and improves VBench Motion Quality by ∼67% (35.51 →59.26).The video experiments use Wan2.1-1.3B and the first stage of the LongLive protocol.
  • Ablation and Mechanisms: Dynamic Score Adaptation and Repulsive Landscape Sharpening jointly provide fine-grained update modulation and stronger correction of pathological regions.The two mechanisms respectively refine gradient magnitude and direction, and sharpen repulsive forces around failure regions.
  • Training Dynamics: Reward and generation quality rise together during training, while AMD shows faster improvement and a more stable trajectory than the baseline.On the 2D multimodal analysis, selective guidance suppresses low-reward modes while global alignment recovers the full distribution.

5. Conclusion

The conclusion identifies Forbidden Zones as a fundamental DMD failure regime and presents AMD as a framework for correcting it. Across image and video benchmarks, AMD improves fidelity and diversity while offering a unified perspective on DMD-style optimization.

  • Conclusion: Forbidden Zones arise when unreliable real-teacher guidance and insufficient fake-teacher repulsion jointly stall DMD optimization.AMD is proposed to address this failure regime across image and video generation benchmarks.
  • Conclusion: AMD provides a unified optimization perspective that interprets prior DMD methods as implicit strategies for navigating corrupted regions.The framework reframes distillation as a latent navigation task and motivates more adaptive, self-correcting approaches.

Impact Statement

The impact statement emphasizes that AMD depends on pretrained teacher and reward models, so the distilled student’s safety and ethical alignment depend on the quality of those foundational signals.

  • Impact Statement: AMD’s safety and ethical alignment are intrinsically tied to the quality of its pretrained teacher and reward-model signals.The authors advocate responsible deployment and content authentication to mitigate misuse risks such as misleading deepfakes.

A. Implementation Details

The paper evaluates AMD across diverse benchmarks, metrics, diffusion backbones, and image/video generation settings. Its optimization update is framed as gradient descent on a contrastive DMD potential.

  • Benchmarks: AMD evaluates image and video generation using benchmarks including MS-COCO, ImageNet, DrawBench, HPD v2, GenEval, VBench, VBench++, and VideoGen-Eval.The evaluation suite covers distributional fidelity, compositional fidelity, human preference, text alignment, temporal quality, and other video dimensions.
  • Evaluation Metrics: HPS v2 predicts user preferences by fine-tuning CLIP on HPD v2 and estimates whether synthesized images will be preferred.It is used as a human-preference benchmark for text-to-image models across diverse image distributions.
  • Optimization Formulation: The DMD latent update follows x_new = x − η_eff∇_xV_DMD(x), characterizing gradient descent on a contrastive potential.The potential gradient is expressed through the difference between real-teacher and fake-teacher directions.

C.2. Theoretical Analysis of Prior Art

The analysis interprets prior DMD variants as mechanisms for reducing or escaping Forbidden Zones, where teacher gradients become unreliable or vanish. These mechanisms modify noise levels, teacher support, or introduce external corrective forces.

  • D-DMD: D-DMD uses a higher teacher noise level to move distorted samples outside the Forbidden Zone and provide a coarse-grained directional guide.The broadened Gaussian kernel increases density around low-quality samples and pulls them toward high-density regions.
  • Support Expansion: Increasing diffusion time expands the teacher distribution’s effective support and reduces the probability that student samples enter the Forbidden Zone.At the student noise level, distorted samples can have vanishing density and undefined or oscillating real-teacher scores.
  • MagicDistillation: MagicDistillation adapts the real teacher on student samples, creating a continuous energy-landscape bridge from the student distribution to the target manifold.The adapted potential remains defined in regions where the original data density is near zero, allowing smooth traversal.
  • External Forces: DMD, DMD2, and DMDR introduce external forces that dominate when intrinsic distillation gradients vanish or become stochastic noise.DMD uses ground-truth tethering, DMD2 uses discriminator boundary repulsion, and DMDR uses reward-gradient steering.

D. Limitation and Future Work

The paper identifies limitations in AMD’s reliance on reward-based Forbidden Zone detection and in the sophistication of its current adaptive operator. It proposes more robust detection and advanced optimization strategies as future directions.

  • Limitations: AMD’s self-correction depends on accurately identifying Forbidden Zones through a reward model’s advantage signal.Noisy or poorly calibrated feedback may fail to trigger adaptation and leave some corrupted regions unrectified.
  • Future Work: Future work could reduce dependence on a single proxy through robust, unsupervised, or ensemble-based Forbidden Zone detection.This direction directly targets the sensitivity of the correction mechanism to proxy quality.
  • Future Work: The current adaptive operator H_AMD leaves room for more sophisticated strategies incorporating momentum, orthogonal gradients, or second-order information.Such refinements are proposed to further stabilize training trajectories and accelerate convergence in complex generative tasks.

F.1. Analytical Experiments

The analytical experiments compare naive adaptation with AMD on a 2D multimodal toy dataset. AMD preserves distributional structure while steering toward high-reward regions, whereas naive adaptation collapses the distribution.

  • Experimental Setup: The experiment compares H_naive and H_AMD under the same reward proxy on a 2D multimodal toy dataset.The real-teacher energy represents a symmetric multimodal distribution, while the reward model favors a specified region.
  • Naive Adaptation: Naive adaptation causes catastrophic distribution collapse during iterations 1000–2000, dissolving the student’s multimodal structure into an uninformative mass.Simple linear scaling of d_real and d_fake fails to resolve the conflict between global teacher guidance and local reward prioritization.
  • AMD: AMD reaches the high-reward regime while maintaining the structural integrity of the generation manifold.Its decoupled distribution-matching and conditional-alignment terms avoid the gradient deadlock seen with naive adaptation.
  • Conclusion: The results validate decoupled adaptation as necessary for stable, reward-aware distribution matching in this toy setting.The figure shows AMD preserving sharp, distinct modes while naive adaptation merges and collapses them.

F.2. Quantitative Results

AMD is evaluated across image and video benchmarks, where it generally surpasses DMD2 and achieves higher aggregate performance on Wan2.1-14B. The results also reveal a motion-quality versus visual-quality and text-alignment trade-off.

  • Image Generation: AMD’s winning rate exceeds DMD2’s 50% baseline across DrawBench and HPDv2.
  • Video Generation: AMD achieves a higher overall score than DMD2 on Wan2.1-14B, improving Total Score from 118.61 to 122.15.
  • Evaluation Setup: The internal image benchmark contains 419 diverse web-collected prompts spanning visual styles, subjects, layouts, and rich Gemini-generated captions.
  • Video Generation: VideoAlign improves Motion Quality by 16.24, sometimes at the expense of static visual fidelity and strict text-image alignment.

F.3. Visualization Results

Visualization results cover SDXL image synthesis across DrawBench and HPDv2 and Wan2.1 video generation. AMD is reported to improve image quality and texture while producing smoother, more logically consistent motion than standard distillation.

  • Image Generation: AMD improves synthesized image quality and texture over standard DMD2 across DrawBench and multiple HPDv2 subsets.
  • Image Generation: The SDXL gallery includes synthesized images from DrawBench and the concept-art, painting, photo, and anime subsets of HPDv2.
  • Video Generation: Standard DMD often exhibits temporal flickering or motion degradation in challenging video scenarios.
  • Video Generation: AMD preserves dynamic fidelity and generates smoother, more logically consistent motion than baseline DMD.
  • Video Generation: The text-to-video comparison reports superior motion smoothness, visual fidelity, and prompt alignment for AMD over the baseline.
Loading 2602.07345v2…