Manifold Drift in Flow Preference Optimization: A Root Cause of Reward Hacking

Technical Summary: Manifold Drift in Flow Preference Optimization

Problem Definition: Manifold Drift
The paper identifies a fundamental failure mode in applying Direct Preference Optimization (DPO) to continuous-time generative models, specifically Flow Matching (FM) and diffusion models. While DPO is effective for discrete models (like LLMs) by reweighting output probabilities, continuous models generate samples by transporting noise along learned trajectories toward a data manifold.

The authors argue that standard preference updates (e.g., FlowDPO) modify these transport dynamics in a way that is not inherently constrained to the pretrained data manifold. This leads to Manifold Drift, a phenomenon where the terminal samples of the fine-tuned model move off the support of the pretrained distribution ($\text{supp}(\mu_\theta) \not\subseteq M_0$).

• Theoretical Cause: Theoretically, optimal Flow Matching recovers the terminal data distribution. However, a preference update introduces a gradient step where the induced endpoint displacement has a non-zero component normal to the pretrained manifold. This "off-manifold" displacement causes the model to generate samples that, while potentially scoring higher on a specific reward function (e.g., OCR accuracy), exhibit visual artifacts, semantic distortions, or degraded fidelity.
• Reward Hacking: The paper posits that this drift is a root cause of reward hacking. Models may achieve high scores on "manifold-unaware" rewards (like text correctness) by generating images that are geometrically invalid or out-of-distribution, effectively bypassing the reward function by leaving the valid data support.

Methodology: THERMODPO
To address manifold drift, the authors propose THERMODPO, a temperature-controlled objective that explicitly anchors preference optimization to the pretrained manifold.

1. Thermodynamic Formulation: The method frames alignment as a balancing act between maximizing reward and preserving the manifold, governed by a time-dependent temperature function $\tau(t)$.
2. Three-State System: The objective minimizes the negative log-probability of a preferred sample within a system of three states:
   • Preferred state ($E_w$)
   • Rejected state ($E_l$)
   • Reference/Baseline state ($E_b$)
     The loss function is formulated as:
     $$L_{\text{ThermoDPO}}(\theta) = \mathbb{E}\left[ -\tau(t) \cdot t^2 \cdot \log \frac{e^{E_w}}{e^{E_w} + e^{E_l} + e^{E_b}} \right]$$
3. Theoretical Properties:
   • Reduction to RFT: As temperature $\tau \to 0$, the objective collapses to Rejection Sampling Fine-Tuning (RFT), which strictly enforces reconstruction of preferred samples on the manifold.
   • Anchored FlowDPO: For $\tau > 0$, the objective decomposes into a temperature-scaled FlowDPO term plus a non-negative winner-side anchoring term. This anchor penalizes deviations from the pretrained manifold for the preferred samples.
   • Geometric Bound: Theoretically, the loss upper-bounds the squared distance of the reconstructed preferred sample to the data manifold, providing a pointwise guarantee against drift.
4. Practical Implementation (THERMODPO-weighted): The standard formulation includes a $t^2$ factor that weakens the manifold anchor near the terminal endpoint ($t=0$), where preservation is most critical. To resolve this, the authors introduce THERMODPO-weighted, which replaces the global $t^2$ with a $(1-t)^2$ term. This dynamically activates the anchor near the terminal time, ensuring robust preservation without sacrificing the theoretical guarantees of the core objective.

Key Contributions

• Formalization of Manifold Drift: The paper formally defines manifold drift as a failure mode where fine-tuned trajectories leave the pretrained terminal support, identifying the non-zero normal component of the update as the sufficient condition for this drift.
• THERMODPO Framework: It introduces a unified objective that bridges Rejection Sampling Fine-Tuning (RFT) and FlowDPO, explicitly controlling the trade-off between reward maximization and manifold preservation.
• Theoretical Guarantees: The authors provide proofs showing that THERMODPO reduces to RFT in the low-temperature limit and decomposes into FlowDPO plus a non-negative anchoring term. They also establish an upper bound on the manifold distance for preferred samples.
• Empirical Validation: Through toy experiments and real-image benchmarks, the paper demonstrates that THERMODPO-weighted achieves a superior balance between preference alignment and sample quality compared to existing methods.

Results

• Toy Benchmark: On a synthetic 3D manifold, THERMODPO-weighted achieved a StrictScore of 0.899, significantly outperforming FlowDPO (0.629) and FlowDPO+RFT (0.857). While vanilla FlowDPO achieved high "Win" rates, it suffered from 0% "OnManifold" validity, confirming the drift hypothesis.
• Real-Image Generation (SD3.5-M): Evaluated on Stable Diffusion 3.5-M using an OCR preference dataset:
  • THERMODPO-weighted improved OCR scores by 47.5% (at CFG=4.5) compared to the baseline.
  • It achieved a 16.0% average improvement across four metrics (GenEval, OCR, HPSv3.0, UniReward).
  • Crucially, unlike FlowDPO variants which showed noticeable quality degradation and lower scores on held-out metrics (GenEval, HPSv3.0), THERMODPO-weighted maintained visual fidelity and performed competitively on metrics not directly optimized for.
• FLUX.2-klein-base-4B: Similar trends were observed on a different model architecture, reinforcing the method's effectiveness across different flow-based generators.

Significance and Claims
The paper claims that preference optimization in continuous-time models fundamentally alters transport dynamics, not just output selection. By formalizing manifold drift, the authors argue that ignoring this structural issue leads to reward hacking and degraded sample quality. THERMODPO is presented as a practical solution that allows for preference alignment without sacrificing the generative prior learned during pretraining. The authors position this as a necessary step for reliable alignment in image and video generation, particularly for tasks where "manifold-unaware" rewards (like text generation) might otherwise encourage out-of-distribution behaviors. The work is framed as a technical contribution for controlled offline research settings, aiming to make aligned generators more reliable rather than claiming immediate safety for unrestricted deployment.