Taming Preference Mode Collapse via Directional Decoupling Alignment in Diffusion Reinforcement Learning

This paper introduces Directional Decoupling Alignment (D$^2$-Align), a novel framework that mitigates Preference Mode Collapse in diffusion reinforcement learning by applying directional corrections to reward signals, thereby preserving generative diversity while achieving superior human preference alignment.

The Gist

Imagine you have a brilliant, creative artist named Diffusion. This artist is amazing at painting whatever you ask for, from "a cat in space" to "a medieval castle." However, you want to teach this artist to paint things that humans really love. So, you hire a Judge (the Reward Model) to give the artist a score based on how much they like the painting.

The Problem: The "Goldilocks" Trap (Preference Mode Collapse)

At first, the artist tries everything. But soon, they figure out a secret trick. They notice that the Judge always gives high scores to paintings that are super bright, have a specific shiny texture, or feature a very specific type of face.

Instead of trying to be creative and diverse, the artist gets lazy. They decide, "Hey, if I just paint only shiny, bright faces, I'll get a perfect score every time!"

This is what the paper calls Preference Mode Collapse (PMC).

• The Result: The artist stops being an artist and becomes a photocopier. Every single painting looks exactly the same: overly bright, slightly plastic-looking, and boring.
• The Irony: The artist is technically "winning" because they have the highest scores, but they have lost their soul (diversity). They are "hacking" the system.

Existing methods tried to fix this by telling the artist, "Hey, don't just paint shiny faces; try to be different too!" But these methods were like trying to stop a runaway train by gently tapping the brakes. They either slowed the train down too much (lowering the quality) or didn't stop it at all.

The Solution: D²-Align (The "Compass" Correction)

The authors of this paper, D²-Align, realized the problem wasn't that the artist was bad, but that the Judge was biased. The Judge had a hidden preference (like loving "shiny" too much) that wasn't actually what humans wanted.

Instead of forcing the artist to change, they decided to fix the Judge's compass.

Here is how they did it, using a simple analogy:

Step 1: Finding the "Bias Vector" (The Correction Compass)

Imagine the Judge's brain is a giant map. On this map, "Shiny" is a direction that points way too far to the right.

• The researchers froze the artist (so they didn't change yet).
• They asked the artist to paint a few things.
• They then calculated a "Directional Vector" (let's call it a Correction Compass). This compass points in the opposite direction of the Judge's bias.
• Analogy: If the Judge is pulling the artist toward "Over-exposed Plastic," the Compass pulls them back toward "Natural and Varied."

Step 2: The Two-Stage Dance

1. Stage 1 (Calibrating the Compass): They teach the computer to find that perfect "Correction Compass" direction. They do this without changing the artist at all. They just figure out: "Okay, to get a true human score, we need to subtract this specific bias."
2. Stage 2 (Guided Painting): Now, they let the artist paint again. But this time, every time the Judge gives a score, they apply the Correction Compass.
   • If the Judge says, "Wow, that shiny face gets a 10!", the Compass says, "Wait, that's just the bias. Let's adjust the score down and encourage variety."
   • This guides the artist to find a sweet spot where the paintings are high quality (humans love them) but also highly diverse (no two look the same).

The Result: Breaking the Trade-off

Before this paper, you had to choose:

• Option A: High scores, but boring, identical images (Mode Collapse).
• Option B: Diverse images, but lower scores (because the Judge didn't like them).

D²-Align breaks this rule. It proves you can have both.

• The Analogy: Imagine a restaurant. Before, the chef only served "Spicy Noodles" because the food critic loved spicy noodles. Everyone got the same dish.
• With D²-Align: The chef realizes the critic actually loves flavor, not just spice. So, the chef starts making a diverse menu (Sushi, Tacos, Pasta) that is all delicious. The critic is happier, and the customers are happier because they aren't eating the same thing every day.

Why This Matters

The paper introduces a new "test" called DivGenBench to measure how boring an AI is. They showed that their method creates images that are not only beautiful but also unique, covering a wide range of styles, faces, and layouts, whereas other methods just churn out the same "plastic" look over and over.

In short: They didn't just tell the AI to "try harder." They fixed the way the AI listens to the feedback, ensuring it learns to be creative rather than just a score-chasing robot.

Technical Summary

Here is a detailed technical summary of the paper "Taming Preference Mode Collapse via Directional Decoupling Alignment in Diffusion Reinforcement Learning."

1. Problem Statement: Preference Mode Collapse (PMC)

Recent advancements in Text-to-Image (T2I) generation have utilized Reinforcement Learning from Human Feedback (RLHF) to align diffusion models with human preferences. However, the authors identify a critical failure mode termed Preference Mode Collapse (PMC).

• The Phenomenon: While existing RL methods (e.g., DanceGRPO, Flow-GRPO) achieve high scores on automated reward metrics (like HPS-v2.1), they inadvertently cause the model to "hack" the reward function. Instead of generating diverse, high-quality images, the model converges on a narrow set of high-scoring outputs.
• Symptoms: The generated images become homogeneous, exhibiting:
  • Monolithic Styles: All images look like the same artistic style.
  • Recurring Visual Features: Identical facial structures, lighting (e.g., pervasive overexposure), or layouts regardless of the prompt.
  • Loss of Diversity: The model fails to follow specific instructions regarding identity, style, layout, or tone, sacrificing generative breadth for a higher reward score.
• Root Cause: The authors posit that reward models possess inherent biases. The optimization process overfits to these biases, driving the generator toward specific "modes" that the reward model favors but which do not reflect true human preference for diversity.
• Evaluation Gap: Existing benchmarks prioritize fidelity (quality) and lack standardized, quantitative metrics to measure generative diversity, making PMC difficult to detect and quantify.

2. Methodology: Directional Decoupling Alignment (D²-Align)

To address PMC, the authors propose D²-Align, a novel two-stage framework that decouples the process of reward signal correction from policy alignment.

Core Concept: Directional Correction

Instead of merely modulating the magnitude of the reward (as done by KL-divergence regularization in other methods), D²-Align corrects the direction of the reward signal within the continuous embedding space of the text encoder.

Stage 1: Learning the Directional Correction Vector ($b_v$)

• Goal: Identify a vector in the text embedding space that counteracts the reward model's intrinsic biases.
• Process:
  1. The generator (diffusion model) is frozen.
  2. A learnable directional vector $b_v$ is introduced.
  3. The text embedding $e_{text}$ is perturbed in two directions: $e^+ = \text{normalize}(e_{text} + b_v)$ and $e^- = \text{normalize}(e_{text} - b_v)$.
  4. A guided embedding $\tilde{e}_{text}$ is constructed by extrapolating from the negative to the positive direction using a guidance scale $\omega$.
  5. The vector $b_v$ is optimized to minimize the guided reward loss, effectively learning a direction that suppresses the reward model's bias (e.g., reducing the score for "overly glossy" images unless "realistic" is explicitly requested).
• Outcome: A frozen, optimized vector $b^*_v$ that represents the "correction direction."

Stage 2: Guided Policy Alignment

• Goal: Align the generator using the corrected reward signal.
• Process:
  1. The generator is unfrozen for training.
  2. The learned vector $b^*_v$ is frozen and applied to the text embeddings during the RL optimization process.
  3. The generator is optimized to maximize the guided reward ($R_{guided}$), which incorporates the bias correction.
• Result: The model is steered away from the narrow, biased modes and toward a broader solution space that balances high fidelity with genuine diversity.

Technical Enablers

• Ground-Truth Noise Prior: To ensure stable reward signals during training, the method uses a ground-truth noise prior technique to reconstruct high-fidelity images from noisy latents for reward evaluation, avoiding the instability of one-step denoising at high noise levels.
• DivGenBench: A new benchmark introduced to quantitatively measure PMC. It uses "keyword-driven" prompts across four dimensions:
  1. ID: Identity diversity (age, ethnicity, features).
  2. Style: Artistic style coverage.
  3. Layout: Spatial arrangement and object counting.
  4. Tonal: Lighting, contrast, and saturation properties.
  • Metrics: Identity Divergence Score (IDS), Artistic Style Coverage (ASC), Spatial Dispersion Index (SDI), and Photographic Variance Score (PVS).

3. Key Contributions

1. Identification of PMC: The paper formally defines and quantifies "Preference Mode Collapse," highlighting the trade-off between automated reward scores and generative diversity.
2. D²-Align Framework: Proposes a novel method that corrects reward signals directionally in the embedding space, breaking the trade-off between quality and diversity without requiring complex hyperparameter tuning of regularization coefficients.
3. DivGenBench: Introduces a comprehensive benchmark with 3,200 prompts and four specialized metrics to evaluate generative diversity, filling a gap in current evaluation standards.
4. State-of-the-Art Performance: Demonstrates that D²-Align achieves superior alignment with human preference while significantly outperforming baselines in diversity metrics.

4. Experimental Results

The authors evaluated D²-Align on the FLUX.1.Dev model against baselines like DanceGRPO, Flow-GRPO, and SRPO.

• Quantitative Results (Quality & Diversity):

  • Human Preference: D²-Align achieved the highest scores on aesthetic metrics (Aesthetic, ImageReward, PickScore) and semantic consistency (CLIP, DeQA, GenEval), often outperforming baselines even when baselines had higher raw HPS-v2.1 scores.
  • Diversity (DivGenBench): D²-Align achieved the best scores across all four diversity metrics (lowest IDS, highest ASC, SDI, and PVS).
  • Baseline Failure: Baselines like Flow-GRPO and DanceGRPO showed severe drops in diversity (e.g., Flow-GRPO had an ASC of 0.044 vs. D²-Align's 0.253), confirming they suffered from PMC.

• Qualitative Results:

  • Baselines generated near-identical faces for different identity prompts and failed to render specific styles (e.g., "low key" or "black and white"), defaulting to a generic "realistic" look.
  • D²-Align successfully generated distinct identities, varied artistic styles, and correct spatial layouts while maintaining high image fidelity.

• Human Evaluation:

  • In user studies, D²-Align won 48.2% of overall preference votes, significantly outperforming the base model and RL baselines.
  • Crucially, in diversity-specific user studies, D²-Align was preferred over baselines that had collapsed, proving it preserves the "flavor" of the base model while improving quality.

• Efficiency: D²-Align is more efficient, achieving higher performance in fewer training steps (20 steps in Stage 2) compared to baselines requiring 250+ steps.

5. Significance

This work addresses a fundamental flaw in current RLHF pipelines for generative AI: the tendency to optimize for a single metric at the expense of the model's creative potential.

• Paradigm Shift: It moves the focus from simply "tuning" reward magnitudes to correcting the direction of the reward signal, offering a more principled approach to bias mitigation.
• Practical Impact: By preventing mode collapse, D²-Align ensures that generative models remain useful for creative tasks, data augmentation, and downstream applications where diversity is as critical as quality.
• Standardization: The introduction of DivGenBench provides the community with a necessary tool to rigorously evaluate and prevent diversity loss in future T2I models.

In summary, D²-Align successfully "tames" the reward hacking phenomenon by decoupling bias correction from policy learning, resulting in models that are both highly aligned with human taste and richly diverse in their outputs.