Optimizing Backward Policies in GFlowNets via Trajectory Likelihood Maximization

## Abstract

Abstract Generative Flow Networks (GFlowNets) are a family of generative models that learn to sample objects with probabilities proportional to a given reward function. The key concept behind GFlowNets is the use of two stochastic policies: a forward policy, which incrementally constructs compositional objects, and a backward policy, which sequentially deconstructs them. Recent results show a close relationship between GFlowNet training and entropy-regularized reinforcement learning (RL) problems with a particular reward design. However, this connection applies only in the setting of a fixed backward policy, which might be a significant limitation. As a remedy to this problem, we introduce a simple backward policy optimization algorithm that involves direct maximization of the value function in an entropy-regularized Markov Decision Process (MDP) over intermediate rewards. We provide an extensive experimental evaluation of the proposed approach across various benchmarks in combination with both RL and GFlowNet algorithms and demonstrate its faster convergence and mode discovery in complex environments.

## 1 Introduction

Generative Flow Networks (GFlowNets, Bengio et al., ) are models designed to sample compositional discrete objects, e.g. graphs, from distributions defined by unnormalized probability mass functions. They operate by constructing an object through a sequence of stochastic transitions defined by a forward policy. This policy is trained to match the marginal distribution over constructed objects with the target distribution of interest. Since this marginal distribution is generally intractable, an auxiliary backward policy is introduced, and the problem is reduced to the one of matching distributions over complete trajectories, bearing similarities with variational inference (Malkin et al., ).

GFlowNets have found success in various areas, such as biological sequence design (Jain et al., ), material discovery (Hernandez-Garcia et al., ), molecular optimization (Zhu et al., ), recommender systems (Liu et al., ), large language model (LLM) and diffusion model fine-tuning (Hu et al., ; Venkatraman et al., ; Uehara et al., ; Zhang et al., ), neural architecture search (Chen & Mauch, ), combinatorial optimization (Zhang et al., ), and causal discovery (Atanackovic et al., ).

Theoretical foundations of GFlowNets have been established in the seminal works of Bengio et al. (; ). Most of the literature has since focused on practical applications of these models, leaving their theoretical properties largely unexplored, with a few exceptions (Krichel et al., ; Silva et al., ). However, a recent line of works has brought attention to connections between GFlowNets and reinforcement learning (Tiapkin et al., ; Mohammadpour et al., ; Deleu et al., ; He et al., ), showing that the GFlowNet learning problem is equivalent to a specific RL problem with entropy regularization (also called soft RL, Neu et al., ; Geist et al., ). This opened a new perspective for understanding GFlowNets. The importance of these findings is supported by empirical evidence, as various RL algorithms have proven useful for improving GFlowNets (Tiapkin et al., ; Lau et al., ; Morozov et al., ).

However, these connections still have a limitation related to GFlowNet backward policies. While GFlowNets can be trained with a fixed backward policy, standard GFlowNet algorithms allow the training of the backward policy together with the forward policy (Bengio et al., ; Malkin et al., ; Madan et al., ), resulting in faster convergence of the optimization process. Other algorithms for optimizing backward policies have been proposed in the literature as well (Mohammadpour et al., ; Jang et al., ), showing benefits for GFlowNet performance. The theory connecting GFlowNets and entropy-regularized RL is based on using the backward policy to add a ”correction” to GFlowNet rewards and shows the equivalence between two problems only when the backward policy is fixed. Thus, understanding the backward policy optimization remains a missing piece of this puzzle. Moreover, Tiapkin et al. () demonstrated that this theoretical gap has practical relevance, as optimizing the backward policy using the same RL objective as the forward policy can either fail to improve or even slow down convergence, highlighting the need for a more refined approach.

In this study, we introduce the trajectory likelihood maximization (TLM) approach for backward policy optimization, which can be integrated with any existing GFlowNet method, including entropy-regularized RL approaches.

To develop this method, we first formulate the GFlowNet training problem as a unified objective involving both forward and backward policies. We then propose an alternating minimization procedure consisting of two steps: (1) maximizing the backward policy likelihood of trajectories sampled from the forward policy and (2) optimizing the forward policy within an entropy-regularized Markov decision process that corresponds to the updated backward policy. The latter step can be achieved by any existing GFlowNet or soft RL algorithm, as it was outlined by Deleu et al. (). By approximating these two steps through a single stochastic gradient update, we derive an adaptive approach for combining backward policy optimization with any GFlowNet method, including soft RL methods.

Our main contributions are as follows:

- •
We derive the trajectory likelihood maximization (TLM) method for backward policy optimization;
- •
The proposed method represents the first unified approach for adaptive backward policy optimization in soft RL-based GFlowNet methods. The method is easy to implement and can be integrated with any existing GFlowNet training algorithm.
- •
We provide extensive experimental evaluation of TLM in four tasks, confirming the findings of Mohammadpour et al. (), which emphasize the benefits of training the backward policy in complex environment with less structure.

Source code: [github.com/tgritsaev/gflownet-tlm](https://github.com/tgritsaev/gflownet-tlm).

## 2 Background

### 2.1 GFlowNets

We aim at sampling from a probability distribution over a finite discrete space $\mathcal{X}$ that is given as an unnormalized probability mass function $\mathcal{R}\colon\mathcal{X}\to\mathbb{R}_{\geq 0}$, which we call the GFlowNet reward. We denote $\mathrm{Z}=\sum_{x\in\mathcal{X}}\mathcal{R}(x)$ to be an (unknown) normalizing constant.

To formally define a generation process in GFlowNets, we introduce a directed acyclic graph (DAG) $\mathcal{G}=(\mathcal{S},\mathcal{E})$, where $\mathcal{S}$ is a state space and $\mathcal{E}\subseteq\mathcal{S}\times\mathcal{S}$ is a set of edges (or transitions). There is exactly one state, $s_{0}$, with no incoming edges, which we refer to as the initial state. All other states can be reached from $s_{0}$, and the set of terminal states with no outgoing edges coincides with the space of interest $\mathcal{X}$. Non-terminal states $s\notin\mathcal{X}$ correspond to “incomplete” objects, and edges $s\to s^{\prime}$ represent adding ”new components” to such objects, transforming $s$ into $s^{\prime}$. Let $\mathcal{T}$ denote the set of all complete trajectories $\tau=\mathopen{}\mathclose{{}\left(s_{0},s_{1},\ldots,s_{n_{\tau}}}\right)$ in the graph, where $\tau$ is a sequence of states such that $(s_{i}\to s_{i+1})\in\mathcal{E}$ and that starts at $s_{0}$ and finishes at some terminal state $s_{n_{\tau}}\in\mathcal{X}$. As a result, any complete trajectory can be viewed as a sequence of actions that constructs the object corresponding to $s_{n_{\tau}}$ starting from the ”empty object” $s_{0}$.

We say that a state $s^{\prime}$ is a child of a state $s$, if there is an edge $(s\to s^{\prime})\in\mathcal{E}$. In this case we also say that $s$ is a parent of $s^{\prime}$. Next, for any state $s$, we introduce the forward policy, denoted by $\mathcal{P}_{\mathrm{F}}(s^{\prime}|s)$ for $(s\to s^{\prime})\in\mathcal{E}$, as an arbitrary probability distribution over the set of children of the state $s$. In a similar fashion, we define the backward policy as an arbitrary probability distribution over the parents of a state $s$ and denote it as $\mathcal{P}_{\mathrm{B}}(s^{\prime}|s)$, where $(s^{\prime}\to s)\in\mathcal{E}$.

Given these two definitions, the main goal of GFlowNet training is a search for a pair of policies such that the induced distributions over complete trajectories in the forward and backward directions coincide:

$$ $\prod_{t=1}^{n_{\tau}}\mathcal{P}_{\mathrm{F}}\mathopen{}\mathclose{{}\left(s_ {t}\mid s_{t-1}}\right)=\frac{\mathcal{R}(s_{n_{\tau}})}{\mathrm{Z}}\prod_{t=1 }^{n_{\tau}}\mathcal{P}_{\mathrm{B}}\mathopen{}\mathclose{{}\left(s_{t-1}\mid s _{t}}\right)\,,\quad\forall\tau\in\mathcal{T}\,.$ (1) $$

The relation (1) is known as the trajectory balance constraint (Malkin et al., ). We refer to the left and right-hand sides of (1) as to the forward and backward trajectory distributions and denote them as

$$ $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}}(\tau):=\prod_{i=1}^{n_{ \tau}}\mathcal{P}_{\mathrm{F}}(s_{i}|s_{i-1})\,,\qquad\mathsf{P}_{\mathcal{T}} ^{\mathcal{P}_{\mathrm{B}}}(\tau):=\frac{\mathcal{R}(s_{n_{\tau}})}{\mathrm{Z} }\cdot\prod_{i=1}^{n_{\tau}}\mathcal{P}_{\mathrm{B}}(s_{i-1}|s_{i})\,,$ (2) $$

where $\tau=(s_{0},s_{1},\ldots,s_{n_{\tau}})\in\mathcal{T}$. If the condition (1) is satisfied for all complete trajectories, sampling a trajectory in the forward direction using $\mathcal{P}_{\mathrm{F}}$ will result in a terminal state being sampled with probability $\mathcal{R}(x)/\mathrm{Z}$. We will call such $\mathcal{P}_{\mathrm{F}}$ a proper GFlowNet sampling policy.

In practice, we train a model (usually a neural network) that parameterizes the forward policy (and possibly other auxiliary functions) to minimize an objective function that enforces the constraint (1) or its equivalent. The main existing objectives are Trajectory Balance (TB, Malkin et al., ), Detailed Balance (DB, Bengio et al., ) and Subtrajectory Balance (SubTB, Madan et al., ). The SubTB objective is defined as

$$ $\textstyle\mathcal{L}_{\mathrm{SubTB}}(\theta;\tau)=\sum\limits_{0\leq j<k\leq n _{\tau}}w_{jk}\mathopen{}\mathclose{{}\left(\log\frac{F_{\theta}(s_{j})\prod_{ t=j+1}^{k}\mathcal{P}_{\mathrm{F}}(s_{t}|s_{t-1},\theta)}{F_{\theta}(s_{k}) \prod_{t=j+1}^{k}\mathcal{P}_{\mathrm{B}}(s_{t-1}|s_{t},\theta)}}\right)^{2}\,,$ (3) $$

where $F_{\theta}(s)$ is a neural network that approximates the flow function of the state $s$, see (Bengio et al., ; Madan et al., ) for more details on the flow-based formalization of the GFlowNet problem. Here $F_{\theta}(s)$ is substituted with $\mathcal{R}(s)$ for terminal states $s$, and $w_{jk}$ is usually taken to be $\lambda^{k-j}$ and then normalized to sum to $1$. TB and DB objectives can be viewed as special cases of (3), which are obtained by only taking the term corresponding to the full trajectory or to individual transitions, respectively. All objectives allow either training the model in an on-policy regime using the trajectories sampled from $\mathcal{P}_{\mathrm{F}}$ or in an off-policy mode using the replay buffer or some exploration techniques. In addition, it is possible to either optimize $\mathcal{P}_{\mathrm{B}}$ along with $\mathcal{P}_{\mathrm{F}}$ or to use a fixed $\mathcal{P}_{\mathrm{B}}$, e.g., the uniform distribution over parents of each state. One can show that given any fixed $\mathcal{P}_{\mathrm{B}}$, there exists a unique $\mathcal{P}_{\mathrm{F}}$ that satisfies (1); see, e.g., (Malkin et al., ).

### 2.2 GFlowNets as Soft RL

In reinforcement learning (Sutton & Barto, ), a typical performance measure of an agent is a value function, which is defined as an expected discounted sum of rewards when acting via a given policy. Entropy-regularized reinforcement learning (RL), also known as soft RL (Neu et al. ; Geist et al. ; Haarnoja et al. ), incorporates Shannon entropy $\mathcal{H}(\mathcal{P}_{\mathrm{F}}(\cdot\mid s))\triangleq\mathbb{E}_{s^{
\prime}\sim\mathcal{P}_{\mathrm{F}}(\cdot\mid s)}[-\log\mathcal{P}_{\mathrm{F}
}(s^{\prime}\mid s)]$ into the value function. This addition encourages the optimal policy to be more exploratory:

$$ $\textstyle V^{\mathcal{P}_{\mathrm{F}}}_{\lambda}(s;r)\triangleq\mathbb{E}_{ \mathcal{P}_{\mathrm{F}}}\bigg{[}\sum\limits_{t=0}^{\infty}\gamma^{t}\big{(}r( s_{t},s_{t+1})+\lambda\mathcal{H}(\mathcal{P}_{\mathrm{F}}(\cdot\mid s_{t})) \big{)}\,\Big{|}\,s_{0}=s\bigg{]},$ (4) $$

where $\lambda\geq 0$ is the regularization coefficient. Note that we use the next state instead of the more conventional action representation, as there is a one-to-one correspondence between the action taken and the resulting next state in DAG environments. Similarly, the regularized Q-values $Q^{\mathcal{P}_{\mathrm{F}}}_{\lambda}(s,s^{\prime})$ are defined as the expected discounted sum of rewards, augmented by Shannon entropy, given an initial state $s_{0}=s$ and the next state $s_{1}=s^{\prime}$. The regularized optimal policy $\mathcal{P}^{\star}_{\mathrm{F},\lambda}$ is the policy that maximizes $V^{\mathcal{P}_{\mathrm{F}}}_{\lambda}(s)$ for any state $s$.

Note. In standard RL notation, a policy is typically denoted as $\pi$. Here, we use $\mathcal{P}_{\mathrm{F}}$, consistently with GFlowNet notation, to avoid notational clutter.

It was proven by Tiapkin et al. () that training a GFlowNet policy $\mathcal{P}_{\mathrm{F}}$ with a fixed $\mathcal{P}_{\mathrm{B}}$ can be formulated as a soft RL problem in a deterministic MDP represented by a directed acyclic graph $\mathcal{G}$, where actions correspond to transitions over edges. For transitions $(s\to x)$ that lead to terminal states, the RL rewards are defined as $r^{\mathcal{P}_{\mathrm{B}}}(s,x)=\log\mathcal{P}_{\mathrm{B}}(s\mid x)+\log
\mathcal{R}(x)$, while for intermediate transitions $(s\to s^{\prime})$, the rewards are $r^{\mathcal{P}_{\mathrm{B}}}(s,s^{\prime})=\log\mathcal{P}_{\mathrm{B}}(s\mid s
^{\prime})$. By setting $\lambda=1$ and $\gamma=1$, it can be shown that the optimal policy $\mathcal{P}^{\star}_{\mathrm{F},\lambda=1}(\cdot\mid s)$ in this regularized MDP coincides with the proper GFlowNet forward policy $\mathcal{P}_{\mathrm{F}}(\cdot\mid s)$, which is uniquely determined by $\mathcal{P}_{\mathrm{B}}$ and $\mathcal{R}$ (Theorem 1, Tiapkin et al., ), thereby inducing the correct marginal distribution over terminal states $\mathcal{R}(x)/\mathrm{Z}$.

In addition, Proposition 1 of Tiapkin et al. () provides a connection between the corresponding regularized value function at the initial state $s_{0}$ for any forward policy $\mathcal{P}_{\mathrm{F}}$ (not necessarily proper) and KL-divergence between the induced trajectory distributions:

$$ $\textstyle V^{\mathcal{P}_{\mathrm{F}}}_{\lambda=1}(s_{0};r^{\mathcal{P}_{ \mathrm{B}}})=\log\mathrm{Z}-\operatorname{KL}(\mathsf{P}_{\mathcal{T}}^{ \mathcal{P}_{\mathrm{F}}}\|\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}} )\,.$ $$

The main practical corollary of these results is the fact that any RL algorithm that works with entropy regularization can be utilized to train GFlowNets when $\mathcal{P}_{\mathrm{B}}$ is fixed. For example, Tiapkin et al. () demonstrated the efficiency of the classical SoftDQN algorithm (Haarnoja et al., ) and its modified variant called MunchausenDQN (Vieillard et al., ). Moreover, it turns out that, under this framework, the existing GFlowNet training algorithms can be derived from existing RL algorithms. Tiapkin et al. () showed that on-policy TB corresponds to policy gradient algorithms, and DB corresponds to a dueling variant of SoftDQN. At the same time, Deleu et al. () showed that TB, DB and SubTB algorithms can be derived from path consistency learning (PCL, Nachum et al., ) under the assumption of a fixed backward policy.

### 2.3 Backward Policies in GFlowNets

The idea of backward policy optimization is essential to understanding the GFlowNets training procedure. In particular, the most straightforward approach used in GFlowNet literature (Malkin et al., ; Bengio et al., ) proposes to optimize the forward and backward policies directly through the same GFlowNet objective (e.g., (3)). This approach can accelerate the speed of convergence Malkin et al. (), at the same time potentially leading to less stable training (Zhang et al., ).

This phenomenon motivates studying the backward policy optimization in the recent works of Mohammadpour et al. () and Jang et al. (). Mohammadpour et al. () suggested using the backward policy with maximum possible trajectory entropy, thus focusing on the *exploration* challenges of GFlowNets. Such policy is proven to be $\mathcal{P}_{\mathrm{B}}(s\mid s^{\prime})=n(s)/n(s^{\prime})$, where $n(s)$ is the number of trajectories which starts at $s_{0}$ and end at $s$. It corresponds to the uniform one if the number of paths to all parent nodes is equal. When $n(s)$ cannot be computed analytically, Mohammadpour et al. () propose to learn $\log{n(s)}$, $s\in\mathcal{S}$ alongside the forward policy using its relation to the value function of the soft Bellman equation in the *inverted MDP* (see Definition 2 in Mohammadpour et al. ()).

The approach of Mohammadpour et al. () showed consistently better results in less structured tasks, like QM9 molecule generation (see Section 4.3 for a detailed description). At the same time, more structured environments with a less challenging exploration counterpart show less benefits of the proposed backward training approach.

In another line of work, Jang et al. () claim that the existing GFlowNets training procedures tend to under-exploit the high-reward objects and propose a Pessimistic Backward Policy approach. Thus, the primary aim of Jang et al. () is to focus on the *exploitation* of the current information about high-reward trajectories. Towards this aim, they focus on maximizing the observed backward flow $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}}(\tau)$ (see (2)) for trajectories leading to high-reward objects, which are stored in a replay buffer. Our approach shares similarities with this method, but differs in theoretical motivation and the choice of trajectories for $\mathcal{P}_{\mathrm{B}}$ optimization. Unfortunately, Jang et al. () do not provide enough details about the choice of trajectories that are stored in their replay buffer and stochastic optimization over trajectories in the buffer. Moreover, related design choices vary between different environments, which limits reproducibility of the results.

As an additional limitation of both Mohammadpour et al. () and Jang et al. (), we mention the fact that only a single GFlowNet training objective is used in both papers (SubTB and TB respectively) to evaluate approaches for backward policy optimization, while we carry out experimental evaluation with various GFlowNet training objectives (see Section 4).

## 3 Trajectory Likelihood Maximization

The objective of our method is to formalize the optimization process for the backward policy for reinforcement learning-based approaches. It is worth mentioning that soft RL methods cannot address the changing of the reward function, except for reward shaping schemes (Ng et al., ) that preserve the total reward of any trajectory. Therefore, we need to return to the underlying GFlowNet problem. Let us look at the following optimization problem:

$$ $\textstyle\min_{\mathcal{P}_{\mathrm{F}}\in\Pi_{\mathrm{F}},\mathcal{P}_{ \mathrm{B}}\in\Pi_{\mathrm{B}}}\operatorname{KL}(\mathsf{P}_{\mathcal{T}}^{ \mathcal{P}_{\mathrm{F}}}\|\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}} )\,,$ (5) $$

where $\Pi_{\mathrm{F}}$ and $\Pi_{\mathrm{B}}$ represent the spaces of forward and backward policies, respectively, and $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}}$ and $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}}$ are defined in (2). It is easy to see that any solution to the problem (5) satisfies the trajectory balance constraint (1) and thus induces a proper GFlowNet sampling policy. Additionally, it is worth mentioning that for any fixed $\mathcal{P}_{\mathrm{B}}$ the problem (5) is equivalent to maximizing value

$$ $\textstyle V^{\mathcal{P}_{\mathrm{F}}}_{\lambda=1}(s_{0};r^{\mathcal{P}_{ \mathrm{B}}})=\log\mathrm{Z}-\operatorname{KL}(\mathsf{P}_{\mathcal{T}}^{ \mathcal{P}_{\mathrm{F}}}\|\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}})$ $$

over $\mathcal{P}_{\mathrm{F}}\in\Pi_{\mathrm{F}}$ in an entropy-regularized MDP, as established in Tiapkin et al. (, Proposition 1). Thus, the joint optimization resembles the RL formulation with non-stationary rewards. To leverage the problem’s block structure, we propose a meta-algorithm consisting of two iterative steps, repeated until convergence:

$$ $\textstyle\mathcal{P}_{\mathrm{B}}^{t+1}\approx\arg\min_{\mathcal{P}_{\mathrm{ B}}}\operatorname{KL}\mathopen{}\mathclose{{}\left(\mathsf{P}_{\mathcal{T}}^{ \mathcal{P}_{\mathrm{F}}^{t}}\big{\|}\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{ \mathrm{B}}}}\right)\,,\qquad\mathcal{P}_{\mathrm{F}}^{t+1}\approx\arg\min_{ \mathcal{P}_{\mathrm{F}}}\operatorname{KL}\mathopen{}\mathclose{{}\left( \mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}}\big{\|}\mathsf{P}_{ \mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{t+1}}}\right)\,.$ (6) $$

It is worth noting that if these optimization problems are solved exactly, the algorithm converges after the first iteration. This occurs because, for every fixed backward policy $\mathcal{P}_{\mathrm{B}}$, there is a unique forward policy $\mathcal{P}_{\mathrm{F}}$, such that $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}}=\mathsf{P}_{\mathcal{T}}^{
\mathcal{P}_{\mathrm{B}}}$, see, e.g., (Malkin et al., ), ensuring that the loss function reaches its global minimum. In the following sections, we provide implementation details on approximating these two steps.

##### First Step: Trajectory Likelihood Maximization.

Using the connection between forward KL divergence minimization and maximum likelihood estimation (MLE), we formulate the following trajectory likelihood maximization objective:

$$ $\textstyle\theta_{\mathrm{B}}^{t+1}\approx\arg\min_{\theta}\mathbb{E}_{\tau \sim\mathcal{P}_{\mathrm{F}}^{t}}\mathopen{}\mathclose{{}\left[\mathcal{L}_{ \mathtt{TLM}}(\theta;\tau)}\right],\qquad\mathcal{L}_{\mathtt{TLM}}(\theta; \tau):=-\sum_{i=1}^{n_{\tau}}\log\mathcal{P}_{\mathrm{B}}(s_{i-1}|s_{i},\theta )\,.$ (7) $$

In this formulation, $\tau=(s_{0},s_{1},\ldots,s_{n_{\tau}})$ denotes a trajectory generated by the forward policy $\mathcal{P}_{\mathrm{F}}^{t}$. This step seeks to update the backward policy by minimizing the negative log-likelihood of trajectories generated from the forward policy. Additionally, instead of solving (7) for exact $\arg\min$ for every $t$, we perform one stochastic gradient update

$$ $\textstyle\theta_{\mathrm{B}}^{t+1}=\theta_{\mathrm{B}}^{t}-\gamma\nabla_{ \theta}\mathcal{L}_{\mathtt{TLM}}(\theta_{\mathrm{B}}^{t};\tau)\,.$ $$

##### Second Step: Non-Stationary Soft RL Problem.

To approximate the second step of (6), we exploit the equivalence between the GFlowNet framework and the entropy-regularized RL problem. This leads to the following expression:

$$ $\textstyle\mathcal{P}_{\mathrm{F}}^{t+1}\approx\arg\min_{\mathcal{P}_{\mathrm{ F}}\in\Pi_{\mathrm{F}}}\operatorname{KL}\mathopen{}\mathclose{{}\left(\mathsf{ P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}}\big{\|}\mathsf{P}_{\mathcal{T}}^{ \mathcal{P}_{\mathrm{B}}^{t+1}}}\right)\iff\mathcal{P}_{\mathrm{F}}^{t+1} \approx\arg\max_{\mathcal{P}_{\mathrm{F}}\in\Pi_{\mathrm{F}}}V^{\mathcal{P}_{ \mathrm{F}}}_{\lambda=1}\mathopen{}\mathclose{{}\left(s_{0};r^{\mathcal{P}_{ \mathrm{B}}^{t+1}}}\right)\,,$ (8) $$

where $r^{\mathcal{P}_{\mathrm{B}}}$ is the RL reward function corresponding to the backward policy $\mathcal{P}_{\mathrm{B}}$. This step can be solved using any soft RL method, such as SoftDQN (Haarnoja et al., ). Additionally, it is noteworthy that all existing GFlowNet algorithms with a fixed backward policy can be viewed as variations of existing RL methods, see, e.g., (Deleu et al., ). Thus, they can be used to solve the optimization problem in (8).

To mitigate the computational overhead of searching for exact $\arg\min$ in (8), we also propose to perform a single stochastic gradient update in the corresponding GFlowNet training algorithm

$$ $\theta_{\mathrm{F}}^{t+1}=\theta_{\mathrm{F}}^{t}-\eta\nabla_{\theta}\mathcal{ L}_{\mathtt{Alg}}(\theta_{\mathrm{F}}^{t};\tau,\mathcal{P}_{\mathrm{B}}^{t+1})\,,$ $$

where $\mathcal{L}_{\mathtt{Alg}}$ represents the loss function associated with a GFlowNet or soft RL method, such as SubTB or SoftDQN. Here, $\tau$ denotes a (possibly off-policy) trajectory.

The complete procedure can be interpreted as a soft RL method with changing rewards. Our suggested method is summarized in Algorithm 1 and can be paired with any GFlowNet training method Alg (e.g., DB, TB, SubTB, or SoftDQN). While for $\mathcal{P}_{\mathrm{B}}$ training our approach uses on-policy trajectories, $\mathcal{P}_{\mathrm{F}}$ can still be trained off-policy, e.g., by utilizing a replay buffer that stores trajectories or transitions.

Figure: Algorithm 1 Trajectory Likelihood Maximization

##### Convergence of the method

Next, we show why this method indeed solves the GFlowNet learning problem. First, we introduce a non-stationary soft reinforcement learning problem of minimizing the worst-case dynamic average regret

$$ $\textstyle\overline{\mathfrak{R}}^{T}:=\frac{1}{T}\sum_{t=1}^{T}V^{\mathcal{P} _{\mathrm{F}}^{\star}}_{\lambda=1}(s_{0};r^{t})-V^{\mathcal{P}_{\mathrm{F}}^{t }}_{\lambda=1}(s_{0};r^{t})\,,$ (9) $$

where $\{r^{t}\}_{t\in[T]}$ is a sequence of reward functions, and $r^{t}$ is revealed to a learner before selecting a policy $\mathcal{P}_{\mathrm{F}}^{t}$.
Following Zahavy et al. (), we conjecture that existing RL algorithms are adaptive to the setting of known but non-stationary reward sequences. The implementation of Sampler player of EntGame algorithm by Tiapkin et al. () is an example of such a regret minimization algorithm. Additionally, we notice that the optimization of dynamic regret is well-studied in the online learning literature, even in a more challenging setting of revealing the corresponding reward function *after* playing a policy (Zinkevich, ; Besbes et al., ).

Next, we provide the convergence result for our two-step procedure, using a stability argument for the first step. The proof is given in Appendix A.1.

###### Theorem 3.1 .

Assume that (1) the backward updates are stable, i.e., $\sup_{t\geq 0}\lVert\mathcal{P}_{\mathrm{B}}^{T}-\mathcal{P}_{\mathrm{B}}^{T+t
}\rVert_{1}\to 0$ as $T\to\infty$, and (2) the forward updates follow a non-stationary regret minimization algorithm, i.e., $\overline{\mathfrak{R}}^{T}\to 0$ as $T\to\infty$. Then, there exists a proper GFlowNet forward policy $\mathcal{P}_{\mathrm{F}}^{\star}$ that induces the marginal distribution $\mathcal{R}(x)/\mathrm{Z}$ over terminal states, such that $\frac{1}{T}\sum_{t=1}^{T}\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}^{t
}}\to\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}^{\star}}$ as $T\to\infty$.

During numerical experiments, we observed that enforcing stability in backward updates, particularly by using a decaying learning rate, significantly improves convergence in practice. Furthermore, as the theorem shows, this stability is essential for theoretical convergence. We discuss stability techniques we utilize alongside our algorithm in Appendix A.2.

###### Theorem 3.1 .

##### Discussion.

We underline our approach’s similarity with a celebrated EM-algorithm (Dempster et al., ) and Hinton’s Wake-Sleep algorithm (Hinton et al., ). Both these methods also attempt to address the minimization problem using the block structure and alternating minimization approaches. Additionally, our approach can be connected to cooperative game theory. Indeed, one may interpret $\mathcal{P}_{\mathrm{F}}$ as a Forward player and $\mathcal{P}_{\mathrm{B}}$ as a Backward player, and both players attempt to minimize KL-divergence between corresponding trajectory distributions.

## 4 Experiments

We conduct our experimental evaluation on hypergrid (Bengio et al., ) and bit sequence (Malkin et al., ) environments, as well as on two molecule design environments: sEH (Bengio et al., ) and QM9 (Jain et al., ). Additional experimental details and hyperparameter choices are provided in Appendix A.3.

We evaluate four GFlowNet training methods: MunchausenDQN (based on the framework of Tiapkin et al. ()), DB (Bengio et al., ), TB (Malkin et al., ), and SubTB (Madan et al., ), which we refer to as GFlowNet algorithms (denoted as $\mathcal{L}_{\mathtt{Alg}}$ in the previous section). On hypergrids, we also include results for SoftDQN (Tiapkin et al., ). Alongside these algorithms, we consider four strategies for learning or selecting the backward policy:

- •
our approach (TLM);
- •
fixed uniform backward policy (uniform);
- •
simultaneously learning the backward policy with $\mathcal{P}_{\mathrm{F}}$ using the same objective (naive);
- •
maximum entropy backward policy (maxent) (Mohammadpour et al., ).

We refer to these strategies as backward approaches. In this section, we denote the distribution induced by a forward policy $\mathcal{P}_{\mathrm{F}}$ over terminal states as $P_{\theta}(x)$ for $x\in\mathcal{X}$, representing the probability of sampling $x$ from our GFlowNet.

### 4.1 Hypergrid

We start experiments with synthetic hypergrid environments introduced by Bengio et al. (). These environments are sufficiently small to compute target distribution in the closed form, allowing us to directly examine the convergence of $P_{\theta}(x)$ to $\mathcal{R}(x)/\mathrm{Z}$.

The environment is a $d$-dimensional hypercube with a side length equal to $H$. The state space is represented as $d$-dimensional vectors $(s_{1},\ldots,s_{d})^{\mathsf{\scriptscriptstyle T}}\in\{0,\ldots,H-1\}^{d}$ with the initial state being $(0,\ldots,0)^{\mathsf{\scriptscriptstyle T}}$. For each state $(s_{1},\ldots,s_{d-1})$, there are at most $d+1$ actions. The first action always corresponds to an exit action that transfers the state to its terminal copy, and the rest of $d$ actions correspond to incrementing one coordinate by $1$ without leaving the grid. The number of terminal states here is $|\mathcal{X}|=H^{d}$.

There are $2^{d}$ regions with high rewards near the corners of the grid, while states outside have much lower rewards. The exact expression for the rewards is given in Section A.3.2.

We explore environments with the reward parameters taken from Malkin et al. (), referred to as “standard case”, and with the reward parameters from Madan et al. (), referred to as a “hard case”. In the second case, background rewards are lower, which makes mode exploration more challenging. We conduct experiments on a 4-dimensional hypercube with a side length of 20. As an evaluation metric, we use $L^{1}$ distance between the true reward distribution and the empirical distribution of the last $2\cdot 10^{5}$ terminal states sampled during training.

Figure 1 presents the results. For SoftDQN, MunchausenDQN, and DB, TLM shows the fastest convergence for both “standard” and “hard” reward designs. For the SubTB algorithm, TLM shows similar performance to naive and outperforms uniform and maxent. TB is known to have difficulties in this environment (Madan et al., ), all approaches fail to converge under the “hard” reward design. At the same time, with the “standard” one, maxent backward shows the best convergence. An important note is that our results reproduce the findings of Tiapkin et al. (): for SoftDQN and MunchausenDQN training with uniform backward converges faster than with naive algorithm, while TLM shows stable improvement over uniform. The results and the ranking of algorithms are almost the same for SoftDQN and MunchausenDQN, so we leave only MunchausenDQN out of two for further experiments.

Figure: Figure 1: $L^{1}$ distance between target and empirical sample distributions over the course of training on the standard (top row) and hard (bottom row) hypergrid environments for each method. *Lower values indicate better performance.*
Refer to caption: x1.png

### 4.2 Bit Sequences

In this section, we consider the bit sequence generation task introduced by Malkin et al. (). Following the experimental setup of (Tiapkin et al., ), we modify the state and action spaces to create a non-tree DAG structure, similar to the approach introduced in Zhang et al. ().

This task is to generate binary sequences of a fixed length $n$, using a vocabulary of $k$-bit blocks. The state space of this environment corresponds to sequences of $n/k$ words, and each word in these sequences is either an empty word $\oslash$ or one of $2^{k}$ possible $k$-bit words. The initial state $s_{0}$ corresponds to a sequence of empty words. The possible actions in each state are replacing an empty word $\oslash$ with one of $2^{k}$ non-empty words in the vocabulary. The set of terminal states $\mathcal{X}$ consists of sequences without empty words and corresponds to binary strings of length $n$. The reward function is defined as ${\mathcal{R}(x)=\exp(-2\cdot\min_{x^{\prime}\in\mathcal{M}}d(x,x^{\prime}))}$, where $\mathcal{M}$ is a set of modes and $d$ is Hamming distance. We fix $n=120$ and $k=8$ for our experiments. The terminal state space size is $|\mathcal{X}|=2^{120}$.

Importantly, for this environment, the uniform backward coincides with maxent, see Proposition 1 of Zhang et al. () and Remark 3 of Theorem 3 of Mohammadpour et al. ().

To evaluate the performance, we use the same metrics as in Malkin et al. () and Tiapkin et al. (): the number of modes found during training (number of sequences from $\mathcal{M}$ for which a terminal state within a distance of $30$ has been sampled) and Spearman correlation on the test set between $\mathcal{R}(x)$ and an estimate of $P_{\theta}$. Since computing the exact probability of sampling a terminal state is intractable due to a large number of paths leading to it, we use a Monte Carlo estimate following the approach of Zhang et al. (); see also Section A.3.

Figure 2 shows the number of modes for different GFlowNet algorithms and backward approaches found over the course of training. We observe that TLM shows a significant improvement for DB and a minor one for MunchausenDQN in comparison to other backward approaches, where in the later case we find all 60 modes. TB and SubTB also find almost all modes, and TLM does not affect the results much. Full plots for modes across varying learning rates are presented in Figure 5 in Appendix. Figure 3 (top) presents Spearman correlation between $\mathcal{R}$ and $P_{\theta}$ on the test set for the same GFlowNet algorithms and varying learning rates. TLM shows better or similar performance to the baselines across all GFlowNet algorithms if the optimal learning rate is chosen. Moreover, for DB and SubTB, TLM shows steady improvement over the baselines for all learning rates.

Figure: Figure 2: Top row: Bit Sequences, the number of discovered modes out of total 60 modes for different methods with learning rate $10^{-3}$. Center row: QM9, the number of Tanimoto-separated modes with reward higher or equal to $1.125$ for different methods with learning rate $5\cdot 10^{-4}$. Bottom row: sEH, the number of Tanimoto-separated modes with reward higher or equal to $0.875$ for different methods with learning rate $5\cdot 10^{-4}$. *Higher values indicate better performance*.
Refer to caption: x3.png

### 4.3 Molecule Design, sEH and QM9

Our final experiments are carried out on molecule design tasks of sEH (Bengio et al., ) and QM9 (Jain et al., ).

In both tasks, the goal is to generate molecular graphs, with reward emphasizing some desirable property. For both problems, we use pre-trained reward proxy neural networks. For the sEH task, the model is trained to predict the binding energy of a molecule to a particular protein target (soluble epoxide hydrolase) (Bengio et al., ). For the QM9 task, the proxy is trained on the QM9 dataset (Ramakrishnan et al., ) to predict the HOMO-LUMO gap (Zhang et al., ).

For the sEH task, we follow the framework proposed by Jin et al. () and generate molecules using a predefined vocabulary of fragments. It is essential to mention that these fragments are explicitly selected for the sEH task to simplify high-quality object generation. The states are represented as trees of fragments. The actions correspond to choosing a new fragment, then choosing an atom to which the fragment will be attached. There is also a special stop action that moves the state to its terminal copy and stops the generation process.

For QM9 task, molecules are generated atom-by-atom and bond-by-bond. Every state is a
connected graph, and actions either add a new node and
edge or set the attribute on edge. Thus, the graph-building environment is much more expressive than the tree-building environment, but it results in a more complex generation task and can lead to construction of invalid molecules.

We use the same evaluation metrics for both tasks as proposed in previous works (Madan et al., ; Tiapkin et al., ). We track the number of Tanimoto-separated modes above a certain reward threshold captured over the course of training, and Pearson correlation on the test set between $\log\mathcal{R}(x)$ and $\log P_{\theta}(x)$. For sEH task we use the same test set as in Bengio et al. (), and for QM9 we use a subset of the dataset introduced in Ramakrishnan et al. ().

Figure: Figure 3: Top row: Bit Sequences, Spearman correlation between $\mathcal{R}$ and $P_{\theta}$ on a test set for different methods and varying learning rate $\in\{5\cdot 10^{-4},10^{-3},2\cdot 10^{-3}\}$. Center row: QM9, Pearson correlation between $\log\mathcal{R}$ and $\log P_{\theta}$ on the fixed subset of the QM9 dataset (Ramakrishnan et al., ) for different methods and varying learning rate $\in\{5\cdot 10^{-5},10^{-4},5\cdot 10^{-4},10^{-3}\}$. Bottom row: sEH, Pearson correlation between $\log\mathcal{R}$ and $\log P_{\theta}$ on the test set from Bengio et al. () for different methods and varying learning rate $\in\{5\cdot 10^{-5},10^{-4},5\cdot 10^{-4},10^{-3}\}$. *Higher values indicate better performance*.
Refer to caption: x6.png

Figure 2 (center and bottom) shows the number of modes for different GFlowNet algorithms and backward approaches found over the course of training. Overall, TLM greatly speeds up mode discovery on QM9 for all GFlowNet algorithms, but shows similar or worse performance when compared to other backward approaches on sEH. However, we note that on sEH no backward approach shows significant improvement over uniform in terms of mode discovery. Full plots for modes across varying learning rates are presented in Figure 6 in Appendix. It is worth noting that on QM9 TLM shows robust improvement over the baselines across most learning rates for all GFlowNet algorithms. Figure 3 (center and bottom) shows Pearson correlation between $\log\mathcal{R}$ and $\log P_{\theta}$ estimate measured on the test set for various learning rates. TLM results in better correlations when paired with MunchausenDQN and SubTB, and shows similar results to the baselines when paired with DB and TB.

Our results support the findings of Mohammadpour et al. (), where it was also observed that the backward policy optimization enhances GFlowNets on more complex and less structured environments like QM9.

## 5 Conclusion

In this paper, we proposed a new method for backward policy optimization that enhances mode exploration and accelerates convergence in complex GFlowNet environments. This approach represents the first principled method for learning a backward policy in soft reinforcement learning-based GFlowNet algorithms, such as SoftDQN and MunchausenDQN.

In line with the findings of Mohammadpour et al. (), we hypothesize that backward policy learning algorithms, such as TLM, may be more advantageous in complex environments with less inherent structure, such as QM9. This is supported by the observation that TLM yields relatively small or negligible improvements in the highly structured sEH task, which has been simplified in prior work (e.g., Jin et al., ). In contrast, the unstructured and complex nature of QM9 presents greater challenges, where our approach achieves the most significant performance gains. A deeper understanding of methods like TLM or maxent could enhance the effectiveness of GFlowNets in tackling less structured and more challenging problems.

Another promising direction is using backward policy for exploration as proposed in Kim et al. (), He et al. (). Indeed, the ability to sample trajectories that start from high-reward terminal states via $\mathcal{P}_{\mathrm{B}}$ provides an opportunity to improve mode exploration. We expect that combining such methods with TLM-like approaches will additionally improve their performance.

## Appendix A Appendix

### A.1 Omitted proofs

###### Proof of Theorem 3.1 .

From the stability of backward updates, the Cauchy criterion implies that there is $\mathcal{P}_{\mathrm{B}}^{\star}\in\Pi_{\mathrm{B}}$ such that $\mathcal{P}_{\mathrm{B}}^{T}\to\mathcal{P}_{\mathrm{B}}^{\star}$. At the same time, by the choice or rewards $r^{t}=r^{\mathcal{P}_{\mathrm{B}}^{t}}$, Proposition 1 by Tiapkin et al. () and joint convexity of KL-divergence

$$ $\overline{\mathfrak{R}}^{T}=\frac{1}{T}\sum_{t=1}^{T}\operatorname{KL} \mathopen{}\mathclose{{}\left(\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F }}^{t}}\big{\|}\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{t}}}\right) \geq\operatorname{KL}\mathopen{}\mathclose{{}\left(\frac{1}{T}\sum_{t=1}^{T} \mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}^{t}}\bigg{\|}\frac{1}{T} \sum_{t=1}^{T}\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{t}}}\right)\,.$ $$

Notice that a mapping $\mathcal{P}_{\mathrm{B}}\mapsto\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{
B}}}$ is continuous, thus $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{T}}\to\mathsf{P}_{\mathcal
{T}}^{\mathcal{P}_{\mathrm{B}}^{\star}}$, and, as a result, averages of $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{t}}$ also converge to $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{\star}}$. Finally, applying Pinkser’s inequality, we have

$$ $\textstyle\big{\|}\frac{1}{T}\sum_{t=1}^{T}\mathsf{P}_{\mathcal{T}}^{\mathcal{ P}_{\mathrm{F}}^{t}}-\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{\star }}\big{\|}_{1}\leq\sqrt{2\overline{\mathfrak{R}}^{T}}+\lVert\frac{1}{T}\sum_{t =1}^{T}\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{t}}-\mathsf{P}_{ \mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{\star}}\rVert_{1}\,.$ $$

The right-hand side of the inequality tends to zero as $T\to+\infty$, thus $\frac{1}{T}\sum_{t=1}^{T}\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}^{t
}}\to\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{\star}}$. Finally, since for any $\mathcal{P}_{\mathrm{B}}^{\star}$ there is $\mathcal{P}_{\mathrm{F}}^{\star}$ such that $\mathsf{P}_{\mathcal{T}}^{\mathcal{P}_{\mathrm{B}}^{\star}}=\mathsf{P}_{
\mathcal{T}}^{\mathcal{P}_{\mathrm{F}}^{\star}}$, we conclude the statement.
∎

###### Proof of Theorem 3.1 .

### A.2 Stability Techniques

In this section, we highlight important practical techniques and design choices motivated by Theorem 3.1 that we use alongside our TLM algorithm to enforce stability into the training of $\mathcal{P}_{\mathrm{B}}$.

First, we found it beneficial to either use a lower learning rate for the backward policy or decay it over the course of training (see the next sections for detailed descriptions).

Second, akin to how the Deep Q-Network algorithm (Mnih et al., ) utilizes a target network to estimate the value of the next state, we utilize target networks for the backward policy when calculating the loss for the forward policy. For example, (3) transforms into

$$ $\textstyle\mathcal{L}_{\mathrm{SubTB}}(\theta;\tau)=\sum\limits_{0\leq j<k\leq n _{\tau}}w_{jk}\mathopen{}\mathclose{{}\left(\log\frac{F_{\theta}(s_{j})\prod_{ t=j+1}^{k}\mathcal{P}_{\mathrm{F}}(s_{t}|s_{t-1},\theta)}{F_{\theta}(s_{k}) \prod_{t=j+1}^{k}{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{ 1,0,0}\mathcal{P}_{\mathrm{B}}(s_{t-1}|s_{t},\bar{\theta})}}}\right)^{2}$ (10) $$

where the parameters $\bar{\theta}$ of $\mathcal{P}_{\mathrm{B}}(s_{t-1}|s_{t},\bar{\theta})$ are updated via exponential moving average (EMA) of the online parameters $\theta$ of $\mathcal{P}_{\mathrm{B}}(s_{t-1}|s_{t},\theta)$. So the loss for the backward policy $\mathcal{L}_{\mathtt{TLM}}$ is computed using an online backward policy $\mathcal{P}_{\mathrm{B}}(s_{t-1}|s_{t},\theta)$, and the loss for the forward policy $\mathcal{L}_{\mathtt{Alg}}$ is computed using a target backward policy $\mathcal{P}_{\mathrm{B}}(s_{t-1}|s_{t},\bar{\theta})$, which is frozen during the gradient update of $\mathcal{P}_{\mathrm{F}}$.

Finally, we find it helpful to initialise $\mathcal{P}_{\mathrm{B}}$ to the uniform distribution at the beggining of training, which is done by zero-initialization of the last linear layer weight and bias.

We ablate the impact of the proposed techniques on QM9, where we try to separately turn off each of the three. Results are presented in Figure 4. We observe that using target model and lower learning rate is crucial, whereas disabling uniform initialization increases variance and shows slightly worse results. For this experiment, we choose DB as the base algorithm because TLM overall shows the greatest impact when applied with it compared to TB, SubTB, and MunchausenDQN.

Figure: Figure 4: Ablation study of stability techniques on QM9. The number of Tanimoto-separated modes with a reward at least $1.125$ is shown. As a base algorithm we use DB with a learning rate of $5\cdot 10^{-4}$.
Refer to caption: x9.png

### A.3 Experimental Details

We utilize PyTorch (Paszke et al., ) in our experiments. For hypergrid and bit sequence environments, we base our implementation upon the published code of Tiapkin et al. (). For molecule design experiments, our implementations are based on the open source library by Recursion Pharma.(^1^11 [https://github.com/recursionpharma/gflownet](https://github.com/recursionpharma/gflownet)) In all our experiments, $\mathcal{P}_{\mathrm{F}}$ and $\mathcal{P}_{\mathrm{B}}$ share the same neural network backbone, predicting logits via two separate linear heads. For all experiments, mean and std values are computed over three random seeds.

#### A.3.1 Hypergrid

The reward at a terminal state $s$ with coordinates $(s^{1},\ldots,s^{D})$ is defined as

$$ $\mathcal{R}(s)=R_{0}+R_{1}\times\prod_{i=1}^{D}\mathbb{I}\mathopen{}\mathclose {{}\left[0.25<\mathopen{}\mathclose{{}\left|\frac{s^{i}}{H-1}-0.5}\right|} \right]+R_{2}\times\prod_{i=1}^{D}\mathbb{I}\mathopen{}\mathclose{{}\left[0.3< \mathopen{}\mathclose{{}\left|\frac{s^{i}}{H-1}-0.5}\right|<0.4}\right].$ $$

Standard reward uses parameters $(R_{0}=10^{-3},R_{1}=0.5,R_{2}=2.0)$ and hard reward uses $(R_{0}=10^{-4},R_{1}=1.0,R_{2}=3.0)$, taken from Bengio et al. () and Madan et al. () respectively.

All models are parameterized using an MLP with 2 hidden layers and 256 hidden units. We use the Adam optimizer with a learning rate of $10^{-3}$ and a batch size of 16 trajectories. For the TLM backward policy, we use the same initial learning rate of $10^{-3}$ with an exponential scheduler, tuning $\gamma$ from $\{0.999,0.9999\}$. For SubTB, we set $\lambda=0.9$, following Madan et al. (). For SoftDQN and MunchausenDQN, we use a prioritized replay buffer (Schaul et al., ) and adopt the same hyperparameters as in Tiapkin et al. (). We sample 256 transitions from the buffer to compute the loss for these methods. The backward policy target network is updated using soft updates with a parameter of $\tau=0.25$ (Silver et al., ). Since the environment is small, we precompute $\log n(s)$ for all states, allowing us to obtain the maxent backward exactly. Hypergrid experiments were performed on CPUs.

The table summarizes the chosen hyperparameters:

| Hyperparatemeter | Value |
| --- | --- |
| Training trajectories | $10^{6}$ |
| Learning rate | $10^{-3}$ |
| $Z$ learning rate for TB | $10^{-1}$ |
| Adam optimizer $\beta,\varepsilon,\lambda$ | $(0.9,0.999),10^{-8},0$ |
| Number of hidden layers | $2$ |
| Hidden embedding size | $256$ |
| SubTB $\lambda$ | $0.9$ |
| MunchausenDQN $\alpha$ | $0.15$ |
| MunchausenDQN $\l_{0}$ | $-100$ |
| Prioritized RB size | $10^{5}$ |
| Prioritized RB $\alpha,\beta$ | $0.5,0.0$ |
| Prioritized RB batch size | $256$ |
| TLM backward learning rate | $10^{-3}$ |
| TLM backward lr decay $\gamma$ | $\{0.999,0.9999\}$ |
| TLM backward target model $\tau$ | $0.25$ |

#### A.3.2 Bit Sequences

The set of modes $M$ is defined as in Malkin et al. (), and we choose the same size, $|M|=60$. We set $H=\{^{\prime}00000000^{\prime},^{\prime}11111111^{\prime},^{\prime}11110000^{
\prime},^{\prime}00001111^{\prime},^{\prime}00111100^{\prime}\}$. Each sequence in $M$ is generated by randomly selecting $n/8$ elements from $H$ with replacement and then concatenating them. The test set for evaluating reward correlations is generated by taking a mode and flipping $i$ random bits in it, where this is repeated for every mode and for each $0\leq i<n$.

We utilize the same Monte Carlo estimate for $P_{\theta}$ as presented in Tiapkin et al. () with $N=10$:

$$ $P_{\theta}(x)=\mathbb{E}_{\mathcal{P}_{\mathrm{B}}(\tau\mid x)}\mathopen{} \mathclose{{}\left[\frac{\mathcal{P}_{\mathrm{F}}(\tau\mid\theta)}{\mathcal{P} _{\mathrm{B}}(\tau\mid x)}}\right]\approx\frac{1}{N}\sum_{i=1}^{N}\frac{ \mathcal{P}_{\mathrm{F}}(\tau^{i}\mid\theta)}{\mathcal{P}_{\mathrm{B}}(\tau^{i }\mid x)},\quad\tau^{i}\sim\mathcal{P}_{\mathrm{B}}(\tau\mid x).$ $$

Notice that any valid $\mathcal{P}_{\mathrm{B}}$ can be used here, but for each model, we take the $\mathcal{P}_{\mathrm{B}}$ that was fixed/trained alongside the corresponding $\mathcal{P}_{\mathrm{F}}$ since such a choice will lead to a lower estimate variance. However, we note that the metric is still very noisy, so for each training run, we compute the metric for all model checkpoints (done every 2000 iterations) and use the maximum value.

All models are parameterized as Transformers Vaswani et al. () with 3 hidden layers, 8 attention heads, and a hidden dimension of 64. Each model is trained for 50,000 iterations and a batch size of 16 with Adam optimizer. We provide results for learning rates from $\{5\times 10^{-4},10^{-3},2\times 10^{-3}\}$. For the TLM backward policy, we use the same initial learning rate as for the forward policy and utilize an exponential scheduler, where $\gamma$ is tuned from $\{0.9997,0.9999\}$. For SubTB we use $\lambda=0.9$. For MunchausenDQN we use a prioritized replay buffer (Schaul et al., ) and take the same hyperparameters as in Tiapkin et al. (). To compute the loss, we sample 256 transitions from the buffer for MunchausenDQN. The backward policy target network uses soft updates with a parameter $\tau=0.25$ (Silver et al., ). Bit sequence experiments were performed on NVIDIA V100 GPUs.

To closely follow the setting of previous works (Malkin et al., ; Madan et al., ; Tiapkin et al., ), we use $\varepsilon$-uniform exploration with $\varepsilon=10^{-3}$. We note that this can introduce a small bias into the gradient estimate of $\nabla_{\theta}\mathcal{L}_{\mathtt{TLM}}(\theta_{\mathrm{B}}^{t};\tau)$ since $\tau$ will not be sampled exactly from $\mathcal{P}_{\mathrm{F}}$.

The table summarizes the chosen hyperparameters:

| Hyperparatemeter | Value |
| --- | --- |
| Training iterations | $5\times 10^{4}$ |
| Learning rate | $\{5\times 10^{-4},10^{-3},2\times 10^{-3}\}$ |
| $Z$ learning rate for TB | $10^{-1}$ |
| Adam optimizer $\beta,\varepsilon,\lambda$ | $(0.9,0.999),10^{-8},10^{-5}$ |
| $\varepsilon$-uniform exploration | $10^{-3}$ |
| Number of transformer layers | $3$ |
| Hidden embedding size | $64$ |
| Dropout | $0.1$ |
| SubTB $\lambda$ | $0.9$ |
| MunchausenDQN $\alpha$ | $0.15$ |
| MunchausenDQN $\l_{0}$ | $-100$ |
| MunchausenDQN leaf coefficient | $5$ |
| Prioritized RB size | $10^{5}$ |
| Prioritized RB $\alpha,\beta$ | $0.9,0.1$ |
| Prioritized RB batch size | $256$ |
| TLM backward learning rate | learning rate |
| TLM backward lr decay $\gamma$ | $0.9999$ |
| TLM backward target model $\tau$ | $0.1$ |

#### A.3.3 Molecules

For sEH, we use the test set from Bengio et al. (). For QM9, we select a subset of 773 molecules from the QM9 dataset (Ramakrishnan et al., ) containing between 3 and 8 atoms. The subset is constructed to ensure an approximately equal representation of different molecule sizes. To compute correlation, we use the same Monte Carlo estimate as in the bit sequence task. It is important to note that Mohammadpour et al. () used a different evaluation approach, computing correlation on sampled molecules rather than on a fixed dataset.

We set $\mathcal{R}(x)=\exp(-75.0)$ for invalid molecules $x$ in QM9. In the sEH task, rewards are divided by a constant of 8, and the reward exponent is set to 10. For QM9, we subtract the 95th percentile from all rewards, resulting in the majority of rewards being distributed between 0 and 1, with 5% of molecules having a reward greater than 1. We track the number of Tanimoto-separated modes as described in Bengio et al. (), using a Tanimoto similarity threshold of 0.7. After normalization, the reward thresholds are $0.875$ and $1.125$ for sEH and QM9, respectively.

We use the graph transformer architecture from Jain et al. () with 8 layers and 256 embedding dimensions for both tasks. Each model is trained for 20,000 iterations using the Adam optimizer. We present results for learning rates in the set $\{5\times 10^{-4},10^{-4},5\times 10^{-4},10^{-3}\}$. For TLM, the learning rate for $\mathcal{P}_{\mathrm{B}}$ is initially set to be 10 times smaller than that for $\mathcal{P}_{\mathrm{F}}$, while other approaches use the same learning rate for $\mathcal{P}_{\mathrm{B}}$ as for $\mathcal{P}_{\mathrm{F}}$. The learning rate for $\mathrm{Z}$ is fixed at $10^{-3}$ for all experiments. We apply exponential schedulers to all learning rates at each training step and refer readers to the hyperparameter table for the exact decay values. The batch sizes are 256 for sEH and 128 for QM9. The backward policy target network is updated using soft updates with a parameter of $\tau=0.05$ (Silver et al., ). For SubTB, we set $\lambda=1.0$, following Madan et al. (). For MunchausenDQN, we do not use a replay buffer, training the model on-policy while using the same other hyperparameters as in Tiapkin et al. (). To learn $\log n(s)$ for the maxent backward policy, we follow the approach from Mohammadpour et al. (), implemented in the Recursion Pharma library. All molecule generation experiments were conducted on NVIDIA A100 GPUs.

In line with previous works (Malkin et al., ; Madan et al., ; Tiapkin et al., ), we use $\varepsilon$-uniform exploration with $\varepsilon=0.05$. To account for the bias introduced into the gradient estimate of $\nabla_{\theta}\mathcal{L}_{\mathtt{TLM}}(\theta_{\mathrm{B}}^{t};\tau)$, we linearly anneal $\varepsilon$ to zero over the course of training. This approach is applied to all GFlowNet algorithms and backward methods to ensure a fair comparison.

The table summarizes chosen hyperparameters:

| Hyperparatemeter | sEH | QM9 |
| --- | --- | --- |
| Reward exponent | $10$ | $1$ |
| Batch size | $256$ | $128$ |
| Training iterations | $2\times 10^{4}$ |  |
| Learning rate | $\{5\times 10^{-4},10^{-4},5\times 10^{-4},10^{-3}\}$ |  |
| Learning rate decay $\gamma$ | $2^{-1/20000}$ |  |
| $Z$ learning rate for TB | $10^{-3}$ |  |
| $Z$ learning rate for TB decay | $2^{-1/50000}$ |  |
| Adam optimizer $\beta,\varepsilon,\lambda$ | $(0.9,0.999),10^{-8},10^{-8}$ |  |
| $\varepsilon$-uniform exploration | $0.05$ linearly annealed to $0.0$ |  |
| Number of transformer layers | $8$ |  |
| Hidden embedding size | $256$ |  |
| SubTB $\lambda$ | $1.0$ |  |
| MunchausenDQN $\alpha$ | $0.15$ |  |
| MunchausenDQN $\l_{0}$ | $-500$ |  |
| MunchausenDQN leaf coefficient | $10$ |  |
| TLM backward learning rate | $0.1\times$ learning rate |  |
| TLM backward lr decay $\gamma$ | $2^{-1/5000}$ | $2^{-1/20000}$ |
| TLM backward target model $\tau$ | $0.05$ |  |

### A.4 Full plots

Figure: Figure 5: Bit Sequences, the number of modes discovered over the course of training for different methods and a learning rate $\in\{5\cdot 10^{-4},10^{-3},2\cdot 10^{-3}\}$. TB results at the learning rate of $2\cdot 10^{-3}$ are not full because of exploding gradients at a certain point in training.
Refer to caption: x10.png

Figure: Figure 6: Top row: QM9, the number of Tanimoto-separated modes discovered over the course of training with reward higher or equal to $1.125$ for different methods and learning rate in $\in\{5\cdot 10^{-5},10^{-4},5\cdot 10^{-4},10^{-3}\}$. Bottom row: sEH, the number of Tanimoto-separated modes discovered over the course of training with reward higher or equal to $0.875$ for different methods and learning rate $\in\{5\cdot 10^{-5},10^{-4},5\cdot 10^{-4},10^{-3}\}$.
Refer to caption: x11.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: NA
soundness: 3
strengths: - This paper is well written. The presentation and structure of this paper are clear.
- The proposed method, which employs an EM-style mechanism to optimize the training objective, is interesting. Theoretical analysis guarantees its convergence, and it can be used in any GFlowNet algorithm.
summary: This paper proposes a novel method to optimize the backward policy of the GFlowNet during training. The proposed method can be combined with any GFlowNet training method, and the experimental results demonstrate that it can bring performance improvement in some tasks.
weaknesses: - The performance improvement is marginal. On the sEH task, the proposed method even harms the performance, as explained in line 485. Moreover, the experiments (e.g., Qm9 and sEH) included in this paper are somewhat toy. I encourage the authors to consider more complex and large-scale benchmarks, such as AMP and RMA generation [1].


[1] Biological Sequence Design with GFlowNets, ICML 2022

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Could the authors provide a direct experimental comparison with the pessimistic backward policy approach proposed by Jang et al. (2024)? This would help establish the relative merits of TLM in similar settings.
2. Can the authors provide an explanation for the performance of TLM with different choices of the GFlowNet training objective (TB, DB, SubTB)? Are there certain objectives where TLM shows particularly strong or weak performance?
soundness: 3
strengths: - The proposed TLM method is designed to be easily implementable and compatible with existing GFlowNet training algorithms.
- The authors validate their approach across four different tasks, demonstrating its effectiveness in various scenarios.
- The paper is well-structured and clearly written, with careful explanations of technical concepts and methodological choices.
summary: This paper introduces a method called Trajectory Likelihood Maximization (TLM) for optimizing backward policies in Generative Flow Networks (GFlowNets). The authors formulate the GFlowNet training problem as a unified objective involving both forward and backward policies and propose an alternating minimization procedure. The method is validated through extensive experimental evaluation across various benchmarks.
weaknesses: - Limited comparative analysis: While the paper acknowledges recent work by Jang et al. (2024) on pessimistic backward policies, it doesn't provide a direct experimental comparison with this highly relevant approach. Such a comparison would help readers understand the relative advantages of TLM. Though the experiments are extensive, it would be valuable to see more detailed analysis of how TLM performs compared to other backward policy optimization methods across different types of environments.
- Ablation studies: The paper could benefit from more detailed ablation studies to understand the contribution of different components of the proposed method.
- Results presentation: It will be better to combine the last three columns in Figure 1 to a single figure as they are just different variants of GFlowNets training objectives. Including the best-performing loss will better present the results (e.g., row 2).

Hyosoon Jang, Yunhui Jang, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Pessimistic backward policy for gflownets. arXiv preprint arXiv:2405.16012, 2024.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: For the related work, I still cannot understand what connections you are trying to make and which research you are trying to use and build upon. More specifically, what limitations are you trying to overcome? What are the similarities between your algorithm and the celebrated EM-algorithm (Dempster et al., 1977) and Hinton’s Wake-Sleep algorithm (Hinton et al., 1995) and how does TLM differ from their algorithm?
The experimental work does not seem well justified, but there is no discussion of the varying performances across the environments; what conclusions are you deriving? Do you have any explanations for the deficiencies your algorithm has which were shown in some parts of the empirical results? Do you have any suggestions in order to overcome those shortages? 
The experiments do not provide convincing evidence of the correctness of the proposed approach or its utility compared to existing approaches. There are so many missing details it is difficult to draw many conclusions:
- Why did you use an online backward policy to compute the loss for the backward policy L_TLM?
- Why did you use a target backward policy to compute the loss for the forward policy L_Alg?
- In Hypergrid, all models are parameterized by MLP with 2 hidden layers, why? There seems no study of the decaying learning rate conducted on the Hypergrid environment, any comments?
- Why did you choose transformers to parameterize your models in the Bit Sequences environment?
- There is no clear rule for choosing your hyperparameters; it appears that you are simply taking the values established in previous related work. Why?
The differences in the “full plots” compared to the original ones are very small when choosing decaying learning rates. Can you provide some context to understand the significance of this difference?
soundness: 3
strengths: The authors study a very important problem in the GFlowNets and entropy-regularized RL literature which is backward policy optimization. They propose the TLM method that represents the first unified approach for adaptive backward policy optimization in soft RL-based GFlowNet methods which enhances mode exploration and accelerates convergence in complex GFlowNet environments. They formulate the GFlowNet training problem as a unified objective involving both forward and backward policies, proposing an alternating minimization procedure. For novelty, they approximate these two steps through a single stochastic gradient update and derive an adaptive approach for combining backward policy optimization with any GFlowNet method, including soft RL methods. The significance of their work arises from working on the limitations of two main previous research; Mohammadpour et al. (2024) and Jang et al. (2024). Furthermore, their approach can be connected to cooperative game theory by interpreting the forward policy PF as a forward player and the backward policy PB as a backward player, with both players attempting to minimize KL-divergence between corresponding trajectory distributions. For theory, they introduce a new theorem that shows that stability is essential for theoretical convergence. The quality of the proposed method is shown by their numerical experiments. TLM greatly speeds up mode discovery on the QM9 environment for all GFlowNet algorithms. Their results support the findings of Mohammadpour et al. (2024), where it was also observed that the backward policy optimization enhances GFlowNets on more complex and less structured environments like QM9. As a last note, the math behind their algorithm is neat and well-explained.

The main strengths are as follows:
- Deriving the trajectory likelihood maximization (TLM) method for backward policy optimization.
- The proposed method represents the first unified approach for adaptive backward policy optimization in soft RL-based GFlowNet methods. The method can be integrated with any existing GFlowNet training algorithm.
- Providing an extensive experimental evaluation of TLM in four tasks, confirming the findings of
Mohammadpour et al. (2024), emphasized the benefits of training the backward policy in complex environments with less structure.
- Providing the basis for implementing and extending their approach in different domains such as cooperative game theory.
summary: This paper studies the connection between GFlowNet training and entropy-regularized reinforcement learning problems–also known as soft RL. It specifically tackles this relationship in the setting of the fixed backward policy which is considered a significant limitation in the literature. For this indeed, they propose a simple backward policy optimization algorithm; Trajectory Likelihood Maximization (TLM)  that primarily involves a direct maximization of the value function in an entropy-regularized Markov Decision Process (MDP) over intermediate rewards. They formulate the GFlowNet training problem as a unified objective involving both forward and backward policies and propose an alternating minimization procedure. Their main contribution is deriving the first unified approach for adaptive backward policy optimization in soft RL-based GFlowNet methods, i.e. the TLM. They carry out an experimental evaluation on Hypergrid and Bit Sequence as well as two molecule design environments: sEH and QM9 to conclude that their algorithm converges faster and that can be integrated with any existing GFlowNet training algorithm. Lastly, there are empirical experiments done to show the performance of TLM under different hyperparameter choices.
In general, this paper could be a good contribution, with the caveat of some clarifications on the theory and experiments. Given these clarifications in an author's response, I would be willing to increase the score.
weaknesses: The paper does not do a great job of demonstrating the deficiencies of the proposed algorithm. No section in this paper clearly explains and shows the difference between the proposed algorithm and other related work. In other words, the reader may not understand what specifically the work you are trying to build your work on. The authors point to references to establish closely related work such as Mohammadpour et al. (2024) and Jang et al. (2024), but do not clearly and plainly show what exactly they are improving, and what previous limitations they are trying to overcome. Furthermore, the authors mention that there are some similarities between their algorithm and a celebrated EM-algorithm (Dempster et al., 1977) and Hinton’s Wake-Sleep algorithm (Hinton et al., 1995), but did not mention what aspects those similarities are.
The experiments do not seem well justified, in the Hypergrid environment, Figure 1, the authors do not clearly explain why their method TLM fails to outperform the naive method in the TB training. In the Bit Sequences environment, the authors again do not show a sufficient explanation for the different results shown in Figure 2-Top row such as the improvement in DB and why TLM does not affect the results much in TB and SubTB. In addition, there seems to be no clear reason for producing results using two different metrics: Found Modes, and the Spearman Correlation between R and P_theta in three different environments, Bit sequences, sEH, and QM9. No explanation for Figure 2 (center and bottom) where TLM greatly speeds up mode discovery on QM9 for all GFlowNet algorithms, but shows similar or worse performance when compared to other backward approaches on sEH. There is no justified explanation of the poor performance of TLM on sEH in terms of mode discovery in all four Gflownets training algorithms.
Lastly, in Figure 3, their conclusion “TLM results in better correlations when paired with MunchausenDQN and SubTB, and shows similar results to the baselines when paired with DB and TB.” is not well justified. In what aspects does that affect your method? It would be much better to explain their method’s performance alongside its deficiencies in a more active discussion section. 


Suggestions:
- Specify related work, similarities, and differences in a separate section.
- Explain your empirical results, i.e. your algorithm’s deficiencies and mention suggestions to enhance its performance.

Minor comments:
1. Appendix. A.1 Omitted Proofs P. 14 Line 2, “by the choice of rewards”

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. The paper highlights the advantages of the TLM method in complex environments, but its performance is suboptimal in simpler environments, such as the sEH task. Can the reasons behind this underperformance in simpler tasks be explored in more depth? Is it due to limitations of the model or characteristics inherent to the environment?

2. The paper notes that TLM struggles to handle non-stationary rewards effectively. Could a variant of TLM be designed specifically for dynamic reward structures to address this issue? How would such an adaptation impact TLM's performance in non-stationary reward environments?

3. Maximum entropy policies are beneficial for enhancing exploration. Could the TLM method be combined with maximum entropy policies to improve exploration capabilities in certain tasks? If so, how might these two approaches be effectively integrated for optimal collaboration?

4. Could the TLM method be integrated with other types of reinforcement learning algorithms, such as hard-constrained RL, to address different task requirements? Would such a combination further improve backward policy optimization and overall performance?
soundness: 3
strengths: 1. The TLM method shares a strong theoretical connection with existing reinforcement learning approaches, such as entropy-regularized RL. It introduces a novel joint optimization framework that enhances the efficiency of learning the backward policy. These theoretical connections are validated through experimental results presented in the paper.

2. The method is grounded in the maximum likelihood estimation problem from reinforcement learning and proposes a unified approach for optimizing the backward policy. This provides a robust theoretical basis for addressing the longstanding challenge of backward policy optimization in GFlowNets.

3. The paper is written in a clear and accessible manner, particularly in its detailed explanations of complex theoretical concepts, making it easier for readers to comprehend.
summary: This paper introduces a novel method called Trajectory Likelihood Maximization (TLM) for optimizing the backward policy in Generative Flow Networks (GFlowNets). GFlowNets are a class of generative models, and the TLM method accelerates convergence in complex tasks by directly maximizing the likelihood of trajectories. Additionally, it significantly enhances exploration capabilities in high-dimensional tasks.
weaknesses: 1. The paper primarily focuses on the introduction and validation of the TLM method, but provides limited discussion on comparisons with other backward policy methods, such as those proposed by Jang et al. and Mohammadpour et al. Notably, the improvement in backward policy optimization is not significant in more structured tasks, such as the sEH task.

2. While the paper demonstrates the advantages of the TLM method in complex environments, it does not sufficiently explore scenarios where TLM might underperform, particularly in simpler or highly structured tasks.

3. Some of the experimental environments, such as the sEH task, are relatively simple. Although the TLM method performs well in these tasks, validation in more complex or realistic environments would better illustrate the practical value of the approach.

4. Although the paper emphasizes the TLM method, exploring combinations with other backward policy approaches, such as maximum entropy policies, could offer more innovative directions for future research.

5. The paper could further discuss the impact of dynamic reward structures on the TLM method, particularly in handling complex, non-stationary reward distributions.

Meta-review: This paper proposes to optimize the backward policy of the GFlowNet. The proposed method is the GFlowNet training method, and the experimental results demonstrate that it can improve performance. This approach has a reasonable theoretical foundation and shows a unified framework for policy optimization. The paper is clearly written and well presented. The paper is above the bar of ICLR.

Decision: Accept (Poster)