Latent Reward Registers for Diffusion Preference Alignment

• URL: http://arxiv.org/abs/2608.03929v2
• Date: Wed, 05 Aug 2026 03:09:27 GMT
• Title: Latent Reward Registers for Diffusion Preference Alignment
• Abstract summary: We propose Latent Reward Registers, a mechanism that estimates terminal preference directly from intermediate noisy latents.<n>This independent readout mechanism extracts latent reward evidence without altering the generator's hidden states or velocity field.<n>For training, Reward-Gradient On-Policy Distillation (RG-OPD) distills reward-guided updates along on-policy trajectories, bypassing the computationally expensive rollouts of standard policy gradients.<n>For inference, Reward-Guided Sampling (RGS) steers trajectories via magnitude-matched reward gradients without parameter updates.
• Score: 46.49072106092868
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Aligning diffusion models with human preferences usually relies on a sparse terminal reward evaluated on the final generated samples, presenting a severe temporal credit-assignment challenge across the multi-step denoising process. We propose Latent Reward Registers, a mechanism that estimates terminal preference directly from intermediate noisy latents by prepending learnable, position-free register tokens to the input sequence of a frozen Diffusion Transformer (DiT). This independent readout mechanism extracts latent reward evidence without altering the generator's hidden states or velocity field. The resulting dense, differentiable reward signal throughout the full denoising process facilitates two alignment strategies. For training, Reward-Gradient On-Policy Distillation (RG-OPD) distills reward-guided updates along on-policy trajectories, bypassing the computationally expensive rollouts of standard policy gradients. For inference, Reward-Guided Sampling (RGS) steers trajectories via magnitude-matched reward gradients without parameter updates. Empirically, at high noise levels (u = 0.8), the registers reach the highest pairwise accuracy among the evaluated latent reward models. Furthermore, RG-OPD outperforms online reinforcement learning baselines while reducing GPU hours by up to 33x, and RGS establishes a new state-of-the-art among training-free methods, strictly enhancing both alignment and perceptual metrics. Code and weights are available at https://github.com/Guanys-dar/latent-reward-register

Related papers

• Survey of Novel Deep Learning Architectures for Denoising Gravitational-wave Signals [0.0]
  We present the first controlled comparison of five neural-network architectures for gravitational-wave denoising.<n>We train identically across the full astrophysically-motivated spinning binary-black-hole parameter space.<n>A unifying principle emerges: matching network structure to the spectral anatomy of a coalescence outperforms brute-force model scaling.
  arXiv  Detail & Related papers  (2026-09-07T18:00:10Z)
• Gradients Know What Outcomes Don't: Unlocking Reinforcement Learning for LLM Reasoning with Gradient-Aligned Rewards [24.636191382674184]
  Gradient Reward (GAR) operates in the policy's own gradient space.<n>C cosine similarity with an expert-anchor gradient yields a dense, reasoning-aware reward.
  arXiv  Detail & Related papers  (2026-09-03T04:00:28Z)
• SATB-VR: Training Few-Step Video Restoration Diffusion Model using SNR-Aware Trajectory Blending [55.764018145256216]
  We present SATB-VR, a few-step paradigm that jump-starts the denoising process via an auxiliary predictor.<n>We show that SATB-VR performs favorably against existing approaches on synthetic, real-world, and AIGC benchmarks.
  arXiv  Detail & Related papers  (2026-06-27T01:32:16Z)
• RODS: Reward-Driven Online Data Synthesis for Multi-Turn Tool-Use Agents [13.362388357375082]
  Multi-turn tool-use RL is bottlenecked by the rapid depletion of informative samples in static datasets.<n>We propose RODS (Reward-driven Online Data Synthesis) to resolve this depletion.
  arXiv  Detail & Related papers  (2026-06-17T13:13:32Z)
• Gradient-Guided Reward Optimization for Inference-time Alignment [22.28081683932275]
  We introduce a lightweight inference-time method that performs targeted, minimal intervention during decoding via gradient guidance.<n>Experiments show that GGRO consistently improves inference-time alignment across safety, helpfulness, and reasoning benchmarks.<n>It also increases coverage of high-quality responses and robustness to reward hacking, with minimal computational overhead.
  arXiv  Detail & Related papers  (2026-06-08T15:33:13Z)
• Precise: SDE-Consistent Stochastic Sampling for RL Post-Training of Flow-Matching Models [56.67321805551389]
  Reinforcement learning (RL) has become an effective way to improve prompt alignment and perceptual quality in diffusion and flow-matching generators.<n>The sampler, controlling the exploration behavior and denoising dynamics, is thus part of the policy.<n>We propose a new sampler that balances effective exploration with stability.
  arXiv  Detail & Related papers  (2026-05-22T11:37:22Z)
• Clipping Bottleneck: Stabilizing RLVR via Stochastic Recovery of Near-Boundary Signals [83.0127582612634]
  Near-boundary Rescue (NSR) is a minimal, plug-and-play modification that retains out-of-bound tokens to recover lost signals.<n>NSR substantially improves training stability and delivers consistent gains over strong baselines such as DAPO and GSPO.
  arXiv  Detail & Related papers  (2026-05-21T16:45:31Z)
• Diff-Instruct with Diffused Reward: Towards Principled One-step Generator RL [16.321536102397733]
  One-step text-to-image generation has enabled real-time synthesis with remarkable efficiency and quality.<n>Previous reinforcement learning methods for one-step generators combine image-space reward optimization with diffusion noisy-space distribution matching.<n>We propose Diff-Instruct with Diffused Reward (DIDR), a data-free trajectory-level alignment framework derived from KL minimization.
  arXiv  Detail & Related papers  (2026-05-18T21:26:33Z)
• Verifier-Free RL for LLMs via Intrinsic Gradient-Norm Reward [69.99652051809737]
  In this work, we propose Verifier-free Intrinsic Gradient-Norm Reward (VIGOR)<n>VIGOR is a simple reward that uses only the policy model itself.<n>It exhibits cross-domain transfer to code benchmarks when trained only on math data.
  arXiv  Detail & Related papers  (2026-05-11T03:15:37Z)
• Beyond Verifiable Rewards: Rubric-Based GRM for Reinforced Fine-Tuning SWE Agents [9.505140329883762]
  We introduce a Generative Reward Model (GRM) that provides richer learning signals.<n>The GRM is equipped with human-designed rubrics that indicate criteria for encouraging or discouraging specific behavioral patterns.<n>When used for Reinforced Fine-Tuning (RFT) on SWE Tasks, our approach outperforms terminal-score-only rejection sampling.
  arXiv  Detail & Related papers  (2026-03-13T02:23:49Z)
• Euphonium: Steering Video Flow Matching via Process Reward Gradient Guided Stochastic Dynamics [49.242224984144904]
  We propose Euphonium, a novel framework that steers generation via process reward gradient guided dynamics.<n>Our key insight is to formulate the sampling process as a theoretically principled algorithm that explicitly incorporates the gradient of a Process Reward Model.<n>We derive a distillation objective that internalizes the guidance signal into the flow network, eliminating inference-time dependency on the reward model.
  arXiv  Detail & Related papers  (2026-02-04T08:59:57Z)
• DenseGRPO: From Sparse to Dense Reward for Flow Matching Model Alignment [49.45064510462232]
  GRPO-based approaches for text-to-image generation suffer from the sparse reward problem.<n>We introduce textbfDenseGRPO, a novel framework that aligns human preference with dense rewards.
  arXiv  Detail & Related papers  (2026-01-28T03:39:05Z)

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.