Beyond Pairwise Preferences: Listwise Reward-Aware Alignment for Diffusion Models

Technical Summary: Diffusion LAIR

Problem Statement
Current preference optimization methods for aligning text-to-image diffusion models, such as Direct Preference Optimization (DPO) and its diffusion-specific variants (e.g., Diffusion-DPO, DSPO), largely rely on pairwise supervision. These methods reduce training data to binary "winner-loser" comparisons for a given prompt. The authors argue this reduction is limiting for two primary reasons:

1. Data Structure: Practical preference datasets (e.g., Pick-a-Pic) often contain multiple candidate images for a single prompt, creating a natural listwise structure that pairwise methods ignore.
2. Information Loss: Modern reward models provide continuous scores. Reducing these to a single binary label discards the relative quality of non-selected candidates and ignores the magnitude of reward gaps between samples.
   Furthermore, existing online reinforcement learning (RL) approaches for diffusion alignment are computationally expensive, difficult to parallelize, and sensitive to distribution shifts due to their on-policy nature. While offline methods are more efficient, they have not fully leveraged the richer signal available in listwise, reward-scored data.

Methodology: Diffusion LAIR
The paper proposes Diffusion LAIR (Listwise Advantage-weighted Implicit Reward), an offline preference optimization method that treats alignment as a listwise, reward-aware problem.

• Data Representation: Instead of independent pairs, the method processes a prompt $c$ with a set of $N_c$ candidate images $\{x^{(i)}_0\}$ and their corresponding continuous reward scores $\{r_i\}$.
• Advantage Weighting: The reward scores are converted into centered advantage weights ($w_i$). First, a softmax probability $p_i$ is computed using a temperature $\tau$. Then, the weights are centered against a uniform baseline:
  $$w_i = p_i - \frac{1}{N_c}$$
  This ensures $\sum w_i = 0$, assigning positive weights to high-reward samples and negative weights to low-reward samples relative to the group average.
• Optimization Objective: The method optimizes an advantage-weighted regression objective on the implicit reward. The implicit reward is defined as the denoising-loss improvement of the current model ($\theta$) over a fixed reference model ($ref$):
  $$s_\theta(x_0, x_t, t, c, \epsilon) = \omega(\lambda_t) (l_{ref} - l_\theta)$$
  The loss function is:
  $$L_{LAIR} = \mathbb{E} \left[ -\sum_{i=1}^{N_c} w_i s^{(i)}_\theta + \frac{\lambda}{N_c} \sum_{i=1}^{N_c} (s^{(i)}_\theta)^2 \right]$$
  The first term encourages the model to increase the implicit reward for high-reward samples and decrease it for low-reward samples. The second term is a quadratic penalty that regularizes the magnitude of the implicit reward, preventing aggressive deviation from the reference model.

Theoretical Analysis
The authors provide a theoretical characterization of the objective:

1. Closed-Form Optimum: The objective admits a unique, bounded closed-form optimum in the implicit-reward space: $s^*_i = \frac{N_c}{2\lambda} w_i$. This clarifies that the regularization strength $\lambda$ directly controls the magnitude of the preference update.
2. Surrogate KL Bound: Under specific assumptions (log-ratio approximation and bounded full-support extension), the authors derive a bound on the Kullback-Leibler (KL) divergence between the induced distribution and the reference distribution: $D_{KL} \le \frac{\Delta_c}{\eta}$. This suggests that $\lambda$ controls the sharpness of the preference tilting, ensuring the update remains conservative.
3. Advantage over Pairwise: Unlike pairwise objectives (e.g., Bradley-Terry) which can drive the implicit reward to infinity to maximize the margin between a winner and loser, LAIR's listwise formulation assigns finite, relative targets based on the entire candidate set, preserving ranking structure and reward spacing.

Experimental Results
The method was evaluated on Stable Diffusion 1.5 (SD1.5) and Stable Diffusion XL (SDXL) using the Pick-a-Pic v2 dataset.

• Setup: Images were scored offline using the PickScore reward model. The models were fine-tuned and evaluated on text-to-image generation (HPD, Parti-prompts), compositional generation (GenEval), and instruction-based image editing (InstructPix2Pix).
• Performance: Diffusion LAIR outperformed strong baselines (Diffusion-DPO, DSPO, Diffusion-KTO, MaPO, InPO) across most metrics.
  • On SD1.5, it achieved the highest scores on PickScore, HPS, Aesthetics, CLIP, and ImageReward for both Parti-prompt and HPD datasets.
  • On SDXL, it similarly outperformed baselines, including MaPO and Diffusion-DPO.
  • Generalization: Notably, despite being trained only with PickScore, the model showed significant improvements in other metrics (e.g., ImageReward), suggesting the listwise objective improves general sample quality beyond the specific reward model used.
  • Compositional & Editing: The method demonstrated consistent gains in GenEval (compositional reasoning) and InstructPix2Pix (image editing) win rates compared to pairwise baselines.
• Efficiency: The training process was significantly more efficient than online RL baselines and comparable or more efficient than other offline methods, requiring approximately 24 GPU hours for SD1.5 and 140 A100 GPU hours for SDXL.

Significance and Claims
The paper claims that Diffusion LAIR offers a simple and effective direction for aligning diffusion models by moving beyond the limitations of pairwise reduction. Its key contributions are:

1. Listwise Learning: It is the first method to directly leverage groups of reward-scored images for diffusion preference optimization, utilizing the full structure of relative sample quality.
2. Theoretical Clarity: It provides a bounded closed-form optimum, offering insight into how regularization controls distribution shift and update magnitude.
3. Empirical Superiority: It demonstrates strong performance across generation, composition, and editing tasks, proving that listwise, reward-aware optimization yields broader improvements than pairwise methods.

The authors note limitations, specifically the reliance on the quality of the reward model used for training and the idealized assumptions in the theoretical KL bound. However, they posit that the method is robust to noisy preference labels common in human-annotated datasets.