Original: L EARNING GF LOW N ETS FROM PARTIAL EPISODES FOR IMPROVED CONVERGENCE AND STABILITY
Revised: L EARNING GF LOW N ETS FROM PARTIAL EPISODES FOR IMPROVED CONVERGENCE AND STABILITY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks.
Revised: Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks.

Original: Existing training objectives for GFlowNets are either local to states or transitions, or propagate a reward signal over an entire sampling trajectory.
Revised: Existing training objectives for GFlowNets are either local to states or transitions, or propagate a reward signal over an entire sampling trajectory.

Original: We argue that these alternatives represent opposite ends of a gradient bias-variance tradeoff and propose a way to exploit this tradeoff to mitigate its harmful effects.
Revised: We argue that these alternatives represent opposite ends of a gradient bias-variance tradeoff and propose a way to exploit this tradeoff to mitigate its harmful effects.

Original: Inspired by the TD( 𝜆 ) algorithm in reinforcement learning, we introduce subtrajectory balance or SubTB( 𝜆 ), a GFlowNet training objective that can learn from partial action subsequences of varying lengths.
Revised: Inspired by the TD( 𝜆 ) algorithm in reinforcement learning, we introduce subtrajectory balance or SubTB( 𝜆 ), a GFlowNet training objective that can learn from partial action subsequences of varying lengths.

Original: We show that SubTB( 𝜆 ) accelerates sampler convergence in previously studied and new environments and enables training GFlowNets in environments with longer action sequences and sparser reward landscapes than what was possible before.
Revised: We show that SubTB( 𝜆 ) accelerates sampler convergence in previously studied and new environments and enables training GFlowNets in environments with longer action sequences and sparser reward landscapes than what was possible before.

Original: We also perform a comparative analysis of stochastic gradient dynamics, shedding light on the bias-variance tradeoff in GFlowNet training and the advantages of subtrajectory balance.
Revised: We also perform a comparative analysis of stochastic gradient dynamics, shedding light on the bias-variance tradeoff in GFlowNet training and the advantages of subtrajectory balance.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Generative flow networks (GFlowNets; Bengio et al., 2021a) are generative models that construct objects lying in a target space X by taking sequences of actions sampled from a learned policy.
Revised: Generative flow networks (GFlowNets; Bengio et al., 2021a) are generative models that construct objects lying in a target space X by taking sequences of actions sampled from a learned policy.

Original: GFlowNets are trained so as to make the probability of sampling an object 𝑥 ∈ X proportional to a given nonnegative reward 𝑅 ( 𝑥 ) .
Revised: GFlowNets are trained so as to make the probability of sampling an object 𝑥 ∈ X proportional to a given nonnegative reward 𝑅 ( 𝑥 ) .

Original: GFlowNets’ use of a parametric policy that can generalize to states not seen during training makes them a competitive alternative to methods based on local exploration in various probabilistic modeling tasks (Bengio et al., 2021a; Malkin et al., 2022; Zhang et al., 2022;
Revised: GFlowNets’ use of a parametric policy that can generalize to states not seen during training makes them a competitive alternative to methods based on local exploration in various probabilistic modeling tasks (Bengio et al., 2021a; Malkin et al., 2022; Zhang et al., 2022;

Original: Jain et al., 2022; Deleu et al., 2022).
Revised: Jain et al., 2022; Deleu et al., 2022).

Original: GFlowNets solve the variational inference problem of approximating a target distribution over X with the distribution induced by the sampling policy, and they are trained by algorithms reminiscent of reinforcement learning (although GFlowNets model the diversity present in the reward distribution, rather than maximizing reward by seeking its mode).
Revised: GFlowNets solve the variational inference problem of approximating a target distribution over X with the distribution induced by the sampling policy, and they are trained by algorithms reminiscent of reinforcement learning (although GFlowNets model the diversity present in the reward distribution, rather than maximizing reward by seeking its mode).

Original: In most past works (Bengio et al., 2021a;
Revised: In most past works (Bengio et al., 2021a;

Original: Malkin et al., 2022; Zhang et al., 2022; Jain et al., 2022), GFlowNets are trained by exploratory sampling from the policy and receive their training signal from the reward of the sampled object.
Revised: Malkin et al., 2022; Zhang et al., 2022; Jain et al., 2022), GFlowNets are trained by exploratory sampling from the policy and receive their training signal from the reward of the sampled object.

Original: The flow matching (FM) and detailed balance (DB) learning objectives for GFlowNets proposed in
Revised: The flow matching (FM) and detailed balance (DB) learning objectives for GFlowNets proposed in

Original: Bengio et al.
Revised: Bengio et al.

Original: (2021a;b) resemble temporal difference learning (Sutton & Barto, 2018).
Revised: (2021a;b) resemble temporal difference learning (Sutton & Barto, 2018).

Original: A third objective, trajectory balance (TB), was proposed in Malkin et al.
Revised: A third objective, trajectory balance (TB), was proposed in Malkin et al.

Original: (2022) to address the problem of slow temporal credit assignment with the FM and DB objectives.
Revised: (2022) to address the problem of slow temporal credit assignment with the FM and DB objectives.

Original: The TB objective propagates learning signals over entire episodes, while the temporal difference-like objectives (FM and DB) make updates local to states or actions.
Revised: The TB objective propagates learning signals over entire episodes, while the temporal difference-like objectives (FM and DB) make updates local to states or actions.

Original: It has been hypothesized by Malkin et al.
Revised: It has been hypothesized by Malkin et al.

Original: (2022) that the improved credit assignment with TB comes at the cost of higher gradient variance, analogous to the bias-variance tradeoff seen in temporal difference learning (TD ( 𝑛 ) or TD ( 𝜆 ) ) with different eligibility trace schemes (Sutton & Barto, 2018; Kearns & Singh, 2000; van Hasselt et al., 2018;
Revised: (2022) that the improved credit assignment with TB comes at the cost of higher gradient variance, analogous to the bias-variance tradeoff seen in temporal difference learning (TD ( 𝑛 ) or TD ( 𝜆 ) ) with different eligibility trace schemes (Sutton & Barto, 2018; Kearns & Singh, 2000; van Hasselt et al., 2018;

Original: Bengio et al., 2020).
Revised: Bengio et al., 2020).

Original: This hypothesis is one of the starting points for the present paper.
Revised: This hypothesis is one of the starting points for the present paper.

Original: In this paper, we propose a new learning objective for GFlowNets, called subtrajectory balance (SubTB, or SubTB( 𝜆 ) when its real-valued hyperparameter 𝜆 is specified).
Revised: In this paper, we propose a new learning objective for GFlowNets, called subtrajectory balance (SubTB, or SubTB( 𝜆 ) when its real-valued hyperparameter 𝜆 is specified).

Original: Building upon theoretical results of Bengio et al.
Revised: Building upon theoretical results of Bengio et al.

Original: (2021b)
Revised: (2021b)

Original: ; Malkin et al.
Revised: ; Malkin et al.

Original: , we show how the SubTB( 𝜆 ) objective allows the flexibility of learning from partial experiences of any length.
Revised: , we show how the SubTB( 𝜆 ) objective allows the flexibility of learning from partial experiences of any length.

Original: Experiments on two synthetic and four real-world domains support the following empirical claims:
Revised: Experiments on two synthetic and four real-world domains support the following empirical claims:

Original: (1) SubTB( 𝜆 ) improves convergence of GFlowNets in previously studied environments: models trained with SubTB( 𝜆 ) approach the target distribution in fewer training iterations and are less sensitive to hyperparameter choices.
Revised: (1) SubTB( 𝜆 ) improves convergence of GFlowNets in previously studied environments: models trained with SubTB( 𝜆 ) approach the target distribution in fewer training iterations and are less sensitive to hyperparameter choices.

Original: (2) SubTB( 𝜆 ) enables training of GFlowNets in environments where past approaches perform poorly due to sparsity of the reward function or length of action sequences.
Revised: (2) SubTB( 𝜆 ) enables training of GFlowNets in environments where past approaches perform poorly due to sparsity of the reward function or length of action sequences.

Original: (3) The benefits of SubTB( 𝜆 ) are explained by lower variance of the stochastic gradient, with the parameter 𝜆 allowing interpolation between the high-bias, low-variance DB objective and the low-bias, high-variance TB objective.
Revised: (3) The benefits of SubTB( 𝜆 ) are explained by lower variance of the stochastic gradient, with the parameter 𝜆 allowing interpolation between the high-bias, low-variance DB objective and the low-bias, high-variance TB objective.

Original:  M ETHOD 2. P RELIMINARIES
Revised:  M ETHOD 2. P RELIMINARIES

Original: In this section, we summarize the necessary preliminaries on GFlowNets.
Revised: In this section, we summarize the necessary preliminaries on GFlowNets.

Original: We follow the notation of
Revised: We follow the notation of

Original: Malkin et al.
Revised: Malkin et al.

Original: (2022), to which the reader is directed for a more thorough exposition written with a view towards motivating the trajectory and subtrajectory balance objectives.
Revised: (2022), to which the reader is directed for a more thorough exposition written with a view towards motivating the trajectory and subtrajectory balance objectives.

Original: A deeper introduction is given in Bengio et al.
Revised: A deeper introduction is given in Bengio et al.

Original: (2021b).
Revised: (2021b).

Original: Let 𝐺 =
Revised: Let 𝐺 =

Original: (S , A) be a directed acyclic graph.
Revised: (S , A) be a directed acyclic graph.

Original: The vertices 𝑠 ∈ S are called states and the directed edges (
Revised: The vertices 𝑠 ∈ S are called states and the directed edges (

Original: 𝑢 → 𝑣 ) ∈
Revised: 𝑢 → 𝑣 ) ∈

Original: A are actions .
Revised: A are actions .

Original: If ( 𝑢 → 𝑣 ) is an edge, we say 𝑣 is a child of 𝑢 and 𝑢 is a parent of 𝑣 .
Revised: If ( 𝑢 → 𝑣 ) is an edge, we say 𝑣 is a child of 𝑢 and 𝑢 is a parent of 𝑣 .

Original: There is a unique initial state 𝑠 0 ∈ S with no parents.
Revised: There is a unique initial state 𝑠 0 ∈ S with no parents.

Original: States with no children are called terminal , and the set of terminal states is denoted by X .
Revised: States with no children are called terminal , and the set of terminal states is denoted by X .

Original: A trajectory or an action sequence is a sequence of states 𝜏 = ( 𝑠 𝑚
Revised: A trajectory or an action sequence is a sequence of states 𝜏 = ( 𝑠 𝑚

Original: →
Revised: →

Original: 𝑠 𝑚 + 1 → . . .
Revised: 𝑠 𝑚 + 1 → . . .

Original: → 𝑠 𝑛 ) , where each ( 𝑠 𝑖 → 𝑠 𝑖 + 1 ) is an action.
Revised: → 𝑠 𝑛 ) , where each ( 𝑠 𝑖 → 𝑠 𝑖 + 1 ) is an action.

Original: The trajectory is complete if 𝑠 𝑚 =
Revised: The trajectory is complete if 𝑠 𝑚 =

Original: 𝑠 0 and 𝑠 𝑛 is terminal.
Revised: 𝑠 0 and 𝑠 𝑛 is terminal.

Original: The set of complete trajectories is denoted by T .
Revised: The set of complete trajectories is denoted by T .

Original: A (forward) policy is a collection of distributions 𝑃 𝐹 (−| 𝑠 ) over the children of every nonterminal state 𝑠 ∈ S .
Revised: A (forward) policy is a collection of distributions 𝑃 𝐹 (−| 𝑠 ) over the children of every nonterminal state 𝑠 ∈ S .

Original: A forward policy determines a distribution over T by
Revised: A forward policy determines a distribution over T by

Original: [Equation]
Revised: [Equation]

Original: Any distribution over complete trajectories that arises from a forward policy satisfies a Markov property: the marginal choice of action out of a state 𝑠 is independent of how 𝑠 was reached.
Revised: Any distribution over complete trajectories that arises from a forward policy satisfies a Markov property: the marginal choice of action out of a state 𝑠 is independent of how 𝑠 was reached.

Original: Conversely, any Markovian distribution over T arises from a forward policy (Bengio et al., 2021b).
Revised: Conversely, any Markovian distribution over T arises from a forward policy (Bengio et al., 2021b).

Original: A forward policy can thus be used to sample terminal states 𝑥 ∈ X by starting at 𝑠 0 and iteratively sampling actions from 𝑃 𝐹 , or, equivalently, taking the terminating state of a complete trajectory 𝜏 ∼ 𝑃 𝐹 ( 𝜏 ) .
Revised: A forward policy can thus be used to sample terminal states 𝑥 ∈ X by starting at 𝑠 0 and iteratively sampling actions from 𝑃 𝐹 , or, equivalently, taking the terminating state of a complete trajectory 𝜏 ∼ 𝑃 𝐹 ( 𝜏 ) .

Original: The marginal likelihood of sampling 𝑥 ∈ X is the sum of likelihoods of all complete trajectories that terminate at 𝑥 .
Revised: The marginal likelihood of sampling 𝑥 ∈ X is the sum of likelihoods of all complete trajectories that terminate at 𝑥 .

Original: Suppose that a nontrivial (not identically 0) nonnegative reward function 𝑅 : X → R ≥ 0 is given.
Revised: Suppose that a nontrivial (not identically 0) nonnegative reward function 𝑅 : X → R ≥ 0 is given.

Original: The learning problem solved by GFlowNets is to estimate a policy 𝑃 𝐹 such that the likelihood of sampling 𝑥 ∈ X is proportional to 𝑅 ( 𝑥 ) .
Revised: The learning problem solved by GFlowNets is to estimate a policy 𝑃 𝐹 such that the likelihood of sampling 𝑥 ∈ X is proportional to 𝑅 ( 𝑥 ) .

Original: That is, there should exist a constant 𝑍 such that
Revised: That is, there should exist a constant 𝑍 such that

Original: [Equation]
Revised: [Equation]

Original: If (2) is satisfied, then 𝑍 = (cid:205)
Revised: If (2) is satisfied, then 𝑍 = (cid:205)

Original: 𝑥 ∈X 𝑅 ( 𝑥 ) .
Revised: 𝑥 ∈X 𝑅 ( 𝑥 ) .

Original: 2. GF LOW N ET TRAINING OBJECTIVES
Revised: 2. GF LOW N ET TRAINING OBJECTIVES

Original: Because the sum in (2) may be intractable to compute, it is in general not possible to directly convert this constraint into a training objective.
Revised: Because the sum in (2) may be intractable to compute, it is in general not possible to directly convert this constraint into a training objective.

Original: To solve this problem, GFlowNet training objectives introduce auxiliary variables in the parametrization in various ways, but all have the property that (2) is satisfied at the global optimum.
Revised: To solve this problem, GFlowNet training objectives introduce auxiliary variables in the parametrization in various ways, but all have the property that (2) is satisfied at the global optimum.

Original: The key properties of these objectives are summarized in Table 1.
Revised: The key properties of these objectives are summarized in Table 1.

Original: Flow matching (FM; Bengio et al., 2021a).
Revised: Flow matching (FM; Bengio et al., 2021a).

Original: Motivating the ‘flow network’ terminology, Bengio et al.
Revised: Motivating the ‘flow network’ terminology, Bengio et al.

Original: (2021a) proved that (2) is satisfied if 𝑃 𝐹 arises from an edge flow function satisfying certain constraints.
Revised: (2021a) proved that (2) is satisfied if 𝑃 𝐹 arises from an edge flow function satisfying certain constraints.

Original: Namely, an assignment 𝐹 :
Revised: Namely, an assignment 𝐹 :

Original: A → R ≥ 0 of a nonnegative number (flow) to each action defines a policy via
Revised: A → R ≥ 0 of a nonnegative number (flow) to each action defines a policy via

Original: [Equation]
Revised: [Equation]

Original: A sufficient condition for the terminating distribution of 𝑃 𝐹 to be proportional to the reward 𝑅 ( 𝑥 ) is that a family of flow-matching (flow in = flow out) conditions is satisfied at all interior states and a
Revised: A sufficient condition for the terminating distribution of 𝑃 𝐹 to be proportional to the reward 𝑅 ( 𝑥 ) is that a family of flow-matching (flow in = flow out) conditions is satisfied at all interior states and a

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: family of reward-matching conditions is satisfied at terminal states:
Revised: family of reward-matching conditions is satisfied at terminal states:

Original: [Equation]
Revised: [Equation]

Original: The flow 𝐹 ( 𝑠 → 𝑡 ) is then proportional to the marginal likelihood that a complete trajectory sampled from 𝑃 𝐹 includes the action 𝑠 → 𝑡 .
Revised: The flow 𝐹 ( 𝑠 → 𝑡 ) is then proportional to the marginal likelihood that a complete trajectory sampled from 𝑃 𝐹 includes the action 𝑠 → 𝑡 .

Original: In Bengio et al.
Revised: In Bengio et al.

Original: (2021a), a GFlowNet is described by a parametric estimate of the edge flow function, 𝐹 ( 𝑢 → 𝑣 ; 𝜃 ) (a neural net with parameters 𝜃 ).
Revised: (2021a), a GFlowNet is described by a parametric estimate of the edge flow function, 𝐹 ( 𝑢 → 𝑣 ; 𝜃 ) (a neural net with parameters 𝜃 ).

Original: These conditions can be converted into objectives that are minimized when (4) is satisfied.
Revised: These conditions can be converted into objectives that are minimized when (4) is satisfied.

Original: For example, the flow-matching objective at a nonterminal state 𝑠 is defined by
Revised: For example, the flow-matching objective at a nonterminal state 𝑠 is defined by

Original: [Equation]
Revised: [Equation]

Original: where 𝜖 is a smoothing constant that can safely be set to 0 if the flows are constrained to be strictly positive, and a similar objective (or a constraint by construction) is defined to force the flow 𝐹 ( 𝑠 → 𝑥 ) into terminal states 𝑥 to match 𝑅 ( 𝑥 ) .
Revised: where 𝜖 is a smoothing constant that can safely be set to 0 if the flows are constrained to be strictly positive, and a similar objective (or a constraint by construction) is defined to force the flow 𝐹 ( 𝑠 → 𝑥 ) into terminal states 𝑥 to match 𝑅 ( 𝑥 ) .

Original: If these objectives are globally minimized for all states 𝑠 , then the policy 𝑃 𝐹 (−|− ; 𝜃 ) defined by 𝐹 (− ; 𝜃 ) via (3) satisfies (2), with 𝑍
Revised: If these objectives are globally minimized for all states 𝑠 , then the policy 𝑃 𝐹 (−|− ; 𝜃 ) defined by 𝐹 (− ; 𝜃 ) via (3) satisfies (2), with 𝑍

Original: = (cid:205) 𝑡 : ( 𝑠 0 → 𝑡 )∈A 𝐹 ( 𝑠 → 𝑡 ; 𝜃 ) = (cid:205)
Revised: = (cid:205) 𝑡 : ( 𝑠 0 → 𝑡 )∈A 𝐹 ( 𝑠 → 𝑡 ; 𝜃 ) = (cid:205)

Original: 𝑥 ∈X 𝑅 ( 𝑥 ) .
Revised: 𝑥 ∈X 𝑅 ( 𝑥 ) .

Original: The question of how to sample states 𝑠 for training is discussed below.
Revised: The question of how to sample states 𝑠 for training is discussed below.

Original: Detailed balance (DB; Bengio et al., 2021b; Malkin et al., 2022).
Revised: Detailed balance (DB; Bengio et al., 2021b; Malkin et al., 2022).

Original: In the DB parametrization, a forward policy model 𝑃 𝐹 (−|− ; 𝜃 ) is learned directly, jointly with two additional objects: a backward policy model 𝑃 𝐵 (−|− ; 𝜃 ) , which can predict a distribution over the parents of any noninitial state, and a state flow function 𝐹 ( 𝑠 ; 𝜃 ) (typically parametrized in the log domain).
Revised: In the DB parametrization, a forward policy model 𝑃 𝐹 (−|− ; 𝜃 ) is learned directly, jointly with two additional objects: a backward policy model 𝑃 𝐵 (−|− ; 𝜃 ) , which can predict a distribution over the parents of any noninitial state, and a state flow function 𝐹 ( 𝑠 ; 𝜃 ) (typically parametrized in the log domain).

Original: The detailed balance conditions state that
Revised: The detailed balance conditions state that

Original: [Equation]
Revised: [Equation]

Original: for all actions ( 𝑠 → 𝑡 ) and 𝐹 ( 𝑥 ; 𝜃 ) =
Revised: for all actions ( 𝑠 → 𝑡 ) and 𝐹 ( 𝑥 ; 𝜃 ) =

Original: 𝑅 ( 𝑥 ) for 𝑥 terminal.
Revised: 𝑅 ( 𝑥 ) for 𝑥 terminal.

Original: Satisfaction of these conditions for all actions ( 𝑠 → 𝑡 ) and 𝑥 ∈ X implies that 𝑃 𝐹 samples proportionally to the reward (i.e., satisfies (2), with 𝑍 = 𝐹 ( 𝑠 0 ) ).
Revised: Satisfaction of these conditions for all actions ( 𝑠 → 𝑡 ) and 𝑥 ∈ X implies that 𝑃 𝐹 samples proportionally to the reward (i.e., satisfies (2), with 𝑍 = 𝐹 ( 𝑠 0 ) ).

Original: The DB condition (6) can be converted into a squared log-ratio objective L DB ( 𝑠 → 𝑡 ) in the same way that (4) yields (5), and L DB ( 𝑠 → 𝑡 ) can be optimized over sampled actions ( 𝑠 → 𝑡 ) .
Revised: The DB condition (6) can be converted into a squared log-ratio objective L DB ( 𝑠 → 𝑡 ) in the same way that (4) yields (5), and L DB ( 𝑠 → 𝑡 ) can be optimized over sampled actions ( 𝑠 → 𝑡 ) .

Original: Trajectory balance (TB; Malkin et al., 2022).
Revised: Trajectory balance (TB; Malkin et al., 2022).

Original: The parametrization required for the TB objective includes forward and backward policy models 𝑃 𝐹 (−|− ; 𝜃 ) and 𝑃 𝐵 (−|− ; 𝜃 ) , as well as an estimate 𝑍 𝜃 of the constant of proportionality in (2).
Revised: The parametrization required for the TB objective includes forward and backward policy models 𝑃 𝐹 (−|− ; 𝜃 ) and 𝑃 𝐵 (−|− ; 𝜃 ) , as well as an estimate 𝑍 𝜃 of the constant of proportionality in (2).

Original: Satisfaction of the following condition for all complete trajectories 𝜏 =
Revised: Satisfaction of the following condition for all complete trajectories 𝜏 =

Original: ( 𝑠 0 → . . .
Revised: ( 𝑠 0 → . . .

Original: → 𝑠 𝑛 ) implies that (2) is satisfied:
Revised: → 𝑠 𝑛 ) implies that (2) is satisfied:

Original: [Equation]
Revised: [Equation]

Original: where we have used the conventions
Revised: where we have used the conventions

Original: [Equation]
Revised: [Equation]

Original: The condition (7) can again be made into a squared log-ratio objective L TB ( 𝜏 ) and optimized for complete trajectories 𝜏 taken from some training policy.
Revised: The condition (7) can again be made into a squared log-ratio objective L TB ( 𝜏 ) and optimized for complete trajectories 𝜏 taken from some training policy.

Original: In Malkin et al.
Revised: In Malkin et al.

Original: (2022), the TB objective was empirically demonstrated to have better convergence properties than FM and DB on various problem domains.
Revised: (2022), the TB objective was empirically demonstrated to have better convergence properties than FM and DB on various problem domains.

Original: Training policy and exploration.
Revised: Training policy and exploration.

Original: Global minimization of the FM, DB, and TB objectives for all values of their respective arguments (states, actions, or complete trajectories) implies satisfaction of (2).
Revised: Global minimization of the FM, DB, and TB objectives for all values of their respective arguments (states, actions, or complete trajectories) implies satisfaction of (2).

Original: Therefore, given a sufficiently expressive model and convergence of the optimization procedure, a GFlowNet policy that samples 𝑥 with likelihood proportional to 𝑅 ( 𝑥 ) can be trained by minimizing any of these losses over a distribution with full support, enabling offline training of GFlowNets.
Revised: Therefore, given a sufficiently expressive model and convergence of the optimization procedure, a GFlowNet policy that samples 𝑥 with likelihood proportional to 𝑅 ( 𝑥 ) can be trained by minimizing any of these losses over a distribution with full support, enabling offline training of GFlowNets.

Original: As in other RL algorithms, the distribution over sampled states, actions, or episodes can be fixed and off-policy, or can vary over the course of training and use available information about terminal states in interesting ways (Zhang et al., 2022; Deleu et al., 2022).
Revised: As in other RL algorithms, the distribution over sampled states, actions, or episodes can be fixed and off-policy, or can vary over the course of training and use available information about terminal states in interesting ways (Zhang et al., 2022; Deleu et al., 2022).

Original: The simplest approach, which is also taken in this paper, is on-policy learning or a very similar off-policy variant that flattens the current policy to ensure exploration.
Revised: The simplest approach, which is also taken in this paper, is on-policy learning or a very similar off-policy variant that flattens the current policy to ensure exploration.

Original: Complete trajectories 𝜏
Revised: Complete trajectories 𝜏

Original: = ( 𝑠 0 → . . .
Revised: = ( 𝑠 0 → . . .

Original: → 𝑠 𝑛 ) are sampled from the forward policy 𝑃 𝐹 (−|− ; 𝜃 ) (tempered or mixed with a uniform policy with a small weight so as to ensure full support and exploration).
Revised: → 𝑠 𝑛 ) are sampled from the forward policy 𝑃 𝐹 (−|− ; 𝜃 ) (tempered or mixed with a uniform policy with a small weight so as to ensure full support and exploration).

Original: One then takes gradient descent steps on L TB ( 𝜏 ) , on L
Revised: One then takes gradient descent steps on L TB ( 𝜏 ) , on L

Original: DB ( 𝑠 𝑖 → 𝑠 𝑖 + 1 ) over all actions in 𝜏 , or on L FM ( 𝑠 𝑖 ) for all intermediate states in 𝜏 .
Revised: DB ( 𝑠 𝑖 → 𝑠 𝑖 + 1 ) over all actions in 𝜏 , or on L FM ( 𝑠 𝑖 ) for all intermediate states in 𝜏 .

Original: The GFlowNets in this paper are trained on-policy, or off-policy with a training policy that is a mixture of 𝑃 𝐹 with a uniform policy: 𝜏 = ( 𝑠 0 → 𝑠 1 → . . .
Revised: The GFlowNets in this paper are trained on-policy, or off-policy with a training policy that is a mixture of 𝑃 𝐹 with a uniform policy: 𝜏 = ( 𝑠 0 → 𝑠 1 → . . .

Original: → 𝑠 𝑛 ) is sampled with 𝑠 𝑖 + 1 ∼ ( 1 − 𝜖 )
Revised: → 𝑠 𝑛 ) is sampled with 𝑠 𝑖 + 1 ∼ ( 1 − 𝜖 )

Original: 𝑃 𝐹 ( 𝑠 𝑖 + 1 | 𝑠 𝑖 ; 𝜃 ) + 𝜖 1 # { 𝑡 : ( 𝑠 → 𝑡 )∈A} .
Revised: 𝑃 𝐹 ( 𝑠 𝑖 + 1 | 𝑠 𝑖 ; 𝜃 ) + 𝜖 1 # { 𝑡 : ( 𝑠 → 𝑡 )∈A} .

Original: Here 𝜖 is the random exploration weight.
Revised: Here 𝜖 is the random exploration weight.

Original: 2. S UBTRAJECTORY BALANCE : L EARNING FROM PARTIAL EPISODES
Revised: 2. S UBTRAJECTORY BALANCE : L EARNING FROM PARTIAL EPISODES

Original: Recall the GFlowNet parametrization used in the DB objective above, with a state flow estimator 𝐹 (−|− ; 𝜃 ) and a pair of policies 𝑃 𝐹 (−|− ; 𝜃 ) , 𝑃 𝐵 (−|− ; 𝜃 ) .
Revised: Recall the GFlowNet parametrization used in the DB objective above, with a state flow estimator 𝐹 (−|− ; 𝜃 ) and a pair of policies 𝑃 𝐹 (−|− ; 𝜃 ) , 𝑃 𝐵 (−|− ; 𝜃 ) .

Original: It is shown in §A.2 of Malkin et al.
Revised: It is shown in §A.2 of Malkin et al.

Original: (2022) that the detailed balance conditions (6) are satisfied for all actions
Revised: (2022) that the detailed balance conditions (6) are satisfied for all actions

Original: if and only if the following subtrajectory balance constraint holds for all (not necessarily complete) trajectories 𝜏 = ( 𝑠 𝑚 → . . .
Revised: if and only if the following subtrajectory balance constraint holds for all (not necessarily complete) trajectories 𝜏 = ( 𝑠 𝑚 → . . .

Original: → 𝑠 𝑛 ) :
Revised: → 𝑠 𝑛 ) :

Original: [Equation]
Revised: [Equation]

Original: where we again enforce that 𝐹 ( 𝑥 ; 𝜃 )
Revised: where we again enforce that 𝐹 ( 𝑥 ; 𝜃 )

Original: =
Revised: =

Original: 𝑅 ( 𝑥 ) if 𝑥 is terminal.
Revised: 𝑅 ( 𝑥 ) if 𝑥 is terminal.

Original: Observe that the DB condition (6) is a special case of (8) when the trajectory consists of one action, and the TB condition (7) is precisely the case when 𝜏 is complete, with the identification 𝑍 𝜃 =
Revised: Observe that the DB condition (6) is a special case of (8) when the trajectory consists of one action, and the TB condition (7) is precisely the case when 𝜏 is complete, with the identification 𝑍 𝜃 =

Original: 𝐹 ( 𝑠 0 ; 𝜃 ) .
Revised: 𝐹 ( 𝑠 0 ; 𝜃 ) .

Original: The above constraint yields the subtrajectory balance objective
Revised: The above constraint yields the subtrajectory balance objective

Original: [Equation]
Revised: [Equation]

Original: If this objective is made equal to 0 for all partial trajectories 𝜏 , where 𝑅 ( 𝑠 𝑛 ) is substituted for 𝐹 ( 𝑠 𝑛 ; 𝜃 ) if 𝑠 𝑛 is terminal, then the policy 𝑃 𝐹 satisfies the desired condition (2).
Revised: If this objective is made equal to 0 for all partial trajectories 𝜏 , where 𝑅 ( 𝑠 𝑛 ) is substituted for 𝐹 ( 𝑠 𝑛 ; 𝜃 ) if 𝑠 𝑛 is terminal, then the policy 𝑃 𝐹 satisfies the desired condition (2).

Original: (Proof: When
Revised: (Proof: When

Original: L SubTB ( 𝜏 ) = 0 , (8) is satisfied, implying satisfaction of both (7) and (6).
Revised: L SubTB ( 𝜏 ) = 0 , (8) is satisfied, implying satisfaction of both (7) and (6).

Original: Either of these conditions is a sufficient condition for (2), as shown by Bengio et al.
Revised: Either of these conditions is a sufficient condition for (2), as shown by Bengio et al.

Original: (2021b)
Revised: (2021b)

Original: ; Malkin et al.
Revised: ; Malkin et al.

Original: Extracting subtrajectories for training.
Revised: Extracting subtrajectories for training.

Original: Suppose that an episode (complete trajectory)
Revised: Suppose that an episode (complete trajectory)

Original: 𝜏
Revised: 𝜏

Original: = ( 𝑠 0 → 𝑠 1 → . . .
Revised: = ( 𝑠 0 → 𝑠 1 → . . .

Original: → 𝑠 𝑛 ) is sampled for training.
Revised: → 𝑠 𝑛 ) is sampled for training.

Original: There are (cid:0) 𝑛 + 12 (cid:1) =
Revised: There are (cid:0) 𝑛 + 12 (cid:1) =

Original: 𝑂 ( 𝑛 2 ) nontrivial subtrajectories:
Revised: 𝑂 ( 𝑛 2 ) nontrivial subtrajectories:

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Having sampled a complete trajectory 𝜏 for training, we make gradient steps on a convex combination of the subtrajectory balance losses L SubTB ( 𝜏 𝑖 : 𝑗 ) : 𝜃 ← 𝜃 − ∇ 𝜃 L , where
Revised: Having sampled a complete trajectory 𝜏 for training, we make gradient steps on a convex combination of the subtrajectory balance losses L SubTB ( 𝜏 𝑖 : 𝑗 ) : 𝜃 ← 𝜃 − ∇ 𝜃 L , where

Original: [Equation]
Revised: [Equation]

Original: Here 𝜆 > 0 is a hyperparameter controlling the weights assigned to subtrajectories of different lengths, and when 𝜆 is set to 1, it leads to a uniform weighting scheme.
Revised: Here 𝜆 > 0 is a hyperparameter controlling the weights assigned to subtrajectories of different lengths, and when 𝜆 is set to 1, it leads to a uniform weighting scheme.

Original: Notice that the 𝜆 → 0 + limit leads precisely to the average detailed balance loss L DB ( 𝑠 𝑖 → 𝑠 𝑖 + 1 ) over all transitions in 𝜏 , while the 𝜆 → +∞ limit gives the trajectory balance objective L TB ( 𝜏 ) .
Revised: Notice that the 𝜆 → 0 + limit leads precisely to the average detailed balance loss L DB ( 𝑠 𝑖 → 𝑠 𝑖 + 1 ) over all transitions in 𝜏 , while the 𝜆 → +∞ limit gives the trajectory balance objective L TB ( 𝜏 ) .

Original: Other schemes for weighting subtrajectories are possible and should be explored in future work.
Revised: Other schemes for weighting subtrajectories are possible and should be explored in future work.

Original: Computational considerations.
Revised: Computational considerations.

Original: It may appear that the optimization of (11) induces a computation cost that is quadratic in the trajectory length.
Revised: It may appear that the optimization of (11) induces a computation cost that is quadratic in the trajectory length.

Original: However, a closer inspection of the gradient of (11) with respect to the state flows log 𝐹 ( 𝑠 𝑖 ; 𝜃 ) and the forward and backward policy logits shows that gradient computation requires only one forward and one backward pass through the neural networks giving log 𝐹 ( 𝑠 ; 𝜃 ) , log 𝑃 𝐹 (−| 𝑠 𝑖 ; 𝜃 ) , and log 𝑃 𝐵 (−| 𝑠 𝑖 ; 𝜃 ) .
Revised: However, a closer inspection of the gradient of (11) with respect to the state flows log 𝐹 ( 𝑠 𝑖 ; 𝜃 ) and the forward and backward policy logits shows that gradient computation requires only one forward and one backward pass through the neural networks giving log 𝐹 ( 𝑠 ; 𝜃 ) , log 𝑃 𝐹 (−| 𝑠 𝑖 ; 𝜃 ) , and log 𝑃 𝐵 (−| 𝑠 𝑖 ; 𝜃 ) .

Original: The quadratic computation cost is incurred only in performing linear operations on these log-flows and policy logits, not in the evaluation of the deep networks.
Revised: The quadratic computation cost is incurred only in performing linear operations on these log-flows and policy logits, not in the evaluation of the deep networks.

Original: Thus the SubTB loss has little computation overhead over DB or TB.
Revised: Thus the SubTB loss has little computation overhead over DB or TB.

Original: Hypothesized benefits.
Revised: Hypothesized benefits.

Original: We hypothesize that SubTB( 𝜆 ) brings two benefits to GFlowNet training:
Revised: We hypothesize that SubTB( 𝜆 ) brings two benefits to GFlowNet training:

Original: V ARIANCE REDUCTION .
Revised: V ARIANCE REDUCTION .

Original: The TB loss terms L TB ( 𝜏 ) for trajectories 𝜏 that take a given sequence of actions until a state 𝑠 , then diverge, share the terms log 𝑍 and the policy logits for all transitions preceding 𝑠 inside the square.
Revised: The TB loss terms L TB ( 𝜏 ) for trajectories 𝜏 that take a given sequence of actions until a state 𝑠 , then diverge, share the terms log 𝑍 and the policy logits for all transitions preceding 𝑠 inside the square.

Original: However, the ‘tail’ of the TB loss, involving the forward and backward policy logits for transitions that appear after 𝑠 in 𝜏 , can be seen as a stochastic least-squares regression target.
Revised: However, the ‘tail’ of the TB loss, involving the forward and backward policy logits for transitions that appear after 𝑠 in 𝜏 , can be seen as a stochastic least-squares regression target.

Original: That is, if 𝑠 = 𝑠 𝑚 in a trajectory 𝜏
Revised: That is, if 𝑠 = 𝑠 𝑚 in a trajectory 𝜏

Original: =
Revised: =

Original: ( 𝑠 0 → 𝑠 1 → . . .
Revised: ( 𝑠 0 → 𝑠 1 → . . .

Original: → 𝑠 𝑛 ) , then
Revised: → 𝑠 𝑛 ) , then

Original: [Equation]
Revised: [Equation]

Original: is regressed to
Revised: is regressed to

Original: [Equation]
Revised: [Equation]

Original: Similarly, for trajectories that share the transitions following 𝑠
Revised: Similarly, for trajectories that share the transitions following 𝑠

Original: but may differ in their initial actions, (12) is a stochastic regression target for (13).
Revised: but may differ in their initial actions, (12) is a stochastic regression target for (13).

Original: The subtrajectory balance loss terms L SubTB ( 𝜏 𝑚 : 𝑗 ) for partial trajectories beginning at 𝑠 regress the log-state flow log 𝐹 ( 𝑠 ) to (parts of) expressions like (13), while loss terms L SubTB ( 𝜏 𝑖 : 𝑚 ) regress (parts of) expressions like (12) to the log-state flow log 𝐹 ( 𝑠 ) .
Revised: The subtrajectory balance loss terms L SubTB ( 𝜏 𝑚 : 𝑗 ) for partial trajectories beginning at 𝑠 regress the log-state flow log 𝐹 ( 𝑠 ) to (parts of) expressions like (13), while loss terms L SubTB ( 𝜏 𝑖 : 𝑚 ) regress (parts of) expressions like (12) to the log-state flow log 𝐹 ( 𝑠 ) .

Original: The learned log 𝐹 ( 𝑠 ) is thus a learned estimate of a stochastic piece of the TB loss for trajectories that contain 𝑠 .
Revised: The learned log 𝐹 ( 𝑠 ) is thus a learned estimate of a stochastic piece of the TB loss for trajectories that contain 𝑠 .

Original: Replacing a stochastic term in the TB loss by a learned estimate of its expectation is guaranteed to introduce bias into the gradient (with respect to the gradient of the TB loss), but is expected to reduce variance.
Revised: Replacing a stochastic term in the TB loss by a learned estimate of its expectation is guaranteed to introduce bias into the gradient (with respect to the gradient of the TB loss), but is expected to reduce variance.

Original: This is akin to the variance-reducing effect of actor-critic methods in RL.
Revised: This is akin to the variance-reducing effect of actor-critic methods in RL.

Original: This hypothesis is studied empirically in our experiments and in particular §4.1.1, where we provide evidence that SubTB( 𝜆 ) is a practically useful interpolation between TB (high variance) and DB (low variance, high bias relative to the true TB gradient) losses.
Revised: This hypothesis is studied empirically in our experiments and in particular §4.1.1, where we provide evidence that SubTB( 𝜆 ) is a practically useful interpolation between TB (high variance) and DB (low variance, high bias relative to the true TB gradient) losses.

Original: F ASTER LEARNING DUE TO GENERALIZATION OF STATE FLOWS .
Revised: F ASTER LEARNING DUE TO GENERALIZATION OF STATE FLOWS .

Original: Another benefit of subtrajectory balance for convergence speed may come from the ability of estimated state flow functions log 𝐹 ( 𝑠 ; 𝜃 ) to be modeled with high precision and generalize between states 𝑠 faster than the often high-dimensional policy logits log 𝑃 𝐹 (−| 𝑠 ; 𝜃 ) , log 𝑃 𝐵 (−| 𝑠 ; 𝜃 ) .
Revised: Another benefit of subtrajectory balance for convergence speed may come from the ability of estimated state flow functions log 𝐹 ( 𝑠 ; 𝜃 ) to be modeled with high precision and generalize between states 𝑠 faster than the often high-dimensional policy logits log 𝑃 𝐹 (−| 𝑠 ; 𝜃 ) , log 𝑃 𝐵 (−| 𝑠 ; 𝜃 ) .

Original: Such generalization is important in problems where the state graph becomes ‘wide’ far from the initial state, making the learning signal sparse at states that are near termination.
Revised: Such generalization is important in problems where the state graph becomes ‘wide’ far from the initial state, making the learning signal sparse at states that are near termination.

Original: Indeed, in all of our experiment domains except the hypergrids in §4.1 – and for the largest hypergrids – the number of terminal states is many orders of magnitude larger than the total number of states seen in training.
Revised: Indeed, in all of our experiment domains except the hypergrids in §4.1 – and for the largest hypergrids – the number of terminal states is many orders of magnitude larger than the total number of states seen in training.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Eligibility traces.
Revised: Eligibility traces.

Original: SubTB( 𝜆 ) draws inspiration from the TD( 𝜆 ) algorithm in RL (Sutton, 1988;
Revised: SubTB( 𝜆 ) draws inspiration from the TD( 𝜆 ) algorithm in RL (Sutton, 1988;

Original: Sutton & Barto, 2018), which forms an estimate of the expected return via a convex combination of 𝑛 -step returns, each weighed by ( 1 − 𝜆 )
Revised: Sutton & Barto, 2018), which forms an estimate of the expected return via a convex combination of 𝑛 -step returns, each weighed by ( 1 − 𝜆 )

Original: 𝜆 𝑛 − 1 .
Revised: 𝜆 𝑛 − 1 .

Original: The parameter 𝜆 ∈ [ 0 , 1 ] enables a bias-variance tradeoff (Kearns & Singh, 2000).
Revised: The parameter 𝜆 ∈ [ 0 , 1 ] enables a bias-variance tradeoff (Kearns & Singh, 2000).

Original: Intuitively, larger 𝜆 leads to lower bias and higher variance, since the estimate of the expected return approaches the single-point Monte Carlo estimate as 𝜆 → 1 .
Revised: Intuitively, larger 𝜆 leads to lower bias and higher variance, since the estimate of the expected return approaches the single-point Monte Carlo estimate as 𝜆 → 1 .

Original: We take inspiration from this idea to mix together different (possibly all) subtrajectories, akin to how 𝑛 -step returns are mixed together.
Revised: We take inspiration from this idea to mix together different (possibly all) subtrajectories, akin to how 𝑛 -step returns are mixed together.

Original: We hypothesize that the right mixing may reduce variance, compared to TB, with the additional benefits of inducing consistency between the flows of intermediate states, and thus of helping propagate credit faster and enable faster convergence.
Revised: We hypothesize that the right mixing may reduce variance, compared to TB, with the additional benefits of inducing consistency between the flows of intermediate states, and thus of helping propagate credit faster and enable faster convergence.

Original: In addition,
Revised: In addition,

Original: GFlowNet training objectives are reminiscent of residual gradient RL methods (Baird, 1995; Zhang et al., 2020) since the “endpoint” (e.g. 𝐹 ( 𝑠 𝑛 ) in (9)) is also considered in the gradient.
Revised: GFlowNet training objectives are reminiscent of residual gradient RL methods (Baird, 1995; Zhang et al., 2020) since the “endpoint” (e.g. 𝐹 ( 𝑠 𝑛 ) in (9)) is also considered in the gradient.

Original: MaxEnt RL.
Revised: MaxEnt RL.

Original: RL has a rich literature on energy-based, or maximum entropy, methods (Ziebart, 2010; Mnih et al., 2016;
Revised: RL has a rich literature on energy-based, or maximum entropy, methods (Ziebart, 2010; Mnih et al., 2016;

Original: Haarnoja et al., 2017; Nachum et al., 2017; Schulman et al., 2017; Haarnoja et al., 2018), which are close or equivalent to the GFlowNet framework in certain settings (in particular when the MDP has a tree structure (Bengio et al., 2021a)).
Revised: Haarnoja et al., 2017; Nachum et al., 2017; Schulman et al., 2017; Haarnoja et al., 2018), which are close or equivalent to the GFlowNet framework in certain settings (in particular when the MDP has a tree structure (Bengio et al., 2021a)).

Original: Also related are methods that maximize entropy not on the policy, but rather on the state visitation distribution (Hazan et al., 2019; Islam et al., 2019; Zhang et al., 2021) or some proxy of it (Eysenbach et al., 2018), which achieve a similar objective to GFlowNet models by flattening the state visitation distribution.
Revised: Also related are methods that maximize entropy not on the policy, but rather on the state visitation distribution (Hazan et al., 2019; Islam et al., 2019; Zhang et al., 2021) or some proxy of it (Eysenbach et al., 2018), which achieve a similar objective to GFlowNet models by flattening the state visitation distribution.

Original: If the state graph of the environment is a directed tree, the loss L SubTB on individual subtrajectories is equivalent to that of path consistency learning (Nachum et al., 2017).
Revised: If the state graph of the environment is a directed tree, the loss L SubTB on individual subtrajectories is equivalent to that of path consistency learning (Nachum et al., 2017).

Original: However, attempts to use path consistency learning in settings without intermediate rewards have only computed the loss on subtrajectories that have length 1 or include a terminal state (Guo et al., 2021).
Revised: However, attempts to use path consistency learning in settings without intermediate rewards have only computed the loss on subtrajectories that have length 1 or include a terminal state (Guo et al., 2021).

Original: [Figure]
Revised: [Figure]

Original:  past work, on all grid sizes.
Revised: = 0 . 9 ) consistently gives faster convergence than TB, the strongest objective from past work, on all grid sizes.

Original: The difference is especially visible for the harder variant of the reward function (last row).
Revised: The difference is especially visible for the harder variant of the reward function (last row).

Original: The 𝑥 -axis is the cumulative number of training trajectories (episodes).
Revised: The 𝑥 -axis is the cumulative number of training trajectories (episodes).

Original:  E XPERIMENTS 4. H YPERGRID : R OBUSTNESS TO SPARSE REWARDS
Revised:  E XPERIMENTS 4. H YPERGRID : R OBUSTNESS TO SPARSE REWARDS

Original: We study the synthetic hypergrid environment introduced in Bengio et al.
Revised: We study the synthetic hypergrid environment introduced in Bengio et al.

Original: (2021a).
Revised: (2021a).

Original: The set of interior states is a 𝑑 -dimensional hypergrid of size
Revised: The set of interior states is a 𝑑 -dimensional hypergrid of size

Original: 𝐻 × 𝐻 × · · · × 𝐻 with a multimodal reward function concentrated near each of the 2 𝑑 corners of the hypergrid (see Bengio et al.
Revised: 𝐻 × 𝐻 × · · · × 𝐻 with a multimodal reward function concentrated near each of the 2 𝑑 corners of the hypergrid (see Bengio et al.

Original: (2021a)
Revised: (2021a)

Original: ; Malkin et al.
Revised: ; Malkin et al.

Original: (2022) and Fig. 1).
Revised: (2022) and Fig. 1).

Original: The initial state is ( 0 , 0 , . . . , 0 ) , and each action is a step that increments one of the 𝑑 coordinates by 1 without leaving the grid.
Revised: The initial state is ( 0 , 0 , . . . , 0 ) , and each action is a step that increments one of the 𝑑 coordinates by 1 without leaving the grid.

Original: A special termination action is also allowed from each state.
Revised: A special termination action is also allowed from each state.

Original: This environment is designed to challenge a learning agent to infer and discover new modes from those that have been already been visited.
Revised: This environment is designed to challenge a learning agent to infer and discover new modes from those that have been already been visited.

Original: We study various sizes of 2-dimensional and 4-dimensional hypergrids, using the hardest variant of the reward function from past work (the minimal reward, away from the corners of the grid, is set to 10 − 3 ).
Revised: We study various sizes of 2-dimensional and 4-dimensional hypergrids, using the hardest variant of the reward function from past work (the minimal reward, away from the corners of the grid, is set to 10 − 3 ).

Original: We train GFlowNets to sample from the target reward functions and plot the evolution of the 𝐿 1 distance between the target distribution and the empirical distribution of the last 2 · 10 5 states seen in training.
Revised: We train GFlowNets to sample from the target reward functions and plot the evolution of the 𝐿 1 distance between the target distribution and the empirical distribution of the last 2 · 10 5 states seen in training.

Original: In all cases, we tune the learning rates for the
Revised: In all cases, we tune the learning rates for the

Original: TB and SubTB( 𝜆 = 0 . 9 ) objectives.
Revised: TB and SubTB( 𝜆 = 0 . 9 ) objectives.

Original: (See §A for details.)
Revised: (See §A for details.)

Original: [Figure]
Revised: [Figure]

Original: The results (mean and standard deviation over three random runs) are shown in the first two rows of Fig.
Revised: The results (mean and standard deviation over three random runs) are shown in the first two rows of Fig.

Original: Models trained with
Revised: Models trained with

Original: SubTB( 𝜆 ) converge faster, and with less variance between random seeds, to the true distribution than with TB for all hypergrid sizes.
Revised: SubTB( 𝜆 ) converge faster, and with less variance between random seeds, to the true distribution than with TB for all hypergrid sizes.

Original: We also study an even sparser variant of the environment, in which the background reward is set to 10 − 4 .
Revised: We also study an even sparser variant of the environment, in which the background reward is set to 10 − 4 .

Original: In this case, SubTB( 𝜆 ) continues to perform strongly (last row of Fig. 3), while models trained with TB do not even discover all modes of the target distribution for grids larger than 8 × 8 (Fig. 2).
Revised: In this case, SubTB( 𝜆 ) continues to perform strongly (last row of Fig. 3), while models trained with TB do not even discover all modes of the target distribution for grids larger than 8 × 8 (Fig. 2).

Original: Additional results are given in §A.1.
Revised: Additional results are given in §A.1.

Original: In particular, SubTB( 𝜆 ) continues to perform strongly when only subtrajectories of less than a certain length are used for training, which can be beneficial in realistic settings where only partial episodes are given.
Revised: In particular, SubTB( 𝜆 ) continues to perform strongly when only subtrajectories of less than a certain length are used for training, which can be beneficial in realistic settings where only partial episodes are given.

Original: We also show the effect of 𝜆 on the convergence rate (Fig.
Revised: We also show the effect of 𝜆 on the convergence rate (Fig.

Original: A.2) and of more exploratory training policies (Fig.
Revised: A.2) and of more exploratory training policies (Fig.

Original: A.3).
Revised: A.3).

Original: [Figure]
Revised: [Figure]

Original: A CLOSER LOOK AT GRADIENT VARIANCE
Revised: A CLOSER LOOK AT GRADIENT VARIANCE

Original: [Figure]
Revised: [Figure]

Original: - ents, showing DB < SubTB( 𝜆 ) <
Revised: - ents, showing DB < SubTB( 𝜆 ) <

Original: TB in gradient variance.
Revised: TB in gradient variance.

Original: Below: Similarity of small-batch
Revised: Below: Similarity of small-batch

Original: DB, SubTB( 𝜆 ), and TB gradients to the large-batch TB gradient, showing that the small-batch
Revised: DB, SubTB( 𝜆 ), and TB gradients to the large-batch TB gradient, showing that the small-batch

Original: SubTB( 𝜆 ) gradient is a good estimator of large-batch TB.
Revised: SubTB( 𝜆 ) gradient is a good estimator of large-batch TB.

Original: We take a closer look at gradient bias and variance to understand the benefits of training GFlowNets with SubTB( 𝜆 ).
Revised: We take a closer look at gradient bias and variance to understand the benefits of training GFlowNets with SubTB( 𝜆 ).

Original: The methodology of these experiments is inspired by Ilyas et al.
Revised: The methodology of these experiments is inspired by Ilyas et al.

Original: We train GFlowNets on the 8 × 8 grid environment using
Revised: We train GFlowNets on the 8 × 8 grid environment using

Original: SubTB( 𝜆 = 0 . 8 ) and monitor various gradient metrics during training.
Revised: SubTB( 𝜆 = 0 . 8 ) and monitor various gradient metrics during training.

Original: To remove the effect of parameter sharing between policies at different states and to isolate the effect of the objective, we use a tabular representation of the GFlowNet, i.e., all flows and policy logits are optimized as independent parameters.
Revised: To remove the effect of parameter sharing between policies at different states and to isolate the effect of the objective, we use a tabular representation of the GFlowNet, i.e., all flows and policy logits are optimized as independent parameters.

Original: Gradient variance.
Revised: Gradient variance.

Original: To measure gradient variance, we use the following procedure for each training objective (DB, TB, or
Revised: To measure gradient variance, we use the following procedure for each training objective (DB, TB, or

Original: SubTB( 𝜆 )).
Revised: SubTB( 𝜆 )).

Original: A large batch of 2 10 = 1024 trajectories is sampled, and the gradient 𝑔 ( 0 ) 𝑗 of the objective with respect to the policy logits at all states is computed for each trajectory 𝜏 𝑗 in the batch.
Revised: A large batch of 2 10 = 1024 trajectories is sampled, and the gradient 𝑔 ( 0 ) 𝑗 of the objective with respect to the policy logits at all states is computed for each trajectory 𝜏 𝑗 in the batch.

Original: Then, for each 𝑘 ∈ { 0 , 1 , . .
Revised: Then, for each 𝑘 ∈ { 0 , 1 , . .

Original: , 9 } , the gradients 𝑔 ( 0 ) 𝑖 are combined into 2 10 − 𝑘 sub-batches, each of size 2 𝑘 .
Revised: , 9 } , the gradients 𝑔 ( 0 ) 𝑖 are combined into 2 10 − 𝑘 sub-batches, each of size 2 𝑘 .

Original: The subbatch gradient 𝑔 ( 𝑘 ) 𝑖 for the 𝑖 -th sub-batch is set to the average of trajectory gradients 𝑔 ( 0 ) 𝑗 contained within the sub-batch and computed for 𝑖 ∈ { 1 , 2 , . . .
Revised: The subbatch gradient 𝑔 ( 𝑘 ) 𝑖 for the 𝑖 -th sub-batch is set to the average of trajectory gradients 𝑔 ( 0 ) 𝑗 contained within the sub-batch and computed for 𝑖 ∈ { 1 , 2 , . . .

Original: , 2 10 − 𝑘 } .
Revised: , 2 10 − 𝑘 } .

Original: We then report the average cosine similarity between the sub-batch and full-batch gradients:
Revised: We then report the average cosine similarity between the sub-batch and full-batch gradients:

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: If this quantity is positive, then gradient steps of infinitesimally small norm along the stochastic sub-batch gradient decrease the full-batch objective in expectation.
Revised: If this quantity is positive, then gradient steps of infinitesimally small norm along the stochastic sub-batch gradient decrease the full-batch objective in expectation.

Original: Fig.
Revised: Fig.

Original: 4 (left) shows the dependence of this metric on 𝑘 at various iterations.
Revised: 4 (left) shows the dependence of this metric on 𝑘 at various iterations.

Original: A steeper curve, such as those of DB and SubTB( 𝜆 ), indicates lower gradient variance.
Revised: A steeper curve, such as those of DB and SubTB( 𝜆 ), indicates lower gradient variance.

Original: Fig.
Revised: Fig.

Original: 5 (top) shows the metric at 𝑘 = 6 (corresponding to the batch size of 64 used for training) over the course of training.
Revised: 5 (top) shows the metric at 𝑘 = 6 (corresponding to the batch size of 64 used for training) over the course of training.

Original: We see that the DB gradient has the highest self-consistency at all iterations,
Revised: We see that the DB gradient has the highest self-consistency at all iterations,

Original: TB has the lowest, and SubTB( 𝜆 = 0 . 8 ) is in between.
Revised: TB has the lowest, and SubTB( 𝜆 = 0 . 8 ) is in between.

Original: Gradient bias.
Revised: Gradient bias.

Original: We next compare the small-batch stochastic gradients with large-batch stochastic gradients, using different objectives for the small and full batches.
Revised: We next compare the small-batch stochastic gradients with large-batch stochastic gradients, using different objectives for the small and full batches.

Original: Specifically, we compare the small-batch DB, SubTB( 𝜆 ), and TB gradients with the full-batch TB gradient.
Revised: Specifically, we compare the small-batch DB, SubTB( 𝜆 ), and TB gradients with the full-batch TB gradient.

Original: (The full-batch TB gradient can be seen as a ‘canonical’ gradient against which bias can be measured, as its expectation equals the gradient of the KL divergence between the distribution over trajectories defined by 𝑃 𝐹 and that defined by the reward 𝑅 and 𝑃 𝐵 ; see §A.3 of Malkin et al.
Revised: (The full-batch TB gradient can be seen as a ‘canonical’ gradient against which bias can be measured, as its expectation equals the gradient of the KL divergence between the distribution over trajectories defined by 𝑃 𝐹 and that defined by the reward 𝑅 and 𝑃 𝐵 ; see §A.3 of Malkin et al.

Original: Fig.
Revised: Fig.

Original: 5 (bottom) shows the cosine similarity at the batch size used for training.
Revised: 5 (bottom) shows the cosine similarity at the batch size used for training.

Original: Notably, at intermediate iterations, the similarity of SubTB( 𝜆 ) with TB is higher than that of TB with TB: despite its bias, the small-batch SubTB( 𝜆 ) gradient estimates the full-batch TB gradient better than the small-batch TB gradient does .
Revised: Notably, at intermediate iterations, the similarity of SubTB( 𝜆 ) with TB is higher than that of TB with TB: despite its bias, the small-batch SubTB( 𝜆 ) gradient estimates the full-batch TB gradient better than the small-batch TB gradient does .

Original: Fig.
Revised: Fig.

Original: 4 (right) shows the dependence of the similarity on 𝑘 at selected iterations and suggests that this effect may be even larger for smaller batch sizes.
Revised: 4 (right) shows the dependence of the similarity on 𝑘 at selected iterations and suggests that this effect may be even larger for smaller batch sizes.

Original: Moreover, at 𝑘 = 10 , the similarity of SubTB( 𝜆 ) vs. TB always lies between DB vs. TB and TB vs. TB, indicating that SubTB( 𝜆 ) interpolates between TB’s unbiased and DB’s biased estimates of the TB gradient.
Revised: Moreover, at 𝑘 = 10 , the similarity of SubTB( 𝜆 ) vs. TB always lies between DB vs. TB and TB vs. TB, indicating that SubTB( 𝜆 ) interpolates between TB’s unbiased and DB’s biased estimates of the TB gradient.

Original: The effect of learned state flows.
Revised: The effect of learned state flows.

Original: For additional experiments, see §A.2.
Revised: For additional experiments, see §A.2.

Original: 16 reward exponent 0.
Revised: 16 reward exponent 0.

Original: l o g p ( x ) t o l o g
Revised: l o g p ( x ) t o l o g

Original: R ( x ) c o
Revised: R ( x ) c o

Original: rr e l a t i o n
Revised: rr e l a t i o n

Original: FMDBTBSubTB(1) 1010 3 learning rate 0(DB) 0.5  (TB) SubTB parameter
Revised: FMDBTBSubTB(1) 1010 3 learning rate 0(DB) 0.5  (TB) SubTB parameter

Original: S MALL MOLECULE SYNTHESIS
Revised: S MALL MOLECULE SYNTHESIS

Original: We use SubTB( 𝜆 ) to train models on the molecule generation task of Bengio et al.
Revised: We use SubTB( 𝜆 ) to train models on the molecule generation task of Bengio et al.

Original: (2021a).
Revised: (2021a).

Original: The task is to generate binders of the sEH (soluble epoxide hydrolase) protein, based on a docking prediction (Trott & Olson, 2010).
Revised: The task is to generate binders of the sEH (soluble epoxide hydrolase) protein, based on a docking prediction (Trott & Olson, 2010).

Original: To be precise, molecules are generated by sequentially joining ‘blocks’ from a fixed library to the partial molecular graph (Jin et al., 2020; Kumar et al., 2012), resulting in a state space of estimated size 10 12 .
Revised: To be precise, molecules are generated by sequentially joining ‘blocks’ from a fixed library to the partial molecular graph (Jin et al., 2020; Kumar et al., 2012), resulting in a state space of estimated size 10 12 .

Original: The reward function 𝑅 is given by a pretrained proxy model made available by Bengio et al.
Revised: The reward function 𝑅 is given by a pretrained proxy model made available by Bengio et al.

Original: (2021a).
Revised: (2021a).

Original: To adjust the greediness of the agent, an inverse temperature hyperparameter 𝛽 is used, i.e., the reward used for training is 𝑅 ( 𝑥 ) =
Revised: To adjust the greediness of the agent, an inverse temperature hyperparameter 𝛽 is used, i.e., the reward used for training is 𝑅 ( 𝑥 ) =

Original: (cid:101) 𝑅 ( 𝑥 ) 𝛽 , where (cid:101) 𝑅 ( 𝑥 ) is the proxy’s prediction.
Revised: (cid:101) 𝑅 ( 𝑥 ) 𝛽 , where (cid:101) 𝑅 ( 𝑥 ) is the proxy’s prediction.

Original: We train models with the DB, TB, and SubTB( 𝜆 ) objectives, with four values each of 𝜆 , 𝛽 , and learning rate, averaging the results over 3 random runs for each setting.
Revised: We train models with the DB, TB, and SubTB( 𝜆 ) objectives, with four values each of 𝜆 , 𝛽 , and learning rate, averaging the results over 3 random runs for each setting.

Original: We measure how well the trained models match the target distribution by the correlation of log 𝑅 ( 𝑥 ) and log 𝑝 𝜃 ( 𝑥 ) , the log-probability assigned to 𝑥 by the GFlowNet, computed on a held-out set of terminal states 𝑥 .
Revised: We measure how well the trained models match the target distribution by the correlation of log 𝑅 ( 𝑥 ) and log 𝑝 𝜃 ( 𝑥 ) , the log-probability assigned to 𝑥 by the GFlowNet, computed on a held-out set of terminal states 𝑥 .

Original: The results are shown in Fig.
Revised: The results are shown in Fig.

Original: SubTB( 𝜆 ), in particular with 𝜆 = 1 , performs better than both DB and TB when the optimal hyperparameters 𝛼, 𝛽 are used (solid lines) and is far more robust to the choice of hyperparameters (dashed lines).
Revised: SubTB( 𝜆 ), in particular with 𝜆 = 1 , performs better than both DB and TB when the optimal hyperparameters 𝛼, 𝛽 are used (solid lines) and is far more robust to the choice of hyperparameters (dashed lines).

Original: Additional details can be found in §B.
Revised: Additional details can be found in §B.

Original: 4. S EQUENCE GENERATION
Revised: 4. S EQUENCE GENERATION

Original: We consider three sequence generation tasks in which sequences are generated left to right, with each action appending one symbol from a vocabulary to a partial sequence: a synthetic task with varying sequence lengths and vocabulary sizes (§4.3.1), a practical biological sequence design task (§4.3.2), and a new protein design task with longer sequences (4.3.3).
Revised: We consider three sequence generation tasks in which sequences are generated left to right, with each action appending one symbol from a vocabulary to a partial sequence: a synthetic task with varying sequence lengths and vocabulary sizes (§4.3.1), a practical biological sequence design task (§4.3.2), and a new protein design task with longer sequences (4.3.3).

Original: For all three tasks, we consider the baselines Soft Actor-Critic (Haarnoja et al., 2018; Christodoulou, 2019), A2C with Entropy regularization (Williams & Peng, 1991;
Revised: For all three tasks, we consider the baselines Soft Actor-Critic (Haarnoja et al., 2018; Christodoulou, 2019), A2C with Entropy regularization (Williams & Peng, 1991;

Original: Mnih et al., 2016) and MARS-like MCMC (Xie et al., 2021) and compare them with three GFlowNet training objectives: TB, FM, and SubTB( 𝜆 ).
Revised: Mnih et al., 2016) and MARS-like MCMC (Xie et al., 2021) and compare them with three GFlowNet training objectives: TB, FM, and SubTB( 𝜆 ).

Original: In §F, we also study a non-autoregressive sequence generation problem (inverse protein folding).
Revised: In §F, we also study a non-autoregressive sequence generation problem (inverse protein folding).

Original: 4.3. B IT SEQUENCES
Revised: 4.3. B IT SEQUENCES

Original: We consider the synthetic sequence generation setting from Malkin et al.
Revised: We consider the synthetic sequence generation setting from Malkin et al.

Original: (2022), where the goal is to generate sequences of bits of fixed length
Revised: (2022), where the goal is to generate sequences of bits of fixed length

Original: 𝑛 = 120 .
Revised: 𝑛 = 120 .

Original: The reward is specified by a set of modes 𝑀 ⊂
Revised: The reward is specified by a set of modes 𝑀 ⊂

Original: X = { 0 , 1 } 𝑛 that is unknown to the learning agent.
Revised: X = { 0 , 1 } 𝑛 that is unknown to the learning agent.

Original: The reward of a generated sequence 𝑥 is defined in terms of Hamming distance 𝑑 from the modes: 𝑅 ( 𝑥 ) = exp (− min 𝑦 ∈
Revised: The reward of a generated sequence 𝑥 is defined in terms of Hamming distance 𝑑 from the modes: 𝑅 ( 𝑥 ) = exp (− min 𝑦 ∈

Original: 𝑀 𝑑
Revised: 𝑀 𝑑

Original: ( 𝑥, 𝑦 )) .
Revised: ( 𝑥, 𝑦 )) .

Original: The vocabulary size can be varied: for any integer 𝑘 dividing 120, we take a vocabulary consisting of words of length 𝑘
Revised: The vocabulary size can be varied: for any integer 𝑘 dividing 120, we take a vocabulary consisting of words of length 𝑘

Original: (so that the vocabulary size is 2 𝑘
Revised: (so that the vocabulary size is 2 𝑘

Original: and the full sequence is generated in 𝑛𝑘 actions).
Revised: and the full sequence is generated in 𝑛𝑘 actions).

Original: By varying the value of 𝑘 and keeping 𝑛 and 𝑀 constant, we study the behavior of learning agents with varying action space sizes and trajectory lengths without changing the underlying modeling problem.
Revised: By varying the value of 𝑘 and keeping 𝑛 and 𝑀 constant, we study the behavior of learning agents with varying action space sizes and trajectory lengths without changing the underlying modeling problem.

Original: Most experiment settings are taken from Malkin et al.
Revised: Most experiment settings are taken from Malkin et al.

Original: ; see §C.
Revised: ; see §C.

Original: [Figure]
Revised: [Figure]

Original: Models are evaluated by computing the Spearman correlation, on a test set of sequences 𝑥 , between the probability of generating 𝑥 and the reward 𝑅 ( 𝑥 ) .
Revised: Models are evaluated by computing the Spearman correlation, on a test set of sequences 𝑥 , between the probability of generating 𝑥 and the reward 𝑅 ( 𝑥 ) .

Original: We also track the number of modes discovered during the training process for all the methods, see Fig.
Revised: We also track the number of modes discovered during the training process for all the methods, see Fig.

Original: We find that models trained with the
Revised: We find that models trained with the

Original: SubTB( 𝜆 ) objective have a higher Spearman correlation at the end of training and discover modes faster compared to the other GFlowNet objectives and non-GFlowNet baselines.
Revised: SubTB( 𝜆 ) objective have a higher Spearman correlation at the end of training and discover modes faster compared to the other GFlowNet objectives and non-GFlowNet baselines.

Original: 4.3. A NTIMICROBIAL PEPTIDE GENERATION
Revised: 4.3. A NTIMICROBIAL PEPTIDE GENERATION

Original: [Table]
Revised: [Table]

Original: Next, we consider the task of generating peptides with antimicrobial properties (AMPs).
Revised: Next, we consider the task of generating peptides with antimicrobial properties (AMPs).

Original: These sequences have maximum length 60 and use a vocabulary of 20 amino acids (and an end-of-sequence token), resulting in a state space of size 21 60 .
Revised: These sequences have maximum length 60 and use a vocabulary of 20 amino acids (and an end-of-sequence token), resulting in a state space of size 21 60 .

Original: The reward function is a pretrained proxy neural network that estimates the antimicrobial activity.
Revised: The reward function is a pretrained proxy neural network that estimates the antimicrobial activity.

Original: (See Jain et al.
Revised: (See Jain et al.

Original: for details on this task
Revised: for details on this task

Original: [Table]
Revised: [Table]

Original: We train GFlowNets with the SubTB( 𝜆 ), TB, and FM losses and compare them with baselines.
Revised: We train GFlowNets with the SubTB( 𝜆 ), TB, and FM losses and compare them with baselines.

Original: To evaluate the trained models, we sample 2048 sequences from the policy, then compute the mean reward and mean pairwise edit distance of the top-100 reward sequences.
Revised: To evaluate the trained models, we sample 2048 sequences from the policy, then compute the mean reward and mean pairwise edit distance of the top-100 reward sequences.

Original: The metrics and model architecture are taken from Malkin et al.
Revised: The metrics and model architecture are taken from Malkin et al.

Original: ; see §D. The results are presented in Table 2.
Revised: ; see §D. The results are presented in Table 2.

Original: SubTB( 𝜆 ) provides significant improvements over all the baselines (including TB, FM, and DB GFlowNets) in both reward and diversity.
Revised: SubTB( 𝜆 ) provides significant improvements over all the baselines (including TB, FM, and DB GFlowNets) in both reward and diversity.

Original: 4.3. F LUORESCENT PROTEIN GENERATION
Revised: 4.3. F LUORESCENT PROTEIN GENERATION

Original: [Table]
Revised: [Table]

Original: We consider the task of generating protein sequences with fluorescence properties (Trabucco et al., 2022) to evaluate SubTB( 𝜆 ) in settings with longer trajectories.
Revised: We consider the task of generating protein sequences with fluorescence properties (Trabucco et al., 2022) to evaluate SubTB( 𝜆 ) in settings with longer trajectories.

Original: In this task, sequences have a fixed length of 237, and the size of the state space is 20 237 .
Revised: In this task, sequences have a fixed length of 237, and the size of the state space is 20 237 .

Original: The proxy reward function 𝑅 ( 𝑥 ) is trained on a dataset of proteins with their fluorescence scores from Sarkisyan et al.
Revised: The proxy reward function 𝑅 ( 𝑥 ) is trained on a dataset of proteins with their fluorescence scores from Sarkisyan et al.

Original: The metrics and models are the same as in §4.3.2; see §E for details.
Revised: The metrics and models are the same as in §4.3.2; see §E for details.

Original: The GFlowNet objectives outperform all other methods in both metrics, finding more diverse and higher-reward sequences (Table 3).
Revised: The GFlowNet objectives outperform all other methods in both metrics, finding more diverse and higher-reward sequences (Table 3).

Original: SubTB( 𝜆 ) significantly outperforms TB, while achieving a similar diversity.
Revised: SubTB( 𝜆 ) significantly outperforms TB, while achieving a similar diversity.

Original: We note that the advantage of SubTB( 𝜆 ) is greater than that in the AMP task (Table 2) and speculate that the benefits of SubTB( 𝜆 ) become more prominent for longer action sequences.
Revised: We note that the advantage of SubTB( 𝜆 ) is greater than that in the AMP task (Table 2) and speculate that the benefits of SubTB( 𝜆 ) become more prominent for longer action sequences.

Original:  D ISCUSSION AND CONCLUSION
Revised:  D ISCUSSION AND CONCLUSION

Original: We have given evidence of a bias-variance tradeoff in GFlowNet training algorithms.
Revised: We have given evidence of a bias-variance tradeoff in GFlowNet training algorithms.

Original: The highvariance stochastic regression objective of TB and the low-variance local consistency objective of
Revised: The highvariance stochastic regression objective of TB and the low-variance local consistency objective of

Original: DB lie at opposite ends of this range.
Revised: DB lie at opposite ends of this range.

Original: We showed that SubTB( 𝜆 ) can harness the variance-reducing effects of local objectives while retaining the fast credit assignment properties of trajectory-level objectives.
Revised: We showed that SubTB( 𝜆 ) can harness the variance-reducing effects of local objectives while retaining the fast credit assignment properties of trajectory-level objectives.

Original: We see learnable strategies for selecting and weighting (sub)trajectories for training – e.g., a dynamic choice of 𝜆 and an active-learning approach to sampling trajectories – as the most interesting questions for future work.
Revised: We see learnable strategies for selecting and weighting (sub)trajectories for training – e.g., a dynamic choice of 𝜆 and an active-learning approach to sampling trajectories – as the most interesting questions for future work.

Original: The ability of subtrajectory objectives to learn from incomplete episodes also makes their application in RL environments an appealing research direction.
Revised: The ability of subtrajectory objectives to learn from incomplete episodes also makes their application in RL environments an appealing research direction.

Original: R EPRODUCIBILITY STATEMENT
Revised: R EPRODUCIBILITY STATEMENT

Original: We provide extensive experiment details, such as learning rates, batch sizes, number of training steps, choices of 𝜆 , description of attempted hyperparameters, and additional clarifying experiments in the Appendices.
Revised: We provide extensive experiment details, such as learning rates, batch sizes, number of training steps, choices of 𝜆 , description of attempted hyperparameters, and additional clarifying experiments in the Appendices.

Original: Code for experiments on the hypergrid domain (§4.1) and on the molecule domain (§4.2) is also provided with the submission
Revised: Code for experiments on the hypergrid domain (§4.1) and on the molecule domain (§4.2) is also provided with the submission

Original: A E XPERIMENT DETAILS : H YPERGRID
Revised: A E XPERIMENT DETAILS : H YPERGRID

Original: The environment is identical to that in Malkin et al.
Revised: The environment is identical to that in Malkin et al.

Original: (2022), with reward function parameters ( 𝑅 0 , 𝑅 1 , 𝑅 2 )
Revised: (2022), with reward function parameters ( 𝑅 0 , 𝑅 1 , 𝑅 2 )

Original: =
Revised: =

Original: ( 10 − 3 , 0 . 5 , 2 ) for the standard variant of the grid and ( 10 − 4 , 1 . 0 , 3 . 0 ) for the harder variant.
Revised: ( 10 − 3 , 0 . 5 , 2 ) for the standard variant of the grid and ( 10 − 4 , 1 . 0 , 3 . 0 ) for the harder variant.

Original: The models giving logits of 𝑃 𝐹 (−| 𝑠 ) and 𝑃 𝐵 (−| 𝑠 ) , as well as log 𝐹 ( 𝑠 ) ,
Revised: The models giving logits of 𝑃 𝐹 (−| 𝑠 ) and 𝑃 𝐵 (−| 𝑠 ) , as well as log 𝐹 ( 𝑠 ) ,

Original: are MLPs of the same architecture as in Bengio et al.
Revised: are MLPs of the same architecture as in Bengio et al.

Original: (2021a), taking a one-hot representation of the coordinates of 𝑠 as input and sharing all layers except the last.
Revised: (2021a), taking a one-hot representation of the coordinates of 𝑠 as input and sharing all layers except the last.

Original: The initial state flow log 𝑍 = log 𝐹 ( 𝑠 0 ) is an independent parameter whose learning rate is set to 10 × the learning rate of other parameters.
Revised: The initial state flow log 𝑍 = log 𝐹 ( 𝑠 0 ) is an independent parameter whose learning rate is set to 10 × the learning rate of other parameters.

Original: All models are trained with the Adam optimizer and a batch size of 16 for a total of 10trajectories (62500 batches).
Revised: All models are trained with the Adam optimizer and a batch size of 16 for a total of 10trajectories (62500 batches).

Original: The optimal learning rate for each experiment is chosen from { 0 . 0005 , 0 . 00075 , 0 . 001 , 0 . 003 , 0 . 005 , 0 . 0075 , 0 . 01 } , and 𝜆 = 0 .
Revised: The optimal learning rate for each experiment is chosen from { 0 . 0005 , 0 . 00075 , 0 . 001 , 0 . 003 , 0 . 005 , 0 . 0075 , 0 . 01 } , and 𝜆 = 0 .

Original: 9 is chosen as the optimal value from the set { 0 . 8 , 0 . 9 , 0 . 99 } .
Revised: 9 is chosen as the optimal value from the set { 0 . 8 , 0 . 9 , 0 . 99 } .

Original: Gradient bias and variance experiments are conducted in the harder variant of the 8 × 8 grid.
Revised: Gradient bias and variance experiments are conducted in the harder variant of the 8 × 8 grid.

Original: The tabular GFlowNet is trained using Adam with a learning rate 0 .
Revised: The tabular GFlowNet is trained using Adam with a learning rate 0 .

Original: 007 and the SubTB ( 𝜆 = 0 . 8 ) objective.
Revised: 007 and the SubTB ( 𝜆 = 0 . 8 ) objective.

Original: A. A DDITIONAL EXPERIMENTS
Revised: A. A DDITIONAL EXPERIMENTS

Original: Fig.
Revised: Fig.

Original: A.1 shows additional results on more difficult grid environments.
Revised: A.1 shows additional results on more difficult grid environments.

Original: We perform another experiment in which only short (up to length 4) subtrajectories are used for training with the SubTB( 𝜆 ) objective (i.e., the sum in (11) is truncated to exclude pairs ( 𝑖, 𝑗 ) with 𝑗 − 𝑖 > 4 ).
Revised: We perform another experiment in which only short (up to length 4) subtrajectories are used for training with the SubTB( 𝜆 ) objective (i.e., the sum in (11) is truncated to exclude pairs ( 𝑖, 𝑗 ) with 𝑗 − 𝑖 > 4 ).

Original: The results, shown in Fig.
Revised: The results, shown in Fig.

Original: A.4, show that SubTB( 𝜆 ) continues to perform strongly in this restricted setting.
Revised: A.4, show that SubTB( 𝜆 ) continues to perform strongly in this restricted setting.

Original: [Figure]
Revised: [Figure]

Original: Fig.
Revised: Fig.

Original: A.2 shows the effect of the SubTB parameter 𝜆 on the training curves, showing a gradual interpolation between DB and TB and fastest convergence at values slightly less than
Revised: A.2 shows the effect of the SubTB parameter 𝜆 on the training curves, showing a gradual interpolation between DB and TB and fastest convergence at values slightly less than

Original: [Figure]
Revised: [Figure]

Original: Fig.
Revised: Fig.

Original: A.3 contains visualizations of the exploration behavior of different training algorithms.
Revised: A.3 contains visualizations of the exploration behavior of different training algorithms.

Original: It shows that TB can perform better with off-policy training and can benefit from a higher temperature of the policy logits
Revised: It shows that TB can perform better with off-policy training and can benefit from a higher temperature of the policy logits

Original: [Figure]
Revised: [Figure]

Original: but still does not learn as fast as SubTB( 𝜆 )
Revised: but still does not learn as fast as SubTB( 𝜆 )

Original: [Figure]
Revised: [Figure]

Original: nor does it find all the modes in the maximum number of training iterations.
Revised: nor does it find all the modes in the maximum number of training iterations.

Original: A. M ORE ON BIAS AND VARIANCE : T HE EFFECT OF LEARNED STATE FLOWS
Revised: A. M ORE ON BIAS AND VARIANCE : T HE EFFECT OF LEARNED STATE FLOWS

Original: To better understand the variance-reducing properties of SubTB( 𝜆 ), we perform the gradient bias experiments with a modified computation of gradients that removes the factor of learning the state flows.
Revised: To better understand the variance-reducing properties of SubTB( 𝜆 ), we perform the gradient bias experiments with a modified computation of gradients that removes the factor of learning the state flows.

Original: Recall from §2.1 that a forward policy 𝑃 𝐹 uniquely determines a distribution over trajectories.
Revised: Recall from §2.1 that a forward policy 𝑃 𝐹 uniquely determines a distribution over trajectories.

Original: If the initial state flow 𝑍 and forward policy 𝑃 𝐹 are fixed, there is a unique state flow function 𝐹 𝐹 and backward policy 𝑃 𝐵 that satisfy the detailed balance conditions (6).
Revised: If the initial state flow 𝑍 and forward policy 𝑃 𝐹 are fixed, there is a unique state flow function 𝐹 𝐹 and backward policy 𝑃 𝐵 that satisfy the detailed balance conditions (6).

Original: This ‘true forward’ flow function, written 𝐹 𝐹 ( 𝑠 )
Revised: This ‘true forward’ flow function, written 𝐹 𝐹 ( 𝑠 )

Original: = 𝑍 (cid:205)
Revised: = 𝑍 (cid:205)

Original: 𝜏 : 𝑠 ∈ 𝜏 𝑃 𝐹 ( 𝜏 ) , is determined by an initial state flow fixed to the true
Revised: 𝜏 : 𝑠 ∈ 𝜏 𝑃 𝐹 ( 𝜏 ) , is determined by an initial state flow fixed to the true

Original: 12× 101010 6 0.
Revised: 12× 101010 6 0.

Original: 16× 101010 6 0.
Revised: 16× 101010 6 0.

Original: 20× 101010 6 0.
Revised: 20× 101010 6 0.

Original: 24× 101010 6 0.
Revised: 24× 101010 6 0.

Original: 28× 101010 6 0.
Revised: 28× 101010 6 0.

Original: 32× temp=1temp=2temp=3temp=4SubTB
Revised: 32× temp=1temp=2temp=3temp=4SubTB

Original: Training with tempered training policy 101010 6 0.
Revised: Training with tempered training policy 101010 6 0.

Original: 12× 101010 6 0.
Revised: 12× 101010 6 0.

Original: 16× 101010 6 0.
Revised: 16× 101010 6 0.

Original: 20× 101010 6 0.
Revised: 20× 101010 6 0.

Original: 24× 

Revised: 24×

Original: [Figure]
Revised: [Figure]

Original: 28× 101010 6 0.
Revised: 28× 101010 6 0.

Original: 32× =0.0=0.0625=0.125=0.25=0.5=
Revised: 32× =0.0=0.0625=0.125=0.25=0.5=

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: partition function 𝑍 = (cid:205) 𝑥 ∈X 𝑅 ( 𝑥 ) and the learned forward policy 𝑃 𝐹 .
Revised: partition function 𝑍 = (cid:205) 𝑥 ∈X 𝑅 ( 𝑥 ) and the learned forward policy 𝑃 𝐹 .

Original: Similarly, the ‘true backward’ flow function, written 𝐹 𝐵 ( 𝑠 ) =
Revised: Similarly, the ‘true backward’ flow function, written 𝐹 𝐵 ( 𝑠 ) =

Original: (cid:205)
Revised: (cid:205)

Original: 𝜏 : 𝑠 ∈ 𝜏 𝑃 𝐵 ( 𝜏 ) 𝑅 ( 𝑥 𝜏 ) where 𝑥 𝜏 is the terminal state of 𝜏 , is uniquely determined by the reward function 𝑅 and the learned backward policy 𝑃 𝐵 .
Revised: 𝜏 : 𝑠 ∈ 𝜏 𝑃 𝐵 ( 𝜏 ) 𝑅 ( 𝑥 𝜏 ) where 𝑥 𝜏 is the terminal state of 𝜏 , is uniquely determined by the reward function 𝑅 and the learned backward policy 𝑃 𝐵 .

Original: In particular, 𝐹 𝐵 ( 𝑠 0 ) =
Revised: In particular, 𝐹 𝐵 ( 𝑠 0 ) =

Original: (cid:205)
Revised: (cid:205)

Original: 𝑥 ∈X 𝑅 ( 𝑥 ) .
Revised: 𝑥 ∈X 𝑅 ( 𝑥 ) .

Original: We repeat the experiments on gradient bias, but by replacing the learned state flows 𝐹 in the losses by either the true forward or the true backward state flows ( 𝐹 𝐹 or 𝐹 𝐵 respectively) computed exactly using the current values of the learned 𝑃 𝐹 and 𝑃 𝐵 .
Revised: We repeat the experiments on gradient bias, but by replacing the learned state flows 𝐹 in the losses by either the true forward or the true backward state flows ( 𝐹 𝐹 or 𝐹 𝐵 respectively) computed exactly using the current values of the learned 𝑃 𝐹 and 𝑃 𝐵 .

Original: (These modifications are not applied in training, but are used only to compute the gradient similarities.
Revised: (These modifications are not applied in training, but are used only to compute the gradient similarities.

Original: The small size of the environment makes computation of the true state flows tractable; this is not possible in general.
Revised: The small size of the environment makes computation of the true state flows tractable; this is not possible in general.

Original: The gradient similarity over the course of training is shown in Fig.
Revised: The gradient similarity over the course of training is shown in Fig.

Original: A.5 (cf.
Revised: A.5 (cf.

Original: Fig.
Revised: Fig.

Original: 5 in the main text).
Revised: 5 in the main text).

Original: The similar behavior of SubTB( 𝜆 ) with learned and true forward state flows suggests that the learned state flows remain close enough to their optimal values and that the variance-reducing benefits of
Revised: The similar behavior of SubTB( 𝜆 ) with learned and true forward state flows suggests that the learned state flows remain close enough to their optimal values and that the variance-reducing benefits of

Original: SubTB( 𝜆 ) with true state flows are retained.
Revised: SubTB( 𝜆 ) with true state flows are retained.

Original: B E XPERIMENT DETAILS : M OLECULES
Revised: B E XPERIMENT DETAILS : M OLECULES

Original: All experiments with SubTB( 𝜆 ) are based upon the published code of Malkin et al.
Revised: All experiments with SubTB( 𝜆 ) are based upon the published code of Malkin et al.

Original: (2022), which extends that of Bengio et al.
Revised: (2022), which extends that of Bengio et al.

Original: (2021a).
Revised: (2021a).

Original: The proxy model giving the reward, the held-out set of molecules used to compute the correlation metric, and the GFlowNet model architecture – a graph neural network – are identical to those in Bengio et al.
Revised: The proxy model giving the reward, the held-out set of molecules used to compute the correlation metric, and the GFlowNet model architecture – a graph neural network – are identical to those in Bengio et al.

Original: (2021a), and the off-policy exploration rate and early stopping likelihood are the same as those tuned for the training with the TB objective in
Revised: (2021a), and the off-policy exploration rate and early stopping likelihood are the same as those tuned for the training with the TB objective in

Original: [Figure]
Revised: [Figure]

Original: Malkin et al.
Revised: Malkin et al.

Original: All models are trained for a maximum of 50000 batches of 4 trajectories each.
Revised: All models are trained for a maximum of 50000 batches of 4 trajectories each.

Original: (Some training runs terminated early because of numerical overflows in the gradients, in which case we report the metric of the last stable model whose cumulative number of batches is a multiple of 5000.)
Revised: (Some training runs terminated early because of numerical overflows in the gradients, in which case we report the metric of the last stable model whose cumulative number of batches is a multiple of 5000.)

Original: C E XPERIMENT DETAILS : B IT SEQUENCES
Revised: C E XPERIMENT DETAILS : B IT SEQUENCES

Original: The modes 𝑀 as well as the test sequences are selected as described in Malkin et al.
Revised: The modes 𝑀 as well as the test sequences are selected as described in Malkin et al.

Original: The policy for all methods is parameterized by a Transformer (Vaswani et al., 2017) with 3 layers, dimension 64, and 8 attention heads.
Revised: The policy for all methods is parameterized by a Transformer (Vaswani et al., 2017) with 3 layers, dimension 64, and 8 attention heads.

Original: All methods are trained for 50,000 iterations with minibatch size of 16 using Adam optimizer.
Revised: All methods are trained for 50,000 iterations with minibatch size of 16 using Adam optimizer.

Original: For GFlowNets with FM objective as well as the baselines, we use the exact same implementation and hyperparameters reported in Malkin et al.
Revised: For GFlowNets with FM objective as well as the baselines, we use the exact same implementation and hyperparameters reported in Malkin et al.

Original: For TB and SubTB( 𝜆 ), we pick the best learning rate from { 0 . 0075 , 0 . 001 , 0 .
Revised: For TB and SubTB( 𝜆 ), we pick the best learning rate from { 0 . 0075 , 0 . 001 , 0 .

Original: 001 , 0 . 003 , 0 . 005 } for forward logits, and for Z, use a learning rate of 10 × the learning rate for forward logits.
Revised: 001 , 0 . 003 , 0 . 005 } for forward logits, and for Z, use a learning rate of 10 × the learning rate for forward logits.

Original: For SubTB( 𝜆 ), we found the best 𝜆 value of 1.9 from the values { 0 . 8 , 0 .
Revised: For SubTB( 𝜆 ), we found the best 𝜆 value of 1.9 from the values { 0 . 8 , 0 .

Original: D E XPERIMENT DETAILS : A NTIMICROBIAL PEPTIDE GENERATION
Revised: D E XPERIMENT DETAILS : A NTIMICROBIAL PEPTIDE GENERATION

Original: Following Malkin et al.
Revised: Following Malkin et al.

Original: we use the following amino acids: [‘A’, ‘C’, ‘D’, ‘E’, ‘F’, ‘G’, ‘H’, ‘I’, ‘K’, ‘L’, ‘M’, ‘N’, ‘P’, ‘Q’, ‘R’, ‘S’, ‘T’, ‘V’, ‘W’, ‘Y’] .
Revised: we use the following amino acids: [‘A’, ‘C’, ‘D’, ‘E’, ‘F’, ‘G’, ‘H’, ‘I’, ‘K’, ‘L’, ‘M’, ‘N’, ‘P’, ‘Q’, ‘R’, ‘S’, ‘T’, ‘V’, ‘W’, ‘Y’] .

Original: We take 6438 known AMP sequences and 9522 non-AMP sequences from the DBAASP database Pirtskhalava et al.
Revised: We take 6438 known AMP sequences and 9522 non-AMP sequences from the DBAASP database Pirtskhalava et al.

Original: The classifier that serves as the proxy reward function is trained on this dataset, using 20% of the data as the validation set.
Revised: The classifier that serves as the proxy reward function is trained on this dataset, using 20% of the data as the validation set.

Original: The reward model is a Transformer, with 4 hidden layers, hidden dimension 64 , and 8 attention heads.
Revised: The reward model is a Transformer, with 4 hidden layers, hidden dimension 64 , and 8 attention heads.

Original: We train it with a minibatch of size 256 , with learning rate 10 − 4 , and with early stopping on the validation set.
Revised: We train it with a minibatch of size 256 , with learning rate 10 − 4 , and with early stopping on the validation set.

Original: We use a Transformer with 3 hidden layers with hidden dimension 64 with 8 attention heads as the architecture of the policy for all methods.
Revised: We use a Transformer with 3 hidden layers with hidden dimension 64 with 8 attention heads as the architecture of the policy for all methods.

Original: All methods are trained for 20 , 000 iterations, with a minibatch size of 16 , using the reported hyperparameters for all the baselines from (Malkin et al., 2022).
Revised: All methods are trained for 20 , 000 iterations, with a minibatch size of 16 , using the reported hyperparameters for all the baselines from (Malkin et al., 2022).

Original: For
Revised: For

Original: TB and SubTB( 𝜆 ), we pick the best learning rates from { 0 . 005 , 0 .
Revised: TB and SubTB( 𝜆 ), we pick the best learning rates from { 0 . 005 , 0 .

Original: 07 } for forward logits and from { 0 . 007 , 0 . 01 , 0 . 03 , 0 .
Revised: 07 } for forward logits and from { 0 . 007 , 0 . 01 , 0 . 03 , 0 .

Original: for log 𝑍 .
Revised: for log 𝑍 .

Original: For SubTB( 𝜆 ), the best performing 𝜆 value of 1 . 9 chosen from { 0 . 9 , 0 . 99 , 1 . 1 , 1 . 2 , 1 . 3 , 1 . 4 , 1 . 6 , 1 . 7 , 1 . 8 , 1 . 9 , 2 . 0 } is used.
Revised: For SubTB( 𝜆 ), the best performing 𝜆 value of 1 . 9 chosen from { 0 . 9 , 0 . 99 , 1 . 1 , 1 . 2 , 1 . 3 , 1 . 4 , 1 . 6 , 1 . 7 , 1 . 8 , 1 . 9 , 2 . 0 } is used.

Original: E E XPERIMENT DETAILS : F LUORESCENT PROTEIN GENERATION
Revised: E E XPERIMENT DETAILS : F LUORESCENT PROTEIN GENERATION

Original: We consider a variant of the GFP task from Trabucco et al.
Revised: We consider a variant of the GFP task from Trabucco et al.

Original: The vocabulary of amino acids is the same as §D: [‘A’, ‘C’, ‘D’, ‘E’, ‘F’, ‘G’, ‘H’, ‘I’, ‘K’, ‘L’, ‘M’, ‘N’, ‘P’, ‘Q’, ‘R’, ‘S’, ‘T’, ‘V’, ‘W’, ‘Y’] .
Revised: The vocabulary of amino acids is the same as §D: [‘A’, ‘C’, ‘D’, ‘E’, ‘F’, ‘G’, ‘H’, ‘I’, ‘K’, ‘L’, ‘M’, ‘N’, ‘P’, ‘Q’, ‘R’, ‘S’, ‘T’, ‘V’, ‘W’, ‘Y’] .

Original: Following Trabucco et al.
Revised: Following Trabucco et al.

Original: (2022), we consider the dataset of 56,086 proteins from Sarkisyan et al.
Revised: (2022), we consider the dataset of 56,086 proteins from Sarkisyan et al.

Original: (2016) processed based on Brookes et al.
Revised: (2016) processed based on Brookes et al.

Original: Each protein is accompanied by a score quantifying its fluorescence.
Revised: Each protein is accompanied by a score quantifying its fluorescence.

Original: As with the AMP data, we keep 20% of the data as a validation set used for early-stopping.
Revised: As with the AMP data, we keep 20% of the data as a validation set used for early-stopping.

Original: The regressor trained with the dataset is a Transformer, with 4 hidden layers, hidden dimension 64 , and 8 attention heads.
Revised: The regressor trained with the dataset is a Transformer, with 4 hidden layers, hidden dimension 64 , and 8 attention heads.

Original: We train it with a minibatch of size 256 , with learning rate 10 − 4 , with early stopping on the validation set.
Revised: We train it with a minibatch of size 256 , with learning rate 10 − 4 , with early stopping on the validation set.

Original: The architecture of the policy for all methods is a Transformer withhidden layers with hidden dimension 64 with 8 attention heads.
Revised: The architecture of the policy for all methods is a Transformer withhidden layers with hidden dimension 64 with 8 attention heads.

Original: All methods are trained for 20 ,iterations, with a minibatch size of 16 .
Revised: All methods are trained for 20 ,iterations, with a minibatch size of 16 .

Original: We use the same implementation for all methods as the ones used in §D.
Revised: We use the same implementation for all methods as the ones used in §D.

Original: To define an exploratory training policy, we set the the random action probability to 0 .
Revised: To define an exploratory training policy, we set the the random action probability to 0 .

Original: 01 selected from { 0 . 0001 , 0 . 0005 , 0 . 001 , 0 . 01 } and the reward exponent 𝛽 (having the same meaning as in §4.2) to 3 selected from { 2 , 3 , 4 } .
Revised: 01 selected from { 0 . 0001 , 0 . 0005 , 0 . 001 , 0 . 01 } and the reward exponent 𝛽 (having the same meaning as in §4.2) to 3 selected from { 2 , 3 , 4 } .

Original: For trajectory balance we use a learning rate of 5 × 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } for the flow parameters and 1 × 10 − 2 for log 𝑍 .
Revised: For trajectory balance we use a learning rate of 5 × 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } for the flow parameters and 1 × 10 − 2 for log 𝑍 .

Original: For SubTB( 𝜆 ), we choose the best 𝜆 from { 0 . 7 , 0 . 8 , 0 . 9 , 0 . 99 } , and found 𝜆 = 0 .
Revised: For SubTB( 𝜆 ), we choose the best 𝜆 from { 0 . 7 , 0 . 8 , 0 . 9 , 0 . 99 } , and found 𝜆 = 0 .

Original: 99 to perform the best.
Revised: 99 to perform the best.

Original: For TB and
Revised: For TB and

Original: SubTB( 𝜆 ), we tune for the best learning rates from { 0 . 0001 , 0 . 0003 , 0 . 0005 , 0 . 00075 , 0 . 001 } for the forward logits.
Revised: SubTB( 𝜆 ), we tune for the best learning rates from { 0 . 0001 , 0 . 0003 , 0 . 0005 , 0 . 00075 , 0 . 001 } for the forward logits.

Original: For log 𝑍 , we use a learning rate of 10 × the learning rate for the forward logits.
Revised: For log 𝑍 , we use a learning rate of 10 × the learning rate for the forward logits.

Original: For FM we use a learning rate of 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } with leaf loss coefficient 𝜆 𝑇 = 30 .
Revised: For FM we use a learning rate of 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } with leaf loss coefficient 𝜆 𝑇 = 30 .

Original: For A2C with entropy regularization we share parameters between the actor and critic networks, and use learning rate of 5 × 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } with entropy regularization coefficient 5 × 10 − 2 selected from { 10 − 4 , 10 − 3 , 5 × 10 − 3 , 10 − 2 , 5 × 10 − 2 } .
Revised: For A2C with entropy regularization we share parameters between the actor and critic networks, and use learning rate of 5 × 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } with entropy regularization coefficient 5 × 10 − 2 selected from { 10 − 4 , 10 − 3 , 5 × 10 − 3 , 10 − 2 , 5 × 10 − 2 } .

Original: For SAC we use the formulation in Christodoulou (2019) with a learning rate of 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } , a target network update frequency of 400 and initial random steps of 200 .
Revised: For SAC we use the formulation in Christodoulou (2019) with a learning rate of 10 − 3 selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } , a target network update frequency of 400 and initial random steps of 200 .

Original: For the MARS baseline, we set the learning rate to 5 × 10 −selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } .
Revised: For the MARS baseline, we set the learning rate to 5 × 10 −selected from { 10 − 5 , 10 − 4 , 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } .

Original: We run the experiments on 3 seeds and report the mean and standard error over the three runs in Table 3.
Revised: We run the experiments on 3 seeds and report the mean and standard error over the three runs in Table 3.

Original: F I NVERSE PROTEIN FOLDING : N ON - AUTOREGRESSIVE SEQUENCE GENERATION
Revised: F I NVERSE PROTEIN FOLDING : N ON - AUTOREGRESSIVE SEQUENCE GENERATION

Original: We consider the inverse protein folding problem suggested in Sinai et al.
Revised: We consider the inverse protein folding problem suggested in Sinai et al.

Original: A target protein 3D backbone conformation is given, and the task is to sample amino acid sequences of a fixed length 𝐿 = 40 from the Boltzmann distribution corresponding to their energy in the target conformation.
Revised: A target protein 3D backbone conformation is given, and the task is to sample amino acid sequences of a fixed length 𝐿 = 40 from the Boltzmann distribution corresponding to their energy in the target conformation.

Original: The energy is provided by a physics model (Rohl et al., 2004; Chaudhury et al., 2010).
Revised: The energy is provided by a physics model (Rohl et al., 2004; Chaudhury et al., 2010).

Original: The policy model is a 3-layer convolutional architecture that closely follows previous work (Sinai et al., 2020).
Revised: The policy model is a 3-layer convolutional architecture that closely follows previous work (Sinai et al., 2020).

Original: Specifically, for the policy function, the convolution size was set to 7 with 32 hidden features and
Revised: Specifically, for the policy function, the convolution size was set to 7 with 32 hidden features and

Original: ReLU activation in each layer.
Revised: ReLU activation in each layer.

Original: The policy network has one additional convolutional layer of size(number of amino acids), and without the activation function.
Revised: The policy network has one additional convolutional layer of size(number of amino acids), and without the activation function.

Original: The flow network has an additional two linear layers of sizes [1280,64], and [64, 1] with ReLU activation in between.
Revised: The flow network has an additional two linear layers of sizes [1280,64], and [64, 1] with ReLU activation in between.

Original: We report mean result over three runs.
Revised: We report mean result over three runs.

Original: For this task, rather than generating sequences from left to right, we consider an action space in which actions modify one letter at a time at arbitrary positions.
Revised: For this task, rather than generating sequences from left to right, we consider an action space in which actions modify one letter at a time at arbitrary positions.

Original: The first action uniformly randomly samples an amino acid sequence.
Revised: The first action uniformly randomly samples an amino acid sequence.

Original: On each subsequent action, the agent selects a position in the sequence and replaces the letter in this position with another letter in the vocabulary.
Revised: On each subsequent action, the agent selects a position in the sequence and replaces the letter in this position with another letter in the vocabulary.

Original: Generation terminates after exactly 𝑁 = 40 replacement steps.
Revised: Generation terminates after exactly 𝑁 = 40 replacement steps.

Original: The forward policy is conditioned on the number of steps taken so far in the trajectory; the backward policy is fixed to be uniform over the 𝑁 · 𝐿 actions.
Revised: The forward policy is conditioned on the number of steps taken so far in the trajectory; the backward policy is fixed to be uniform over the 𝑁 · 𝐿 actions.

Original: As a metric of how well the learned model matches the target distribution, we measure the correlation between log 𝑅 ( 𝑥 ) and the marginal sampling likelihood log 𝑝 𝜃 ( 𝑥 ) on a held-out set of terminal states.
Revised: As a metric of how well the learned model matches the target distribution, we measure the correlation between log 𝑅 ( 𝑥 ) and the marginal sampling likelihood log 𝑝 𝜃 ( 𝑥 ) on a held-out set of terminal states.

Original: The results are presented in Fig.
Revised: The results are presented in Fig.

Original: F.1.
Revised: F.1.

Original: We observe that intermediate values of lambda lead to the best fit to the target distribution.
Revised: We observe that intermediate values of lambda lead to the best fit to the target distribution.