Paper: Local Search GFlowNets (ICLR-2024)

Reviewer comment (Evaluation):
The proposed method adds some computational overhead

Author response:
**W1: About computational overhead**

LS-GFN does not introduce any additional computational overhead. Below are the wall clock times for each training round in the QM9 task.

|  | $I$ |$M$ | Wall clock time per round|
| -------- | -------- | -------- | -------- |
| TB   | 0     | 32     |   1.93 ± 0.13 seconds  | 
| TB + LS-GFN    | 7     | 4     |  1.91 ± 0.07 seconds |

The wall clock time is evaluated with a single RTX 3090 GPU and Intel(R) Xeon(R) Gold 5317 CPU @ 3.00GHz and 256 GB RAM.




This restriction is in place to ensure that the sample complexity remains consistent with that of the baseline GFN methods.

The calculation of sampling complexity follows this formula: $(I+1) \times M$, where $I$ signifies the number of local search iterations, and $M$ denotes the batch size per training round.

| | $I$ |$M$ | Sampling Complexity per Training Round |
| -------- |-------- |-------- |-------- |
| GFN        | 0     | 32     | $1 \times 32 = 32$     |
| LS-GFN         | 7     | 4     |$(7+1) \times 4 = 32$      |

Reviewer comment (Experiments):
How is the number of steps K picked for backtracking? If it's fixed, is there a way to pick it automatically?

Author response:
**Q1: How is the number of steps K picked for backtracking? If it's fixed, is there a way to pick it automatically?**

We used $K=\lfloor (L+1) / 2\rfloor$, where $L$ is length of trajecotry. We also provided a hyperparameter analysis in the original manuscript in Appendix B.5. As you mentioned, we can automatically pick $K$ to adaptively balance the wideness of the local search region, which could be a great avenue for future work.

Reviewer comment (Experiments):
The paper does not declare the sampling complexity of the proposed method. The local search may require more sampling, which can lead to an unfair comparison.

Author response:
Thank you for highlighting this for clarification. All our experiments were conducted under a fair setting, as we used exactly the same number of samples across all experiments. Note this is already outlined in Section 5.3.

The calculation of sampling complexity follows this formula: $(I+1) \times M$, where $I$ stands for the number of local search iterations, and $M$ denotes the batch size per training round. 



| | $I$ |$M$ | Sampling Complexity per Training Rounds |
| --------  | -------- |-------- |-------- |
| GFN      | 0     | 32     | $1 \times 32 = 32$     |
| LS-GFN        | 7     | 4     |$(7+1) \times 4 = 32$      |

Reviewer comment (Evaluation):
The limitations of the proposed algorithm and potential drawbacks are not discussed in detail.

Author response:
A limitation of LS-GFN lies in the potential impact of the quality of the backward policy on its performance, particularly when the acceptance rate of the local search becomes excessively low. One immediate remedy is to introduce an exploratory element into the backward policy, utilizing techniques like $\epsilon$-greedy or even employing a uniform distribution to foster exploration within the local search.

A promising avenue for future research could involve fine-tuning backward policy to enhance the acceptance rate of the local search. We plan to incorporate this as a topic in our paper's "Limitations and Further Work" after we get one additional page for the final paper (we now put this into the appendix), and we appreciate your valuable suggestion.

Reviewer comment (Experiments):
What about the sampling complexity of local search? Can you conduct experiments under the same sampling complexity or compare your algorithm with an upper bound sampling times?

Author response:
Thank you for highlighting this for clarification. All our experiments were conducted under a fair setting, as we used exactly the same number of samples across all experiments. Note this is already outlined in Section 5.3.

The calculation of sampling complexity follows this formula: $(I+1) \times M$, where $I$ stands for the number of local search iterations, and $M$ denotes the batch size per training round. 



| | $I$ |$M$ | Sampling Complexity per Training Rounds |
| --------  | -------- |-------- |-------- |
| GFN      | 0     | 32     | $1 \times 32 = 32$     |
| LS-GFN        | 7     | 4     |$(7+1) \times 4 = 32$      |

Reviewer comment (Experiments):
How to decide the local search interaction I? What about the results with different I?

Author response:
We have conducted an ablation study on the hyperparameter $I$, as now detailed in Appendix B.3. 

It's worth noting that across various hyperparameter candidates, namely $I \in \{1, 3, 7, 15, 31\}$, LS-GFN ($I > 0$) consistently outperforms GFN ($I = 0$). This observation suggests that choosing the appropriate hyperparameter $I$ for LS-GFN is relatively straightforward, as it consistently yields superior results.

Reviewer comment (Experiments):
As PRT is adopted in LS-GFN training, is PER used for RL methods in the experiments?

Author response:
Taking into consideration your feedback, we implemented reward-prioritized experience replay training in conjunction with the off-policy RL baseline, SQL. It's worth noting that PPO and A2C are on-policy RL methods in which replay training is not directly utilized. Here are new experiment results that compare replay-trained SQL and LS-GFN:



| | Number of Modes | TopK reward |
| -------- | -------- | -------- |
| SQL    | 232 ± 8     | 0.56 ± 0.01     |
|   SQL (replay trained)   | 235 ± 6     | 0.56 ± 0.01     |
| TB + LS-GFN   | 793 ± 4     | 0.67 ± 0.00     |

Reviewer comment (Experiments):
In addition to the hyper-parameter $I$ (number of revisions with local searches), I wished the number of backtracking steps and the acceptance rate were also studied — how much tuning did they require in order to get these good results?

Author response:
**W2: In addition to the hyper-parameter 
 (number of revisions with local searches), I wished the number of backtracking steps and the acceptance rate were also studied — how much tuning did they require in order to get these good results?**
 
We've already conducted experiments on hyperparameter $I$ in Appendix B.3, as well as explored the acceptance rate in Appendix B.6.

It's worth noting that across various hyperparameter candidates, namely $I \in \{1, 3, 7, 15, 31\}$, LS-GFN ($I > 0$) consistently outperforms GFN ($I = 0$). This observation suggests that choosing the appropriate hyperparameter $I$ for LS-GFN is relatively straightforward, as it consistently yields superior results.

Reviewer comment (Evaluation):
Wall-clock time is never mentioned — how much overhead does the refining process of LS-GFNs incur? How about in terms of sampled states (instead of training rounds)?

Author response:
**W3: Wall-clock time is never mentioned — how much overhead does the refining process of LS-GFNs incur? How about in terms of sampled states (instead of training rounds)?**


There's negligible additional wall clock time overhead.

Here are the wall clock times for each training round in the QM9 task. Wall clock time is computed as the mean of three independent seeds. To remove the bias from warmup, we compute the wall clock time per round after 10 initial training rounds. There is no meaningful wall time difference between the two methods: 

|  | $I$ |$M$ | Wall clock time per round|
| -------- | -------- | -------- | -------- |
| TB   | 0     | 32     |   1.93 ± 0.13 seconds  | 
| TB + LS-GFN    | 7     | 4     |  1.91 ± 0.07 seconds | 


The wall clock time is evaluated with a single RTX 3090 GPU and Intel(R) Xeon(R) Gold 5317 CPU @ 3.00GHz and 256 GB RAM.


The wall clock time remains consistent because we always use the same number of samples per training round ($M \times (I+1) = 32$, $M$: batch size, $I$: number of local searches) for every method across all experiments. This choice is particularly crucial considering that evaluating the reward per sample is a significant bottleneck (in some applications, one could spend a day for reward computation), especially in real-world scenarios like bio and chemical discovery tasks. By keeping the sample complexity uniform, we ensure that the computational overhead remains stable.

Reviewer comment (Experiments):
One baseline I wished were included: given a first trajectory, resample the last K steps and only keeping the best candidate. This is akin to beam search in LLMs or go-explore in RL. Other nice-to-have baselines include top-p and top-k sampling.

Author response:
**W4: One baseline I wished was included: given a first trajectory, resample the last K steps, and only keeping the best candidate. This is akin to beam search in LLMs or go-explore in RL. Other nice-to-have baselines include top-p and top-k sampling.**

Thank you for recommending baselines that can effectively showcase the LS-GFN method. We update a new section in Appendix B.8, including new experiments and analysis of results. Here is a summarized version of Appendix B.8.

We conducted new experiments based on your recommendations. In our evaluation, we compared LS-GFN to the baseline methods you suggested, specifically focusing on the Trajectory Balance (TB) objective function in the QM9 task.

In the context of resampling, we initiate a trajectory and subsequently resample the last $K$ steps, repeating this process four times. We concurrently evaluate and select the best trajectory among these resampled trajectories within batches of size $M=8$. Consequently, the total sampling complexity amounts to $8 \times 4 = 32$.

In the case of Top K filtering, we generate a total of $M=32$ trajectories and apply a filtering process to identify the Top K  trajectories (Top2 among 32 trajectories), which are then utilized for training purposes.


**< QM9 >**
| | Number of modes|Top K Reward| 
| -------- | -------- | -------- |
| TB    | 699 ± 14     | 0.66 ± 0.00   | 
| TB + Resampling     |  778 ± 6   | **0.67 ± 0.00**     |
| TB + Top K filtering   | 768 ± 6  | **0.67 ± 0.00**      |
| TB + LS-GFN | **793 ± 4**   | **0.67 ± 0.00**   |

**< L14_RNA1 >**
| | Number of modes|Top K Reward| 
| -------- | -------- | -------- |
| TB    |   6 ± 1  | 0.87 ± 0.01  | 
| TB + Resampling     |  4 ± 0    | 0.87 ± 0.01   |  
| TB + Top K filtering   | 4 ± 2  |  0.87 ± 0.01  |  
| TB + LS-GFN | **18 ± 0**   |  **0.97 ± 0.00** |

In the QM9 task, LS-GFN outperforms baselines, having superior performances in terms of number modes. For the RNA task, which contains a larger search space than QM9, an effective search for reward improvement is crucial. The result shows that simple resampling and filtering method is not enough as they are too focused on the exploitation of existing trajectories and are unable to find novel modes effectively.

However, LS-GFN harnesses a backward policy for backtracking, resulting in a diverse set of backtracked states that are subsequently reconstructed by the forward policy. This approach has proven to be notably more effective in navigating the extensive RNA search space compared to the resampling method, which merely samples the last $K$ steps without taking advantage of a backward policy.

Reviewer comment (Writing):
Why call it “destroy” since the original trajectory isn’t always discarded? A more intuitive name could be “backtrack”, “rewind”, or anything that doesn’t suggest destruction.

Author response:
**Q1: Why call it “destroy” since the original trajectory isn’t always discarded? A more intuitive name could be “backtrack,” “rewind,” or anything that doesn’t suggest destruction.**

We are in agreement on this matter. We have updated our manuscript to incorporate the term "backtrack" instead of "destroy."

Reviewer comment (Presentation):
The top plots in Figure 5 are strange: it looks like some curves go beyond 100% accuracy. Could you either fix them so we can still see the curves on the plot, or explain what is happening?

Author response:
**Q2: The top plots in Figure 5 are strange: it looks like some curves go beyond 100% accuracy. Could you either fix them so we can still see the curves on the plot or explain what is happening?**

Please note that Figure 5 is plotted accurately in accordance with the baseline from the paper by Shen et al., 2023, where we have set our accuracy metrics following their guidelines:


$\text{Acc}(p(x;\theta)) = 100 \times \text{min}(\frac{E_{p(x;\theta)}[R(x)]}{E_{p^{*}(x)}[R(x)]},1)$


This provides an upper bound of 100% accuracy. In response to your suggestion, we have included a graph illustrating accuracy beyond the 100% mark in Appendix B.7.

Reviewer comment (Theory):
How can LS-GFNs recover from a biased backward policy? In other words, assume the forward policy is fine but the backward policy always backtracks to states which yield the same (high reward) candidate objects — how can LS-GFNs overcome this lack of exploration?

Author response:
**Q3: How can LS-GFNs recover from a biased backward policy? In other words, assume the forward policy is fine but the backward policy always backtracks to states that yield the same (high reward) candidate objects — how can LS-GFNs overcome this lack of exploration?**

Biased backward policies can indeed pose challenges. In such scenarios, we can employ an $\epsilon$-greedy method to introduce exploratory back-sampling into the biased backward policy. It's worth noting that LS-GFN performs exceptionally well in MaxEnt GFN, even when the backward policy is set to a uniform distribution (i.e., $\epsilon = 1$).

Building upon your valuable feedback, one potential avenue for future research could involve fine-tuning the backward policy to maximize the local search acceptance rate. This approach would aim to ensure that the backward policy identifies diverse intermediate states that lead to improved rewards without exhibiting a bias towards yielding the same candidate objects repeatedly. We appreciate your insightful comment.

Reviewer comment (Experiments):
Please confirm that lines 7 and 8 in Algorithm 1 aren’t swapped. If they aren’t (which partially addresses my question above), wouldn’t swapping them and extending $\mathcal{D}$ with $\tau_m$ further improve exploitation? Maybe this should be added as an ablation as well.

Author response:
**Q4: Please confirm that lines 7 and 8 in Algorithm 1 aren’t swapped. If they aren’t (which partially addresses my question above), wouldn’t swapping them and extending them further improve exploitation? Maybe this should be added as an ablation as well.**
 
No swapping is involved here. We update every sample to the dataset, irrespective of whether it is accepted or not. The rationale behind this approach is that every sample where the reward is calculated is considered valuable. During the refinement phase, our primary objective is to explore the local region in order to discover high-reward samples rather than determining whether the found sample should be used in training.

In the training process, we employ PRT to sample trajectories in proportion to their rewards from the dataset. Consequently, accepted samples have a higher likelihood of being included during training.

Regarding your suggestion to swap lines 7 and 8, it would entail updating the data only with accepted samples. However, as you pointed out, while this could enhance exploitation, it might excessively hinder exploration. To shed more light on this, here are the new results of our ablation study:



| Column 1 | Number of modes| TopK reward |
| -------- | -------- | -------- |
| Swapped version     | 754 ± 6    | 0.66 ± 0.00    |
| Original LS-GFN    | 793 ± 4     | 0.67 ± 0.00    |

Based on the above results, the original LS-GFN gives higher performances than the swapped version.