Step-level Denoising-time Diffusion Alignment with Multiple Objectives

1. Problem Statement

Diffusion models are typically pre-trained on large-scale datasets to match data distributions but lack alignment with specific human preferences (e.g., aesthetic quality, text-image consistency, safety). While Reinforcement Learning (RL) fine-tuning has been successful for single-objective alignment, real-world applications require balancing multiple, often conflicting objectives.

Existing approaches to multi-objective alignment face significant limitations:

• Multi-objective RL Fine-tuning: Requires training a massive number of models (exponential in the number of objectives) to cover the space of preference weights, or solving complex conflict-avoiding update directions, which is computationally prohibitive.
• Denoising-time Alignment (Training-free): Existing methods that fuse pre-aligned models at inference time generally fall into two categories:
  1. Reward-dependent: Require access to reward functions and their gradients/values, necessitating repeated sampling and estimation (e.g., via Tweedie estimators), which is computationally expensive.
  2. Reward-free: Fuse models without reward access (e.g., Reward Soup, DB-MPA, DERADIFF). However, these methods often rely on approximations derived from intractable optimal policies, introducing unquantified approximation errors in the resulting denoising objectives.

Core Question: Can we design a retraining-free, denoising-time alignment method for multiple objectives that requires no access to individual reward functions and introduces no approximation error?

2. Methodology

The authors propose MSDDA (Multi-objective Step-level Denoising-time Diffusion Alignment), a framework that achieves exact multi-objective alignment without retraining. The methodology consists of three key theoretical pillars:

A. Step-level RL Formulation

To overcome the intractability of tracking the optimal policy in standard RL formulations (which depend on circular policy-advantage estimates), the authors reformulate the problem as a Step-level RL task.

• Inspired by Trust Region Policy Optimization (TRPO), they decompose the reward and KL regularization terms into step-wise advantages.
• They derive a Step-level DPO (Direct Preference Optimization) objective. Unlike standard diffusion DPO, this loss includes an explicit step-wise regularization term ($\Delta_{diff}$) that encourages the model to stay closer to the pre-trained model on less-preferred trajectories while allowing deviation on preferred ones.
• Key Insight: This formulation resolves the circular dependence, making the optimal policy tractable and dependent only on the pre-trained policy's advantages.

B. Closed-Form Multi-Objective Fusion

Building on the step-level formulation, the authors prove that for a weighted reward $r_w = \sum w_i r_i$, the optimal reverse denoising distribution $p_w$ can be constructed purely by fusing the distributions of single-objective aligned models $\{p_i\}$.

• Theorem 1: If each single-objective model $p_i$ is a Gaussian distribution $N(\mu_i, \sigma_i^2 I)$, the optimal multi-objective distribution $p_w$ is also a Gaussian $N(\mu_w, \sigma_w^2 I)$.
• Closed-Form Parameters: The mean and variance are explicit functions of the base models:
  $$\sigma_w^2 = \left( \sum_{i=1}^M \frac{w_i}{\sigma_i^2} \right)^{-1}$$
  $$\mu_w = \sigma_w^2 \sum_{i=1}^M \frac{w_i}{\sigma_i^2} \mu_i$$
• This allows for exact fusion at inference time without any additional training or reward function access.

C. Theoretical Equivalence

The paper proves that the denoising-time objective derived via this fusion is exactly equivalent to the solution of the step-level RL fine-tuning formulation.

• Unlike previous reward-free methods (e.g., DB-MPA, DERADIFF) that drop terms in the diffusion dynamics to make the problem tractable (introducing errors), MSDDA drops no terms.
• It generalizes DB-MPA (which assumes identical variances for all objectives) to handle heterogeneous variances ($\sigma_i^2$).

3. Key Contributions

1. Step-level RL Framework: Introduced a novel step-level RL formulation for diffusion models that avoids the intractability of marginal KL optimization and enables the derivation of a step-level DPO objective.
2. Exact Closed-Form Solution: Proposed MSDDA, a retraining-free algorithm that computes the optimal reverse denoising distribution for any preference weight vector in closed form.
3. Zero Approximation Error: Proved theoretically that the MSDDA fusion is mathematically equivalent to the optimal step-level RL policy, eliminating the approximation errors inherent in prior denoising-time methods.
4. Reward-Free Operation: The method requires no access to reward functions or gradients during inference, relying solely on the parameters of pre-aligned single-objective models.

4. Experimental Results

The authors evaluated MSDDA using Stable Diffusion v1.5 with two reward objectives:

• ImageReward: For text-image alignment.
• VILA: For aesthetic quality.

Baselines:

• Stable Diffusion (Pre-trained).
• Reward Soup (RS): Linear interpolation of model parameters.
• CoDe & RGG: Reward gradient/value-based guidance methods.
• DB-MPA / DERADIFF: Other reward-free fusion methods.

Findings:

• Performance: MSDDA consistently outperformed all baselines across various preference weights ($w$). It achieved the highest Pareto front scores, balancing both aesthetic quality and text consistency better than Reward Soup or gradient-based methods.
• Qualitative Alignment: Visual results showed MSDDA generated images that strictly followed complex prompts (e.g., "a purple colored dog," "a pink banana") where the base SD model failed and Reward Soup often produced artifacts or incorrect colors.
• Efficiency:
  • MSDDA is significantly faster than reward-gradient methods (CoDe, RGG) because it avoids repeated reward evaluations during denoising.
  • It is slightly slower than SD/RS (roughly 2x runtime) because it runs two diffusion models in parallel to compute the fused mean/variance, but this is a negligible cost compared to the training overhead of multi-objective RL.

5. Significance

This paper represents a significant theoretical and practical advancement in diffusion model alignment:

• Theoretical Rigor: It bridges the gap between RL theory and practical denoising-time alignment by providing the first exact closed-form solution for multi-objective fusion, removing the "black box" nature of previous approximation-based methods.
• Practical Deployment: By eliminating the need for reward access during inference and avoiding retraining, MSDDA offers a highly efficient, plug-and-play solution for deploying diffusion models in dynamic environments where user preferences (weights) change frequently.
• Generalizability: The framework is not limited to image generation; the step-level RL formulation and closed-form fusion logic could potentially be adapted to other sequential generation tasks with multi-objective constraints.