[Paper Review] Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models

Masatoshi Uehara, Yulai Zhao+6 more|arXiv (Cornell University)|May 30, 2024

Advanced Mathematical Modeling in Engineering4 citations

TL;DR

This paper proposes BRAID, a conservative fine-tuning framework that combines diffusion models with reward modeling to improve offline design optimization. By introducing a doubly conservative approach—penalizing out-of-distribution regions in reward modeling and enforcing KL divergence constraints—it enables diffusion models to surpass the best offline designs while avoiding invalid or adversarial generations.

ABSTRACT

AI-driven design problems, such as DNA/protein sequence design, are commonly tackled from two angles: generative modeling, which efficiently captures the feasible design space (e.g., natural images or biological sequences), and model-based optimization, which utilizes reward models for extrapolation. To combine the strengths of both approaches, we adopt a hybrid method that fine-tunes cutting-edge diffusion models by optimizing reward models through RL. Although prior work has explored similar avenues, they primarily focus on scenarios where accurate reward models are accessible. In contrast, we concentrate on an offline setting where a reward model is unknown, and we must learn from static offline datasets, a common scenario in scientific domains. In offline scenarios, existing approaches tend to suffer from overoptimization, as they may be misled by the reward model in out-of-distribution regions. To address this, we introduce a conservative fine-tuning approach, BRAID, by optimizing a conservative reward model, which includes additional penalization outside of offline data distributions. Through empirical and theoretical analysis, we demonstrate the capability of our approach to outperform the best designs in offline data, leveraging the extrapolation capabilities of reward models while avoiding the generation of invalid designs through pre-trained diffusion models.

Motivation & Objective

• Address the challenge of overoptimization in offline model-based optimization, where reward models mislead models in out-of-distribution regions.
• Reconcile generative modeling and model-based optimization by fine-tuning pre-trained diffusion models using reward functions learned from static offline datasets.
• Develop a method that maintains fidelity to the valid design space while enabling extrapolation beyond the best observed designs in offline data.
• Ensure robustness against distributional shift by incorporating uncertainty-aware penalization and KL regularization during fine-tuning.

Proposed method

• Train a conservative reward model using offline data, incorporating an uncertainty quantification term that increases penalties for out-of-distribution inputs.
• Use a soft-entropy regularized Markov decision process framework to theoretically justify the inclusion of conservative regularization in the fine-tuning objective.
• Fine-tune pre-trained diffusion models by optimizing the conservative reward model via reinforcement learning, with an added KL penalty to constrain generation within the valid design space.
• Decompose the training process into two stages: (1) offline reward model training with conservative regularization, and (2) diffusion model fine-tuning using the conservative reward.
• Define a composite reward function that combines the learned reward with a consistency signal from vision-language models (e.g., LLaVA) to ensure prompt adherence in image generation.
• Apply the method to both biological sequence design (DNA/UTR/enhancer) and image generation tasks, using Enformer and Stable Diffusion as pre-trained models.

Experimental results

Research questions

• RQ1 Can a conservative fine-tuning approach mitigate overoptimization in offline model-based optimization when reward models are trained on limited static datasets?
• RQ2 How can the extrapolative power of reward models be harnessed without compromising the validity of generated designs?
• RQ3 To what extent does incorporating uncertainty-aware penalization in reward modeling improve generalization and performance beyond the best offline designs?
• RQ4 Does the combination of conservative reward modeling and KL regularization lead to more robust and faithful generation in both sequence and image generation tasks?
• RQ5 How does the proposed method compare to conditional fine-tuning and standard RL fine-tuning in terms of reward performance and distributional fidelity?

Key findings

• BRAID outperforms the best designs in offline data across both DNA sequence and image generation tasks, demonstrating the ability to generalize beyond the observed data distribution.
• In image generation, BRAID-generated images achieve higher aesthetic scores (via LLaVA and LAION predictor) while maintaining prompt adherence and visual quality, unlike STRL which generates images that deviate from prompts.
• The training curve in terms of the learned reward ($\hat{r}$) shows stable and consistent improvement in BRAID, whereas STRL exhibits rapid reward gain but poor generalization and distributional drift.
• Empirical results show that BRAID-boot and BRAID-bonus significantly outperform STRL in reward performance, confirming the benefit of conservative regularization.
• Theoretical analysis frames the method in a soft-entropy regularized MDP setting, showing that conservative fine-tuning reduces regret and enables superior performance over offline baselines.
• Ablation studies confirm that both conservative reward modeling and KL penalization are essential—removing either leads to performance degradation and increased out-of-distribution generation.