[论文解读] 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 Engineering被引用 4

一句话总结

本文提出BRAID，一种保守微调框架，通过结合扩散模型与奖励建模，提升离线设计优化性能。通过引入双重保守策略——在奖励建模中惩罚分布外区域，并施加KL散度约束，该方法使扩散模型能够超越最佳离线设计，同时避免生成无效或对抗性样本。

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.

研究动机与目标

• 解决离线基于模型的优化中过度优化的问题，即奖励模型在分布外区域误导模型。
• 通过使用从静态离线数据集学习到的奖励函数微调预训练扩散模型，弥合生成建模与基于模型的优化之间的差距。
• 开发一种方法，在保持对有效设计空间忠实性的同时，实现对离线数据中最佳观测设计的外推。
• 通过在微调过程中引入不确定性感知惩罚和KL正则化，提升对分布偏移的鲁棒性。

提出的方法

• 使用离线数据训练保守奖励模型，引入不确定性量化项，使分布外输入的惩罚增加。
• 采用软熵正则化马尔可夫决策过程框架，从理论上证明在微调目标中引入保守正则化的合理性。
• 通过强化学习优化保守奖励模型，对预训练扩散模型进行微调，并增加KL惩罚项以约束生成结果在有效设计空间内。
• 将训练过程分解为两个阶段：(1) 使用保守正则化进行离线奖励模型训练；(2) 使用保守奖励模型对扩散模型进行微调。
• 定义复合奖励函数，将学习到的奖励与视觉-语言模型（如LLaVA）的一致性信号结合，以确保图像生成中对提示的遵循。
• 将该方法应用于生物序列设计（DNA/UTR/增强子）和图像生成任务，分别采用Enformer和Stable Diffusion作为预训练模型。

实验结果

研究问题

• RQ1 当奖励模型在有限静态数据集上训练时，保守微调方法是否能缓解离线基于模型优化中的过度优化问题？
• RQ2 如何在不损害生成设计有效性的前提下，利用奖励模型的外推能力？
• RQ3 在奖励建模中引入不确定性感知惩罚在多大程度上能提升对最佳离线设计之外的泛化能力与性能？
• RQ4 保守奖励建模与KL正则化的结合是否能在序列和图像生成任务中带来更鲁棒、更忠实的生成结果？
• RQ5 与条件微调和标准强化学习微调相比，所提出方法在奖励性能和分布保真度方面表现如何？

主要发现

• BRAID在DNA序列和图像生成任务中均超越了离线数据中的最佳设计，展现出对观测数据分布外的泛化能力。
• 在图像生成中，BRAID生成的图像在美学评分（通过LLaVA和LAION预测器评估）方面更高，同时保持了对提示的遵循与视觉质量，而STRM则生成与提示偏离的图像。
• 在学习到的奖励（$\hat{r}$）方面，BRAID的训练曲线表现出稳定且持续的提升，而STRM则呈现快速奖励增长但泛化能力差且出现分布漂移。
• 实证结果表明，BRAID-boot和BRAID-bonus在奖励性能上显著优于STRM，证实了保守正则化的有效性。
• 理论分析将该方法置于软熵正则化MDP框架中，表明保守微调可降低遗憾值，并在性能上超越离线基线方法。
• 消融研究证实，保守奖励建模与KL惩罚均不可或缺；若任一缺失，性能将下降且分布外生成增加。