Reward-Aware Trajectory Shaping for Few-step Visual Generation

• URL: http://arxiv.org/abs/2604.14910v2
• Date: Tue, 21 Apr 2026 03:58:54 GMT
• Title: Reward-Aware Trajectory Shaping for Few-step Visual Generation
• Abstract summary: We propose textbfReward-Aware Trajectory Shaping (RATS), a lightweight framework for preference-aligned few-step generation.<n>RATS effectively transfers preference-relevant knowledge from high-step generators without incurring additional test-time computational overhead.<n> Experimental results demonstrate that RATS substantially improves the efficiency--quality trade-off in few-step visual generation.
• Score: 48.211558783057086
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Achieving high-fidelity generation in extremely few sampling steps has long been a central goal of generative modeling. Existing approaches largely rely on distillation-based frameworks to compress the original multi-step denoising process into a few-step generator. However, such methods inherently constrain the student to imitate a stronger multi-step teacher, imposing the teacher as an upper bound on student performance. We argue that introducing \textbf{preference alignment awareness} enables the student to optimize toward reward-preferred generation quality, potentially surpassing the teacher instead of being restricted to rigid teacher imitation. To this end, we propose \textbf{Reward-Aware Trajectory Shaping (RATS)}, a lightweight framework for preference-aligned few-step generation. Specifically, teacher and student latent trajectories are aligned at key denoising stages through horizon matching, while a \textbf{reward-aware gate} is introduced to adaptively regulate teacher guidance based on their relative reward performance. Trajectory shaping is strengthened when the teacher achieves higher rewards, and relaxed when the student matches or surpasses the teacher, thereby enabling continued reward-driven improvement. By seamlessly integrating trajectory distillation, reward-aware gating, and preference alignment, RATS effectively transfers preference-relevant knowledge from high-step generators without incurring additional test-time computational overhead. Experimental results demonstrate that RATS substantially improves the efficiency--quality trade-off in few-step visual generation, significantly narrowing the gap between few-step students and stronger multi-step generators.

Related papers

• Eliciting Weak-to-Strong Generalization with On-Policy Reverse Distillation [41.854144563456046]
  Weak-to-strong generalization asks whether stronger models can learn from weaker supervisors and surpass them.<n>We introduce On-Policy Reverse Distillation, which evaluates the teacher's policy shift and amplifies the student's verifier-driven policy gradient.<n>In both successive model transfer and multi-teacher distillation, OPRD achieves prohibitively higher performance with fewer student updates than existing RL and distillation approaches.
  arXiv  Detail & Related papers  (2026-09-08T14:26:35Z)
• On-Policy Adversarial Flow Distillation for Autoregressive Video Generation [34.70887671905755]
  Adversarial Flow Distillation (AFD) is an on-policy framework for heterogeneous black-box video distillation.<n>AFD queries the teacher and rolls out the current student on the same prompts.<n>AFD provides dense velocity-field supervision while requiring no teacher scores, latents, denoising trajectories, step alignment, or reverse-chain reinforcement learning.
  arXiv  Detail & Related papers  (2026-05-25T17:58:12Z)
• On-Policy Distillation with Best-of-N Teacher Rollout Selection [54.91780727674628]
  We propose BRTS, a Best-of-N Rollout Teacher Selection framework for on-policy distillation.<n>BRTS augments standard student-context OPD with a teacher-context supervision branch constructed from the teacher trajectory.<n>BRTS improves over standard OPD on challenging reasoning benchmarks, with the largest gains on harder datasets.
  arXiv  Detail & Related papers  (2026-05-10T19:49:00Z)
• Prune-OPD: Efficient and Reliable On-Policy Distillation for Long-Horizon Reasoning [66.52232008796294]
  Prune-OPD dynamically aligns training budgets with supervision quality.<n>It reduces training time by 37.6%--68.0% while preserving, and often improving, performance on challenging benchmarks.
  arXiv  Detail & Related papers  (2026-05-08T14:38:53Z)
• B-DENSE: Branching For Dense Ensemble Network Learning [2.972979055997652]
  We propose B-DENSE, a novel framework that leverages multi-branch trajectory alignment.<n>B-DENSE learns to navigate the solution space from the earliest stages of training, demonstrating superior image generation quality.
  arXiv  Detail & Related papers  (2026-02-17T19:40:58Z)
• Towards One-step Causal Video Generation via Adversarial Self-Distillation [71.30373662465648]
  Recent hybrid video generation models combine autoregressive temporal dynamics with diffusion-based spatial denoising.<n>Our framework produces a single distilled model that flexibly supports multiple inference-step settings.
  arXiv  Detail & Related papers  (2025-11-03T10:12:47Z)
• AdaSwitch: Adaptive Switching Generation for Knowledge Distillation [58.647880811071495]
  Small language models (SLMs) are crucial for applications with strict latency and computational constraints.<n>We propose AdaSwitch, a novel approach that combines on-policy and off-policy generation at the token level.<n>AdaSwitch consistently improves accuracy, offering a practical and effective method for distilling SLMs with acceptable additional overhead.
  arXiv  Detail & Related papers  (2025-10-09T06:38:37Z)
• Training-Free Motion Customization for Distilled Video Generators with Adaptive Test-Time Distillation [53.877572078307935]
  Distilled video generation models offer fast and efficient but struggle with motion customization when guided by reference videos.<n>We propose MotionEcho, a training-free test-time distillation framework that enables motion customization by leveraging diffusion teacher forcing.
  arXiv  Detail & Related papers  (2025-06-24T06:20:15Z)
• Distill Not Only Data but Also Rewards: Can Smaller Language Models Surpass Larger Ones? [58.80794196076336]
  Distilling large language models (LLMs) typically involves transferring the teacher model's responses through supervised fine-tuning (SFT)<n>We propose a novel distillation pipeline that transfers both responses and rewards.<n>Our method generates pseudo-rewards through a self-supervised mechanism that leverages the inherent structure of both teacher and student responses.
  arXiv  Detail & Related papers  (2025-02-26T20:50:11Z)
• DOLLAR: Few-Step Video Generation via Distillation and Latent Reward Optimization [50.30051934609654]
  We introduce a distillation method that combines variational score distillation and consistency distillation to achieve few-step video generation.<n>Our method demonstrates state-of-the-art performance in few-step generation for 10-second videos (128 frames at 12 FPS)<n>One-step distillation accelerates the teacher model's diffusion sampling by up to 278.6 times, enabling near real-time generation.
  arXiv  Detail & Related papers  (2024-12-20T09:07:36Z)
• Efficient Diffusion as Low Light Enhancer [63.789138528062225]
  Reflectance-Aware Trajectory Refinement (RATR) is a simple yet effective module to refine the teacher trajectory using the reflectance component of images. textbfReflectance-aware textbfDiffusion with textbfDistilled textbfTrajectory (textbfReDDiT) is an efficient and flexible distillation framework tailored for Low-Light Image Enhancement (LLIE)
  arXiv  Detail & Related papers  (2024-10-16T08:07:18Z)
• Gap Preserving Distillation by Building Bidirectional Mappings with A Dynamic Teacher [43.678380057638016]
  Gap Preserving Distillation (GPD) method trains an additional dynamic teacher model from scratch along with training the student to bridge this gap. In experiments, GPD significantly outperforms existing distillation methods on top of both CNNs and transformers architectures. GPD also generalizes well to the scenarios without a pre-trained teacher, including training from scratch and fine-tuning, yielding a large improvement of 1.80% and 0.89% on ResNet18.
  arXiv  Detail & Related papers  (2024-10-05T12:29:51Z)
• Distribution Backtracking Builds A Faster Convergence Trajectory for Diffusion Distillation [19.88187051373436]
  We propose Distribution Backtracking Distillation (DisBack) to speed up the sampling speed of diffusion models.<n>DisBack achieves faster and better convergence than the existing distillation method, with FID score of 1.38 on ImageNet 64x64 dataset.
  arXiv  Detail & Related papers  (2024-08-28T17:58:17Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.