Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models

Authors: masatoshi  uehara
yulai  zhao
ehsan  hajiramezanali
gabriele  scalia
gokcen  eraslan
avantika  lal
sergey  levine
tommaso  biancalani
  genentech
princeton  university
u c berkeley
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 00:28:50

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.

Main paper: 
 introduction: 
 Computational design involves synthesizing designs that optimize a particular reward function. This approach finds applications in various scientific domains, including DNA/RNA/protein design (Sample et al., 2019; Gosai et al., 2023; Wu et al., 2024) . While physical simulations are often used in design problems, lacking extensive knowledge of underlying physical processes necessitates solutions that solely rely on experimental data. In these scenarios, we need an algorithm that synthesizes an improved design by utilizing a dataset of past experiments (i.e., static offline dataset). Existing research has addressed computational design from two primary angles. The first angle is generative modeling such as diffusion models (Ho et al., 2020) , which aim to directly model the distribution of valid designs by emulating the offline data. This approach allows us to model the space of "valid" designs (e.g., natural images, natural DNA sequences, foldable protein sequences (Avdeyev et al., 2023) ). The second angle is offline model-based optimization (MBO), which entails learning the reward model from static offline data and optimizing it with respect to design inputs (Brookes et al., 2019; Trabucco et al., 2021; Angermueller et al., 2019; Linder and Seelig, 2021; Fannjiang and Figure 1 : The left figure illustrates our setup with a pre-trained generative model and offline data. On the right, the motivation of the algorithm is depicted. The region surrounded by the green line is the original entire design space, with the colored region indicating the valid design space (e.g., natural images, human-like DNA sequences). The red region denotes areas with more offline data available, while the blue region indicates areas with less data available. We aim to add penalties to the blue regions using conservative reward modeling to prevent overoptimization while imposing a stricter KL penalty on the non-colored regions to prevent the generation of invalid designs.Listgarten, 2020; Chen et al., 2022) . This class of methods potentially enables us to surpass the best design observed in the offline data by harnessing the extrapolative capabilities of reward models.In our work, we explore how the generative modeling and MBO perspectives could be reconciled, inspired by recent work on RL-based fine-tuning of diffusion models (e.g., Black et al. (2023) ; Fan et al. (2023) ), which aims to finetune diffusion models by optimizing down-stream reward functions. Although these studies do not originally address computational design, we can potentially leverage the strengths of both perspectives. However, these existing studies often focus on scenarios where online reward feedback can be queried or accurate reward functions are available. Such approaches are not well-suited for the typical offline setting, where we lack access to true reward functions and need to rely solely on static offline data (Levine et al., 2020; Kidambi et al., 2020; Yu et al., 2020) . In scientific fields, this offline scenario is common due to the high cost of acquiring feedback data. In such contexts, existing works for fine-tuning diffusion models may easily lead to overoptimization, where optimized designs are misled by the trained reward model from the offline data, resulting in out-of-distribution adversarial designs instead of genuinely high-quality designs.To mitigate overoptimization, we develop a conservative fine-tuning approach for generate models aimed at computational design. Specifically, we consider a critical scenario where we have offline data (with feedback) and a pre-trained diffusion model capable of capturing the space of "valid" designs, and propose a two-stage method (Figure 1 ). In the initial stage, we train a conservative reward model using offline data, incorporating an uncertainty quantification term that assigns higher penalties to out-of-distribution regions. Subsequently, we finetune pre-trained diffusion models by optimizing the conservative reward model to obtain high-quality designs and prevent the generation of out-of-distribution designs. In the fine-tuning process, we also introduce a KL penalization term to ensure that the generated designs remain within the valid design space.Our primary contribution lies in the introduction of a novel framework, BRAID (douBly conseRvAtive fine-tuning diffusIon moDels). The term "doubly conservative" reflects the incorporation of two types of conservative terms, both in reward modeling and KL penalization. By properly penalizing the fine-tuned diffusion model when it deviates significantly from the offline data distribution, we effectively address overoptimization. Additionally, by framing our fine-tuning procedure within the context of soft-entropy regularized Markov Decision Processes, we offer theoretical justification for the inclusion of these conservative terms in terms of regret. This theoretical result shows that finetuned generative models outperform the best designs in the offline data, leveraging the extrapolation capabilities of reward models while avoiding the generation of invalid designs. Furthermore, through empirical evaluations, we showcase the efficacy of our approach across diverse domains, such as DNA/RNA sequences and images. 
related works: 
 We summarize related works. For additional works such as fine-tuning on LLMs, refer to Section A.Fine-tuning diffusion models via reward functions. Several previous studies have aimed to improve diffusion models by optimizing reward functions using various methods, including supervised learning (Lee et al., 2023; Wu et al., 2023) , RL (Black et al., 2023; Fan et al., 2023) and control-based techniques (Clark et al., 2023; Xu et al., 2023; Prabhudesai et al., 2023) . In contrast to our work, their emphasis is not on an offline setting, i.e., their setting assumes online reward feedback is available or accurate reward functions are known. Additionally, while Fan et al. (2023) include the KL term in their algorithms, our innovation lies in integrating conservative reward modeling to mitigate overoptimization and formal statistical guarantees in terms of regret (Theorem 1, 2).Conditional diffusion models. Conditional diffusion models, which learn conditional distributions of designs given the rewards, have been extensively studied (Ho and Salimans, 2022; Dhariwal and Nichol, 2021; Song et al., 2020; Bansal et al., 2023) . However, for the purpose of MBO, these approaches require that the offline data has good coverage on values we want to condition on (Brandfonbrener et al., 2022) . Compared to conditional diffusion models, our approach aims to obtain designs that can surpass the best design in offline data by leveraging the extrapolation capabilities of reward models. We compare these approaches with our work in Section 7.Offline MBO is also known as offline black-box optimization and is closely related to offline contextual bandits and offline RL (Levine et al., 2020) . While conservative approaches have been studied there (e.g., Kidambi et al. (2020) ; Yu et al. (2020) and more in Section A); most of the works are not designed to incorporate a diffusion model, unlike our approach. Hence, it remains unclear how these methods can generate designs that remain within intricate valid design spaces (e.g., generating natural images).It is worth noting a few exceptions (Yuan et al., 2023; Krishnamoorthy et al., 2023) that attempt to integrate diffusion models into MBO. However, the crucial distinctions lie in the fact that we directly optimize rewards with diffusion models, whereas these prior works focus on using conditional diffusion models. Additionally, we delve into the incorporation of conservative terms, an aspect not explored in their works. We compare these methods with ours empirically in Section 7. 
preliminaries: 
 We outline our framework for offline model-based optimization with a pre-trained generative model. Subsequently, we highlight the challenges arising from distributional shift. Additionally, we provide an overview of diffusion models, as we will employ them as pre-trained generative models.Our objective is to find a high-quality design within a design space, X . Each design x ∈ X is associated with a reward, r(x), where r : X → [0, 1] is an unknown reward function. Then, our aim is to find a high-quality generative model p ∈ ∆(X ), that yields a high r(x). It is formulated asargmax p∈∆(X ) E x∼p [r(x)].(1)Avoiding invalid designs. In MBO, the design space X is typically huge. However, in practice, the valid design space denoted by X pre is effectively contained within this extensive X as a potentially lower-dimensional manifold. For instance, in biology, our focus often centers around discovering highly bioactive protein sequences. While the raw search space might encompass |20| B possibilities (where B is the length), the actual design space corresponding to valid proteins is significantly more constrained. Consequently, our problem can be formulated as:EQUATIONNote supposing that a reward r(•) is 0 outside of X pre , this is actually still equivalent to (1).Offline data with a pre-trained generative model. Based on the above motivation, we consider scenarios where we have an offline dataset D off , used for learning the reward function. More specifically, the dataset, D off = {x (j) , y (j) } n off j=1 contains pairs of designs x ∼ p off (•) and their associated noisy reward feedbacks y = r(x) + ϵ, where ϵ is noise.Compared to settings in many existing papers on MBO, we also assume access to a pre-trained generative model (diffusion model) trained on a large dataset comprising valid designs, in addition to the offline data D off . For example, in biology, this is expected to capture the valid design space X pre such as human DNA sequences or physically feasible proteins (Avdeyev et al., 2023; Li et al., 2024; Sarkar et al., 2024; Stark et al., 2024; Campbell et al., 2024) . These pre-trained generative models are anticipated to be beneficial for narrowing down the raw search space X to the design space X pre . In our work, denoting the distribution induced by the pre-trained model by p pre , we regard the support of p pre as X pre .To understand our challenges, let's first examine a simple approach for MBO with a pre-trained generative model. For instance, we can adapt methods from Clark et al. (2023) ; Prabhudesai et al. (2023) to our scenario. This approach involves two steps. In the first step, we perform reward learning: r = argmin r∈F n off i=1 {r(x (i) ) -y (i) } 2 , where F represents a function class that includes mappings from X to [0, 1], aiming to capture the true reward function r(•). Then, in the second step, we fine-tune a pre-trained diffusion model to optimize r.Despite its simplicity, this approach faces two types of distributional shifts. Firstly, the fine-tuned generative model might produce invalid designs outside of X pre . As discussed in Section 3.1, we aim to prevent this situation. Secondly, the fine-tuned generative model may over-optimize r, exploiting uncertain regions of the learned model r. Indeed, in regions not covered by offline data distribution p off , the learned reward r can easily have higher values, while the actual reward values in terms of r might be lower due to the higher uncertainty. We aim to avoid situations where we are misled by out-to-distribution adversarial designs.We present an overview of denoising diffusion probabilistic models (DDPM) (Song et al., 2020; Ho et al., 2020; Sohl-Dickstein et al., 2015) . Note while the original diffusion model was initially introduced in Euclidean spaces, it has since been extended to simplex spaces for biological sequences (Avdeyev et al., 2023) , which we will use in Section 7. In diffusion models, the goal is to develop a generative model that accurately emulates the data distribution from the dataset. Specifically, denoting the data distribution by p pre ∈ ∆(X ), a DDPM aims to approximate using a parametric model structured as p(x 0 ; θ) = p(x 0:T ; θ)dx 1:T , where p(x 0:T ; θ) = p T +1 (x T ; θ)1 t=T p t (x t-1 |x t ; θ).Here, each p t is considered as a policy, which is a mapping from a design space X to a distribution over X . By optimizing the variational bound on the negative log-likelihood, we can obtain a set of policies {p t } 1 t=T +1 such that p(x 0 ; θ) ≈ p pre (x 0 ). For simplicity, in this work, assuming that pre-trained diffusion models are accurate, we denote the pre-trained policy as{p pre t (•|•)} 1 t=T +1, and the generated distribution by the pre-trained diffusion model at x 0 by p pre . With slight abuse of notation, we often denotep pre T +1 (•) by p pre T +1 (•|•) 
doubly conservative generative models: 
 We've discussed how naïve approaches for computational design may yield invalid designs or overoptimize reward functions, with both challenges stemming from distributional shift. Our goal in this section is to develop doubly conservative generative models to mitigate this distributional shift.To avoid invalid designs, we begin by considering the following generative model:EQUATIONwhereKL(p∥p pre ) = E x∼p [log(p(x)/p pre (x))].In this formulation, the generative model is designed as an optimizer of a loss function composed of two parts: the first component encourages designs with high rewards, while the second component acts as a regularizer penalizing the generative model for generating invalid designs. This formulation is inspired by our initial objective in (2), where we substitute an indicator function with log(p/p pre ). This regularizer takes ∞ when p is not covered by p pre , and α governs the strength of the regularizer.Next, we address the issue of overoptimization. This occurs when we are fooled by the learned reward model in uncertain regions. Therefore, a natural approach is to penalize generative models when they produce designs in uncertain regions.As a first step, let's consider having an uncertainty oracle ĝ : X → [0, 1], which is a random variable of D off . This oracle is expected to quantify the uncertainty of the learned reward function r.Assumption 1 (Uncertanity oracle). With probability 1 -δ, we haveEQUATIONThese calibrated oracles are well-established when using a variety of models such as linear models, Gaussian processes, and neural networks. We will provide detailed examples of such calibrated oracles in Section 4.3. Essentially, as long as the reward model is well-specified (i.e., there exists r ∈ F such that ∀x ∈ X pre : r(x) = r(x)), we can create such a calibrated oracle.Doubly Conservative Generative Models. Utilizing the uncertainty oracle defined in Assumption 1, we present our proposal:πα (•) = exp ((r -ĝ)(•)/α) p pre (•) exp ((r -ĝ)(x)/α) p pre (x)dx (:= argmax p∈∆(X ) E x∼p [(r -ĝ)(x)]Penalized reward-αKL(p∥p pre )).(5)Here, to combat overoptimization, we introduce an additional penalty term ĝ(x). This penalty term is expected to prevent πα from venturing into regions with high uncertainty because it would take a higher value in such regions. We refer to πα as a doubly conservative generative model due to the incorporation of two conservative terms.An attentive reader might question the necessity of simultaneously introducing two conservative terms. Specifically, the first natural question is whether KL penalties, intended to prevent invalid designs, can replace uncertainty-oracle-based penalties. However, this may not hold true because even if we can entirely avoid venturing outside of X pre (support of p pre ), we may still output designs on uncertain regions not covered by p off . The second question is whether uncertainty-oracle-based penalties can substitute KL penalties. While it is partly true in situations where the support of p off is contained within that of p pre , uncertainty-oracle-based penalties, lacking leverage on pre-trained generative models, are ineffective in preventing invalid designs. In contrast, KL penalties are considered a more direct approach to stringently avoid invalid designs by leveraging pre-trained generative models.Example 1 (Gaussian processes.). When we use an RKHS as F (a.k.a. GPs) associated with a kernel Srinivas et al., 2009) , a typical construction of r and ĝ isk(•, •) : X × X → R (r(•) = Y(K + λI) -1 k(•), ĝ(•) = c(δ) k(•, •), where c(δ) ∈ R >0 , λ ∈ R >0 , k(x) = [k(x (1) , x), • • • , k(x (n off ) , x)] ⊤ , Y = [y (1) , • • • , y (n off ) ], {K} p,q = k(x (p) , x (q) ), k(x, x ′ ) = k(x, x ′ ) -k(x) ⊤ {K + λI} -1 k(x ′ ).Note that when using deep neural networks, by considering the last layer as a feature map, we can still create a kernel (Zhang et al., 2022; Qiu et al., 2022) .Example 2 (Bootstrap). When we use neural networks as F, it is common to use a statistical bootstrap method. Note many variants have been proposed (Chua et al., 2018) , and its theory has been analyzed (Kveton et al., 2019) . Generally, in our context, we generate multiple models r1 , • • • , rM by resampling datasets, and then consider argmin i ri as r -ĝ.Algorithm 1 BRAID (douBly conseRvAtive fIne-tuning Diffusion models)1: Require: Parameter α ∈ R + , a set of policy classes {Π t } where Π t ⊂ [X → ∆(X )], pre-trained diffusion model {p pre t } 1 t=T +1 . 2:Train a conservative reward model r -ĝ using an offline data D off . 3: Update a diffusion model as {p t } t by solving the planning problem:{p t } t = argmax {pt∈Πt} 1 t=T +1 E {pt} [r(x 0 ) -ĝ(x 0 )] Penalized reward -α Σ 1 t=T +1 E {pt} [KL(p t (•|x t )∥p pre t (•|x t ))] KL penalty (6)where the expectation E {pt} [•] is taken with respect to1 t=T +1 p t (x t-1 |x t ). 4: Output: A policy {p t } t 
conservative fine-tuning of diffusion models: 
 In this section, we consider how to sample from a doubly conservative generative model πα , using diffusion models as pre-trained generative models. Our algorithm is outlined in Algorithm 1. Initially, we learn a penalized reward r -ĝ from the offline data and set it as a target to prevent overoptimization in (6). Additionally, we integrate a KL regularization term to prevent invalid designs. The parameter α governs the intensity of this regularization term.Formally, this phase can be conceptualized as a planning problem in soft-entropy-regularized MDPs (Neu et al., 2017; Geist et al., 2019) . In this MDP formulation:• The state space S and action space A correspond to the design space X .• The reward at timet ∈ [0, • • • , T ] (∈ S × A → R) is provided only at T as r -ĝ. • The transition dynamics at time t (∈ [S × A → ∆(S)]) is an identity δ(s t+1 = a t ).• The policy at time t (∈ S → ∆(A)) corresponds to p T +1-t : X → ∆(X ).• The reference policy at t is a policy in the pre-trained model p pre T +1-tIn these entropy-regularized MDPs, the soft optimal policy corresponds to {p t }. Importantly, we can analytically derive the fine-tuned distribution in Algorithm 1 and show that it simplifies to a doubly conservative generative model πα , from which we aim to sample. Theorem 1. Let pα (•) be an induced distribution from optimal policies {p t } 1 t=T +1 in (6), i.e., p(x 0 ) = { 1 t=T +1 pt (x t-1 |x t )}dx 1:T when {Π t } is a global policy class (Π t = {X → ∆(X )}). Then, pα (x) = πα (x).We have deferred to the proof in Section 5.1. While similar results are known in the context of standard entropy regularized RL (Levine, 2018) , our theorem is novel because previous studies did not consider pre-trained diffusion models.Training algorithms. Based on Theorem 1, to sample from πα , what we need to is to solve Equation ( 6). We can employ any off-the-shelf RL algorithms to solve this planning problem. Given that the transition dynamics are known, and differentiable reward models are constructed in our scenario, a straightforward approach to optimize (6) is to directly optimize differentiable loss functions with respect to parameters of neural networks in policies, as detailed in Appendix B. Indeed, this approach has recently been used in fine-tuning diffusion models (Clark et al., 2023; Prabhudesai et al., 2023) , demonstrating its stability and computational efficiency. Remark 1 (Novelty of Theorem 1). A theorem similar to Theorem 1 has been proven for continuoustime diffusion models in Euclidean space (Uehara et al., 2024, Theorem 1) . However, the primary distinction lies in the fact that while their findings are restricted to Euclidean space, where diffusion policies take Gaussian polices, our results are not constrained to any specific domain. Hence, for example, our Theorem 1 can handle scenarios where the domain is discrete or lies on the simplex space (Avdeyev et al., 2023) in order to model biological sequences as we do in Section 7.We explain the sketch of the proof of Theorem 1. The detail is deferred to Theorem C.1.By induction from t = 0 to t = T + 1, we can first showEQUATIONHere, v t (x t ) is a soft optimal value function:E p[r(x 0 ) -ĝ(x 0 ) -α 1 k=t KL(p k (•|x k )∥p pre k (•|x k ))|x t ],which satisfies an equation analogous to the soft Bellman equation:EQUATIONNow, we aim to calculate a marginal distribution at t defined by pt (x t ) = { t k=t+1 pk (x k-1 |x k )}dx t+1:T . Then, by induction, we can show that pt (x t ) = exp(v t (x t )/α)p pre t (x t )/C (9)where C is a normalizing constant. Indeed, supposing that the above (9) hold at t, the equation ( 9) also holds for t -1 as follows:pt-1 (x t-1 |x t )p t (x t )dx t = exp(v t-1 (x t-1 )/α)p pre t-1 (x t-1 )/C = pt-1 (x t-1 ).Finally, by setting t = 0, the statement in Theorem 1 is concluded. 
regret guarantee: 
 In this section, our objective is to demonstrate that a policy pα from Algorithm 1 can provably outperform designs in offline data by establishing the regret guarantee.To assess the performance of our fine-tuned generative model, we introduce the soft-value metric:J α (p) := E x∼p [r(x)] -αKL(p∥p pre ).This metric comprises two components: the expected reward and a penalty term applied when p produces invalid outputs, as we see in (3). Now, in terms of soft-value J α (p), our proposal pα offers the following guarantee.Theorem 2 (Per-step regret). Suppose Assumption 1. Then, with probability 1 -δ, we have∀π ∈ ∆(X ); J α (π) -J α (p α ) Per step regret ≤ 2 C π × E x∼p off [ĝ(x) 2 ] 1/2 Stat , C π := max x∈X off π(x) p off (x) ,where X off = {x ∈ X : p off (x) > 0}. As an immediate corollary,E x∼π [r(x)] -E x∼ pα [r(x)] ≤ αKL(π∥p pre ) + 2 C π × E x∼p off [ĝ(x) 2 ] 1/2 .In the theorem above, we establish that the per-step regret against a generative model π we aim to compete with is small as long as the generative model π falls within X off and the learned model r is calibrated as in Assumption 1. First, the term (Stat) corresponds to the statistical error associated with r over the offline data distribution p off . When the model is well-specified, it is upper-bounded by d/n, where d represents the effective dimension of F, as we will discuss shortly. Secondly, the term C π corresponds to the coverage between a comparator generative model π and our generative model pα . Hence, it indicates that the performance of our learned pα is at least as good as that of a comparator generative model π covered by p off . While this original coverage term C π diverges when π goes outside of X off , we can refine it using the extrapolation capabilities of a function class F, as we will discuss shortly. This refined version ensures that we can achieve high-quality designs that outperform designs in the offline data (i.e., best designs in X off ).Example 3. We consider a scenario where an RKHS is used for F. Let F be a model represented by an infinite-dimensional feature ϕ(•). Let d denote the effective dimension of F (Valko et al., 2013) .Corollary 1 (Informal: Formal characterization is in Section D ). Assuming that the model is well-specified, with probability 1 -δ, we have:J α (π) -J α (p α ) ≤ Cπ × Õ d3 n , Cπ := sup κ:∥κ∥2=1 κ ⊤ E x∼π [ϕ(x)ϕ ⊤ (x)]κ κ ⊤ E x∼p off [ϕ(x)ϕ ⊤ (x)]κ .The refinement of the coverage term in Cπ is characterized as the relative condition number between covariance matrices on a generative model π and an offline data distribution p off , which is smaller than C π . This Cπ could still be finite even if C π is infinite. In this regard, Corollary 1 illustrates that the trained generative model can outperform the best design in the offline data by harnessing the extrapolation capabilities of reward models. 
experiments: 
 We perform experiments to evaluate (a) the effectiveness of conservative methods for fine-tuning diffusion models and (b) the comparison of our approach between existing methods for MBO with diffusion models (Krishnamoorthy et al., 2023; Yuan et al., 2023) . We will start by outlining the baselines and explaining the experimental setups. Regarding more detailed setups, hyperparameters, architecture of neural networks, and ablation studies, refer to Appendix E.Methods to compare. We compare the following methods in our evaluation. For a fair comparison, we always use the same α in BRAID and STRL.3 .• BRAID (proposed method): We consider two approaches: (1) Bonus, as in Example 1 by setting a last layer as a feature map and constructing a kernel, (2) Bootstrap, as in Example 2.• Standard RL (STRL): RL-fine-tuning that optimizes the standard r without any conservative term, following existing works on fine-tuning (Clark et al., 2023; Prabhudesai et al., 2023) .• DDOM (Krishnamoorthy et al., 2023) : We train with weighted classifier-free guidance (Ho and Salimans, 2022) using offline data, conditioning on a class with high y values (top 5%) during inference. Note that this method is training from scratch rather than fine-tuning.• Offline Guidance (Yuan et al., 2023) : After training a classifier using offline data, we use guidance (conditional diffusion models) (Dhariwal and Nichol, 2021) on top of pre-trained diffusion models and condition on classes with high y values (top 5%) at inference time.Evaluation. We assess the performance of each generative model primarily by visualizing the histogram of true rewards r(x) obtained from the generated samples. For completeness, we include similar histograms for both the pre-trained model (Pretrained) and the offline dataset (Offline).As for hyperparameter selection, such as determining the strengths of conservative terms/epochs, we adhere to conventional practice in offline RL (e. 2020)) and choose the best one through a limited number of online interactions. Remark 2. We omit comparisons with pure MBO methods for two reasons: (i) DDOM, which we compare against, already demonstrates a good performance across multiple datasets, and (ii) these methods are unable to model complex valid spaces since they do not incorporate state-of-the-art generative models (e.g., stable diffusion), thereby lacking the capability to generate valid designs (e.g., natural images) as we show in Section 7.2.We examine two publicly available large datasets consisting of enhancers (n ≈ 700k) (Gosai et al., 2023) and UTRs (n ≈ 300k) (Sample et al., 2019) with activity levels collected by massively parallel reporter assays (MPRA) (Inoue et al., 2019) . These datasets have been extensively used in sequence therapy (Castillo-Hair and Seelig, 2021; Ghari et al., 2023; Lal et al., 2024; Ferreira DaSilva et al., 2024) . In the Enhancers dataset, each x is a DNA sequence with a length of 200, while y ∈ R is the measured activity in cell lines. For the UTRs dataset, x is a 5'UTR RNA sequence with a length of 50 while y ∈ R is the mean ribosomal load (MRL) measured by polysome profiling.Setting of oracles and offline data. We aim to explore a scenario where we have a pre-trained model and an offline dataset. Since the true reward function r(•) is typically unknown, we initially divide the original dataset D = {x (i) , y (i) } randomly into two subsets: D ora and D ′ . Then, from D ′ , we select datasets below 95% quantiles for enhancers and 60% quantiles for UTRs and define them as offline datasets D off . Subsequently, we construct an oracle r(•) by training a neural network on D ora and use it for testing purposes. Here, we use an Enformer-based model, which is a state-of-the-art model for DNA sequences (Avsec et al., 2021) . Regarding pre-trained diffusion models, we use ones customized for sequences over simplex space (Avdeyev et al., 2023) . In the subsequent analysis, each algorithm solely has access to the offline data D off and a pre-trained diffusion model, but not r(•).Results. The performance results in terms of r(•) are depicted in Fig 2a and b . It is seen that fine-tuned generative models via RL outperform conditioning-based methods: DDOM and Guidance. This is expected because conditional models themselves are not originally intended to surpass the conditioned value (≈ best value in the offline data). Conversely, fine-tuned generative models via RL are capable of exceeding the best value in offline data by harnessing the extrapolation capabilities of reward modeling, as also theoretically supported in Corollary 1. Secondly, both BRAID-boot and BRAID-bonus demonstrate superior performance compared to STRL.This suggests that conservatism aids in achieving fine-tuned generative models with enhanced rewards while mitigating overoptimization.We consider the task of generating aesthetically pleasing images, following prior works (Fan et al., 2023; Black et al., 2023) . We use Stable Diffusion v1.5 as our pretrained diffusion model, which can generate high-quality images conditioned on prompts such as "cat" and "dog". We use the AVA dataset (Murray et al., 2012) as our offline data and employ a linear MLP on top of CLIP embeddings to train reward models (r and r -ĝ) from offline data for fine-tuning.Setting of oracles. To construct r(x), following existing works, we use the LAION Aesthetic Predictor V2 (Schuhmann, 2022) , already pre-trained on a large-scale image dataset. However, this LAION predictor gives high scores even if generated images are almost identical regardless of prompts, as in Figure 3b . These situations are undesirable because it means fine-tuned models are too far away from pre-trained models. Hence, for our evaluation, we define r(x) as follows: (1) asking vision language models (e.g., LLaVA (Liu et al., 2024) ) whether images contain objects in the original prompts4 (e.g., dog, cat), (2) if Yes, outputting the LAION predictor, and (3) if No, assigning 0. This evaluation ensures that high r(x) still indicates capturing the space of the original stable diffusion.Results. We show that our proposed approach outperforms the baselines in terms of r(x), as in Fig 2c 5 . We also show the generated images in Figure 3cc . Additionally, we plot the training curve during the fine-tuning process in terms of the mean of r(x) of generated samples in Fig 3cb . The results indicate that in STRL, while the learning curve based on the learned reward quickly grows, fine-tuned models no longer necessarily remain within the space of pre-trained models (STRL in Fig 2c) . In contrast, in our proposal, by carefully regularizing on regions outside of the offline data, we can generate more aesthetically pleasing images than STRL, which remain within the space of pre-trained models. For more images/ ablation studies, refer to Appendix E. 
summary: 
 For the purpose of fine-tuning from offline data, we introduced a conservative fine-tuning approach 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. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We discussed potential positive societal impact in the introduction. Regarding negative impact, we have regonized unintended uses is possible. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [No] Justification: Due to the page limit, we haven't included it. However, we recognize diffusion models could be misused to make fake images. We aim to address this problem in the future. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: In this section, we summarize additional related works.Conservative approaches in offline RL/offline contextual bandits. Conservative approaches have been explored in offline RL and contextual bandits. Firstly, in both model-free and model-based RL, one prevalent method involves incorporating an additional penalty on top of the reward functions (Yu et al., 2020; Chang et al., 2021) . Secondly, in model-based RL, a common strategy is to train transition dynamics in a conservative manner (Kidambi et al., 2020; Rigter et al., 2022; Uehara and Sun, 2021) . Thirdly, in model-free RL, typical approaches include conservative learning of q-functions (Kumar et al., 2020; Xie et al., 2021) or the inclusion of KL penalties against behavioral policies (Wu et al., 2019; Fakoor et al., 2021) .However, these works are not designed to incorporate a diffusion model, unlike our approach. Hence, it remains unclear how their works can generate designs that remain within intricate valid design space, such as high-quality images using stable diffusion.Design with generative models. Many works are focusing on design problems with generative models. However, these works are typically limited to the usage of VAEs. (Notin et al., 2021; Gómez-Bombarelli et al., 2018) Our work is still significantly different because we focus on generative models.Fine-tuning in LLMs from human feedbacks. A closely related area of research involves finetuning LLMs through the optimization of reward functions using human feedback (Touvron et al., 2023; Ouyang et al., 2022) . Especially, following works such as Zhan et al. (2023) , from a theoretical viewpoint, Xiong et al. (2023) explores the effectiveness of pessimism in offline scenarios and its theoretical aspect. However, our theoretical findings are more specifically tailored to diffusion models. Indeed, our main result in Theorem 1 is novel, and our algorithm differs significantly, as fine-tuning methods in offline settings in the literature on LLMs typically rely on policy gradient or PPO, whereas we use more direct backpropagation approaches. Furthermore, the meaning of step size is different as well. Hence, in (Xiong et al., 2023) , they do not use soft entropy regularized MDPs.Typically, in the above works, human feedback is considered to be given in the form of preferences.Similarly, in the context of diffusion models, Yang et al. (2023) ; Wallace et al. (2023) discusses fine-tuning of diffusion models using preference-based feedback. However, these works focus on online settings but not offline settings.Sampling from unnormalized distributions. In our approach, we use an RL method to sample from our target distribution that is proportional to exp((r -ĝ(x))p pre (x). While MCMC has traditionally been prevalent in sampling from unnormalized Boltzmann distributions exp(r(x)), recent discussions (Zhang and Chen, 2021; Vargas et al., 2023) , have explored RL approaches similar to our approach. However, their focus differs from ours, as they do not address the fine-tuning of diffusion models (i.e., no p pre ) or sample efficiency in offline settings.Another relevant literature discusses sampling from unnormalized Boltzmann distributions when pre-trained diffusion models are available (Kong et al., 2024) . However, their algorithm closely resembles classifier-based guidance (Dhariwal and Nichol, 2021) , rather than a fine-tuning algorithm. Additionally, they do not examine conservatism in an offline setting.Our planning algorithm has been summarized in Algorithm 2. Here, we parametrize each policy by neural networks.For the sake of explanation, we also add typical cases where the domain is Euclidean in Algorithm 3.Algorithm 2 Direct Back Propagation (General case) Update θ s to θ s+1 by adding the gradient of the following loss L(θ) with respect to θ at θ s : L(θ) = 1 m m i=1 r(x (i) 0 (θ)) -ĝ(x (i) 0 (θ)) -α 1 t=T +1; θ pre ), σ 2 t )} 1 t=T +1 , batch size m, a parameter α ∈ R + . 2: Initialize: θ 1 = θ pre 3: for s ∈ [1, • • • , S] do 4: Set θ = θ s . 5: Collect m samples {x (i) t (θ)} 0t=T +1 from a current diffusion model (i.e., generating by sequentially running polices {N (ρ(t, x t ; θ), σ 2 t )} 1 t=T +1 from t = T + 1 to t = 1)6:Update θ s to θ s+1 by adding the gradient of the following loss L(θ) with respect to θ at θ s :L(θ) = 1 m m i=1 r(x (i) 0 (θ)) -ĝ(x (i) 0 (θ)) -α 1 t=T +1 ∥ρ(x (i) t (θ), t; θ) -ρ(x (i) t (θ), t; θ pre )∥ 2 2σ 2 (t).(11) Proof. To simplify the notation, we let f (x) = r(x) -ĝ(x). As a first step, by using induction, we aim to obtain an analytical form of the optimal policy {p t (•|x t-1 )}.First, we define the soft-optimal value function as follows:v t-1 (x t-1 ) = E { pt} f (x) -α 1 k=t-1 KL(p k (•|x k )∥p pre k (•|x k ))|x t-1 .Then, by induction, we havept (x t-1 |x t ) = argmax pt∈∆(X ) E { pt} [v t-1 (x t-1 ) -αKL(p t (•|x t )∥p pre t (•|x t ))|x t ] .With simple algebra, we obtainEQUATIONHere, notingv t (x t ) = max pt∈∆(X ) E { pt} [v t-1 (x t-1 ) -αKL(p t (•|x t )∥p pre t (•|x t ))|x t ],we get the soft Bellman equation:EQUATIONTherefore, by plugging ( 13) into ( 12), we actually haveEQUATIONFinally, with the above preparation, we calculate the marginal distribution:pt (x t ) := t s=T +1 ps (x s-1 |x s ) dx t+1:T +1 .Now, by using induction, we aim to provept (x t ) = exp v t (x t ) α p pre t (x t ).Indeed, when t = T + 1, from ( 14), this hold as follows:pT +1 (x +1 ) = 1 C exp v T (x T +1 ) α p pre T +1 (x T +1 ).Now, suppose the above holds at t. Then, this also holds for t -1:pt-1 (x t-1 ) = pt (x t-1 |x t )p t (x t )dx t = exp v t-1 (x t-1 ) α {p pre t (x t-1 |x t )}p pre t (x t )dx t (Use Equation 14) = exp v t-1 (x t-1 ) α p pre t-1 (x t-1 ).By invoking the above when t = 0, the statement is concluded.In the following, we condition on the event where ∀x ∈ X pre ; |r(x) -r(x)| ≤ g(x).holds.First, we defineĴα (π) := E x∼π [r(x) -ĝ(x)] -αKL(π∥p un ), J α (π) := E x∼π [r(x)] -αKL(π∥p un ).We note that, πα maximizes Ĵα (π). Therefore, we haveEQUATIONUsing a result in Srinivas et al. (2009) . we showr(x) -r(x) ≤ C(δ) k(x, x)whereC(δ) = c 1 1 + log 3 (n/δ)I n , I n = log(det(I + K)).Then, with probability 1 -δ, we have Recall from the proof of Theorem 3,EQUATIONJ α (π) -J α (π α ) ≤ 2E x∼π [ĝ(x)] = 2C(δ)E x∼π [ k(x, x)]. Now, first, to upper-bound E x∼π [ k(x, x)],we borrow Theorem 25 in Chang et al. (2021) , which showsE x∼π [ k(x, x)] ≤ c 1 Cπ d ′ {d ′ + log(c 2 /δ)} n .whereCπ := sup κ:∥κ∥2=1 κ ⊤ E x∼π [ϕ(x)ϕ ⊤ (x)]κ κ ⊤ E x∼psp [ϕ(x)ϕ ⊤ (x)]κNext, in order to upper-bound C(δ), we borrow Theorem 24 in Chang et al. (2021) , which showsI n ≤ c 1 {d ′ + log(c 2 /δ)}d ′ log(1 + n).The statement in Corollary 1 is immediately concluded.E.1 DNA/RNA sequencesIn this subsection, we add the details of experiments in Section 7.Diffusion models. Regarding diffusion models for sequences, we adopt the architecture and algorithm tailored to biological sequences over the simple space (Avdeyev et al., 2023) . Its architecture is described in Table 1 .Oracles. We use the architecture in Avsec et al. (2021) , which is a state-of-the-art model in sequence modeling. We just change the last layer so that it is tailored to a regression problem as in Lal et al. (2024) .In all experiments, we use A100 GPUs. The important hyperparameters are summarized in the following table ( Hyperparameter selection. The process of selecting hyperparameters in offline RL is known to be a challenging task (Rigter et al., 2022; Paine et al., 2020) in general. A common practice in existing literature is determining crucial hyperparameters with a limited number of online interactions. In our case, the key hyperparameters include the strength of the LCB parameter (utilized online in BRAID-Bonus) and the termination criteria during training (applied to all fine-tuning algorithms such as STRL). To ensure a fair comparison, we operate within a framework where we can utilize 120 × 20 online samples. This implies that, for instance, in STRL and BRAID-Bonus, we conduct an online evaluation using 120 samples for 20 pre-defined epochs. However, in BRAID-Bonus, given the additional hyperparameters to be tested (strengths of the bonus term 0.01, 0.1, 1.0), we use 40 samples for each 20 pre-defined epoch.We performed ablation studies by varying the strength of the bonus parameter C in Figure 5 . We chose the one with the best performance in the main text (See the previous section to see the validity of this procedure). Regarding the effectiveness of KL-regularization, it has been discussed inFan et al. (2023);Uehara et al. (2024). Hence, in our work, we focus on the effectiveness of conservatism in reward modeling.The F1 score of LLaVA for detecting objects was 1.0 as detailed in Appendix E.2.3. Hence, this part is considered to be accurate.Note we omitted DDOM because it is not a fine-tuning method. 
 
 Appendix:  


Decision: The paper proposes a conservative fine-tuning approach to tackle overoptimization in diffusion models, which is common in settings where task-specific reward models are unknown and offline data is used instead for fine-tuning diffusion models. The reviewers highlighted the originality of the proposed method and found that the empirical evaluation satisfactorily demonstrated the method's usefulness. All reviewers agree that the paper should be presented at NeurIPS, and I agree that the paper would be of interest to the NeurIPS community and would be a valuable addition to the program.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: N/A
limitations: 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.)
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.

2. 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.
questions: questions: STRL 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?
limitations: limitations: See my comments on the weakness above.
paper_summary: paper_summary: The 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.

The 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.
strength_weakness: strength_weakness: 1. 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.
Very hard to address and not a must have for ML conferences:
1. 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?
questions: questions: 1. 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?
2. 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?
limitations: limitations: The 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).
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: * Difficulty in tuning hyperparameters without online data interaction.
* Reliance on accurate reward and diffusion models for effective performance.
* Theoretical results depend on certain idealized assumptions that may not hold in all cases.
* 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.

References：

[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.

[2] Liu, Tenglong, et al. "Adaptive Advantage-Guided Policy Regularization for Offline Reinforcement Learning." In International Conference on Machine Learning (ICML). PMLR, 2024.

[3] Zhang, Jing, et al. "Constrained policy optimization with explicit behavior density for offline reinforcement learning." Advances in Neural Information Processing Systems. 2023

[4] Shimizu, Yutaka, et al. "Strategically Conservative Q-Learning." arXiv preprint arXiv:2406.04534 (2024).
questions: questions: * 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? 
* Have the authors considered alternative generative models such as GANs or VAEs, and can they provide a comparative analysis of performance and resource usage?
* How to obtain the $\hat{g}$ in the pseudo-code of Algorithm 1 BRAID? 
* This methods is related to offline reinforcement learning. So can you give me some experiment comparisions with the SOTA offline RL methods ？
limitations: limitations: * The method requires careful selection of hyperparameters, which can be challenging in a purely offline setting without access to additional online data.
* The pseudo-code of the paper does not illustates the algorithm explicitly. The authors can improve it further.
* 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.

Comments: title: Official Comment by Reviewer HnkA
Thanks for your rebuttal
comment: Thanks for your careful answers and clear explanation. To some extend, my concerns have been addressed. So I expect you to discuss SOTA offline RL as mentioned earlier in your paper and conduct comparison experiments with GANs or VAEs. I would like to keep my score, but it is possible to improve my score.
1. I appreciate the additional suggestions for experiments and am sure they could be interesting as well. I did not state that "[your] computational evaluation did not provide enough confidence". I think the computational evaluations are well crafted and can provide a fine degree of confidence. I just pointed this out as a minor weakness prepended with "not a must have" (which was maybe overread in the formatting).\
Anyways, the additional promised experiments and the qualitative recovery of experimental observations do not change my score. I think it is a good paper and should be accepted. It is outstanding in its clarity of writing (thanks for that effort) but not outstanding overall in my current estimation (I hope that is not rude to say :)).

2. Thanks for taking the time to answer - that was a confused question.

3. "Yes" -> Sad, but expected.
After careful consideration, I decided to raise my score to 5.