Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control (2409.08861v5)

Abstract: Dynamical generative models that produce samples through an iterative process, such as Flow Matching and denoising diffusion models, have seen widespread use, but there have not been many theoretically-sound methods for improving these models with reward fine-tuning. In this work, we cast reward fine-tuning as stochastic optimal control (SOC). Critically, we prove that a very specific memoryless noise schedule must be enforced during fine-tuning, in order to account for the dependency between the noise variable and the generated samples. We also propose a new algorithm named Adjoint Matching which outperforms existing SOC algorithms, by casting SOC problems as a regression problem. We find that our approach significantly improves over existing methods for reward fine-tuning, achieving better consistency, realism, and generalization to unseen human preference reward models, while retaining sample diversity.

• The paper reframes reward fine-tuning as a stochastic optimal control problem and introduces the Adjoint Matching algorithm to overcome value function bias.
• It defines a memoryless noise schedule that ensures convergence to the desired tilted distribution in flow and diffusion models.
• Extensive experiments in text-to-image generation show that Adjoint Matching significantly improves metrics such as ClipScore and Human Preference Score v2 while preserving sample diversity.

Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control

This paper focuses on improving dynamical generative models, such as Flow Matching and denoising diffusion models, by leveraging memoryless stochastic optimal control (SOC) for reward fine-tuning. The authors introduce a novel algorithm named Adjoint Matching to tackle the challenges inherent in SOC-based fine-tuning, demonstrating its efficacy through theoretical analysis and extensive experimentation.

Main Contributions

1. SOC Formulation for Reward Fine-tuning:
   • The paper reframes reward fine-tuning as a stochastic optimal control problem. This involves optimizing a stochastic differential equation (SDE) to generate samples that align with a desired tilted distribution proportional to the base model's distribution, adjusted by the human preference reward model.
   • The control problem formulation addresses the value function bias problem that previous approaches encountered. This bias stems from the dependency of the initial noise variable on the generated samples.
2. Memoryless Noise Schedule:
   • The authors introduce the concept of a memoryless noise schedule, which is critical for ensuring the fine-tuning leads to the correct distribution. This schedule mandates a specific diffusion coefficient that removes the dependency between the noise variables and generated samples.
   • They rigorously prove that using this noise schedule is not just sufficient but necessary for converging to the desired tilted distribution.
3. Adjoint Matching Algorithm:
   • The paper presents Adjoint Matching, an innovative algorithm that circumvents the instability issues of traditional gradient-based SOC methods. It combines continuous adjoint state methods with least-squares regression principles to stabilize the training process.
   • Adjoint Matching formulates an objective that matches the gradient of the optimal control, ensuring the scalability and efficiency of the proposed method.
4. Empirical Validation:
   • The authors validate their approach through extensive experiments in text-to-image generation, comparing against various baseline methods. Adjoint Matching shows significant improvements across multiple dimensions, including text-to-image consistency, realism, and generalization to unseen human preference models, while maintaining sample diversity.

Key Numerical Results

• Adjoint Matching achieves superior scores in metrics such as PickScore, ClipScore, and Human Preference Score v2 (HPS v2). For instance, with $\lambda = 12500$, the method obtains a ClipScore of 31.65 and an HPS v2 score of 24.49 when evaluated using a zero-guided sampling procedure.
• The proposed algorithm also displays robust performance across different values of the guidance weight $w$, indicating its flexibility in various fine-tuning scenarios.

Implications and Future Directions

The implications of this research are multifaceted. Practically, the methodology offers a more controlled and theoretically sound approach to fine-tuning generative models, significantly enhancing their output quality in human preference-aligned tasks. Theoretically, the introduction of the memoryless noise schedule sets a new standard for considering noise characteristics in SOC problems, potentially influencing broader applications in machine learning and control theory.

Future research could explore extending Adjoint Matching and memoryless SOC to other types of generative models beyond text-to-image generation, such as text-to-video or text-to-audio tasks. Additionally, the concept of memoryless noise schedules could be investigated further to optimize other stochastic processes in machine learning.

Conclusion

The paper "Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control" presents a significant advance in the fine-tuning of generative models. By addressing the foundational biases in existing methods and proposing a theoretically robust optimization algorithm, the authors provide a highly effective and scalable solution to reward-based model fine-tuning. The empirical results substantiate the practical benefits, making this a noteworthy contribution to the field of generative modeling.