We thank the reviewer for their helpful comments on our work.

What do the reference in-fills look like for the other distributions? It only shows the GFlowNet examples in Table B.3.

We have added samples generated by each method on 4 randomly selected examples to Appendix C in the updated manuscript.

What’s the speed of learning/convergence relative to, e.g. supervised fine tuning? Is the inference unstable/need many restarts, etc? How about compared to PPO training?

In terms of wall time, GFlowNet fine-tuning is slower compared to supervised fine-tuning as it has to explore the space of sequences and has a similar runtime to PPO. We did not encounter instabilities during training that required special treatment like restarts, etc.

Did the authors fine-tune separately for every temperature, or was this done in the style of an amortized sampler where you can dynamically specify the target temperature at run time?

We do not amortize over the temperature, though it has been studied in some prior work on GFlowNets. We linearly anneal the reward over the course of training.

We hope we have addressed all of your concerns. Please don't hesitate to let us know if we can clarify anything else!

We thank the reviewer for their helpful comments on our work.

GFlowNets start with an empty string and add one token at a time in a left-to-right manner. Depending on different tasks, $Z$ should be generated conditional on $X$ or $X, Y$? Here, $X$ or $X, Y$ is omitted?

In the subjectivity and arithmetic experiments, the GFlowNet depends only on $X$ so we can use it on unseen problems at test time when $Y$ is not available. $Y$ is available at test time for the infilling experiment, so we the GFlowNet is conditioned on both $X$ and $Y$.

Besides that 1) the ability to avoid estimating the flow function $F$; 2) SubTB can have a better bias-variance trade-off in GFlowNet training, are there any other benefits to use a modified version? Also, did you try the conventional SubTB objective with the flow function considered?

The primary motivation for the modification in our approach is to avoid learning a state flow function, limiting the learnable objects to just the forward policy (i.e., the same data that is output by an autoregressive LM). We did not experiment with the standard SubTB objectives since it would add additional complexity (adding an extra head to the LM to output the flow) not central to our main contribution.

Did you consider the hyper-parameter $\lambda^{j-i}$ over incomplete trajectories with variable lengths $0 \leq i < j \leq n+1$, like SubTB($\lambda$)?

We did not tune the SubTB hyperparamter $\lambda$ and left it to the default value of 1 for all experiments following prior work on GFlowNets, such as [Hu et al., 2023].

Given that the generation order is fixed (i.e., left-to-right), it results in $P_{B} = 1$. For readers who might be unfamiliar with GFlowNets, it might be helpful to include an explanation or mention $P_B=1$ somewhere in the paper to ensure accessibility for all readers?

Thank you for the feedback! We have added a note to clarify this in the updated draft.

$R(Z) = p_{LM}(XZY) \propto p_{LM}(Z | X, Y)$ --> should be $p_{LM}(Z | X, Y) \propto R(Z) = p_{LM}(XZY)$?

We may be misunderstanding your comment: those two are equivalent, if we are talking about proportionality in $Z$. The expression is intended to define the reward for $Z$ as the likelihood $p_{\text{LM}}(XZY)$ which is the quantity proportional to the desired posterior. Thus we state it as $R(Z) = p_{\text{LM}}(XZY) \propto p_{\text{LM}}(Z | X, Y)$.

How to understand GFlowNet fine-tuning and supervised fine-tuning solely? The former is to train the LM with Eq.3, while the later is to train the LM by maximizing $\log p_{LM}(XZY)$ with $Z$. Supervised fine-tuning corresponds to the variational EM - first update GFlowNet polices and then LM pamemters with $Z$? Thus, supervised fine-tuning should already include GFlowNet fine-tuning?

GFlowNet fine-tuning corresponds to only the E-step in EM. Supervised fine-tuning on its own directly maximizes $\log p_{LM}(Y|X)$ without a latent variable $Z$.

The supervised fine-tuning on top of GFlowNet fine-tuning, however, maximizes $\log p_{LM}(Y|XZ)$ with $Z$ drawn from the GFlowNet (this is the M-step in EM).

In Table 3, GFlowNet fine-tuning + supervised fine-tuning was considered. Then why not to consider it as well in Table 2 & 4?

In story infilling (Table 2), the goal is to generate $Z$ given both $X$ and $Y$, so there is no need to finetune $p_{LM}(Y|X, Z)$. For the tool use problem (Table 4), the $Z$ already contains the solution and consequently, there isn't much improvement to be expected with subsequent supervised fine-tuning of the LM. Thus, we do not consider GFlowNet fine-tuning + supervised fine-tuning in those experiments.

4.1 Sentence continuation - task description $R(Z) = p_{LM}(Z | X)^{\frac{1}{T}}$ --> should be $R(Z) = p_{LM}(XZ)^{\frac{1}{T}}$?

These two rewards are equal up to a multiplicative constant that would only depend on $X$, i.e., $p_{LM}(XZ)^{\frac{1}{T}} = p_{LM}(Z|X)^{\frac{1}{T}}p_{LM}(X)^{\frac{1}{T}}$. The conditioning variable $X$ is given as input to the GFlowNet policy. This means that both rewards describe the same posterior distribution. Therefore, in theory, an optimally-trained GFlowNet would converge to the same solution.

However, in practice, using $R(Z) = p_{LM}(XZ)^{\frac{1}{T}}$ could be more difficult to optimize because the scale of the reward could differ significantly for different $X$'s (e.g., sentences following very long prompts would generally have much lower rewards than sentences following shorter prompts). The choice of using $R(Z) = p_{LM}(Z | X)^{\frac{1}{T}}$ corresponds to substracting a baseline from the log-reward that depends only on $X$, a common variance reduction technique in RL.

We thank the reviewer for the strongly positive assessment and insightful comments on our work.

Limited discussion of the training objective and its relationship to possible alternatives. The training objective is introduced very briefly and without much intuition. I realize that there is extensive literature on training GFlowNets and there is not enough space to go into full detail here. But are there reasons that this objective (among many other GFlowNet objectives) was particularly well-suited to the language modeling case?

We use the subtrajectory balance loss as it has favorable properties such as low gradient variance which enables stable training [Madan et al., 2023]. The modification to account for each state being a valid terminal state comes from [Deleu et al., 2022] and consists in replacing the flow function $F(z_{1:n})$ by $R(z_{1:n}\top)/q_{\text{policy}}(\top\mid z_{1:n})$. We make this choice because it allows us to instantiate the GFlowNet without explicitly parameterizing the flow function, only the policy itself.

Why GFlowNets at all instead of e.g. reweighted wake-sleep (a common method for amortizing intractable posterior inference)?

GFlowNets have an advantage over other variational approaches -- the ability to train on off-policy trajectories without resorting to importance sampling. Approaches such as reweighted wake-sleep rely on importance sampling to utilize off-policy trajectories, which can lead to high-variance and biased estimates of the gradient. This finding was established by [Malkin et al., 2023].

How sensitive is performance to the distribution you use to generate training trajectories?

The distribution of training trajectories we used is a mixture of on-policy trajectories, trajectories from a tempered policy, and trajectories from the replay buffer. The performance indeed depends on the choice of this distribution. Due to the compute constraints, we prioritized ablating factors such as seeding the replay buffer with good rationales over this aspect.

How important is the replay buffer, and how is it populated? In fairness, I am not sure how many of these questions need to be addressed in a short conference paper.

The replay buffer is populated using trajectories sampled from the policy (and a tempered version of it) during training and some seed rationales at the beginning. Trajectories are added to the buffer based on a diversity threshold, i.e., a trajectory is added to the buffer only if it is a distance $\delta$ from every example in the buffer or it has a higher reward than the closest element, and the buffer is instantiated as a priority queue. Our analysis in Table D.4 and Table E.3 illustrates the importance of the number of examples used to seed the buffer. Additionally, we observe that the off-policy trajectories provided by the buffer are critical for reliable training.

For sentence continuation, do you see interesting pathologies at lower reward-model temperatures (e.g., bias toward very short completions, or very repetitive completions)?

Indeed, we do see pathologies at lower reward temperatures, which is why the lowest temperature we provided results for was $T = 0.8$. For instance, when we train with $T = 0.6$ the average sentence length becomes only $1.58$ tokens, with a large proportion of sentences consisting of just an empty space followed by a period, or other short, common, and generic phrases (e.g., "We know."). Furthermore, the diversity decreases to $\sim 0.4$, compared to $\sim 0.75$ at $T = 0.8$. Despite this, the model samples sentences with high log-likelihood values that are comparable to what is achieved at higher temperatures ($-9.18$ maximum log-likelihood). This means that the pathologies are due to a combination of two factors: (1) the LM used for the reward has a bias for short and generic sentences, as is expected, and (2) training at lower temperatures harms exploration, as evidenced by the lower diversity and comparable log-likelihood of the samples. While the second problem can potentially be addressed with small modifications to our training scheme (e.g., more gradual reward temperature annealing), the first problem suggests that LM sentence log-likelihood on its own is a suboptimal reward signal for generating naturalistic sentences and that alternatives should be explored.

We thank the reviewer for their detailed comments. We address each of them below. We have also updated the paper, with the changes colored red.

As a general comment, we are aware that problem-specific models have been proposed for some of the tasks we consider, notably text infilling. (Thank you for suggesting the references.) However, our goal is not to train new models to achieve superior performance on any of these tasks but to extract knowledge from a pretrained LLM by efficiently fine-tuning it. To this end, we propose GFlowNet fine-tuning as a tool to solve general intractable inference problems in LLMs.

Overall the paper is hard to follow: the authors provide little background on reinforcement learning and GFlowNet training. In particular, the authors use many terminologies without/before defining them clearly, examples include “policy”, “reward”, “matching” a target distribution, “rewarding all valid integers equally leads to an expected gradient of zero for policy gradient methods.”

We appreciate the feedback. While we cannot provide a comprehensive background on RL and GFlowNets due to space constraints, we have added pointers to the relevant literature and a glossary in Appendix A.

In section 2, by looking at the problem of using LLMs to generate random numbers between 0 - 100, the authors try to motivate the use of GFlowNet instead of PPO training. PPO training does not resolve the distribution skew because the reward function only considers whether the number lies between 0 - 100.

We would like to clarify that the problem considered in section 2 serves as a simple demonstration of the underlying problem of amortized inference, where one has access to a likelihood and the problem is to sample from the desired distribution. RL algorithms such as PPO are formulated to maximize the reward, and thus fail in this scenario. Note that the reward used is the same for GFlowNets and PPO, and the only thing that differs is the fine-tuning algorithm. This example serves precisely to illustrate that reward maximization (PPO) is less appropriate than distribution matching (GFlowNet) in some settings, motivating our approach.

One correct way to do it could be asking the LLM to generate a sequence of numbers sampled from 0 - 100 uniformly and assign a positive reward only if the frequency of the numbers is close to uniform.

Our goal here is not to sample random numbers from LLMs but to demonstrate a critical shortcoming with existing fine-tuning paradigms. The approach of rewarding a sequence of numbers is undesirable for many reasons: 1) PPO can collapse to deterministically generating a sequence of uniformly distributed but not independent numbers (e.g., just listing the numbers from 0 to 100 in order), which would not yield a reliable sampler; 2) it does not scale to scenarios where we want to sample long reasoning chains, because we will need to generate many of these long chains before receiving a reward; 3) the reward is hard to design and may have a high variance.

A major part of the introduction focuses on intractable posterior inference/conditional probabilities and the fact that Section 2 mentions nothing about them makes it hard to follow.

We have made the transition more clear. Section 2, through a minimal problem of using language models to sample from a distribution given an unnormalized density, demonstrates the shortcomings of reward-maximizing RL for posterior inference and introduces GFlowNet fine-tuning as the appropriate tool.

In section 3 the authors introduced some related problems in NLP that could potentially be solved by GFlowNet and in section 3.3 on page 5 that the authors finally describes GFlowNet and their training objective. What is the original subtrajectory balance objective? How do you modify it? What is the semantics of your objective function? Answers to these questions can help distinguish GFlowNet from other approaches from the methodology perspective.

We limited our exposition of GFlowNets in general to maintain the focus on the setting studied in the paper and direct the reader to relevant prior work. Specifically, the original subtrajectory balance term for a subtrajectory $\tau_{m:n} = s_m\rightarrow\dots\rightarrow s_n$ is:

$ℒ_{\text{subTB} }(\tau_{m:n}) = \bigg[\log\frac{F(s_m)\prod_{i=m}^{n-1}P_F(s_{i+1}\mid s_i)}{F(s_n)\prod_{i=m}^{n-1}P_B(s_i \mid s_{i+1})}\bigg]^2$

where $P_F$ is the sampling policy -- called $q_{\text{policy}}$ in our paper -- and $P_B$ is a "backward policy" (we refer to past GFlowNet work for discussion of what this means, but note that in our setting the $P_B$ terms are always 1 and can be ignored).

Our modification leverages an important aspect of the problem setting: During generation, the trajectory can terminate at each state. To account for this, we incorporate the modification proposed by [Deleu et al., 2022], which involves incorporating the termination likelihood and the reward at each state. Specifically, using the fact that at convergence we have $R(s_n^\top)=F(s_n)P_F(\top\mid s_n)$, we simply substitute $R(s_n^\top)/P_F(\top\mid s_n)$ for $F(s_n)$ in the loss above. Rearrangement of the terms yields exactly the term being summed in our loss (3).

These modifications together allow us to parameterize the GFlowNet through only the forward policy, avoiding the need to train additional estimators.

GFlowNet objectives have been studied extensively in prior work. Semantically, these learning objectives aim to satisfy constraints on the distribution over the trajectories given by the policy to sample proportionally to the reward.

Besides, some important related works of the field are missing from Section 3

Thanks for pointing to the additional related work! We have added the missing references.

for temperature scaling: [1] leverages importance sampling to fine-tune LM p(x) such that it approximates the desired distribution p(x)^{1/T}. Their approach suffers from various problems such that the variance of loss is high due to the exponent 1/T. Given that the authors study this empirically, does the GFlowNet objective also suffer from this issue? If so, how is it resolved?

A key advantage of our method is that we do not rely on importance sampling for off-policy learning, an advantageous property of GFlowNets that was studied by [Malkin et al., 2023]. The variance of our loss is therefore not prohibitive for stable training. Independent from importance sampling, the variance of our loss does increase with lower temperatures (especially early in training) simply because this increases the difference between the target distribution and the initial LM distribution that we are fine-tuning. We mitigate this by slowly annealing the temperature from $1$ down to its final value throughout training. As a result, at any given point during training, the GFlowNet distribution is close to the current target distribution and the variance (and magnitude) of the loss is small. Empirically, we found that this greatly reduces the variance of our loss to the point where it is not a significant concern.

Overall, we would like to reiterate that the key goal of the paper is not to achieve state-of-the-art performance in each of the settings we consider. Rather, the GFlowNet fine-tuning paradigm, as you point out in your review, provides a unified view for all of these intractable inference problems and proposes a single general approach to tackle all of these problems. The breadth of our experiments intends to demonstrate the problem-agnostic nature of GFlowNet fine-tuning. Different problems simply correspond to different reward functions.

We appreciate your engagement and additional comments.

We see the advantages of our approach to intractable inference with LLMs as the following:

• Principled objectives: Ours is the first attempt at general-purpose amortized inference in LLMs that has a guarantee of matching the target distribution when trained to zero loss. As such, it has relatively few "moving parts" beyond those present in the design of any RL algorithm. This is in contrast to prior work that has used Monte Carlo approaches or specialized wake-sleep algorithms for specific intractable inference problems, which, respectively, do not perform amortization and do not feature a loss that can be optimized to zero to yield an exact sampler.
• Bayesian formulation: We offer a clean Bayesian inference perspective on infilling and chain-of-thought reasoning and present an algorithm to perform this inference. This is already a major advantage over the standard approach to chain-of-thought reasoning through judicious prompting and in-context learning, in which a Bayesian formulation emerges only in post-hoc analysis (cf. the recent literature on the Bayesian interpretations of ICL).
• Versatility and breadth of applications: As you noted, we show that our approach can be used to fine-tune LLMs to solve a wide variety of intractable inference problems, namely infilling, chain-of-thought / reasoning chain inference, tool use, and even the fundamental problem of tempered autoregressive sampling. The versatility of GFlowNet fine-tuning is a key advantage over approaches that focus on specific inference problems.

Thus, our general answer is that "if you can afford to GFlowNet-fine-tune your LLM for your inference problem, then you should try to do so".

This leads naturally to the question of limitations, some of which we have already discussed some in our paper and responses. The three main ones we see are:

• Compute cost: GFlowNet fine-tuning is more expensive than supervised fine-tuning, as it involves exploration. (This can in part be mitigated by seeding of the replay buffer with high-quality samples obtained using a different algorithm.)
• Sensitivity to training parameters: GFlowNets are reinforcement learning algorithms and as such require choices of exploration parameters in addition to those present in any fine-tuning setting. (The use of the replay buffer and the temperature annealing schedule are especially important, as we have found.) The need to search for good parameters can be costly, and we have not explored the full range of possible settings and tricks.
• Reliance on reward model and possible misalignment: GFlowNet fine-tuning requires a fixed reward model (or a lightly varying one in the case of an EM loop, but in any case defined by a base LM). While formulating the reward as an unnormalized Bayesian posterior is straightforward in cases such as infilling, it may be difficult to balance quality and diversity in general constrained generation settings. Additionally, as we allude to in the conclusion and in Appendix E, high-reward sequences are not necessarily of high quality. To summarize, while GFlowNet fine-tuning seems to be quite good at learning to sample a posterior, it does not address the question of what that posterior should be, nor does it address the failures of the LLM reward model to capture the desired properties.

We hope that this discussion helps put our contributions into perspective.