Title: Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Diffusion models

Reinforcement learning
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a conservative fine-tuning method called BRAID, which integrates the strengths of diffusion models and model-based optimization (MBO) to improve the performance of pre-trained diffusion models on offline datasets. BRAID optimizes a conservative reward model that includes penalties outside the offline data distribution to prevent overoptimization and generate valid designs. The approach is validated through empirical and theoretical analyses, demonstrating its ability to outperform the best designs in offline data while avoiding the generation of invalid designs. The paper also discusses the method's effectiveness compared to existing conditional diffusion models and traditional MBO techniques, with experiments showcasing its superiority in biological sequence and image generation. The authors acknowledge the limitations of their study, particularly in model selection and hyperparameter tuning, and suggest future research directions.\n\nSTRENGTHS:\n* BRAID incorporates a conservative approach to fine-tuning diffusion models, which includes penalization terms that discourage the model from generating designs outside the distribution of the offline data. This conservative strategy is effective in preventing overoptimization and ensuring the validity of the generated designs.\n\n* The method is supported by both theoretical analysis and empirical results. Theoretically, it provides a regret guarantee, ensuring that the fine-tuned models can outperform the best designs in the offline data. Empirically, it has been validated through experiments across various domains, such as biological sequences and images, demonstrating its ability to generate high-quality designs.\n\nWEAKNESSES:\n* Difficulty in tuning hyperparameters without online data interaction.\n* Reliance on accurate reward and diffusion models for effective performance.\n* Theoretical results depend on certain idealized assumptions that may not hold in all cases.\n* Can you compare the methods with other SOTA offline RL methods to illustrate your proposed augmented methods more effective than the SOTA offline RL methods? I think this paper is very relevant to some offline RL methods, such as ReDS[1], A2PR[2], CPED[3], SCQ[4]. It is not required that experimental comparisons must be given, but at least add some discussion with these methods to the paper.\n\nReferences：\n\n[1] Singh, Anikait, et al. \"ReDS: offline reinforcement learning with heteroskedastic datasets via support constraints.\" Proceedings of the 37th International Conference on Neural Information Processing Systems. 2023.\n\n[2] Liu, Tenglong, et al. \"Adaptive Advantage-Guided Policy Regularization for Offline Reinforcement Learning.\" In International Conference on Machine Learning (ICML). PMLR, 2024.\n\n[3] Zhang, Jing, et al. \"Constrained policy optimization with explicit behavior density for offline reinforcement learning.\" Advances in Neural Information Processing Systems. 2023\n\n[4] Shimizu, Yutaka, et al. \"Strategically Conservative Q-Learning.\" arXiv preprint arXiv:2406.04534 (2024).\n\nQUESTIONS:\n* The methods will need more samples, as shown in the pseudo-code of the algorithms 2 Direct Back Propagation (General case), which may bring more computational burden. Meanwhile, the methods train use two diffusion model to obtain the policy. Can you give me some experiments to show the computational burden? \n* Have the authors considered alternative generative models such as GANs or VAEs, and can they provide a comparative analysis of performance and resource usage?\n* How to obtain the $\\hat{g}$ in the pseudo-code of Algorithm 1 BRAID? \n* This methods is related to offline reinforcement learning. So can you give me some experiment comparisions with the SOTA offline RL methods ？\n\nLIMITATIONS:\n* The method requires careful selection of hyperparameters, which can be challenging in a purely offline setting without access to additional online data.\n* The pseudo-code of the paper does not illustates the algorithm explicitly. The authors can improve it further.\n* This paper lack strong baselines comparisions and enough related works, which is related to some offline reinforcement learning methods. So you can add more related offline reinforcement learning works.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the task of black box optimization in an offline setting. Given a pretrained diffusion model, they first train a surrogate model on the offline data and use it to tilt the diffusion model distribution via finetuning it. The authors distinctly focus on an uncertainty quantification based procedure to bias the diffusion model tilting toward regions where the reward is high and the reward uncertainty is low while not tilting toward regions with high uncertainty. Experiments are carried out on a reasonable set of diverse tasks.\n\nThe paper introduces a small specific challenge and address it well with a reasonable approach and good motivations. The potential impact of the method may be small but it is neat, educational, and should be useful in many cases. I recommend acceptance.\n\nSTRENGTHS:\n1. Identifying a crucial non-obvious overlooked challenge and specifically pinpointing it. The authors identify that previous work on tilting diffusion models with reward models misses to incorporate tilting less toward regions of the reward function in which it has high uncertainty but high reward. Instead, we should only tilt to the regions that have high reward and high certainty to avoid optimizing toward adversarial examples. (The fact that the finetuned diffusion model will steer away from the pretraining distribution seems like a less relevant insight)\n2. The authors identify a relevant overlooked problem in finetuning diffusion models and bring standard techniques from uncertainty quantification into the field to address it in a reasonable fashion. They do not overcomplicate things and their technique could be valuable to several researchers in the area. \n3. The authors prove that the training procedure yields the desired distribution. I have no comments regarding the value/insightfulness of the proof. Maybe other reviewers have a stronger opinion about the relevance.\n4. The authors evaluate their method on a very diverse set of experiments that includes discrete DNA sequence generation, and image generation. The results are convincing and demonstrate the central empirical claim that out of distribution generation is a problem and is effectively avoided with the proposed conservative reward model fine tuning.\nMinor:\n1. Interesting snippets of insights. The authors point out interesting relationships and connections along the way which are non-obvious and well placed for putting their motivations into context.\n2. Exceptional clarity in writing. The paper lays out the task in its precise specification and covers required concepts and related work in equal clarity.\n\nWEAKNESSES:\n1. I would say that the insights in terms of methodological novelty are on the moderate side. The ideas are simple and good, which is appreciated, but the level on which the conceptural changes operate are low level (a tweak to diffusion model tilting) and thus limited in impact. However, it is certainly a good thing to have.\nVery hard to address and not a must have for ML conferences:\n1. Evaluations are inherently limited in their computational nature and the conclusions that can be drawn for the procedures effectiveness in biological sequence optimization is small. Do you authors disagree with this in any way?\n\nQUESTIONS:\n1. Do I understand correctly that theorem 2 only states that the finetuned diffusion model will incur lower regret than the pretrained diffusion model? Why is that useful and would we not rather be interested in relations between using your additional uncertainty quantification bias for the reward model based finetuning and not using it?\n2. It seems to me that reward models are often very poor in scientific applications (more so than the generative models that can be trained on a lot more data). Does this mean that their uncertainty estimates are also bad and your method might not provide any improvements in these cases?\n\nLIMITATIONS:\nThe authors point out the key inherent limitation of their work in the fact that the reward models that are often trained on limited data are likely suboptimal instead of just mentioning useless small limitations that are beside the point (which is the more common practice it seems).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a conservative approach for fine-tuning diffusion models with a reward model learned from offline data. Specifically, the ideas are two-fold: The first idea is to replace the reward model with a conservative estimate based on classical generalization bounds. The second idea is to leverage the KL divergence to force the optimized distribution to not deviate too far from the pretrained model. Experiments and theoretical results show the efficacy of the proposed method in fine-tuning diffusion models without over-optimization.\n\nSTRENGTHS:\n1. The proposed method is well-presented, and the motivation behind the algorithm is interesting. The over-optimization problem is indeed critical when fine-tuning diffusion models with learned rewards.\n\n2. Extensive experimental results show the efficacy of the proposed method in improving the reward model while avoiding reward over-optimization.\n\nWEAKNESSES:\n1. Leveraging generalization bounds via kernel RKHS and bootstrap is interesting, but I doubt their practicality for real applications. Firstly, the RKHS bound is usually too conservative to be useful, while the computational cost for the bootstrap method is pretty high since one has to train the model from scratch multiple times. As far as I can tell, the reward models used in the experiments are mainly single-layer MLPs, and it is doubtful whether this approach is useful when the reward model needs to be a larger model.\n\n2. Another problem with the conservative estimator of the reward models is that it is unclear whether it is useful given the current experimental results. On one hand, KL regularization is a widely-known technique for preventing over-optimization in diffusion models and is thoroughly studied in existing works, so it is certain that the KL regularization term will help. On the other hand, the proposed algorithm mixes both the conservative reward estimator and the KL regularization term together, making it unclear which part is playing the role in avoiding over-optimization. My guess is that, for the most part, only the KL regularization term is effective in the end.\n\nQUESTIONS:\nSTRL methods like AlignProp and DRaFT can work with ODE samplers, which are more commonly used in practice than SDE samplers. However, the method proposed in this work, due to the use of entropy regularization, can only adopt SDE samplers. I wonder if it is possible to design a regularization term for ODE samplers. Could the authors share some insights on this point?\n\nLIMITATIONS:\nSee my comments on the weakness above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

comments: SUMMARY:
1) This paper analyzed the two mainstream angles of computational design.
2) Proposed a hybrid one that offline fine-tunes generative models.
3) Conduct experiments on two tasks to show the performance of their method.

STRENGTHS:
1) Sufficient theoretical analysis and detailed preliminaries.
2) The idea is straightforward.
3) The method is comprehensive.

WEAKNESSES:
1) In the introduction, the advantages and disadvantages of the two mainstream issues are not fully analyzed.
2) Insufficient metrics evaluation for image generation task.

QUESTIONS:
N/A

LIMITATIONS:
1) Lack of gradual guidance from analyzing the advantages and disadvantages of existing methods to proposing the hybrid method.
2) Lack of multi-metric quantitative results analysis. (e.g. in the image generation task, the fidelity and diversity should also be reported by like LPIPS Score/CLIP score/FID/IS, etc.)
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12