Improved Off-policy Reinforcement Learning in Biological Sequence Design

W1: Complexity in implementation

We want to emphasize that the hyperparameters are not excessively tuned. First, the reweighting coefficient $k$ is usually set to 0.01 or 0.001 [1, 2]; the results below (and more in Appendix B.8) show that both settings outperform the baselines. The adaptation is not critical for $\delta$-CS, as performance is decent with just a constant hyperparameter and is robust to its selection. $\lambda$ (which is only required for the adaptive version of $\delta$-CS) is also easy to tune because it is just a scaling factor. We can use a tuned $\lambda$ for similar scale tasks (e.g., we use identical $\lambda$ for short-scale tasks like DNA and RNA, and long-scale tasks like AAV and GFP), which can be adjusted with a simple validation experiment.

To demonstrate hyperparameter robustness, our sensitivity analysis results in Appendix B.2 and the following tables show that our method performs well with various values of these parameters, highlighting its robustness.

[1] Kim, Minsu, et al. "Bootstrapped training of score-conditioned generator for offline design of biological sequences." NeurIPS. 2023.
[2] Kim, Hyeonah, et al. "Genetic-guided GFlowNets for sample efficient molecular optimization." NeurIPS. 2024.

Ablation for $k$

RNA-A

GFP

It is noteworthy that we have not excessively searched these parameters, and the results are not cherry-picked. As evidences, the better combinations for each task are found in our additional results above e.g., $\delta=0.025$ or $k=0.001$ for GFP.

W2: Uncertainty Estimate

Good point. Theoretically, uncertainty can be measured even when prediction accuracy is low, as they have different characteristics. The low prediction accuracy of the proxy arises from out-of-distribution problems—data points not seen during training. Uncertainty is measured based on how such a new data point affects the uncertainty of the proxy function. This can be quantified through posterior inference over the proxy model's parameters, given the current in-distribution training dataset. Although theoretically valid, such posterior inference is difficult in practice. We use ensemble methods to estimate uncertainty, which are sub-optimal but still practical and widely used. Developing better Bayesian posterior inference methods for active learning could be valuable future work. Our method is orthogonal to this, as $\delta$-CS can be used with such improved uncertainty estimation methods.

W3: Choice of GFlowNets

1. Performance gain of GFN-AL + $\delta$-CS over GFN-AL.

We argue that the performance improvement achieved by $\delta$-CS is significant, particularly evident in the large margins observed for RNA-C, GFP, and AAV. Notably, as shown in Table 2, while GFN-AL underperforms in terms of average score compared to the RL baseline (DyNA PPO), our approach (with EI) outperforms it in all metrics, mean and max score, diversity, and novelty.

2. About [1,4,5]

We agree with this suggestion. Such off-policy advancements in the GFlowNet community can potentially improve active learning performance. While prior methods design off-policy exploration [1,4,5] under the assumption that the reward model is accurate, our delta-CS method is built on a different assumption and purpose: conducting conservative search to mitigate the risks from proxy misspecification. Among these prior methods, local search GFlowNets (LS-GFN) [1] seem to be very related to our method and can also be applied in conservative search. We directly compare our method with LS-GFN as follows:

Results (MAXIMIUM)

As shown in the table, our method clearly outperforms LS-GFN. We suspect this is because the search flexibility of delta-CS is better than that of LS-GFN. The algorithm of LS-GFN is based on back-and-forth local search, where they use a backward policy to partially destroy a solution and a forward policy to reconstruct a new solution. However, in the auto-regressive sequence generation setting, such a backward policy must be unidirectional; therefore, their local search region is bounded in the leaf space of the sequences. On the other hand, delta-CS randomly destroys tokens independently, regardless of whether the token is located at the head or leaf node, allowing it to search more flexibly in the sequence space.

We have included experimental results in Appendix B.5 and have added prior training methods of GFlowNets to the related works section.

3. About [2,3]

Regarding other related works, references [2] and [3] have clearly different purposes and are orthogonal to ours, as they focus on improving the training scheme of GFlowNet but we focus to improve exploration which provides experiences for training. Specifically, [2] suggests a new loss function for better credit assignment, while [3] proposes a new method for multi-objective settings. We may utilize [2] with $\delta$-CS when the trajectory length is too large and better credit assignment is needed. Similarly, we might employ [3] with delta-CS when multi-objective active learning is required. These could be exciting directions for future work. We have added this to the discussion section of our manuscript.

Thank you for pointing this out. We can greatly improve our manuscript.

[1] Local Search GFlowNets. Minsu Kim, et al. ICLR 2024 spotlight.

[2] Learning Energy Decompositions for Partial Inference of GFlowNets. Hyosoon Jang, et al. ICLR 2024 oral.

[3] Order-Preserving GFlowNets. Yihang Chen, et al. ICLR 2024.

[4] Learning to Scale Logits for Temperature-Conditional GFlowNets. Minsu Kim, et al. ICML 2024.

[5] QGFN: Controllable Greediness with Action Values. Elaine Lau, et al. NeurIPS 2024.

Thanks for the feedback. We are pleased to know that your concerns except for regarding uncertainty have been addressed.

We would like to emphasize that considering uncertainty in $\delta$ represents an additional component—the adaptive version of $\delta$. Furthermore, in the adaptive $\delta$, the uncertainty is measured for observed data sets, where the proxy is relatively reliable compared to OOD data points, with the properly chosen sailing factor $\lambda$. Our main message is that employing a conservative search using $\delta$ itself is beneficial for active learning in off-policy reinforcement learning for biological sequence design, as it effectively limits the search space of the generative policy. The subsequent experimental results demonstrate that even a constant $\delta$ without uncertainty measuring gives competitive results, outperforming other baselines.

maximum scores

median scores

W1. proxy model misspecification. the hypothesis (*) is not quantitatively validated well enough when the search space is scaled up

Q1. Can we use $\delta$-Conservative Search with any other proxy model?

Thank you for pointing this out. To validate our hypothesis on larger-scale tasks, we conducted additional experiments on AAV and GFP tasks (see tables below), which we have updated in the revised manuscript.

AAV with CNN

GFP with CNN

W2. Another major concern is the lack of proxy model ablation

Thanks for pointing that out. To verify our hypothesis on various types of proxy models, we conducted additional experiments using a different proxy, MuFacNet [1], in AAV, GFP, and RNA tasks (the default proxy model is based on 1d-CNN). With another proxy model, the trends are consistent, validating our hypothesis. See the results below (we have also updated this in the revised manuscript).

[1] Ren, Zhizhou, et al. "Proximal exploration for model-guided protein sequence design." International Conference on Machine Learning. PMLR, 2022.

AAV with MuFacNet

GFP with MuFacNet

Maximum scores

W3. how does δ-CS differ from existing off-policy search methods for GFlowNets such as [7,8]? (Especially, LS-GFN [7])

There are many off-policy methods in GFN, as you mentioned (we have updated them in the related works section). Among them, the most relevant method is Local Search GFN (LS-GFN) [7], where both our methods focus on restricted search in local regions for reward exploitation (the other suggested method, Thompson Sampling [8] is for diverse exploration rather than conservatism).

We made a direct comparison with LS-GFN:

Results (MAXIMIUM)

The reason that LS-GFN performs worse than ours is that they utilize back-and-forth search using backward and forward policies, where their search space in unidirectional sequences is limited to only the leaf part of sequences. On the other hand, our local neighbors are defined by a Hamming ball of sequence space; we distribute the search not only on the leaf part but also equally on the every tokens in sequences (from head to leaf), which makes for a more flexible local search. More results are provided in the updated manuscript (Appendix B.5).

W4.More BO baselines (especially with trust region)

Thank you for your suggestion. We agree that comparison with state-of-the-art trust-region based BO methods can provide stronger evidence of the effectivenss of our method. We conduct experiment with TuRBO [1], a widely used trust-region based BO method for our setting. As shown in the table, while TuRBO exhibits generally higher score than classical BO, our method surpass TuRBO across various tasks, exhibiting the superiority of our $\delta$-CS constraints.

maximum

median

[1] David Eriksson, Michael Pearce, Jacob Gardner, Ryan D Turner, and Matthias Poloczek. Scalable Global Optimization via Local Bayesian Optimization. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, pages 5496–5507, 2019.

W5. A quantitative analysis showing how the level of conservativeness/noise injection amount differs throughout the active learning rounds

Q6. How does noise injection/level of conservativeness differ throughout the active learning rounds? Have you analyzed the behavior with respect to the proxy uncertainty?

The level of conservativeness varies depending on the data point. As active learning progresses through more rounds, it continues to seek new data points with high predictive uncertainty. Therefore, as active learning actively discovers new data points, the conservativeness level associated with these points is also dynamic. No empirical distributional tendencies regarging on the number of active rounds have been found, so the model must adapt in every active learning round and for each data point $x$.

W6. A part that requires a bit of clear discussion is that, if the proxy model is not of good quality then exploration helps to recover from getting stuck at local optima and collect diverse points, whereas over-exploration yields unreliable results as explained for using GFlowNets case. I think this trade-off could have been argued more clearly while motivating the method

Thanks for opening insightful discussion. Such diverse exploration can help recover from bad local optima, yet it can produce unreliable results if the exploration is too extensive. Therefore, we need to balance diversity to focus on reliable regions without being so local that we cannot escape from bad local optima. Our $\delta$ factor can be interpreted as such a balancing parameter.

W7. batch size ablation

Thank you for your suggestion. We conduct additional experiments with query batch size 32 and 512 (default setting is 128). The same ablation is also applied for AdaLead as a baseline. We report the results in the tables below, and we also added Figure 15,16 in Appendix B.6 to the manuscript. The results demonstrate that our method achieves superior performance compared to AdaLead across different batch size, exhibiting robustness.

maximum (bs=32)

maximum (bs=512)

W8. theoretical analysis

Thanks for pointing out. We provided theoretical analysis on exploration range and sample complexity:

1. Exploration range:

Let $\delta \in [0,1]$ be the exploration parameter serving as the success probability in a Bernoulli distribution for each token in a sequence of length $L$. Each token independently has a probability $\delta$ of being flipped (changed) and $1 - \delta$ of remaining the same.

Define the random variable $H$ as the Hamming distance—the number of differing tokens between two sequences. The probability distribution of $H$ is:

where $k = 0, 1, \dots, L$. This is a binomial distribution with parameters $n = L$ and $p = \delta$.

The expected value and variance of $H$ are:

A higher $\delta$ increases $\mathbb{E}[H]$, expanding the exploration region in the sequence space and promoting diversity.

2. Sample Complexity: We do not introduce additional meaningful sample complexity as a scale because the noising process involves sampling from a tractable trivial distribution (Bernoulli). The major bottleneck is sampling sequences proportional to the reward, which is intractable. However, the GFlowNet (GFN) amortizes this process using a neural network, reducing the sampling cost to the neural network's forward pass with token-by-token $O(L)$ complexity, where $L$ is the sequence length. This approach is consistent with every sequence-based decoding method in deep learning.

3.Convergence Throughout: Proving the convergence of deep learning models is notoriously difficult without making substantial assumptions, even on fixed datasets. Therefore, conducting convergence analysis on active learning that includes an inner deep-learning training process is even more challenging. Instead, this paper focused to demonstrate convergence through empirical validation across six tasks.

W9. There is no clear discussion on the limitations and potential drawbacks of the method.

Thank you for your feedback. We added a discussion section (Section 7), including a few notes on the limitations, in the updated manuscript.

Q2. For the baseline methods that require surrogate/proxy (e.g. BO), is the same proxy model of $\delta$-CS used, with the same initial dataset?

We use the same 1dCNN-based proxy (from Sinai et al., 2020) for all algorithms except GFN-AL. We implement two versions and report the better results between them. The first version is the original implementation (Jain et al. 2022) with MLP-based proxy, and the second is our own implementation with 1dCNN-based proxy. Theses are described in Implementation details in Section 6 and Appendix A.3.

Q3.How diversity/novelty is calculated? Is it the diversity of the last top-k batches? Presenting this metric across active learning rounds would have enhanced clarity and provided a more comprehensive understanding of the diversity achieved by $\delta$-CS.

Following GFN-AL (Jain et al. 2022), the diversity and novelty is calculated with top-128 at the final rounds. We added the definition of these metrics in Appendix B.8.

As you suggested, we report the progress of diversity and novelty across active learning rounds in Appendix B.8.

Q4. In Figure 5, AAV task, there is a substantial difference in diversity between the proposed method and baselines. Given that the search space is medium-sized—somewhere between RNA-A and GFP—what might be the reason for this significant disparity?

You must mean the Figure 3. The low divrsity and novelty for the baselines is because:

1. The initial dataset is not very diverse. The diversity of top-128 samples in the initial dataset is around 5.
2. The queried candidates from GFN-AL or GFNSeqEditor have (relatively) low oracle scores, and thus the top-128 batch couldn't be populated by newly queried diverse and novel candidates.

Compared to these baselines, our method generated both diverse and novel high-score candidates through out the active learning rounds.

Thanks a lot to the authors, with the addition of ablation studies and discussions on the updated manuscript. I have the following additional comment about your response:

• I am still concerned about the first key evaluation point I have stated in the weaknesses section. The paper proposes $\delta$-CS as "novel off-policy search method for training GFlowNets designed to improve robustness against proxy misspecification". I appreciate experimenting with other proxy models, which was one of my questions, however, the key missing evaluation is in what level $\delta$-CS helps with this problem? I mean, under poor to good-quality proxy models, what's the actual contribution of $\delta$-CS addition to the GFlowNets (or to GFN-AL) regarding the empirical performance? As I have stated previously, a clear empirical demonstration for this would be to initialize proxy models of different quality (e.g using varying numbers of initial training points or specifically crafting it to be worse) and evaluating different search methods from restricted ($\delta$-CS, LS-GFN) to explorative (GFN-AL and Thompson Sampling [8]) ones. I have asked for the initial dataset size ablation. The provided ablation study (also I wonder why it is not added to the updated manuscript) compares against AdaLead, however, from those results, I cannot properly evaluate why the addition of $\delta$-CS would help. As also pointed out the reviewer 6RU1, the main improvement might stem from the use of GFlowNets rather than from the search method itself?

I believe the paper has been improved, however, still missing key points. Overall, I cannot vote for acceptance here.

Thanks for pointing out key points that can still make our paper improved.

a clear empirical demonstration for this would be to initialize proxy models of different quality (e.g using varying numbers of initial training points or specifically crafting it to be worse) and evaluating different search methods from restricted ($\delta$-CS, LS-GFN) to explorative (GFN-AL and Thompson Sampling [8]) ones.

Following your suggestion, we have conducted additional experiments to clarify the contribution of the $\delta$-CS component, and we have included these results in Appendix B.11 of the revised manuscript.

To assess the impact of $\delta$-CS and address concerns that improvements might stem solely from using GFlowNets, we intentionally degraded the proxy model's reliability. We did this by only using the lower percentiles— 50%, 25%, and 10%— of the initial dataset based on scores. Training the proxy model on datasets with lower percentile cutoffs increases misspecification, particularly affecting higher-reward data points.

Under these conditions, we compared our method with GFN-AL (which employs an explorative search strategy) and LS-GFN (which uses a restricted search approach), as you suggested. This comparison isolates the specific contribution of the $\delta$-CS component over the base GFN-AL method.

The results, also presented in Appendix B.11, demonstrate that all methods experience performance degradation as the proxy model becomes less reliable, which is expected. Notably, explorative search methods like GFN-AL perform significantly worse under proxy misspecification than restricted search methods (LS-GFN and $\delta$-CS). Unfortunately, we could not include Thompson Sampling GFN (explorative method) in our experiments due to the unavailability of their source code.

Table. Maximum values with varying initial dataset quality (AAV)

The results show that our method consistently outperforms the other baselines across all levels of initial dataset percentiles (various proxy reliability). This provides empirical evidence of the clear contribution of the $\delta$-CS component.

These findings suggest that the performance improvements are not solely due to the use of GFlowNets but are significantly enhanced by our proposed search method. We believe this addresses your concern by demonstrating that the addition of $\delta$-CS indeed helps, especially when dealing with unreliable proxy models.

The provided ablation study (also I wonder why it is not added to the updated manuscript) compares against AdaLead, however, from those results, I cannot properly evaluate why the addition of $\delta$-CS would help. As also pointed out the reviewer 6RU1, the main improvement might stem from the use of GFlowNets rather than from the search method itself?

Thank you for pointing this out. In response to your suggestion, we have added GFN baselines—specifically, GFN-AL and LS-GFN—to our ablation study on different initial dataset sizes, setting $|\mathcal{D}_0| = 1000$, which are randomly collected.

We present the results below:

As shown in the table, the inclusion of the $\delta$-CS component leads to performance improvements over both GFN-AL and LS-GFN, which are representative search methods. This indicates that the main improvement stems from the search capability of $\delta$-CS. By restricting the search to a Hamming ball in the sequence space adjusted by the parameter $\delta$, $\delta$-CS enables a more robust search even when the proxy model is inaccurate.

These results have been included in Appendix B.10 of the revised manuscript.

Thanks for your careful and detailed feedback. Please let me know if any additional concerns have not been resolved.

Dear Reviewer WGsn,

We wanted to gently remind you that our discussion period is coming to an end soon. If you feel that your concerns have been fully addressed in our previous responses, we kindly ask you to consider revising the score. Should there be any remaining issues or questions, please do not hesitate to let us know at your earliest convenience. Thank you for your time and consideration.

Warm regards,

The Authors

W1: Novelty

We think a proper way to combine search methods ($\delta$-CS) and learning methods (GFN) together is not trivial, as many works in biological sequence design that use exploration-exploitation balancing methods like UCB or evolutionary search-based methods like GFNSeqEditor did not get satisfactory performance compared to our method. That shows our novel approach, which controls conservativeness parameters in terms of the level of the Hamming ball of local sequence space, was more effective than others, and such integration with GFN-AL and uncertainty adaptation was critical to performance; novel combination of the techniques is clearly categorized as novelty in conference review guidelines.

W2: Evaluation and experiments are testing out-of-distribution sequences, which rely on the accuracy of $f$

Q3: It seems that model $f$ is fixed in training, and it provides reward information for $\delta$-CS; how can we ensure $f$ can provide safe and correct reward information?

It seems there is a slight misunderstanding in this generative active learning benchmark for biological sequence design. In this setting, the oracle function $f$, representing the ground truth, can only be accessed a limited number of times. As a result, researchers focus on introducing a proxy model $f_{\theta}$ that learns to imitate the oracle function using the available dataset.

Thus, the setting does not require consideration of the out-of-distribution (OOD) capability of the oracle $f$. However, we do need to account for the OOD generalization capability of the proxy $f_{\theta}$, which is highly challenging. This motivates the use of conservative search strategies when the proxy exhibits high uncertainty. Consequently, the proxy function's limited accuracy further strengthens the rationale for using the $\delta$-CS off-policy method. This method enables us to compare with existing approaches that rely more heavily on the proxy function’s performance, highlighting the novelty of our approach.

Q1: Regarding the choice of $\delta$, although the authors reported a control experiment on $\delta$ in the appendix, are all the experiments in Table 1 based on the same $\delta$?

For the short sequence task with (e.g., DNA, RNA), we set $\delta = 0.5$, and for the long sequence task (AAV, GFP), we set $\delta = 0.05$. These parameter values were roughly tuned using some validation experiments. Specifically, we conducted tests with a proxy model trained on a portion of the initial dataset and a validation model trained on the entire initial dataset.

Q2: What is the initial dataset $D_0$? How is it initialized? What is the cost to acquire the reward for all sequences in $D_t$ in line 3 of Alg.1?

As we noted in experimental section, the initial dataset $D_0$ for TFbind, RNA, we follows exisiting benchmark [1,2]. We collect $D_0$ for AAV and GPF based on [3], by using evolutionary search on wild type sequences given from [3].

The cost to acquire the reward for all sequences in $D_t$ is not required in line 3 (the section for rank-based sampling on dataset) because they are already precomputed for each data point and included in the dataset.

[1] Kim, Minsu, et al. "Bootstrapped training of score-conditioned generator for offline design of biological sequences." Neural Information Processing Systems (NeurIPS), 2024.

[2] Trabucco, Brandon, et al. "Design-bench: Benchmarks for data-driven offline model-based optimization." International Conference on Machine Learning. PMLR, 2022.

[3] Sinai, Sam, et al. "AdaLead: A simple and robust adaptive greedy search algorithm for sequence design." arXiv preprint arXiv:2010.02141 (2020).

W1: The adaptive $\delta$ seems not to be adaptive to uncertainty, just constant.

We acknowledge that our initial expression may cause confusion. Our method adapts to each datapoint $x$ (a given sequence to edit) based on the prediction uncertainty for $x$, leveraging $\sigma(x)$ to adjust $\delta$. The misleading part arose from our handwavy expression $\lambda \sigma \approx 1/L$. Therefore, we revised it in the manuscript as

$\lambda E_{P_{\mathcal{D}_0}(x)}[\sigma(x)] \approx \frac{1}{L},$

which means we set the scaling hyperparameter $\lambda$ based on the average value of $\sigma(x)$ during the first round. Consequently, $\delta \approx \delta_{\text{const.}} - 1/L$ is not true because: 1) $\sigma$ varies for each datapoint $x$'s, and 2) even for the same $x$, $\sigma(x)$ evolves throughout the active learning rounds.

W2: It would be nice to see the evaluation of various $\delta$ from Appendix B.1 extended to multiple datasets. It seems that the choice of $\delta$ is critical for the achieved performance and it would be interesting what ranges of $\delta$-values outperform the strongest baseline in each of the tasks to develop an understanding for the robustness of $\delta$-CS.

Thank you for your suggestion. We conduct additional experiments with different $\delta$ values to validate the robustness of $\delta$-CS. The summarized results are described in the tables below, and we also include more comprehensive results in Appendix B.1 of the updated manuscript.

RNA-A

GFP

Q1: Can the authors elaborate how $\delta$ is chosen in the experiments? Is this a choice that can be reproduced in practice, i.e., does it require running $\delta$-CS with multiple (all?) values of $\delta$, or does it require determining $\delta$ from a (small) validation experiment, or something else?

We set $\delta = 0.5$ for short sequence tasks (e.g, DNA, RNA), and $\delta = 0.05$ for long sequence tasks (e.g., GFP, AAV). This hyperparameter was very roughly tuned with small validation experiments; we conducted tests with a proxy model trained on just part of the initial dataset and a validation model trained on the entire initial dataset as oracle. These experiments were only done in the short DNA sequence task and the long protein sequence task. This tuning can be further improved if we finely tune task-by-task using such validation experiments.