Paper: Pre-Training and Fine-Tuning Generative Flow Networks (ICLR-2024)

Reviewer comment (Experiments):
Clarifications on the experimental setting (see Q2)

Author response:
*Q2: Not sure I completely understand how the number of discovered modes (normalised and unnormalised) is calculated (see Figures 10 and 11)?*

We use the same evaluation metrics as in the literature in GFlowNets [1-2]. The number of modes is calculated using a sphere-exclusion procedure. Specifically, the generated candidates are sorted by reward, and then scanning down this sorted list, a candidate is added to the list of modes if it is above a certain reward threshold and is further away than some distance threshold from all other modes. For the normalized number of modes, we normalize the number of modes discovered in each task by min-max normalization, and average this value across different tasks. This procedure is adapted from [1-2] and we have added this description in Appendix C.1.

Reviewer comment (Evaluation):
Not sure I completely understand how the number of discovered modes (normalised and unnormalised) is calculated (see Figures 10 and 11)?

Author response:
*Q2: Not sure I completely understand how the number of discovered modes (normalised and unnormalised) is calculated (see Figures 10 and 11)?*

We use the same evaluation metrics as in the literature in GFlowNets [1-2]. The number of modes is calculated using a sphere-exclusion procedure. Specifically, the generated candidates are sorted by reward, and then scanning down this sorted list, a candidate is added to the list of modes if it is above a certain reward threshold and is further away than some distance threshold from all other modes. For the normalized number of modes, we normalize the number of modes discovered in each task by min-max normalization, and average this value across different tasks. This procedure is adapted from [1-2] and we have added this description in Appendix C.1.

Reviewer comment (Writing):
“A remarkable aspect of GFlowNets is the ability to demonstrate the reasoning potential in generating a task-specific policy.“ Not sure the word reasoning would be the right way to describe it as I am not sure it is objective; despite not necessarily agreeing with the wording, I do think the meaning behind it, i.e. transferring to the task without re-training (Section 4.2), is a valuable aspect.

Author response:
*Q3: “A remarkable aspect of GFlowNets is the ability to demonstrate the reasoning potential in generating a task-specific policy.” Not sure the word reasoning would be the right way to describe it as I am not sure it is objective; despite not necessarily agreeing with the wording, I do think the meaning behind it, i.e. transferring to the task without re-training (Section 4.2), is a valuable aspect.*

We would like to thank the reviewer for recognizing the value of Section 4.2. Our goal was to emphasize the brain-inspired reasoning potential [3] that GFlowNets exhibit, while our work takes a step further by enabling efficient fine-tuning on downstream tasks. We acknowledge that the term "reasoning" may not be the most precise choice, and we have revised our wording to avoid confusion.

Reviewer comment (Theory):
There should be a discussions of assumptions behind the OC-GFNs pretraining. Namely, that transfer is only possible when the reward function changes but not if the action-space or the state-space change. Moreover, the goal-conditioning requires a well specified set of outcomes Y — presumably not all states s are terminal states — which makes the proposed method not truly unsupervised. These limitations (together with the applicability mentioned at the end of A.2) could be stated explicitly in the main text, and left to future work.

Author response:
Thanks for the suggestion. For the definition of the terminal states, as in most GFlowNets benchmarks, the agent is allowed to take a "stop" action at every state, so every state can be a terminal state. We explicitly include this in Section 4 in the main text and Appendix D due to space limitation. We hope that our general framework can induce many fruitful follow-ups in the pre-training and fine-tuning paradigm of GFlowNets, and combine the recent success of this paradigm in vision, language, and RL domains.

Reviewer comment (Experiments):
While there are enough benchmarks, I believe none include continuous action/state spaces. Moreover, the experiments only one GFN variant — the detailed-balance one, which is also used for OC-GFN. It would help validate the generality of OC if we had experiments showing it worked on these different settings. Moreover, I’d be curious to know how other pretrained amortized sampling baselines (eg, VAEs, normalizing flows) fare against OC-GFN — and what about pretraining a GFN on task A (without OC) and fine-tuning it on task B?

Author response:
(Continuous state/action spaces) On the one hand, while continuous GFlowNets have recently been studied theoretically [1], in practice they remain limited in terms of applications. As noted in [2], continuous GFlowNets can be quite tricky to train and are challenging to scale to realistic tasks (which is also an open problem in GFlowNets). Therefore, we focus on the discrete case which is much more reliable (and most of the success benchmarks for GFlowNets focus on discrete space), and consider the practical and challenging tasks in the biological sequence generation domain. On the other hand, we do note that our method, in theory, is still applicable to the continuous case. 

(Applicability on other GFlowNets objectives) We include an additional study of OC-GFN when built upon the sub-trajectory balance [3] objective (OC-SubTB) in Section 5.1, where a detailed discussion can also be found in Appendix B.2. The experiments validate the versatility of our proposed methodology, which is general and can be applied to different GFlowNets learning objectives that learn state flow functions. 

*Q3: Moreover, I’d be curious to know how other pretrained amortized sampling baselines (eg, VAEs, normalizing flows) fare against OC-GFN.* 

Approaches such as VAEs and normalizing flows differ from GFlowNets in a key way -- they are learned using samples from some distribution and typically cannot be trained directly with a reward function. Although methods like MCMC can be used to sample proportionally to a given function, they typically do not take advantage of machine learning, and we are not aware of any such methods exploiting unsupervised pre-training for MCMC. While we are not aware of any such approaches that are applicable to the domains we study, we would be happy to include any such baselines for comparison. 

*Q4: What about pretraining a GFN on task A (without OC) and fine-tuning it on task B?*

We include an additional comparison of a direct pre-training and fine-tuning scheme for GFlownets in Appendix C.3 (by pre-training a GFlowNet on task A, and then fine-tuning it on task B), which is also a motivation for our work. Results show that this simple scheme for pre-training GFlowNets does not lead to a universally helpful strategy for efficient transfer. This is because the GFlowNet learns to sample proportionally from the reward function of task A during the "pre-training" stage, and thus, if there is not enough shared structure between the rewards for task A and task B, the GFlowNet model pre-trained on task A may not be beneficial for the "fine-tuning" stage on task B, and could potentially even have a negative impact. On the contrary, our insight for unsupervised pre-training is that we cast this challenging problem as a self-supervised problem by learning to reach any goal for learning a *functional* understanding of the task.

Reviewer comment (Writing):
(minor) The second and fourth paragraphs of Section 4.2 mention the “reasoning potential” of GFNs, and that intractable marginalization leads to “slow thinking”. Are these anthropomorphisms really needed for this paper?

Author response:
We discuss this aspect to distinguish the pre-trained GFlowNet model from traditional RL policies, which typically learn a reward-maximizing policy and do not have this potential. Our goal was to emphasize the brain-inspired reasoning potential [4] that GFlowNets exhibit, while our work takes a step further by enabling efficient fine-tuning on downstream tasks. We acknowledge that the term "reasoning" may not be the most precise choice, and we have revised our wording to avoid confusion.

Reviewer comment (Writing):
(minor) I wished the preliminaries (Section 2) included a training objective like Eq. 5 & 9, and that these more clearly specified which are the optimization variables.

Author response:
Thanks for the suggestion, we did not include the corresponding loss function for the learning objective due to space limitation. We have now included the loss functions and optimization variables for Section 2 in Appendix A.1 for improved readability.

Reviewer comment (Experiments):
The paper lacks a comparison with existing approaches for pre-trained models or goal-conditioned RL methods.

Author response:
We would like to highlight that we propose the *first* pre-training paradigm for GFlowNets, and there are no existing methods for pre-training GFlowNets models to the best of our knowledge. However, as a point of comparison, we also include a discussion of a simple scheme (pre-train a GFlowNet on task A and then fine-tune it on task B) in Appendix C.3, which also serves as a motivation for our work. Results show that this simple scheme for pre-training GFlowNets does not lead to a universally applicable strategy for efficient transfer. This is because the GFlowNet learns to sample proportionally from the reward function of task A during the "pre-training" stage, and thus, if there is not enough shared structure between the rewards for task A and task B, the GFlowNet model pre-trained on task A may not be beneficial for the "fine-tuning" stage on task B, and could potentially even have a negative impact. On the contrary, our insight for unsupervised pre-training is that we cast this challenging problem as a self-supervised problem by learning to reach any outcome for learning a "functional" understanding of the task.

Reviewer comment (Experiments):
How does the proposed approach perform compared to existing methods for pre-trained models or RL methods besides DQN? and what about the computation cost of these methods?

Author response:
We would like to highlight that we propose the *first* pre-training paradigm for GFlowNets, and there are no existing methods for pre-training GFlowNets models to the best of our knowledge. However, as a point of comparison, we also include a discussion of a simple scheme (pre-train a GFlowNet on task A and then fine-tune it on task B) in Appendix C.3, which also serves as a motivation for our work. Results show that this simple scheme for pre-training GFlowNets does not lead to a universally applicable strategy for efficient transfer. This is because the GFlowNet learns to sample proportionally from the reward function of task A during the "pre-training" stage, and thus, if there is not enough shared structure between the rewards for task A and task B, the GFlowNet model pre-trained on task A may not be beneficial for the "fine-tuning" stage on task B, and could potentially even have a negative impact. On the contrary, our insight for unsupervised pre-training is that we cast this challenging problem as a self-supervised problem by learning to reach any outcome for learning a "functional" understanding of the task.

Regarding other RL baselines, we include additional comparisons with two other strong RL baselines (PPO and SAC) besides DQN in Appendix C.4 on two downstream tasks in TF Bind generation due to time limitation for the rebuttal period. As shown, the reward-maximization nature of RL approaches limits the resulting solution diversity [1], where OC-GFN continues to outperform these baselines consistently.

We also include computation costs for each method in Appendix C.4 (Figure 18), where we observe that OC-GFN only induces marginal additional computation cost.

Reviewer comment (Experiments):
Is the trained GAFlowNet necessary? What about its performance and how does it influence the results?

Author response:
Thank you for the question. The GAFlowNet model plays a critical role in the unsupervised pre-training stage. Please find our detailed discussion and experiments below.

(Intuition) GAFlowNet [2] enables the learning of GFlowNet driven by intrinsic rewards (e.g., curiosity about states [3] -- which gives a higher reward for less familiar/unseen states and vice versa). This makes it very suitable in the unsupervised pre-training stage [4] where the agent generates the outcomes autotelically, which requires training it with full support over the outcomes. GAFlowNet improves the efficiency of generating diverse outcomes, which is particularly important in higher-dimensional spaces. 

(Experiments) We include an ablation study for the effect of the GAFlowNet model by replacing it with a random agent and a GFlowNet which is trained with task-irrelevant rewards ($R(x)=1$) in the unsupervised pre-training stage. (Please also note that the effect of completely removing this part has been studied in Section 5.1 (Figure 4), which is a key to the success of the unsupervised pre-training stage.)

(Empirical results) We compare the success rate of each method for reaching given outcomes following Section 5.1. As shown in Figure 14 in Appendix C.2, OC-GFN (GAFN) learns much more efficiently than the other two variants, where the GAFN model is replaced with a GFN model trained with task-irrelevant rewards (OC-GFN (GFN)) or a random agent (OC-GFN (random)), with a more significant margin in larger maps. In addition,  we also study the generated outcome distribution by ablating this model design to gain a better understanding during the learning process. Figure 15 in Appendix C.2 demonstrates the outcome distribution generated by each method at the same iteration in a large map. As shown, GAFN generates outcomes that allow the agent to explore unfamiliar regions, resulting in a more diverse outcome distribution for training the OC-GFN and improving its learning efficiency, which also demonstrates the importance of the GAFN model.