1 Introduction↩︎

Diffusion models [1] have become the leading framework for high-quality image generation. However, because they are trained to reproduce the distribution of their training data, they do not inherently reflect human preferences. Recent work [2], [3] addresses this limitation through reinforcement learning from human feedback (RLHF) [4], which fine-tunes diffusion models using scalar feedback from human or reward models to explicitly optimize for preference alignment.

One major challenge with RLHF in diffusion models is the credit assignment problem: the difficulty of determining which intermediate timesteps in the diffusion process actually contributed to the final reward. Because human feedback is only given on the final image, methods such as DDPO [2] simply assign the same loss to every timestep. This setup ignores the structure of the denoising process and how different timesteps edit the image at different levels of granularity [5]. As a result, the model optimizes transitions uniformly without considering the non-uniform nature of the denoising trajectory, leading to inefficient training.

Prior work attempts to address this issue by contrasting paired denoising trajectories generated from the same initial noise [6], [7]. In these approaches, trajectories remain identical until a designated branching timestep and diverge afterward. While contrasting two trajectories with shared initial noise can highlight the effect of the timestep at which they diverge, the approach does not contrast other timesteps in the sequence. Earlier steps remain identical and therefore uninformative, and later steps still inherit the same uniform advantage, leaving the broader credit assignment problem unresolved. Moreover, these methods require multiple reward evaluations per starting noise, reducing sample efficiency and increasing computational cost.

In this paper, we argue that reward information in diffusion trajectories is inherently unevenly distributed across timesteps and trajectories. Some denoising steps and trajectories contribute significantly more learning signal than others. To exploit this observation, we propose two complementary strategies that emphasize informative training data.

First, we introduce a per-timestep weighting scheme that selectively reweights denoising steps during optimization. We theoretically motivate this weighting by analyzing the relationship between Group Relative Policy Optimization (GRPO) and Proximal Policy Optimization (PPO), showing that each timestep should ideally be weighted proportionally to the variance of the TD-error PPO advantage. Although this quantity is infeasible to estimate efficiently during training, we analyze its behavior and approximate the trend using a simple heuristic based on the squared magnitude of per-timestep latent change. Empirically, this weighting improves training sample efficiency.

Second, we introduce a trajectory replay mechanism inspired by replay buffers commonly used in robotics-based reinforcement learning [8]–[10]. Instead of discarding previously sampled trajectories, we reuse informative trajectories during training. Drawing inspiration from prioritized replay buffers [8], we find that hard-mining trajectories with the largest advantages [11] provides the most sample-efficient learning. By emphasizing informative trajectories, the model can continue improving without repeatedly sampling new trajectories and reward evaluations.

Overall, our method improves the efficiency of the baseline by 2–6\(\times\) while remaining simple and compatible with existing diffusion RLHF pipelines.

In summary, our main contributions are:

• A computationally practical and mathematically motivated per-timestep weighting scheme that mitigates the credit assignment problem in diffusion RLHF.

• A replay buffer mechanism that retrieves informative past trajectories during training, reducing the need for repeated reward queries.

• A simple, plug-and-play method that integrates seamlessly into existing diffusion RLHF pipelines without architectural changes.

• Extensive experiments demonstrating consistent gains across diverse reward functions, underscoring the generality of our approach.

2 Related Work↩︎

2.1 Replay Buffers↩︎

A replay buffer stores previous environment interactions for a policy to train on. It is standard in off-policy methods such as DDPG [12], Soft-Actor Critic [13], and Deep Q-Networks [14], but not in on-policy methods like PPO [15] and TRPO [16], or in current RLHF methods, since these rely on sampling from the most recent weights to stay in-distribution.

Many replay-buffer variants exist. Hindsight experience replay [9] relabels failed trajectories with new goals to synthesize successful ones. Others prioritize transitions under certain criteria. The approach [8] samples by high temporal-difference (TD) error as a proxy for how undertrained a network is at a given point, and Energy-Based Hindsight [17] prioritizes high-kinetic-energy trajectories. Synthetic trajectories generated by diffusion models were proposed by [10], and retrieval was shown by [18] to enhance robot imitation learning. Finally, [19] hardmines previous PPO trajectories using cumulative TD error.

2.2 RLHF↩︎

RLHF has become the central paradigm for aligning large language models (LLMs). Modern RLHF was introduced by [20], who scaled it to pretrained LLMs via a three-step pipeline of supervised fine-tuning, reward modeling, and reinforcement learning. Subsequent approaches such as Group Relative Policy Optimization [21] remove the reward model entirely in favor of group-normalized advantage, and [22] shows GRPO prevents catastrophic forgetting better than supervised fine-tuning. Preference-based approaches such as [23] have a human rank one sample over another rather than give a numerical reward. More recently, [24] segments the LLM token stream before computing segment-level reward, and [25] applies GRPO with prompt hardmining.

RLHF has also been explored in robotics: [26], [27] train policies on preference-based evaluation of robot trajectories, [28] uses reward-guided evolutionary search over a pretrained trajectory diffusion model, and [29] finetunes pretrained imitation-learning diffusion policies such as diffusion policy [30].

2.3 Diffusion RLHF↩︎

RLHF is also effective for aligning diffusion models with human preference. Diffusion RLHF began with DPOK and DDPO [2], [3], which applied RL to diffusion and outperformed reward-weighted regression (RWR) [31]. Later approaches such as B2-DiffuRL [6] and Branch-GRPO [32] branch from shared input noise to sharpen contrast between trajectories, and another line of work demonstrates GRPO can be applied in flow matching and diffusion [32]–[35]. In contrast to GRPO, preference-based methods [7], [36], [37] have a human pick between generated image pairs rather than assign a numerical reward and train on this feedback. Diffusion RLHF has been shown effective in video diffusion [38] and 2-step diffusion [39] settings.

In contrast, some methods [40]–[42] assume a differentiable reward (e.g.a classifier or VLM) and backpropagate through it directly, while others localize spatial regions of reward importance [42], [43]. The work SIPO [44] reweights trajectory timesteps for more stable training under non-numerical pair-wise preference feedback, and a recent work TempFlow [45] explores timestep weighting in flow matching rather than DDPM-style models, but without hardmining. Moreover, they evaluate on a single reward function, whereas our evaluation across 5 reward functions shows their weighting scheme underperforms ours.

3 Preliminaries↩︎

3.1 Markov Decision Process↩︎

Formally, a MDP  [46] can be represented as a tuple \((S, A, P, R,\gamma)\), where \(S\) is the state space, \(A\) is the action space, \(P(s' | s,a)\) is the probability transition function from current state \(s\) and action \(a\) to new state \(s'\), \(R(s,a)\) is the reward function, and \(\gamma\) is the discount factor. At each timestep \(t\), the agent has access to state \(s\) and chooses an action \(a\). The result is a new state \(s'\) sampled from \(\mathcal{P}(s'| s,a)\) and a reward, denoted as \(R(s,a)\).

The goal in reinforcement learning is to maximize the sum of decayed rewards, formally denoted as \(J(\pi) = \mathbb{E}_{\tau\sim p_\pi}\left[\sum_{i=0}^n\gamma^iR(s_i,a_i)\right]\).

3.2 Diffusion Models↩︎

For a given data sample \(x_0 \in \mathbb{R}^d\), the forward diffusion process gradually adds noise over \(T\) steps to attain the noised trajectory \((x_0, x_1, ..., x_T)\). Mathematically, the forward process can be defined as \(q(x_t | x_{t-1}) = \mathcal{N}(x_t ; \sqrt{1 -\beta_t}x_{t-1}; \beta_tI)\), where \(\beta_t\) is the monotonic noise schedule. The reverse diffusion process consists of a trained model \(p_\theta\) where \(p_\theta(x_{t-1} | x_t)\) learns to approximate the posterior distribution \(q(x_{t-1}| x_t, x_0)\), which is also a normal distribution with monotonic variance schedule \(\tilde{\beta}_t\).

With reward feedback, the reverse diffusion process can be interpreted as a MDP with the action being the predicted noise from the diffusion model, the state being the latent space of images, and the transition function being a sampler such as DDIM  [47] and DDPM  [48] that inputs the current latent and predicted noise.

3.3 PPO and Diffusion RLHF↩︎

PPO [15] is a commonly used method in classical reinforcement learning. Formally, the PPO uses the loss function: \[L_{\text{PPO}} = \min\big( A_t \tfrac{\pi_\theta(a|s)}{\pi_{\theta_{\text{old}}}(a|s)}, \operatorname{clip}(\tfrac{\pi_\theta(a|s)}{\pi_{\theta_{\text{old}}}(a|s)},1-\epsilon,1+\epsilon)A_t \big)\] where \(A_t\) is the advantage at timestep \(t\), \(\epsilon\) is the clipping value to prevent overly large weight shifts, and \(\pi_{\theta_\text{old}}(a|s)\) is the probability of taking action \(a\) with the old parameters in \(\theta_\text{old}\) as opposed to the current parameters \(\theta\). In classical robotics RL, the advantage is given by the TD-error \(A_t = \gamma V(s_{t+1}) + R_{t+1} - V(s_t)\), where \(V\) is the value function. In the context of diffusion RLHF, prior methods do not use the TD-error but rather compute the uniform advantage \(A=\frac{r - r_\text{mean}}{\text{std}(r)}\), where \(r\) is the reward at the end of the trajectory, \(r_\text{mean}\) is the average reward in a batch of samples, and \(\text{std}(r)\) is the standard deviation of the rewards in a batch. The same advantage value \(A\) is used for all timesteps in the trajectory.

4 Methodology↩︎

In existing diffusion RLHF pipelines, informative training signals are often obscured by uninformative or low-reward data. Our approach addresses this issue at two levels: (1) timestep-level credit assignment and (2) trajectory-level data selection. We first analyze the credit assignment problem and derive a timestep-dependent weighting scheme, and then introduce a replay-based hard-mining strategy to prioritize informative trajectories.

First, at the timestep level (Sec. 4.2), we try to address the credit assignment problem. While standard diffusion RLHF methods assign uniform credit to all denoising steps, we introduce a practical approximation to a theoretically motivated weighting scheme that assigns greater importance to diffusion steps that contribute most to the final reward (Fig. 1 (left)). Second, at the trajectory level (Sec. 4.4), we introduce a hard-mining mechanism that selectively reuses past trajectories with high absolute advantage, ensuring that the most informative experiences are emphasized during learning (Fig. 1 (right)).

4.1 Non-Uniform Credit Assignment↩︎

Traditional diffusion RLHF uses GRPO math to compute a single advantage \(A = \frac{R - R_\text{mean}}{\text{std}(R)}\). Depending on the diffusion RLHF method, \(R_\text{mean}\) and \(R_\text{std}\) either refer to the group or batch mean and standard deviation. This advantage constant \(A\) is applied evenly in its loss function. The same constant is used in each timestep: \[\mathcal{L}_{\text{ddpo}} = \sum_{t=1}^n -A \cdot\log\left(P_\theta(x_{t-1} | x_t, c)\right)\] This loss function is problematic because it does not account for how much each step contributes to the final reward. To better understand the credit assignment issue, we first analyze how reward evolves along a diffusion trajectory.

4.1.1 Case Study.↩︎

Let’s look at the reward corresponding to the predicted clean image \(\hat{x}_0(x_t)\) at timestep \(t\) be defined as \(R_t = R(\hat{x}_0(x_t))\) for some reward function \(R\). For brevity, we denote this as \(R(x_t)\) since \(\hat{x}_0\) varies smoothly with \(x_t\). Consider the quantity \(\Delta R_t = R(x_{t+1}) - R(x_t)\). For a small diffusion update \(x_{t+1} = x_t + h_t\), a first order Taylor series expansion gives \(R(x_{t+1}) \approx R(x_t)+h_t^\top\nabla R(x_t)\). Summing over timesteps (\(t\)) yields a telescoping decomposition \(R(x_T)-R(x_0)\approx\sum_t \Delta R(x_t)\) under the assumption of local linearity. This provides an approximate decomposition of the final reward in terms of \(\Delta R_t\) and initial reward. If we disregard the initial reward \(R(x_0)\) (which should average to the same quantity given random initial noise), \(\Delta R_t\) can be viewed as a crude proxy for the contribution of timestep \(t\) to the final reward. In Fig. 3, we observe that \(\Delta R_t\) varies significantly across timesteps, suggesting that different parts of the denoising trajectory contribute unequally to the final reward signal. We emphasize that \(\Delta R_t\) should not be used directly for credit assignment, since it captures only the immediate reward change and does not account for the long-term effect of the state \(x_t\) on future denoising steps.

4.1.2 Advantage Variance Proportional Weighting as a Natural Consequence.↩︎

Motivated by this observation, we next connect diffusion RLHF to PPO to derive a principled timestep-level credit assignment scheme.

A key difference between PPO and GRPO is that PPO does per-timestep advantages based on TD-error, thereby reducing the credit assignment issue that occurs with GRPO. We derive an equation to reconstruct the PPO timestep-level advantages from a singular final advantage that is computed using GRPO. We show that the PPO timestep advantages are proportional to the trajectory level advantage \(A_\text{final}\) for some timestep-dependent constants \(w(t)\). By using these PPO advantage values over the uniform GRPO advantages, we can attain better convergence.

The original PPO uses TD-error advantages as: \[A_t = \gamma V(s_{t+1}) + R_{t+1} - V(s_t)\] where \(A_t\) is the timestep-level advantage and \(V\) is the value function from reinforcement learning.

For simplicity of notation, we assume a 20 step diffusion process. Let state \(s_k = \{x_k, T_k\}\) where \(T_k\) is the diffusion timestep and \(x_k\) is the noisy image. Consider a diffusion trajectory \((s_{0}, s_1, ..., s_{20})\) where state \(s_0\) corresponds to pure noise, \(s_{20}\) corresponds to the final image, and \(s_k\) corresponds to intermediate latents. Consider the series of actions \(a_0, ..., a_{19}\). In a RLHF scenario, there is only a single reward from \(s_{20}\), so denote the final reward as \(R_\text{final}\) and all intermediate rewards as 0. Finally, we assume that \(\gamma=1\), as is typical in RLHF.

We have that for \(t = 0, ..., 19\) our equation for \(A_t\) simplifies to: \[A_t = \gamma V(s_{t+1}) + R_{t+1} - V(s_t) = V(s_{t+1}) - V(s_t)\] We have that \(V(s_{20}) = R_\text{final}\). This leads us with \[A_0 + A_1 +... + A_{19}\] \[= \left[V(s_1)-V(s_0)\right] + \left[V(s_2)-V(s_1)\right] + ... + \left[V(s_{20})-V(s_{19})\right]\] \[= V(s_{20}) - V(s_0)= R_\text{final} - V(s_0)\]

Since \(s_0\) is pure noise sampled from \(\mathcal{N}(0,1)\), there is a very high diversity of denoising trajectory that can originate from the same start noise due to the high resampling noise of existing samplers. Thus, we assume that \(V(s_0) \approx \mathbb{E}_{s_0\sim N(0,1)}(R_\text{final}) = R_\text{mean}\). Thus, \[A_0+... + A_{19} = R_\text{final} - V(s_0) \approx R_\text{final} - R_\text{mean} = A_\text{final}\cdot \text{std}(R)\] Under the assumption that we only know the final advantage \(A_\text{final}\) and nothing else, we can estimate each \(A_t\). We have that \(A_t = V(s_{t+1})-V(s_t)\) is zero-meaned over all trajectories because \[V(s_{t}) = \mathbb{E}_{s_{t+1}|s_t}\left[\gamma V(s_{t+1}) + R_t\right] = \mathbb{E}_{s_{t+1}|s_t}\left[V(s_{t+1})\right]\] \[\implies \mathbb{E}_{s_{t+1}|s_t}\left[A_t\right] = \mathbb{E}_{s_{t+1}|s_t}\left[V(s_{t+1})-V(s_{t})\right] = 0\]

Let \(\sigma_t\) be the standard deviation of \(A_t\) over the distribution of trajectories generated by \(\pi\). We further approximate each \(A_t\) as coming from a gaussian \(\mathcal{N}(0,\sigma_t^2)\).

We want to find the conditional distribution for \(P\!\left(A_k \mid \sum_{i=0}^{19}\!A_i = A_{\text{final}}\!\cdot\!\operatorname{std}(R)\right)\). This is the conditional distribution of a multivariate normal under linear constraint. This gives us the distribution:

\[A_k \Big | \sum_{i=0}^{19}A_i = A_\text{final} \cdot \text{std(R)} \sim \mathcal{N}\!\left(\mu, \sigma^2\right),\] \[\text{where } \mu = \frac{\sigma_k^2 \, \text{std}(R)}{\sum_{i=0}^n \sigma_i^2}\, A_\text{final},\quad \sigma^2 = \sigma_k^2\!\left(1 - \frac{\sigma_k^2}{\sum_{i=0}^n \sigma_i^2}\right)\]

Refer to appendix for more details on how this distribution was derived. Thus, the expected value/mean of \(A_k\) is a constant times \(A_\text{final}\) where the constant is: \[w(t) = \frac{\sigma_t^2\cdot \text{std}(R)}{\sum_{i=0}^n \sigma_i^2}\] Note that \(w(t)\) is larger when \(\sigma_k\) is larger, which happens in timesteps that change the final reward the most. Thus, if we align a schedule \(w(t)\) correctly with \(\sigma_t\), then we can construct accurate estimates of the PPO timestep-level advantages \(A_t\). In our formulation, \(\sigma_t\) does not depend on \(s_t\) or \(a_t\) and solely depends on the timestep \(t\). This allows us to use the same \(w(t)\) universally for different trajectories.

4.1.3 Empirical Measurement of \(\text{Var}(A_t)\)↩︎

Intuitively, the advantage variance \(\text{Var}(A_t)=\sigma_t^2\) should vary across diffusion timesteps, since the forward process adds noise with timestep-dependent variance \(\beta_t\). In Fig. 4 we estimate \(\sigma_t^2\) by denoising a trajectory partially till timestep \(t\), and branching off into multiple denoising trajectories and computing variance of their final rewards. This is undesirable during training due to the need of large number of branches needed to accurately compute the variance. However, this clearly indicates the need for a non-uniform weighting scheme and provides an initial indication of the expected weighting trend. We see here a monotonically decreasing trend (except for a slight increase in the final steps, which might be due to image generation stochasticity).

4.2 Strategy 1: Per-Timestep Weighting↩︎

Previously, we motivated the case for using a non-uniform weighting scheme using a case study. We then demonstrated that the PPO timestep-level advantage can be derived from the original GRPO advantage under certain assumptions, and that it depends directly on the variance of the advantage. While this derivation provides theoretical guidance, estimating these quantities exactly during training is impractical. We therefore seek a simple, practical weighting scheme that approximates this behavior.

We experiment with multiple weighting schemes as shown in Fig. 5. An obvious choice for such a monotonically decreasing weighting scheme is using the reverse diffusion standard deviation \(\sqrt{\tilde{\beta}_t}\) used for adding randomness during sampling. (A recent work [45] tried to use a similar strategy of weighting using the standard deviation of noise in Flow-GRPO SDE.) We additionally consider the change in diffusion model’s mean absolute latent change \(|z_{t}-z_{t-1}|\) and mean squared latent change \(|z_{t}-z_{t-1}|^2\) as other choices for weighting schemes, which share the same monotonicity. Here, \(z_t\) is the diffusion state in the latent space of Stable Diffusion VAE. As opposed to the previous static weighing schemes, this last scheme is dynamic, as it is dependent on the latents which in turn are dependent on how the current model denoises. Finally, we also try gaussian weighting schemes with early, middle, and late focus w.r.t.. to the sampling steps. These are not monotonic but can signal towards which sampling step should be focused on more.

We replace the GRPO advantage \(A_\text{final}\) with our new advantage \(w(t) A_\text{final}\). Our new loss is: \[\mathcal{L}_{\text{reweighted}} = \sum_{t=1}^n -A_\text{final} \cdot w(t)\cdot\log\left(P_\theta(x_{t-1} | x_t, c)\right)\] where \(w(t)\) is a scalar that only depends on the timestep of the transition.

4.3 Replay Buffers↩︎

So far, we focused on improving credit assignment at the timestep level. We now address the second source of inefficiency: how trajectories are reused during training. In classical RL for robotics, methods [12]–[14] often make use of a replay buffer that stores previous trajectory from previous epochs. During training, episodes are drawn from the replay buffer and used to train the model alongside data from the current epoch. This allows the policy to sample from a larger dataset than just the trajectories from the current epoch. An extension of replay buffer is the prioritized replay buffer  [8] which replays transitions with a high TD training error and has been shown to speed up training.

In contrast to replay buffers, GRPO disregards diffusion trajectories after the current epoch is done. While this online sampling allows the method to train on trajectories that are generated by the most up-to-date weights, we argue that this is still highly sample inefficient and storing previous trajectories is more effective at increasing sample efficiency. TD error is not available in RLHF settings because rewards are provided only at the end. This prevents us from implementing a prioritized replay buffer. Instead, we implement a buffer that prioritizes high absolute advantage values.

4.4 Strategy 2: Historical Hardmining↩︎

Building on this replay buffer formulation, we propose a hard-mining strategy that prioritizes trajectories with high absolute advantage values. Our intuition is that a trajectory with high absolute advantage, whether strongly positive or strongly negative, contains one or more timesteps that substantially changed the reward, making it highly informative training data. We perform two different methods of trajectory replay: (1) Trajectory-level hardmining, where we find the highest \(k\) trajectories based on \(|A|\) in the previous few epochs, and (2) Random sampling of previous trajectories. We find that keeping a replay buffer of only the last few epochs of trajectories was the most advantageous because sampling from an epoch too long ago resulted in out-of-distribution trajectories.

Overall, our method combines these two strategies: timestep-level reweighting and trajectory-level hard mining. Both strategies are motivated by the idea of reweighting the training data so that more informative samples are emphasized during training. See Algorithm [Our95algorithm] for details.

5 Experiments↩︎

Since our method is highly applicable to existing diffusion RLHF frameworks, we choose to incorporate our method with existing baselines to show the significant increase in sample efficiency. For all experiments, we use the same list of animal training prompts as [2]. We use LoRA [49] weights with rank 4 matrices that were inserted into RunwayML Stable Diffusion v1.5 [50]. For our experiments, we test our method applied to DDPO [2], DPOK [3], and B2-DiffuRL [6]. See full list of hyperparameters in the appendix.

5.1 Reward Functions↩︎

We test our method on 5 reward functions. We select 3 rewards from the DDPO [2] paper. We also experiment with HPS v2 [51] and Image Reward [52]. The first reward function we test on is JPEG compressibility, which prefers compressed images and favors smoother textures that are easier to compress. Conversely, our second reward function is JPEG incompressibility, which favors highly detailed images with diverse textures. Third, we test on the reward function given by the Aesthetic Score classifier from [53] that is trained on human aesthetic preferences. Finally, both HPS v2 and Image Reward simulate human preference by training on aesthetic and prompt-adherence human feedback on generated images.

5.2 Sample Efficiency↩︎

In Figure 10 (a), we see that our modifications are able to achieve higher reward than all existing methods within the same budget of reward queries. For fair comparison, each baseline and its augmented counterpart share identical hyperparameters. For qualitative examples of training images, see Figures 7 and 8.

We also show that our method is capable of generalizing to novel prompts with higher sample efficiency than existing methods. For the list of unseen prompts, we prompted ChatGPT [54] for animals not present in the prompt dataset. As seen in Table \(\ref{tab:Generalization95Table}\), we find that our method still allows for more generalization than previous methods on the same budget of reward queries.

5.3 Ablations↩︎

Since we do both per-timestep weighting and historical hardmining, we show that both of these methods independently make diffusion RLHF more sample efficient. We consider the 4 possibilities: (1) Neither hardmining nor timestep weighting (baseline), (2) Hardmining but not timestep weighting, (3) Timestep weighting but not hardmining, and (4) Both timestep weighting and hardmining. Our results show that each of our techniques contributes significantly to the overall efficiency of the finetuning method. See Figure 10 (b) where, for all the rewards, using hardmining and weighting together is more effective than using only one.

5.4 Timestep Weighting↩︎

Empirically, we find that timestep weighting consistently provides the most sample-efficient results, sometimes up to 6\(\times\) efficiency. Despite our reward functions focusing on different aspects of an image, our timestep weighting boosts sample efficiency for all cases, as seen in Figures 10 (b) and 10 (a).

Choice of Timestep Weighting. Using our Gaussian weighting schemes, we find that weighting earlier timesteps more heavily works better than weighting middle timesteps, which in turn works better than weighting later timesteps. This suggests that a monotonically decreasing weighting scheme is preferable. For reference, uniform (unweighted) weighting falls somewhere between the mid and early focus approaches in sample efficiency.

In our experiments, the most consistently well-performing scheme across all reward functions is the mean squared latent change, \(|z_{t}-z_{t-1}|^2\) (see Fig. 10 (c)). This even outperforms the DDIM sampling standard deviation proposed in the recent work TempFlow [45]. We believe this scheme works well because it assigns higher weight to the timesteps where the model makes the largest changes to the latent representation. By prioritizing these high-change timesteps during reward-alignment training — rather than using fixed coefficients — the model is guided more effectively toward better rewards.

5.5 Historical Hardmining↩︎

We compare to historical sampling, where samples are randomly selected, to hardmining. As seen in Figure 10, we see that for all reward functions, hardmining is more effective than simply rerunning random samples from the previous epochs. This validates our hypothesis that large advantage trajectories are more informative than other trajectories.

5.6 CLIP Score↩︎

In addition to testing sample efficiency, we also test CLIP score [55] on the finetuned models to measure prompt-image adherence. CLIP score measures the embedding similarity between an image and its prompt, serving as a proxy for how faithfully the image reflects the text. We generated 1k images using training prompts and computed the average CLIP score, as seen in Table [tab:CLIP95Table]. We find that there is no visible degradation of prompt adherence compared to the baseline DDPO.

5.7 Generalization Score↩︎

Following experiments [2], we also test generalization score. For our experiments, we first finetune a model based on a reward function for a given number of queries. Afterwards, we find a new list of prompts that were not in the list of training prompts and test how well our model keeps the reward high in those scenarios. As seen in Table 9, we find that our method also generalizes well to other prompts.

6 Conclusion↩︎

In this paper, we provide a simple but effective method to enhance the sample efficiency of diffusion RLHF. Our method is easily adaptable to existing RLHF platforms and has the potential to speed up training up to 6 times while retaining prompt-adherence and the ability to generalize to novel prompts. Our method departs from the common RLHF assumption that each timestep affects the output reward to the same extent. We ground our method on the intuition that both certain timesteps and certain trajectories are more informative than others during training. Future directions could analyze how different schedulers and \(\alpha\)-schedules affect the training sample efficiency or extend the method preference-based models such as D3PO [7].

7 Appendix↩︎

7.1 Hyperparameters For Baselines↩︎

We chose to use the same hyperparameters for baselines as those chosen for all baselines in the D3PO [7] paper. The only difference is that we chose a smaller number samples per epoch as we found that sampling more often led to better performance. We also changed to AdamW over Adam, as consistent with most diffusion RLHF methods. Between baseline and baseline + our method, we kept all hyperparameters the same. All results were obtained with 4 RTX A5000 GPUs or 4 RTX A4000 GPUs. See a comprehensive list of hyperparameters in Tables 1 and 2.

7.2 Prompt List↩︎

We used the same training prompt list as in  [2]. The list of animals is listed below:

For generalization experiments, we used the following list of animals generated by ChatGPT:

7.3 Comparisons with Hardmining↩︎

When using hardmining, even though we use lesser number of reward calls, we use more backpropagation through the network. Hence, an obvious question might be whether using hardmining helps just because there are more samples whose gradients are being backpropagrated or is hardmining helping because of the selection of more important trajectories. In this section, we compare hardmining with two more methods. First, we compare to simply increasing the learning rate proportional to how many more gradient steps we take. Second, we compare to randomly replaying entries from the current epoch. These two settings correspond to using faster learning rate and more gradient calls respectively, both of which are alternatives to hardmining. Note that this is different from Figure. 10 where we compared to random retrieval from previous epochs as opposed to only the current epoch. As seen in Figure 11, our hardmining based on the absolute advantages does better than both increasing the learning rate and repeating samples from the current epoch, showing that hardmining helps not just because of increased samples.

7.4 Proof That \(A_k\) has mean \(\displaystyle\frac{\sigma_k^2\cdot \text{std}(R)}{\sum_{i=1}^n \sigma_i^2}A_\text{final}\)↩︎

We want to find the probability distribution of \(A_k\) where we assume that each \(A_i\) has mean 0 and standard deviation \(\sigma_i\) for \(i \in \{0,1,2,...,n\}\). Furthermore, we have the condition that for a given constant \(C\), \[A_0 +A_1 + A_2 +...+A_n = C\] For simplicity of notation, assume that \(k=0\). The proof can easily generalize to other values of \(k\).

Let \(B = A_1 + A_2+...+A_n\). The mean of \(B\) is the sum of the means of \(A_1+A_2,...A_n\) which is 0. The variance of \(B\) is the sum of the variances of \(A_1+A_2,...A_n\). This means that \(\sigma_B^2 = \sigma_1^2 + \sigma_2^2 + ... + \sigma_n^2\). Finally, we have that from the summation of Gaussian distributions is still a Gaussian distribution. Thus, \[B = A_1+A_2 + A_3 + ...+A_n \sim \mathcal{N}\left(0, \sigma_B^2\right)\] We can rewrite our question as finding the distribution \(A_0 \;|\; A_0+B=C\). We have that the for a given value \(a\) such that \(A_0=a\),

\[\begin{align} &P(A_0=a \;|\;A_0+...+A_n=C)\nonumber\\ &= P(A_0=a \;|\;A_0+B=C) \nonumber\\ &= P(A_0=a \;\&\;A_0+B=C) / P(A_0+B=C) \nonumber\\ &= P(A_0=a \;\&\;B=C-a) / P(A_0+B=C) \nonumber\\ &\propto P(A_0=a \;\&\;B=C-a) \nonumber\\ &= P(A_0 = a)P(B=C-a) \nonumber\\ &= \frac{1}{\sqrt{2\pi\sigma_0^2}}\exp\left({-\frac{a^2}{2\sigma_0^2}}\right)\cdot \frac{1}{\sqrt{2\pi\sigma_B^2}}\exp\left({-\frac{(C-a)^2}{2\sigma_B^2}}\right) \nonumber\\ &= \frac{1}{2\pi \sigma_0\sigma_B}\exp\left({-\frac{a^2}{2\sigma_0^2}-\frac{(C-a)^2}{2\sigma_B^2}}\right) \nonumber\\ &\propto \exp\left({-\frac{a^2}{2\sigma_0^2}-\frac{(C-a)^2}{2\sigma_B^2}}\right)\nonumber\\ &= \text{exp}\left({-\frac{a^2}{2\sigma_0^2}-\frac{(C-a)^2}{2\sigma_B^2}}\right) \nonumber\\ &= \text{exp}\left({\frac{-a^2\sigma_B^2-(C-a)^2\sigma_0^2}{2\sigma_0^2\sigma_B^2}}\right)\nonumber\\ &= \text{exp}\left({\frac{-a^2(\sigma_0^2+\sigma_B^2)+2C\sigma_0^2a - C^2\sigma_0^2}{2\sigma_0^2\sigma_B^2}}\right) \nonumber\\ &= \text{exp}\left({\frac{-a^2(\sigma_0^2+\sigma_B^2)+2C\sigma_0^2a - C^2\sigma_0^2}{2\sigma_0^2\sigma_B^2}}\right)\nonumber\\ &= \exp\Bigg( \frac{ -(\sigma_0^2+\sigma_B^2) \left(a - C\frac{\sigma_0^2}{\sigma_0^2+\sigma_B^2}\right)^2 }{ 2\sigma_0^2\sigma_B^2 } \Bigg) \nonumber\\ &\times \exp\Bigg( \frac{ C^2\frac{\sigma_0^4}{\sigma_0^2+\sigma_B^2} - C^2\sigma_0^2 }{ 2\sigma_0^2\sigma_B^2 } \Bigg) \nonumber\\ &= \underbrace{ \exp\!\left( \frac{ C^2\frac{\sigma_0^4}{\sigma_0^2+\sigma_B^2} - C^2\sigma_0^2 }{ 2\sigma_0^2\sigma_B^2 } \right) }_{\text{constant}} \nonumber\\ &\times \exp\!\left( \frac{ -(\sigma_0^2+\sigma_B^2) \left(a - C\frac{\sigma_0^2}{\sigma_0^2+\sigma_B^2}\right)^2 }{ 2\sigma_0^2\sigma_B^2 } \right) \nonumber\\ &\propto \text{exp}\left({\displaystyle \frac{-(\sigma_0^2+\sigma_B^2)\left(a - C\frac{\sigma_0^2}{\sigma_0^2+\sigma_B^2}\right)^2}{2\sigma_0^2\sigma_B^2}}\right) \nonumber\\ &= \text{exp}\left({\displaystyle \frac{-\left(a - C\frac{\sigma_0^2}{\sigma_0^2+\sigma_B^2}\right)^2}{2\cdot \frac{\sigma_0^2\sigma_B^2}{\sigma_0^2+\sigma_B^2}}}\right)\nonumber\\ \nonumber \end{align}\]

We have that the end result is a gaussian with mean \[\mu = C\frac{\sigma_0^2}{\sigma_0^2+\sigma_B^2} = C\frac{\sigma_0^2}{\sigma_0^2+\sigma_1^2+ ... + \sigma_n^2}\] and variance with value \[\frac{\sigma_0^2\sigma_B^2}{\sigma_0^2 + \sigma_B^2} = \sigma_0^2 \left(1 - \frac{\sigma_0^2}{\sum_{i=0}^n\sigma_i^2}\right)\] Letting \(C = \text{std}(R)A_\text{final}\), we get that the distribution of \(A_0\) under the assumption that \(A_0 + A_1+...+A_n=\text{std}(R)A_\text{final}\) has mean \(\mu = \frac{\sigma_0^2}{\sigma_0^2+\sigma_1^2 + ... + \sigma_n^2}\text{std}(R)A_\text{final}\)

7.5 More Image Samples↩︎

In figures 12, 13, 14, 15, and 16, we give a randomly selected sample of DDPO vs DDPO with our method. We use 30 random seeds in the grid for the 30 images and denoise using finetuned weights.

References↩︎