marginparsep has been altered.
topmargin has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

Investigating Generalization Behaviours of Generative Flow Networks

Anonymous Authors1

††footnotetext:  1Anonymous Institution, Anonymous City, Anonymous Region,
Anonymous Country. Correspondence to: Anonymous Author
<EMAIL_ADDRESS>
Preliminary work. Under review by the International Conference on Machine
Learning (ICML). Do not distribute.

###### Abstract

Generative Flow Networks (GFlowNets, GFNs) are a generative framework for
learning unnormalized probability mass functions over discrete spaces. Since
their inception, GFlowNets have proven to be useful for learning generative
models in applications where the majority of the discrete space is unvisited
during training. This has inspired some to hypothesize that GFlowNets, when
paired with deep neural networks (DNNs), have favourable _generalization_
properties. In this work, we empirically verify some of the hypothesized
_mechanisms_ of generalization of GFlowNets. In particular, we find that the
functions that GFlowNets learn to approximate have an implicit underlying
structure which facilitate generalization. We also find that GFlowNets are
sensitive to being trained offline and off-policy; however, the reward
implicitly learned by GFlowNets is robust to changes in the training
distribution.

##  1 Introduction

Generative Flow Networks (GFlowNets, or GFNs) have emerged as a generative
modelling framework for learning unnormalized probability mass functions of
discrete objects such as graphs, sequences, or sets Bengio et al. (2021;
2023). They have especially shown promise for a wide range of problems and
applications where optimization is conducted over very large combinatorial
spaces. To name a few, discrete probabilistic modelling (Zhang et al., 2022),
molecular discovery (Bengio et al., 2021; Jain et al., 2023), biological
sequence design (Jain et al., 2022), and causal discovery (Deleu et al., 2022;
2023; Atanackovic et al., 2023) are areas where the discrete spaces may be
intractably large. In these applications, the model will realistically only
visit a small fraction of the overall state space. Therefore, understanding
how well GFlowNets model and assign probability mass to the unvisited areas of
the state space is a critical step. This is fundamentally a question of
_generalization_.

Generalization of deep neural networks (DNNs) in supervised learning and deep
learning has been extensively studied (Leshno et al., 1993; Pascanu et al.,
2013; Zhang et al., 2017; 2021; Kawaguchi et al., 2017). Although it is still
not entirely understood, many mechanisms and intuitive notions for
generalization have emerged as the favored schools of thought (Arpit et al.,
2017) and inspire this work. Due to their close relationship with GFlowNets,
works studying generalization in Reinforcement Learning also inspire this work
(Zhang et al., 2018; Packer et al., 2018; Cobbe et al., 2019; Bengio et al.,
2020).

It has been hypothesized that GFlowNets work well because they synergistically
leverage the generalization potential of DNNs to assign probability mass in
unvisited areas of the state space Bengio et al. (2021), and while the works
of Nica et al. (2022) and Shen et al. (2023) probe at this question, they only
do so superficially. Indeed, questions regarding the _mechanisms_ of
generalization in GFlowNets have yet to be systematically investigated and
wholly understood. With this comes a natural question – _how well_ do
GFlowNets generalize to _unvisited_ areas of the state space? Thus,
unravelling some intuitive notions on _why_ and _how_ GFlowNets generalize
forms the basis of this work.

We center our investigation around three primary hypotheses for generalization
in GFlowNets:

  1. 1.

GFlowNets generalize well when their self induced training _distributions_ are
nice, i.e. sampling on-policy from $P_{F}(s^{\prime}|s;\theta)$, or
proportionally to $R(s)$, is ideal.

  2. 2.

GFlowNets generalize well because the objects they are learning have
_structure_ ; $P_{F}(s^{\prime}|s)$ and $F(s)$ are not “arbitrary” functions.

  3. 3.

The difficulty for GFlowNets to generalize is modulated more by the
_complexity_ of the reward (as perceived by a DNN) than the properties of the
distribution induced by the reward.

To investigate these hypotheses, we devise an empirical investigation to
_disentangle_ some of the _mechanisms_ that lead to generalization when
training GFlowNets. We structure our study on the basis of making simplifying
assumptions such that we reduce the complexities/moving parts of training
GFlowNets. These “moving parts” may otherwise act as confounding variables
when investigating mechanisms. This approach helps us isolate specific factors
of variation which may play a role in _why_ GFlowNets generalize.

We describe the details of our experimental setup in §3, and to this end,
construct comprehensive benchmark tasks with well defined and tractably
computable $p(x;\theta),\forall x\in{\mathcal{X}}$, presented in §2.

We note that our proposed empirical study helps us disentangle some of the
mechanisms of generalization in GFlowNets, but does not necessarily determine
the true underlying causal order of these mechanisms. Nonetheless, we believe
this work can pave a road-map for how to approach testing generalization of
GFlowNets and view this as a key step towards understanding important
properties for this class of algorithms.

Our main contributions are summarized as follows:

  * •

We propose a set of benchmark graph generation tasks of varying difficulty,
useful for evaluating GFlowNets’ generalization performance.

  * •

We reify and validate some hypothesized characteristics of GFlowNet
generalization behaviour over discrete spaces. We accomplish this empirically
using benchmark tasks.

  * •

We identify and present a set of observations and empirical findings that form
a basis towards disentangling some of the _mechanisms_ for generalization of
GFlowNets.

For the sake of clarity, we define some terms that we want to make as least
ambiguous as possible in §A.

###  1.1 Generative Flow Networks

Generative Flow Networks (GFlowNets, GFNs) are a generative modeling framework
used to learn to sample from an unnormalized probability distribution. We will
refer to this unnormalized function as a positive _reward_ , $R(s)>0$.
Introduced in the discrete setting (Bengio et al., 2021), but since then
extended to continuous settings (Lahlou et al., 2023), GFlowNets work by
learning a sequential, constructive sampling policy. This policy, $P_{F}$, is
used to sample trajectories $\tau=(s_{0},...,s_{t})$ where states
$s\in{\mathcal{S}}$ are partially constructed objects, making the state space
a pointed directed acyclic graph (DAG)
$\mathcal{G}=({\mathcal{S}},{\mathcal{A}})$ where $(s\to
s^{\prime})\in{\mathcal{A}}\subset{\mathcal{S}}\times{\mathcal{S}}$ is a valid
constructive step. There is a unique initial state $s_{0}$.

While they can be expressed in multiple equivalent ways, we will think of
GFlowNets in this work through three main objects: the flow of a state
$F(s)>0$, and the forward and backward policies, $P_{F}(s^{\prime}|s)$ and
$P_{B}(s|s^{\prime})$. In a perfect GFlowNet, these functions are such that
for any valid partial trajectory $(s_{n},..,s_{m})$:

$F(s_{n})\prod_{i=n}^{m-1}P_{F}(s_{i+1}|s_{i})=F(s_{m})\prod_{i=n}^{m-1}P_{B}(s_{i}|s_{i+1})$
(1)

where for terminal (leaf) states, $F(s)=R(s)$ by construction. If the above
equation is satisfied, then starting at $s_{0}$ and sampling from $P_{F}$
guarantees to reach a leaf state $s_{t}\equiv x$ with probability $p(x)\propto
R(x),x\in{\mathcal{X}}$. We use ${\mathcal{X}}$ to denote the set of terminal
(leaf) states. Note that it may be useful to think of $P_{F}$ and $P_{B}$ as
representing fractions of flows going forward and backward through the (DAG)
network. It may also be useful to think of the flow going through edges:
$F(s\to s^{\prime})=F(s)P_{F}(s^{\prime}|s)$.

The so-called _balance condition_ in (1) leads to a variety of learning
objectives. In this work we primarily use the Sub-trajectory Balance
(specifically SubTB(1)) introduced by Madan et al. (2022), which takes the
above conditions, parameterizes the $\log$ flow and logits of policies, taking
the squared error over all possible subtrajectories:

$\mathcal{L}_{\mathrm{SubTB}}(\tau)=\\!\\!\\!\sum_{n<m\leq
T}\\!\left(\log\frac{F(s_{n})\prod_{i=n}^{m-1}P_{F}(s_{i+1}|s_{i})}{F(s_{m})\prod_{i=n}^{m-1}P_{B}(s_{i}|s_{i+1})}\right)^{2}.$
(2)

We note that trajectory balance (TB) is a special case of the above, where
only $n=0$ and $m=T$ are used. We also note from (1) that
$P_{F}(s^{\prime}|s)=P_{B}(s|s^{\prime})F(s^{\prime})/F(s)\propto
F(s^{\prime})$, and note that $F(s)$ is upper bounded (when $P_{B}$ of all its
descendant edges is 1) by the sum of the rewards of all its descendant leaves.
For a complete overview of GFlowNets, we refer readers to Bengio et al.
(2023).

##  2 A Benchmark Task for Investigating Generalization Behaviours of
GFlowNets

As a first step, we present a novel benchmark environment built on a series of
graph-based tasks. We use graphs as the foundation of the benchmark tasks as
they are a natural choice for compositional discrete objects, and a wide range
of combinatorial problems can be expressed as graph generation. Thus, we use a
graph environment to define tasks of varying difficulty. This helps us
standardize benchmarks for GFlowNets when investigating generalization and
learning $p(x)$ since the environment can be held constant while the reward
difficulty is varied. However, for completion, we also conduct experiments on
two commonly used environments and tasks in GFlowNet literature: the hypergrid
and sequence environments. In the following, we present our graph-based
benchmark environment and the corresponding tasks considered in this work.

###  2.1 A Small Graph Environment

We setup an environment with all possible graphs of size 7 and less, where the
nodes can be one of two colors. This environment has a total of 72296 states,
which allows us to compute the true distribution $p(x)$ and the learned
$p(x;\theta)$ for $x\in{\mathcal{X}}$ relatively quickly (in the order of 10
seconds on a GPU; see §C.1.2). We denote the set of graphs/states as
${\mathcal{X}}$. We define a number of reward functions that induce a variety
of challenges for models.

Reward complexity: We define three different reward functions, which we hope
to be of varying difficulty. The hardest function, cliques, requires the model
to identify subgraphs in the state which are 4-cliques of at least 3 nodes of
the same color. The second function, neighbors, requires the model to verify
whether nodes have an even number of neighbors of the opposite color. Finally,
the third function, counting, simply requires the model to count the number of
nodes of each color in the state. We fully define and show the distribution of
$\log R(x)$ of the respective tasks in §C.1. We can observe that the
distribution of $\log R(x)$ increases in complexity relative to the hardness
of the underlying task.111For example, the relative quantity of high reward
states (“sparseness” of rewards over states) and the general distribution of
$\log R(x)$ (“discontinuity” of reward distribution) are factors that can
drive learning difficulty.

We verify that our intuitive notion of hardness for these tasks translates to
graph neural networks (GNNs). To do so we train standard graph attention
networks (Veličković et al., 2017) with varying capacity to regress to the
(log) reward functions. We use a 90%-10% train-test split, and show the
resulting test error in Figure 1. We see that the ordering seems consistent
with the difficulty of the task.

Figure 1: Training a GNN on 3 different tasks with models of varying capacity.
Most of the variance comes from varying capacity. Dashed lines are for the
highest capacity models. Figure 2: Training a GFlowNet (online and on-policy)
on 4 different tasks. While ordering is mostly preserved, apparent difficulty
depends on the choice of metric.

Structural generalization: The so-called _structural generalization
hypothesis_ posits that function approximators like deep neural networks
(DNNs) capture regular patterns in the data (Arpit et al., 2017; Zhang et al.,
2021), i.e. structure. This enables models to make accurate predictions on
similarly structured inputs not present in the training set. In our
investigation, we reify this intuition by introducing a structured test set.
We generate a test set randomly (with seeds of course) by picking states with
at least 6 nodes (this choice is not arbitrary, see §C.1.3) and adding _all
their descendants_ to the test set. We do so iteratively until the test set is
of the desired size. We include further details in §C.1.3. Generalizing to a
“structured” test set this way poses a more challenging problem than simply
selecting test samples i.i.d. from the data distribution (Tossou et al.,
2023).

Metrics for Performance: We aim to evaluate how well GFlowNets model and learn
the distribution $p(x)$ over all states $x\in{\mathcal{X}}$. To do this, we
consider 2 distributional metrics: Jensen-Shannon (JS) divergence and mean
absolute error (MAE) between $\log p(x)$ and the learned $\log p(x;\theta)$.
We describe details for these metrics in §C.4. We also note that when
generating graphs, properly taking into account isomorphic actions is
essential to learning the right $p(x;\theta)$ (see Ma et al., 2023, for
reference).

###  2.2 Difficulty of Tasks is Preserved When Training GFlowNets Online and
On-policy

We now verify that training GFlowNets on our proposed graph environment and
defined tasks yields the task ranking that we’d expect. We train GFlowNets
online with SubTB$(1)$ and a uniform $P_{B}$.

To lower bound the modeling complexity, we add a fourth reward, one where
$R(x)=1,\forall x\in{\mathcal{X}}$.222Note that, hypothetically, this may be
harder than other rewards. For a constant reward, the model has to learn to
put _equal probability mass everywhere_ , which means being able to model the
entire state space. In contrast, more complex but “peakier” reward functions
may be in some sense easier to “get right” if one cares more about modeling
high-reward states, since there are presumably fewer of them. Note that for
this specific experiment there are no states hidden from the model, i.e. no
test set, since the model is able to sample from its policy and explore the
entirety of the state space. We will start using a test set in §4.1. Here, we
are only verifying the difficulty of the “usual” GFlowNet setup in the
respective tasks.

Figure 2 shows interesting results. The ordering of tasks is more or less
conserved, but depends on the metric we use to measure the discrepancy between
$p(x)$ and $p(x;\theta)$. Using a constant reward indeed lower bounds other
rewards, but not by much. This suggests that there is, unsurprisingly,
inherent difficulty in modeling the dynamics of the environment.

We find that there are two main axes of difficulty with respect to the task
(aside of course from the _scale_ of the problem, which here is kept
constant): (1) how _difficult_ the reward is to model, and (2) how the reward
is _distributed_. Hopefully, Figure 1 and Figure 2 are convincing evidence of
(1). We will come back to (2) later in §4.2.

Readers are also hopefully now on board that this benchmark is sufficiently
complex, and has interesting characteristics for empirical validation and
hypothesis testing purposes. We also consider 2 non-graph environments, a
hypergrid environment and sequence environment, with reasonably sized states
spaces such that we can tractably compute $p(x;\theta)$ exactly. This allows
us to investigate GFlowNet generalization behaviours in contexts other than
graphs. Since the graph, hypergrid, and sequence environments are all discrete
spaces, we expect general findings to be consistent across all tasks. We use
tasks and reward functions used in prior GFlowNet work (Bengio et al., 2021;
Malkin et al., 2022; Jain et al., 2023). We describe the setups and
experimental details of these environments and tasks in §C.2.

##  3 A Method for Disentangling Generalization Mechanisms of GFlowNets

We consider 3 main experimental settings: (1) distilling (regressing to) flow
functions, (2) memorization gaps in GFlowNets, and (3) offline and off-policy
training regimes. Each experimental setup is built on a series of simplifying
assumptions such that we can control for different moving parts and
complexities that are present when training GFlowNets. We now present this
experimental protocol, and then report and discuss our findings in §4.

###  3.1 Distilling (Regressing to) Flow Functions

Just as we can compute the distribution $p(x;\theta)$ over ${\mathcal{X}}$
exactly in this environment, we also compute edge flows $F(s\to s^{\prime})$
and $P_{F}(s^{\prime}|s)$ exactly. We refer to these as the _true_ flow and
forward policies, in opposition to the approximated $F(s\to
s^{\prime};\theta)$ and $P_{F}(s^{\prime}|s;\theta)$. Note that we
parameterize $F(s\to s^{\prime};\theta)$ and $P_{F}(s^{\prime}|s;\theta)$ as
mappings from ${\mathcal{S}}\to{\mathbb{R}}^{|n(s)|}$ where $n(s)$ is the
number of children of $s$.

In this set of experiments, we train DNNs by regressing to the _true_
flow333In this sense, we are “distilling” the _true_ flows into flow functions
parameterized by a DNN. $F(s\to s^{\prime})$ and forward policy
$P_{F}(s^{\prime}|s)$. This removes the use of the GFlowNet training objective
and the use of trajectories sampled from $P_{F}(s^{\prime}|s;\theta)$ as the
training data distribution, drastically simplifying the training procedure. As
such, we basically control for any non-ideal factors introduced by the
GFlowNet training objectives, for example, possible shortcomings due to
temporal credit assignment (Malkin et al., 2022).

We can now train models to regress to $\log F(s\to s^{\prime})$ or $\log
P_{F}(s^{\prime}|s)$. These are both vectors, so we regress to each value
independently and minimize the mean squared error. Note that the true
$P_{F}(s^{\prime}|s)=\mathrm{softmax}(\log F(s\to s^{\prime}))$. Regressing to
$\log F(s\to s^{\prime})$ is the same general problem as regressing to $\log
P_{F}(s^{\prime}|s)$, but with one caveat. When learning $F(s\to
s^{\prime};\theta)$, the model has to not only get the relative magnitudes of
each logit right, but also the absolute magnitudes right. In contrast, when
learning $P_{F}(s^{\prime}|s;\theta)$, because of the normalization within the
softmax, the model only has to get the relative magnitude of the logits right.
Because of this, we expect that learning $P_{F}(s^{\prime}|s)$ is easier,
although learning flow magnitudes could help with generalization. For these
experiments we set $P_{B}$ to be the uniform policy to get unique $F$ and
$P_{F}$.

Using this setup we are able to asses how fundamentally difficult these flow
functions are to learn. Furthermore, we can use our test set of unseen states.
This allows us to asses generalization performance for learning $p(x)$ (which
we are able to tractably compute) by directly computing distributional errors.
Because of the use of a test set, this is in some sense more challenging than
when training GFlowNets directly, i.e. online and on-policy, since standard
training of GFlowNets allows the model to potentially explore the entire space
and “see” the entire dataset444Readers may notice that in this task, because
we’ve computed the regression targets $F$ and $P_{F}$ exactly using the entire
state space, some information of the test set is leaking into the targets.
Since the purpose of this experiment is to asses how hard these functions
fundamentally are to learn, we knowingly allow the model to cheat..

Overall, in this setting we can probe the question of “ _do_ ” GFlowNets
generalize when learning $F$ or $P_{F}$ and assess Hypotheses 2 & 3. In the
following sub-section, we describe a second approach that can help us
investigate some mechanistic intuitions of _why_ flow functions might induce
generalization in GFlowNets, probing the question of “ _why do_ ” GFlowNets
generalize.

###  3.2 Memorization Gaps in GFlowNets

Problem | Data Coupling |  | Reward (data)
---
Structure
| Flow (environment)
---
Structure
Regress $R$ | $s\not\\!\perp\\!\\!\\!\\!\perp R(s)$ | ✓ | ✗
Regress $R$ | $s\\!\perp\\!\\!\\!\\!\perp\tilde{R}(s)$ | ✗ | ✗
Regress $P_{F}$ | $(s,s^{\prime})\not\\!\perp\\!\\!\\!\\!\perp P_{F}(s^{\prime}|s)$ | ✓ | ✓
Regress $P_{F}$ | $(s,s^{\prime})\not\\!\perp\\!\\!\\!\\!\perp\tilde{P}_{F}(s^{\prime}|s)$ | ✗ | ✓
Regress $P_{F}$ | $(s,s^{\prime})\\!\perp\\!\\!\\!\\!\perp P_{F}^{\mathrm{random}}$ | ✗ | ✗

Table 1: High level summary for the memorization gap experiments. Each row
lists the setup of an individual experiment showing the corresponding data
pair coupling and structure of learning problem.

We would like to probe our hypothesis on learning _true_ flows and their
contribution to generalization in GFlowNets. Here, we take on the perspective
of Zhang et al. (2017; 2021) and consider generalization as the act of _not
memorizing_. In other words, we can assess whether a model is generalizing or
not by measuring the gap (in _training_ performance) between training it on
structured data and training on random unstructured data. We use this notion
of _memorization gap_ to try and reason about learning structured flows and
their implications on generalization.

We devise an experiment inspired from Zhang et al. (2017; 2021), where we
train supervised models by, similarly to the previous section, regressing to
$R(s)$ and $P_{F}(s^{\prime}|s)$, but with various degrees of
“de-”structuring.

Consider simply learning to predict $R(s)$ with $R(s;\theta)$. If instead of
regressing from $s$ to $R(s)$ we shuffle the reward labels of each state $s$,
we end up with new $(s,\tilde{R}(s))$ pairings. This induces independence
between the data pairings, i.e. $s\\!\perp\\!\\!\\!\\!\perp\tilde{R}(s)$.
Intuitively, training a DNN with $s\\!\perp\\!\\!\\!\\!\perp\tilde{R}(s)$ will
force the DNN to memorize since we have removed any structured information
that may exist in the problem. We compare the training curve gap between
regressing to $R$ and $\tilde{R}$.

Now consider regressing to $P_{F}(s^{\prime}|s)$; there are two ways to
destructure this function. Recall that $P_{F}(s^{\prime}|s)\propto
F(s^{\prime})$, and that $F(s)$ is a function of the _reward_ of all its
descendants as well as of the _transition structure_ of the state space (the
ways to get to those descendants). Let ${\mathcal{D}}(s)$ be the set of
descendants of $s$. By shuffling $R$ into $\tilde{R}$, we thus remove the
dependence between $s$ and $\tilde{R}(s_{d})\forall s_{d}\in{\mathcal{D}}(s)$,
but keep the dependence between $s$ and _how to get to_ ${\mathcal{D}}(s)$.
The second way to de-structure $P_{F}$ is to simply regress to random logits,
which we will denote $P_{F}^{\mathrm{random}}$.

To recap, we train models on both the coupled and permuted or random data. We
consider 5 overall experiments (see Table 1), regressing to: $R(s)$ using
paired data, $R(s)$ using shuffled data, $P_{F}(s^{\prime}|s)$ using paired
data, $P_{F}(s^{\prime}|s)$ using shuffled data, and $P_{F}(s^{\prime}|s)$
using random logits $P_{F}^{\mathrm{random}}$ sampled (once before training)
uniformly at random from a magnitude-preserving range. Table 1 summarizes the
high level setup of each experiment and the corresponding data/flow structure
of the problem.

Through this setup, we can assess the _memorization gap_ that occurs when
training models on the respective settings and tasks. In particular, removing
the structure of the data but maintaining flow structure in the learning
problem allows us to assess the contribution that learning _true_ flows has on
generalization (i.e. _not memorization_) when training GFlowNets. Following
from Hypothesis 2, if state flows have interesting structure and learning
structured flows reduces the degree of memorization, perhaps this structure
inherently acts as a mechanism for generalization in GFlowNets.

###  3.3 Offline and Off-policy Training Regimes

| ${\mathbb{P}}_{\mathcal{X}}$ | Resemblance to Sampling $x$
---|---|---
Uniform | $\propto{\mathcal{U}}(x)$ | i.i.d.
Log-rewards | $\propto R(x)$ | from an ideal policy
Proxy for on-policy | $\propto p(x;\theta)$ |  $\propto$ to $P_{F}(s^{\prime}|s;\theta)$
Absolute error | $\propto|p(x;\theta)-p(x)|$ |  $\propto$ absolute loss
Squared log-error | $\propto(\log p(x;\theta)-\log p(x))^{2}$ |  $\propto$ squared log-loss (e.g. SubTB(1))

Table 2: Details for different training distributions
${\mathbb{P}}_{\mathcal{X}}$ for offline and off-policy experiments.

Lastly, we would like to investigate the effects of deviating from the self-
induced training distributions of GFlowNets on generalization. To do this, we
consider the setting of training GFlowNets _offline_ and off-policy given a
known dataset of final states ${\mathcal{X}}$. We sample
$x\sim{\mathbb{P}}_{\mathcal{X}}$, where ${\mathbb{P}}_{\mathcal{X}}$ is some
training distribution over ${\mathcal{X}}$. We consider different
distributions for ${\mathbb{P}}_{\mathcal{X}}$ that resemble different
practical approaches and techniques that are used for training GFlowNets (see
Table 2). This setting controls for the effect of sampling from
$P_{F}(s^{\prime}|s;\theta)$ in training GFlowNets.

Given a terminal $x$, we sample a trajectory $\tau=(s_{0},\dots,s_{f})$
backwards using a uniform $P_{B}$555Investigating the use of different
variants/distributions for $P_{B}$ is another interesting direction that may
give insight into the effects of the backward policy on the training dynamics
of GFlowNets (see Shen et al., 2023). We leave this for future work.. Hence,
we term this as _offline_ training, since we are avoiding sampling $x$ using
$P_{F}(s^{\prime}|s;\theta)$. In this way, we are still able to use a standard
GFlowNet objective, e.g. SubTB(1), to train $P_{F}(s^{\prime}|s;\theta)$ using
the backward sampled trajectories. Note that this setup can remove the data
non-stationarity that exists during normal training when using on-policy
samples from $P(s^{\prime}|s;\theta)$.

We also wish to assess to what degree deviating from
$P_{F}(s^{\prime}|s;\theta)$ during training will affect learning $p(x)$.
Understanding this may be useful to develop novel GFlowNet algorithms,
considering that in practice it seems hard to train models that deviate from
on-policy training (Rector-Brooks et al., 2023). To achieve this, we consider
a policy interpolation experiment where we sample trajectories from
$P_{\alpha}(s^{\prime}|s)=(1-\alpha)P_{F}(s^{\prime}|s;\theta)+\alpha
P_{{\mathcal{U}}}(s^{\prime}|s)$, where $0\leq\alpha\leq 1$ and
$P_{{\mathcal{U}}}(s^{\prime}|s)$ denotes a $P_{F}$ that goes to every
terminal state with equal probability, i.e. such that $p(x)\propto 1$. We
approximate $P_{{\mathcal{U}}}(s^{\prime}|s)$ using a policy pre-trained with
a constant reward (thus approximating a uniform $p(x)$).

We then train $P_{F}(s^{\prime}|s;\theta)$ _online_ instead of _offline_. Note
that this is a slightly different setup from what we discussed previously
since we re-introduce some of the effects of self-induced training
distributions of $P_{F}(s^{\prime}|s;\theta)$. We do this because, as we will
see, any degree of deviation from $P_{F}(s^{\prime}|s;\theta)$ during
_offline_ training (i.e. sampling $x\sim{\mathbb{P}}_{\mathcal{X}}$ and using
$P_{B}$ to construct trajectories) appears to have some negative effect. We
observe this even for values of $\alpha$ close to 0. We discuss this in
further detail in section 4.4.

This experimental setup helps us investigate the isolated effect of different
data distributions ${\mathbb{P}}_{\mathcal{X}}$ on training dynamics and
generalization by changing the data non-stationarity typically induced by
$P(s^{\prime}|s;\theta)$. To observe the effects on training dynamics, we
consider the setting of no test set (the model sees all of ${\mathcal{X}}$).
In contrast, to observe the effects on generalization, we consider performance
on unseen test set examples. This allows us to assess Hypothesis 1 by
observing the effects of various deviations from the normal self-induced
training distribution coming from $P_{F}(s^{\prime}|s;\theta)$.

##  4 Experiments

In this section, we report and discuss our observations and findings of the
experimental methodology defined in §3. We conduct all experiments reported in
this section over 3 random seeds. For training online and offline GFlowNets,
we use SubTB(1) and a uniform $P_{B}$.

Figure 3: Training a model to distill edge flows and policies on 4 different
tasks. It generally seems that (a) doing so recovers the intended distribution
fairly well, and (b) modeling $P_{F}$ appears easier than modeling $F$,
insofar as it recovers $p(x)$ better. Figure 4: JS-divergence and MAE gaps
between SubTB(1) trained GFlowNet model and $P_{F}$ distilled model.
Distilling to $P_{F}$ generally appears to yield lower distributional error,
with the exception of observing JS-divergence on the cliques task. Figure 5:
Training models distilled to $P_{F}$ and GFlowNet models trained online/on-
policy for a range of monotonic skew values $\gamma$. Transforming the
distribution of the reward to contain a larger proportion of high reward
values generally does not yield significant influence on performance for
approximating $p(x)$. Figure 6: Memorization gap training curves for counting,
neighbors, and cliques tasks. Maintaining _flow_ structure in the learning
problem (learning $P_{F}$ under shuffled $\tilde{R}$) induces a smaller
_memorization gap_ relative to the fully de-structured setting.

###  4.1 Distilling Flow Functions

We run the setup explained in §3.1. To evaluate the models, we look at the
distributional errors on $p(x;\theta)$. We use a 90%-10% train-test split. We
report results in Figure 3 and Figure 4.

First, we can see that learning $P_{F}$ and learning $F$ yield similar
difficulty in the sense that the distributional errors are systematically
close across tasks. Second, we see that the model is generalizing, getting
sometimes even better scores than a model trained online on the entire states
space (i.e. with no test set) as seen in Figure 4. The exception appears to be
when evaluating the performance of the cliques task using the JS divergence
metric, where the online GFlowNet yields marginally better performance than
the distilled models. While expected, this second result is relevant, because
it suggests that the difficulty of training a GFlowNet comes from both
learning to model the flow functions themselves _and_ performing temporal
credit assignment.

######  Observation 4.1.

Flow functions and flow policies are learnable and standard neural networks
generalize when predicting them.

We confirm these observation in sequence and grid environments (see §E.1). We
include further results in §D.1 where we assess how well GFlowNets assign
probability mass to unvisited states in the graph generation tasks.

###  4.2 Reward complexity and transformations

Consider a monotonic transform
${\mathcal{H}}_{\gamma}:\mathbb{R}\rightarrow\mathbb{R}$ that maps $\log
R^{\prime}(s)={\mathcal{H}}_{\gamma}(\log R(s)),\forall s\in{\mathcal{S}}$
such that $\log R^{\prime}(s)$ is skewed towards higher reward states as a
function of some parameter $\gamma$ (see §C.3 for further details). We use the
setup in §3.1 and also train online GFlowNets in the setting of skewed reward
distributions transformed by ${\mathcal{H}}_{\gamma}$. Results are shown in
Figure 5. We observe that as we increase $\gamma$ (increasing the skewness of
the distribution of $\log R$ towards higher reward values), both distilled and
online trained models are not significantly affected by the augmented reward
distribution.

######  Observation 4.2.

Monotonic transformations of $R(x)$ do not have detrimental impact on
generalization when learning $p(x)$ with GFlowNets.

This implies that applying minor reward transformations, which is a standard
technique in GFlowNet training, is not detrimental to _generalization_ about
$p(x)$ for a given task (of course, e.g. sparsifying a reward makes
exploration harder), supporting Hypothesis 3. We repeat this experiment on the
sequence environment and find consistent results (see §E.2).

###  4.3 Memorization Gaps in GFlowNets

We run the setup explained in §3.2 and Table 1; results are shown in Figures
6, with results involving noise corruption666This is another setting
considered by Zhang et al. (2017; 2021) to achieve de-structuring of the form
$s\\!\perp\\!\\!\\!\\!\perp\tilde{R}(s)$. (corrupted constant reward) and
online trained GFlowNets included in §D.3. For readers not familiar with the
results of Zhang et al. (2021)’s work, a typical observation in the
_memorization_ setup is for the training loss to initially plateau, until some
phase change occurs (one can imagine parameters have self-arranged to create
separate linear regions around each input point) and the training loss goes to
0. We observe a similar phenomenon here, regressing to the original $R$ and
$P_{F}$ induce curves with no plateau, while regressing to $\tilde{R}$ and
$P_{F}^{\mathrm{random}}$ initially have very flat plateaus until a phase
change occurs. This shows a very clear memorization gap, which was expected.

The most interesting result here is where $P_{F}$ is only partially de-
structured, and we regress to the _true_ $P_{F}$ of a shuffled reward
$\tilde{R}$ and retain the environment structure information. We observe that
while there is no flat plateau, (a) $P_{F}$ becomes harder to fit and (b)
there remains a phase change whereby learning accelerates after some slow
period. These results support Hypothesis 2; $P_{F}$ captures structure of both
the reward function and the environment, since removing the structure of
either induces _memorization gaps_.

######  Observation 4.3.

The existence of clear memorization gaps, when removing reward structure but
maintaining environment structure, suggests that there is informative
structure in flow functions and flow policies.

This finding implies that state flows may play a mechanistic role in _why_
GFlowNets generalize to unseen states. In §D.3 we show an extension of these
results where we consider rewards corrupted with noise (rather than shuffled).
We also show similar results for sequences and grids (see §E.3).

###  4.4 Generalization in Offline and Off-policy Training of GFlowNets

(a)  (b)

Figure 7: Evaluation curves for offline and off-policy trained GFlowNets on
the neighbors task for different choices of ${\mathbb{P}}_{{\mathcal{X}}}$.
(a) When training using the full dataset (no test set). (b) When training
using a subset of the full dataset (90%-10% train-test split). Complete
experiments for all graph generation tasks and evaluation metrics are shown in
§D.4.

We run the setup described in §3.3 and Table 2; results on the neighbors task
are shown in Figure 7(a) and Figure 7(b). We include full results for all
tasks in §D.4. First, we observe that GFlowNets trained in this regime are
dependent on environment/task difficulty, the choice of
${\mathbb{P}}_{{\mathcal{X}}}$, and on the evaluation metric. For example, the
best performing choice for ${\mathbb{P}}_{{\mathcal{X}}}$ on the counting task
is typically $x\sim{\mathcal{U}}$ or $x\sim\log R(x)$, while this is not
necessarily the case for neighbors and cliques (see §D.4 Figure 14).

Second, we observe that when training $P_{F}(s^{\prime}|s;\theta)$ offline and
off-policy in the full data (no test set) regime, convergence rate depends on
the choice of ${\mathbb{P}}_{{\mathcal{X}}}$ (see Figure 7(a)), sometimes to a
catastrophic degree. We see this for our usual distributional metrics for
approximating $p(x)$777It is possible that the drawbacks of this training
regime overwhelm our ability to make meaningful conclusions on the choice of
${\mathbb{P}}_{{\mathcal{X}}}$ for offline training on the neighbors and
cliques tasks.. However, when considering the MAE between $\log R(s)$ the
approximated $\log\hat{R}(s)$, we do not observe this lack of convergence. We
can measure this because GFlowNets learn to implicitly predict the reward:
$\log\hat{R}(s)=\log F(s;\theta)+\log P_{F}(s_{f}|s;\theta)$. This suggests
that although training $P_{F}(s^{\prime}|s;\theta)$ offline and off-policy may
be tricky, learning $F(s;\theta)$ jointly with $P_{F}(s^{\prime}|s;\theta)$
might be beneficial. This is supported by observing a reasonable estimate for
$\log\hat{R}(s)$ even when $P_{F}(s^{\prime}|s;\theta)$ fails to adequately
model $p(x)$ (the model is learning something).

We investigate this finding deeper by considering the off-policy interpolation
setup and find that for the more difficult tasks (neighbors and cliques) there
exists an apparent trade-off between on-policy and off-policy training when
estimating $p(x)$ versus $R$ (see §D.4 Figure 15). In particular, when
sampling on-policy (from $P_{F}(s^{\prime}|s;\theta)$, $\alpha=0$),
approximating $p(x)$ is reasonable. In contrast, approximating $R$ is best
when sampling uniformly off-policy (from $P_{{\mathcal{U}}}(s^{\prime}|s)$,
$\alpha=1$) even though error for $p(x)$ high.

######  Observation 4.4.

Deviating from the on-policy training distribution induced by
$P_{F}(s^{\prime}|s;\theta)$ can affect the ability of GFlowNets’ to learn
$p(x)$.

This finding implies that training GFlowNets offline and off-policy may come
with challenges depending on the task. However, we observe that the model
could be learning some informative structure in $F(s)$ that could help it
generalize (seen by reasonable estimates for $\log\hat{R}(s)$). We test this
hypothesis in our second offline and off-policy training setup, and we
consider performance on an unseen test (see Figure 7(b)). We observe that
indeed offline and off-policy trained GFlowNets reasonably approximate $R$ on
unseen states even when failing to adequately model $p(x)$.

######  Observation 4.5.

Learning $F(s)$ is informative to models, who implicitly predict reasonable
rewards, even when $P_{F}$ is unreasonable.

These results support Hypothesis 1; training GFlowNets online and on-policy is
ideal, where $P_{F}(s^{\prime}|s;\theta)$ converges to sampling trajectories
proportional to $R$. We repeat this experiment (with a left out test set) for
the sequence and hypergrid environments. Interestingly, we observe that in
these environments $p(x)$ is adequately modelled in the offline and off-policy
training regime (see §E.4). This could in-part be due to the reduced
complexity of these environments, which we establish by comparing the
difficulty of our tasks to the difficulty of fitting the constant reward (see
§E.1).

##  5 Conclusion

In this work, we conducted an empirical investigation into the generalization
behaviours of GFlowNets. We introduced a set of graph generation tasks of
varying difficulty to benchmark and measure GFlowNet generalization
performance. Our findings support existing hypothesized _mechanisms_ of
generalization of GFlowNets as well as add to our understanding of _why_
GFlowNets generalize. We found that GFlowNets inherently learn to approximate
functions which contain structure favourable for generalization. In addition,
we found that the reward implicitly learned by GFlowNets is robust to changes
in the training distribution, but that GFlowNets are sensitive to being
trained off-policy. We discuss limitations and future directions of this work
in §B.

## Impact Statement

The outcomes of the research conducted in this work have practical
implications for generative modelling in applications that include drug
discovery, material design, and general combinatorial optimization in
commercial settings. Because of this, we recognize the importance of
considering safety and alignment in these closely related domains. We believe
that advancements in GFlowNet research could result in models with improved
generalization, in turn leading to models that are easier to align. Another
important and useful application of GFlowNets is in the causal and reasoning
domains. We believe that advancements in these areas may lead to more
interpretable, easier to understand, and safer models. Lastly, we highlight
that the applications of our findings for the advancements of GFlowNets may
rely on building upon the hypotheses put forward in this work.

## Acknowledgments

The bulk of this research was done at Valence Labs as part of an internship,
using computational resources there. The authors are grateful to Yoshua
Bengio, Berton Earnshaw, Dinghuai Zhang, Johnny Xi, Jason Hartford, Philip
Fradkin, Cristian Gabellini, Julien Roy, and the Valence Labs team for
fruitful discussions and feedback. In addition, we acknowledge funding from
the Natural Sciences and Engineering Research Council of Canada and the Vector
Institute.

## References

  * Arpit et al. (2017) Arpit, D., Jastrzebski, S., Ballas, N., Krueger, D., Bengio, E., Kanwal, M. S., Maharaj, T., Fischer, A., Courville, A., Bengio, Y., et al.  A closer look at memorization in deep networks.  In _International Conference on Machine Learning_. PMLR, 2017.
  * Atanackovic et al. (2023) Atanackovic, L., Tong, A., Wang, B., Lee, L. J., Bengio, Y., and Hartford, J.  DynGFN: Towards bayesian inference of gene regulatory networks with gflownets.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Bengio et al. (2020) Bengio, E., Pineau, J., and Precup, D.  Interference and generalization in temporal difference learning.  In _International Conference on Machine Learning_. PMLR, 2020.
  * Bengio et al. (2021) Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y.  Flow network based generative models for non-iterative diverse candidate generation.  _Advances in Neural Information Processing Systems_ , 34:27381–27394, 2021.
  * Bengio et al. (2023) Bengio, Y., Lahlou, S., Deleu, T., Hu, E. J., Tiwari, M., and Bengio, E.  Gflownet foundations.  _Journal of Machine Learning Research_ , 24(210):1–55, 2023.
  * Cobbe et al. (2019) Cobbe, K., Klimov, O., Hesse, C., Kim, T., and Schulman, J.  Quantifying generalization in reinforcement learning.  In _International Conference on Machine Learning_ , pp. 1282–1289. PMLR, 2019.
  * Deleu et al. (2022) Deleu, T., Góis, A., Emezue, C., Rankawat, M., Lacoste-Julien, S., Bauer, S., and Bengio, Y.  Bayesian structure learning with generative flow networks.  In _Uncertainty in Artificial Intelligence_ , pp. 518–528. PMLR, 2022.
  * Deleu et al. (2023) Deleu, T., Nishikawa-Toomey, M., Subramanian, J., Malkin, N., Charlin, L., and Bengio, Y.  Joint bayesian inference of graphical structure and parameters with a single generative flow network.  _arXiv preprint arXiv:2305.19366_ , 2023.
  * Jain et al. (2022) Jain, M., Bengio, E., Hernandez-Garcia, A., Rector-Brooks, J., Dossou, B. F., Ekbote, C. A., Fu, J., Zhang, T., Kilgour, M., Zhang, D., et al.  Biological sequence design with gflownets.  In _International Conference on Machine Learning_ , pp. 9786–9801. PMLR, 2022.
  * Jain et al. (2023) Jain, M., Raparthy, S. C., Hernández-Garcia, A., Rector-Brooks, J., Bengio, Y., Miret, S., and Bengio, E.  Multi-objective gflownets.  In _International Conference on Machine Learning_ , pp. 14631–14653. PMLR, 2023.
  * Kawaguchi et al. (2017) Kawaguchi, K., Kaelbling, L. P., and Bengio, Y.  Generalization in deep learning.  _arXiv preprint arXiv:1710.05468_ , 1(8), 2017.
  * Lahlou et al. (2023) Lahlou, S., Deleu, T., Lemos, P., Zhang, D., Volokhova, A., Hernández-Garcıa, A., Ezzine, L. N., Bengio, Y., and Malkin, N.  A theory of continuous generative flow networks.  In _International Conference on Machine Learning_ , pp. 18269–18300. PMLR, 2023.
  * Leshno et al. (1993) Leshno, M., Lin, V. Y., Pinkus, A., and Schocken, S.  Multilayer feedforward networks with a nonpolynomial activation function can approximate any function.  _Neural networks_ , 6(6):861–867, 1993.
  * Li et al. (2020) Li, G., Xiong, C., Thabet, A., and Ghanem, B.  Deepergcn: All you need to train deeper gcns.  _arXiv preprint arXiv:2006.07739_ , 2020.
  * Ma et al. (2023) Ma, G., Bengio, E., Bengio, Y., and Zhang, D.  Baking symmetry into gflownets.  In _NeurIPS 2023 AI for Science Workshop_ , 2023.
  * Madan et al. (2022) Madan, K., Rector-Brooks, J., Korablyov, M., Bengio, E., Jain, M., Nica, A., Bosc, T., Bengio, Y., and Malkin, N.  Learning gflownets from partial episodes for improved convergence and stability. arxiv e-prints, page.  _arXiv preprint arXiv:2209.12782_ , 2022.
  * Malkin et al. (2022) Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Y.  Trajectory balance: Improved credit assignment in gflownets.  _Advances in Neural Information Processing Systems_ , 35:5955–5967, 2022.
  * Nica et al. (2022) Nica, A. C., Jain, M., Bengio, E., Liu, C.-H., Korablyov, M., Bronstein, M. M., and Bengio, Y.  Evaluating generalization in gflownets for molecule design.  In _ICLR2022 Machine Learning for Drug Discovery_ , 2022.
  * Packer et al. (2018) Packer, C., Gao, K., Kos, J., Krähenbühl, P., Koltun, V., and Song, D.  Assessing generalization in deep reinforcement learning.  _arXiv preprint arXiv:1810.12282_ , 2018.
  * Pan et al. (2023) Pan, L., Malkin, N., Zhang, D., and Bengio, Y.  Better training of gflownets with local credit and incomplete trajectories.  In _International Conference on Machine Learning_ , 2023.
  * Pascanu et al. (2013) Pascanu, R., Montufar, G., and Bengio, Y.  On the number of response regions of deep feed forward networks with piece-wise linear activations.  _arXiv preprint arXiv:1312.6098_ , 2013.
  * Rector-Brooks et al. (2023) Rector-Brooks, J., Madan, K., Jain, M., Korablyov, M., Liu, C.-H., Chandar, S., Malkin, N., and Bengio, Y.  Thompson sampling for improved exploration in gflownets.  _arXiv preprint arXiv:2306.17693_ , 2023.
  * Shen et al. (2023) Shen, M. W., Bengio, E., Hajiramezanali, E., Loukas, A., Cho, K., and Biancalani, T.  Towards understanding and improving gflownet training.  _arXiv preprint arXiv:2305.07170_ , 2023.
  * Shi et al. (2020) Shi, Y., Huang, Z., Feng, S., Zhong, H., Wang, W., and Sun, Y.  Masked label prediction: Unified message passing model for semi-supervised classification.  _arXiv preprint arXiv:2009.03509_ , 2020.
  * Tossou et al. (2023) Tossou, P., Wognum, C., Craig, M., Mary, H., and Noutahi, E.  Real-world molecular out-of-distribution: Specification and investigation.  2023\.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Veličković et al. (2017) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y.  Graph attention networks.  _arXiv preprint arXiv:1710.10903_ , 2017.
  * Zhang et al. (2018) Zhang, A., Ballas, N., and Pineau, J.  A dissection of overfitting and generalization in continuous reinforcement learning.  _arXiv preprint arXiv:1806.07937_ , 2018.
  * Zhang et al. (2017) Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O.  Understanding deep learning requires rethinking generalization.  In _International Conference on Learning Representations_ , 2017.
  * Zhang et al. (2021) Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O.  Understanding deep learning (still) requires rethinking generalization.  _Communications of the ACM_ , 64(3):107–115, 2021.
  * Zhang et al. (2022) Zhang, D., Malkin, N., Liu, Z., Volokhova, A., Courville, A., and Bengio, Y.  Generative flow networks for discrete probabilistic modeling.  In _International Conference on Machine Learning_ , pp. 26412–26428. PMLR, 2022.

##  Appendix A Vocabulary

  * •

The $p(x;\theta)$ distribution, when using GFNs we never explicitly model
$p(x;\theta)$; instead this measure is induced by the parameterization
($F(s\to s^{\prime};\theta)$ or $P_{F}(s^{\prime}|s;\theta)$) of a sequential
constructive policy. We use $p(x;\theta)$ as the shorthand for _the
distribution over $X$ induced by the chosen parameterization with parameters
$\theta$_. See also §C.1.2.

  * •

Online vs offline, a model is trained online when it is trained from data
generated during the training process – typically according to its own
parameters, i.e. for some model $p(x;\theta)$, we may use $X\sim p(x;\theta)$
to train $p(x;\theta)$. Conversely a model trained offline is trained from a
(usually fixed) data set, $X\sim\mathcal{D}$. It is possible to form a mixture
of online and offline data. GFNs are compatible with this paradigm (unlike RL
methods such as policy gradient methods).

  * •

On-policy vs off-policy, a model is trained on-policy if it is trained from
data generated according to the unperturbed distribution $p(x;\theta)$,
whereas it is trained off-policy if it comes from any other distribution. For
example, taking some actions at random would be considered off-policy
(although a mild version), but so would taking samples from an entirely
different distribution $p(x;\theta^{\prime})$ or from a dataset.

  * •

Self-induced distribution, when training a model we refer to the distribution
$p(x;\theta)$ as “self-induced”. This is mainly to emphasize that this
distribution changes as $\theta$ changes, which in online on-policy contexts
is due to the model generating its _own_ training samples (thus the
“self”-induced).

##  Appendix B Limitations and Future Work

##### Limitations:

We have shown that GFlowNets tend to generalize when learning unnormalized
probability mass functions for approximating $p(x)$ over discrete spaces. In
particular, we have investigate the generalization behaviours of GFlowNets in
the context of distributional errors for $p(x)$. Because of the combinatorial
nature of computing $p(x)$ exactly, and likewise computing $p(x;\theta)$ from
$P(s^{\prime}|s;\theta)$, we are limited in the fact that we need to work
within reasonably sized discrete spaces (as those we have proposed) to study
GFlowNet generalization for approximating $p(x)$. Secondly, we don’t
explicitly test generalization in online and on-policy trained GFlowNets. This
in part a consequence of requiring an environment and state space large enough
such that a sufficient quantity of unseen states can be produced, while also
being able to purposely hide the visited states from parameters updates of the
GFlowNet. Lastly, although our proposed problems and environments have
structure that induces difficult and interesting generalization tasks, they
still may remain far from the structure of the real world.

##### Future Directions:

We hope to have formed a sound set of findings and observations that may lead
to future research in understanding and formalizing generalization in
GFlowNets. For instance, our set of findings could be used as starting points
for formalizing some of our notions and intuitions for GFlowNet generalization
(and the corresponding mechanisms) into mathematical theory. Another direction
that can stem from our work is to further empirically investigate GFlowNet
generalization in the online and on-policy setting. Since we don’t explicitly
look at GFlowNet generalization in the online and on-policy training regime
(i.e. since we don’t hide any states from online trained GFlowNets); it would
be interesting in future work to explicitly investigate this. However, it is
important to note that this is not necessarily a trivial task (as described in
limitations).

##  Appendix C Experimental Details

###  C.1 Details for constructing graph benchmark task

####  C.1.1 Reward functions

Here are the exact log-reward functions we use, as implemented in Python using
the `networkx` (`nx`) and `numpy` (`np`) libraries.

cliques counts the number of 4-cliques in a graph such that at least 3 of the
4 nodes of the cliques share the same color. We then subtract to that number
the total number of cliques in the graph, but add the number of nodes. This is
so that the maximal log-reward is 0. We clip log-rewards below -10.

    def cliques(g, n=4):
        cliques = list(nx.algorithms.clique.find_cliques(g))
        # The number of cliques each node belongs to
        num_cliques = np.bincount(sum(cliques, []))
        colors = {i: g.nodes[i]["v"] for i in g.nodes}

        def color_match(c):
            return np.bincount([colors[i] for i in c]).max() >= n - 1

        cliques_match = [float(len(i) == n) * (1 if color_match(i) else 0.5) for i in cliques]
        return np.maximum(np.sum(cliques_match) - np.sum(num_cliques) + len(g) - 1, -10)

neighbors looks at all the neighbors of every node, counting the number of
nodes with an even number of neighbors of the opposite color. This total is
modified as a function of the number of nodes in order to produce a nice log-
reward distribution between 0 and -10.

    def neighbors(g):
        total_correct = 0
        for n in g:
            num_diff_colr = 0
            c = g.nodes[n]["v"]
            for i in g.neighbors(n):
                num_diff_colr += int(g.nodes[i]["v"] != c)
            total_correct += int(num_diff_colr % 2 == 0) - (1 if num_diff_colr == 0 else 0)
        return np.float32((total_correct - len(g.nodes) if len(g.nodes) > 3 else -5) * 10 / 7)

counting simply counts the number of red and green nodes, red nodes being
”worth” more, and again modifies this count in order to produce a nice 0 to
-10 log-reward distribution.

    def counting(g):
        ncols = np.bincount([g.nodes[i]["v"] for i in g], minlength=2)
        return np.float32(-abs(ncols[0] + ncols[1] / 2 - 3) / 4 * 10)

Figure 8: Distribution of log rewards for the given tasks: (left) counting,
(middle) neighbors, and (right) cliques.

####  C.1.2 Computing $p(x;\theta)$, $F(s)$, $F(s\to s^{\prime})$, and $P_{F}$

We compute the probability of a model with parameters $\theta$ sampling $x$,
$p(x;\theta)$ via what is essentially Dynamic Programming. We visit states $s$
in topological order starting from $s_{0}$. For every child of $s$,
$s^{\prime}$, (valid transition $s\to s^{\prime}$) we add
$p_{v}(s)P_{F}(s^{\prime}|s;\theta)$ to the probability of visiting
$p_{v}(s^{\prime})$, eventually accumulating visitation probability from all
the parents of $s^{\prime}$. This is possible because we are visiting states
in topological order. Note that we start with $p_{v}(s_{0})=1$ and $p_{v}$ set
to 0 for every other state. $p(s;\theta)$ is computed as
$p_{v}(s)P_{F}(\mathrm{stop}|s)$.

The above is fairly easy to batch, requires one forward pass per state.
Computing $P_{F}$ is the most expensive operation and can be batched, as long
as values are accumulated by respecting a topological order afterwards. We
also actually store values on a $\log$ scale, using logaddexp operations for
numerical stability.

To compute _true_ flow functions, we do the reverse, visiting the DAG in
reverse topological order. For every state $s$, we add
$F(s)P_{B}(s^{\prime}|s)$ units of flow to every parent $s^{\prime}$ of $s$,
starting at leaves (due to the reverse topological order) where $F(s)=R(s)$.
We similarly set the value of the edge flow $F(s\to
s^{\prime})=F(s)P_{B}(s^{\prime}|s)$. $P_{F}$ is simply the softmax of edge
flows. Note that we use a uniform $P_{B}$, as explained in the main text.

####  C.1.3 Constructing test set

Why 6 nodes? Considering the maximal number of nodes is 7, 6 may seem “too
close”. On the other hand, picking a graph and excluding its subtree means
that this subgraph will never appear in the training set. This is quite
aggressive, for example creating a test set of 10% of the data only requires
35 such graphs whose subtrees are excluded. These graphs have an average of
$272\pm 205$ descendants. Choosing graphs of 5 nodes would exclude an average
of $\sim 4700$ graphs per pick, meaning that a test set of 10% of the data
would only stem from the exclusion of 2 or 3 graphs. This may not have the
diversity we desire, thus our choice of 6 nodes.

Finally, we believe this is an interesting choice that relates to real-world
applications of GFlowNets. In drug-discovery of small molecules, given the
enormous size of the state space, it is likely for most subgraphs of a
sufficient size to never appear in the training set (because a model can only
train for so many iterations in practice). In a graph generation context we
are thus interested in how a model generalizes to subgraphs it’s never seen.
In this proposed benchmark, 67994 of the 72296 states are graphs of 7 nodes;
only excluding graphs of 6 or 7 nodes thus covers most of state space.

###  C.2 Details for hypergrid and sequence environments and tasks

####  C.2.1 ($M\times M$) Hypergrid:

Grid environments have been used in GFlowNet papers since their inception
(Bengio et al., 2021) as a sanity check environment, presumably inheriting
this custom from Reinforcement Learning. We continue the tradition and report
results on a 2D grid environment. In it, the agent has 3 actions, move in the
$x+$ direction, move in the $y+$ direction, or stop. The agent navigates on a
grid of size $M$ (64 in our experiments), and is forced to stop if any of the
coordinates reaches $M-1$.

As a reward signal we use functions used in past work (Bengio et al., 2021;
Jain et al., 2023).

####  C.2.2 Bit Sequences:

Auto-regressive sequence environments have been used in a variety of past
GFlowNet work Malkin et al. (2022); Madan et al. (2022); Pan et al. (2023). In
this work, we consider a bit sequence environment akin to that originally
introduced by Malkin et al. (2022). We consider a reward of the form
$R(x)=\exp(\frac{-\min_{m\in{\mathcal{M}}}d(x,m))}{l}\times 10)$, where $d$ is
the Levenshtein edit distance between an input sequence $x$ and the closest
“mode” sequence $m$, ${\mathcal{M}}\subset{\mathcal{X}}$ is the set of mode-
sequences, and $l$ is the max sequence length. We select $|{\mathcal{M}}|=60$
sequences from ${\mathcal{M}}$ uniformly at random given the set
${\mathcal{X}}$ of all possible bit sequences up to length $l=15$ to be the
“mode” sequences (we do this once at the start of each run). This yields a
discrete state space of $65,535$ states, comparable in size to our graph
generation environment.

###  C.3 Monotonic reward transformation

Figure 9: Example of monotonic transform ${\mathcal{H}}_{\gamma}$. (left)
distributions of $\log R(x)$ and $\log R^{\prime}(x)$ for true log-rewards and
skewed log-rewards (using $\gamma=0.2$). (right) plot of $\log R(x)$ versus
$\log R^{\prime}(x)$ for various values of $\gamma$. Points above the
$\gamma=0$ line yield skew towards higher log-reward values.

We consider a monotonic transform of the form ${\mathcal{H}}_{\gamma}(\log
R(x))=\log R(x)e^{(-\gamma\log R(x))}$. As described earlier, the transform
${\mathcal{H}}_{\gamma}$ maps $\log R^{\prime}(s)={\mathcal{H}}_{\gamma}(\log
R(s)),\forall s\in{\mathcal{S}}$ such that $\log R^{\prime}(s)$ is skewed
towards higher reward states as a function of some parameter $\gamma$ (see
§C.3 for further details. Note, ${\mathcal{H}}_{\gamma}$ is monotonic for
values of $\log R(x)\leq 0$ and $\gamma\geq 0$. In our graph and sequence
generations tasks, $\log R(x)\leq 0,\forall x\in{\mathcal{X}}$, hence we are
able to use ${\mathcal{H}}_{\gamma}$ as a monotonic transform for $\log R(x)$.
We show an example of ${\mathcal{H}}_{\gamma}$ applied to the cliques task in
Figure 9 for various value sof $\gamma$.

###  C.4 Evaluation metrics

To measure the generative modelling performance of GFlowNets we consider the
Jensen-Shannon (JS) divergence and the MAE between the learned $\log
p(x;\theta)$ and the true $\log p(x)$. For JS divergence, we compute:

$\text{JS}(p(x),p(x;\theta))=\frac{1}{2}\text{KL}(p(x)\|Q)+\frac{1}{2}\text{KL}(p(x;\theta)\|Q),$
(3)

where $Q=\frac{1}{2}(p(x)+p(x;\theta))$. This reduces to

$\text{JS}(p(x),p(x;\theta))=\frac{1}{2}\sum_{x\in{\mathcal{X}}}\left(p(x;\theta)(\log
p(x;\theta)-\log p(x))+p(x)(\log p(x)-\log p(x;\theta))\right).$ (4)

For the MAE, we simply take the absolute error $|\log p(x;\theta)-\log p(x)|$
and average over the cardinality of the state space.

###  C.5 Model architectures

For all graph experiments we use a modified graph transformer (Veličković et
al., 2017; Shi et al., 2020). On top of the normal attention mechanism, we
augment the input of each layer with the output of one round of message
passing (using layers from the work of Li et al., 2020) with a sum
aggregation–we found that this was useful in tasks where counting was
required. We use 8 layers with 128-dimensional embeddings and 4 attention
heads; we use this architecture after having tried different numbers of layers
and embeddings in order to validate our task; this is coincidentally
represented in Fig. 1.

For sequence tasks, we use a vanilla transformer (Vaswani et al., 2017) with 4
layers of 64 embeddings and 2 attention heads. We did not search for
hyperparameters, since this setup has been used in prior work (Malkin et al.,
2022).

For grid tasks we use a LeakyReLU MLP with 3 layers of 128 units. The input is
a one hot representation of each coordinate.

###  C.6 Implementation details

Our experiments are implemented in Pytorch and Pytorch Geometric. Our code is
available at https://github.com/lazaratan/gflownet-generalization.

##  Appendix D Additional Graph Task Experiments

###  D.1 GFlowNets learn rank preservation of unseen states

Task | Supervised |  Distilled $F$ |  Distilled $P_{F}$ |  offline TB |  offline subTB |  offline FM
---|---|---|---|---|---|---
| Test set Spearman correlation
cliques | 0.92 $\pm$ 0.02 | 0.92 $\pm$ 0.01 |  0.93 $\pm$ 0.00 | 0.88 $\pm$ 0.01 | 0.90 $\pm$ 0.01 | 0.86 $\pm$ 0.03
neighbors |  0.98 $\pm$ 0.00 |  0.98 $\pm$ 0.00 |  0.98 $\pm$ 0.00 | 0.96 $\pm$ 0.00 | 0.96 $\pm$ 0.00 | 0.96 $\pm$ 0.00
count |  0.98 $\pm$ 0.00 |  0.98 $\pm$ 0.00 |  0.98 $\pm$ 0.00 | 0.95 $\pm$ 0.01 | 0.95 $\pm$ 0.00 | 0.96 $\pm$ 0.00
| Test set top-100 Spearman correlation
cliques | 0.66 $\pm$ 0.02 | 0.74 $\pm$ 0.10 |  0.83 $\pm$ 0.04 | 0.67 $\pm$ 0.10 | 0.69 $\pm$ 0.11 | 0.53 $\pm$ 0.13
neighbors |  0.87 $\pm$ 0.00 |  0.87 $\pm$ 0.00 |  0.87 $\pm$ 0.00 | 0.86 $\pm$ 0.01 | 0.86 $\pm$ 0.00 | 0.86 $\pm$ 0.01
count |  0.44 $\pm$ 0.00 |  0.44 $\pm$ 0.00 |  0.44 $\pm$ 0.00 | 0.24 $\pm$ 0.06 | 0.24 $\pm$ 0.06 | 0.37 $\pm$ 0.02
| Avg Rank of optimal objects
cliques (15) | 21.5 $\pm$ 14.5 | 39.3 $\pm$ 30.2 |  13.3 $\pm$ 5.3 |  18.1 $\pm$ 4.9 |  12.1 $\pm$ 2.0 |  17.7 $\pm$ 5.2
neighbors (48) |  23.5 $\pm$ 0.0 |  23.5 $\pm$ 0.0 |  23.5 $\pm$ 0.0 | 43.4 $\pm$ 5.5 | 43.6 $\pm$ 9.4 | 40.7 $\pm$ 2.2
count (7) |  3.0 $\pm$ 0.0 |  3.0 $\pm$ 0.0 |  3.0 $\pm$ 0.0 | 62.7 $\pm$ 12.3 | 61.1 $\pm$ 15.1 | 25.6 $\pm$ 4.2
Table 3: Comparing the rankings given by different ways of training a model on
a dataset. Standard deviations are over 4 runs. In the last three rows the
number in brackets is the number of optimal objects in the test set.

We want to assess more precisely where GFNs put probability mass, in
particular in states they’ve never visited. More specifically though, GFNs are
often used to find _likely_ hypotheses, i.e. generate objects “close enough”
to the argmax(es).

Consider the following scenario, which is a common way to use GFNs. Some
dataset of objects and scores $\mathcal{D}=\\{(x_{i},y_{i})\\}$ is given to
us. We train a reward proxy to regress to $R(x_{i};\theta)=y_{i}$, and then
train a GFN on $R(x;\theta)$ in order to generate $x$s (and commonly, find the
most “interesting” $x$). If we assume that supervised learning is “as good as
it gets” to approximate $R$, then the task of finding
$\operatorname*{arg\,max}_{x}R(x)$ by finding
$\operatorname*{arg\,max}_{x}R(x;\theta)$ should also be as good as it gets.
This prompts us to ask, how close are GFNs to this ideal?

In the following experiment, we look at three measures. First, the Spearman
rank correlation on the probabilities $p(x;\theta)$ (or $R(x;\theta)$ for the
supervised model) of the test set, second we do the same for the top-100
graphs in the test set, and finally we look at the mean rank of the set of
optimal $x$s in the test set (e.g. if there are 7 optimal graphs in the test
set and the model perfectly predicts their rank to be $[0,1,2,3,4,5,6]$ then
the mean rank is $(0+1+2..+6)/7=3.0$).

These measures are intended to be proxies for how likely it is that a model
trained on some dataset would be able to generate objects close enough to the
optimal objects.

This third measure is also inspired by the following fact: it has been
observed in a few setups that even though trajectory balance objectives
appeared to fit the distribution much better than flow matching, flow matching
can be more efficient at producing high-reward samples (which seems counter-
intuitive).

We thus compare distilling edge flows to policies, as well as TB objectives to
the FM objective. Specifically for GFN objectives (TB, subTB, FM), we use an
offline paradigm where we sample trajectories by sampling an $x$ uniformly at
random from $\mathcal{D}$ and going backwards with $P_{B}$ uniformly at random
(this is presumably a somewhat ideal training condition).

We report the results in Table 3. The results are quite dense, so let us make
a few observations:

  * •

When considering the entire test set, the supervised model is indeed the gold
standard. Offline GFN models are slightly worse. This isn’t too surprising,
and is in fact reminiscent of results related to the difficulty of training
value functions via TD in RL (Bengio et al., 2020).

  * •

When considering either the top-100 or the optimal objects, we see surprising
results for the cliques task (the hardest task)

    * –

for the top-100 objects, the distilled models do better, and the offline TB
models are on-par with (and occasionally better than) the supervised model

    * –

for the rank of optimal objects, not only does distilled $P_{F}$ do better,
the offline GFN methods do better as well.

Recall that some test set information is leaking into training; it may explain
the advantage of distilled models, but not the advantage of offline GFN
models.

  * •

TB and subTB are generally better at fitting _ranking_ than FM, _except_ when
it comes to optimal objects. This reproduces past observations, suggesting
flow matching over-weights high-reward states, even outside of its training
data; we unfortunately do not have a coherent explanation:

    * –

$F(s,s^{\prime})$ distillation being worse than $P_{F}$ distillation means
modeling $F$ is probably not the advantage;

    * –

This phenomenon is consistent across training, although not on every task.

    * –

This phenomenon is amplified when _not_ correcting for idempotent actions.
This is expected (Ma et al., 2023), but useful to measure.

We repeat these three measures with different train-test ratios, shown in
Figure 19, where we see that the results are as expected, and consistent with
Table 3. To summarize: _GFlowNets implicitly learn to put more mass on high-
reward objects than low-reward objects outside their training set._

The above is interesting, but is still just a statement about _ranking_
between unseen states. Obviously this says nothing about the total probability
mass given to unseen states.

###  D.2 Distilled flow experiments over only unseen states

(a)  (b)

Figure 10: MAE error for only unseen states on constant, counting, neighbors,
and clique graph generations tasks. (a) evaluation for model trained to
distill edge flows and policies. (b) evaluation an online trained GFlowNet via
SubTB(1). We see that considering only the unseen states for MAE metric leads
to comparable observation to the case of evaluating MAE over all states. This
is expected since error on the unseen states should be driving the model’s
overall evaluation performance when evaluating on the entire state space.
Note, for the online GFlowNet, there is no guarantee that the model has not
“seen” the left out unseen states in this evaluation. Figure 11: MAE gap
between online trained GFlowNet via SubTB(1) and distillation trained model
regressing to $P_{F}$ on only unseen states. We observe that when considering
only unseen states, results are consistent with evaluation over the entire
state space. This is also expected since error on the unseen states should be
driving the model’s overall evaluation performance when evaluating on the
entire state space.

###  D.3 Additional memorization experiments

Figure 12: Memorization gap training curves for the constant reward on the
graph generation task. In this experiment, we consider the case where Gaussian
noise corrupts the constant reward signal of the form
$\tilde{R}(x)=R(x)+\epsilon,\epsilon\sim{\mathcal{N}}(0,\sigma)$, thus
inducing de-structuring. Here we consider $\sigma=2$. We observe behaviour
consistent with that found in §4.3. Figure 13: Memorization gap training
curves for counting, neighbors, and cliques tasks for distilled (regressing to
$P_{F}$) and online trained GFlowNet. We observe results are consistent with
findings in §4.3 for the online trained GFlowNet.

###  D.4 Full Experimental results for offline and off-policy training of
GFlowNets

Figure 14: Evaluation curves for offline and off-policy trained GFlowNets on
counting, neighbors, and cliques graph generations tasks for different choices
of ${\mathbb{P}}_{{\mathcal{X}}}$ when training using the full dataset (no
test set). Model performance is dependent on the choice of
${\mathbb{P}}_{{\mathcal{X}}}$. When considering evaluation on the JS
divergence and MAE metrics, depending on task and some choices of
${\mathbb{P}}_{{\mathcal{X}}}$, $p(x)$ is not adequately approximated.
However, regardless of task, offline and off-policy trained GFlowNets appear
to robustly learn $R$. Figure 15: Results for policy mixing interpolation
experiments. Here, we mix
$P_{\alpha}(s^{\prime}|s)=(1-\alpha)P_{F}(s^{\prime}|s;\theta)+\alpha
P_{{\mathcal{U}}}(s^{\prime}|s)$, where $0\leq\alpha\leq 1$. For $\alpha=0$,
the model is fully on-policy, sampling trajectories from
$P_{F}(s^{\prime}|s;\theta)$. For $\alpha=1$, the model is sampling fully off-
policy using $P_{{\mathcal{U}}}(s^{\prime}|s)$. This approximates the
behaviour of training offline and off-policy when using
$x\sim{\mathcal{U}}(x)$ (i.e. uniform sampling of $x$). Here we can observe
the effects of deviating from on-policy samples during training. We see that
at times, specifically for the more difficult tasks (neighbors and cliques),
even small degrees of deviation from $P_{F}(s^{\prime}|s;\theta)$ can lead to
worsened performance when approximating $p(x)$ (seen by sensitivity in
distributional JS divergence and MAE metrics). In contrast, approximating $R$
appears to improve as the model samples more uniformly over ${\mathcal{X}}$
(if not entirely uniform, i.e. when $\alpha=1$). Interestingly, the counting
task appears generally invariant to minor changes in $\alpha$. Figure 16:
Evaluation curves for offline and off-policy trained GFlowNets on counting,
neighbors, and cliques graph generations tasks for different choices of
${\mathbb{P}}_{{\mathcal{X}}}$ when training using a subset of the full
dataset (90%-10% train-test split). In this setting, agnostic to the choice of
${\mathbb{P}}_{{\mathcal{X}}}$ and graph generation task, we observe the
GFlowNet models struggles to converge when considering JS divergence and MAE
distributional metrics for approximating $p(x)$. However, we observe the
GFlowNet models remains robust when learning $R$.  Figure 17: Continuation of
Figure 16, but evaluated over only the unseen sates (test states). We observe
performance is consistent with the evaluations over the entire state space
(Figure 16).

###  D.5 Offline GFlowNets inadequately assign probability mass

Figure 18: The total probability of the test set (dotted) and of the top-100
objects within the test set (full lines); means of 4 seeds. Each color is a
different training/test split ratio. As the model has access to more data, it
seems to overfit _more_ ; it gives more total probability to very few states,
but particularly, to high reward states. We also plot in the bottom right
(stars) what the “true” total probability of the top-100 objects should be (it
varies because of the random seeds and test set size, we show the seed
average).

Let’s dig into the experiment of §D.1. We’ve established that offline GFNs
manage to rank unseen states fairly well, but this says nothing about how much
probability it gives those states, only that the relative probabilities are
well behaved.

In Figure 18 we show how much _total probability_ (i.e.
$\sum_{i\in\mathrm{test}}p(x_{i};\theta)$) is allocated to the test set. In
particular, we show the total probability of the whole test set, as well as of
the top-100 objects within the test set.

We believe the following result is unexpected: as the training set grows
larger, high-reward unseen objects take up more total probability.

This seems to be counter to our intuition on generalization; which suggests
that as a model gets more training data, it should extrapolate better (which
for a probability model should mean that $p(x;\theta)$ should get closer to
$p(x)$).

This result seems to be both terrible news and great news. First the bad news;
this result suggests that naively applying GFN objectives to an offline
training regime simply doesn’t work. Indeed, intuitively because $Z$ is a free
variable in GFNs, one of two things seems likely to happen: either the model
ends up learning $Z=\sum_{i\in\mathrm{train}}R(x_{i})$ and giving 0
probability to the test set, or the model ends up learning an arbitrary large
$Z$ and giving 0 probability to the training set. The good news; while the
latter seems to happen (the training set ends up with fairly low total
probability), there seems to be _structure_ to which states receive the most
probability. This again seems particularly useful in the scenario where we are
looking for the most interesting (high-reward) $x$.

Figure 19: Rank-related metrics during offline GFN training and supervised
regression as a function of the size of the training set. Averages are over 4
seeds (which influence the construction of the test).

##  Appendix E Sequence and Hypergrid Experiments

###  E.1 Online and distilled training

####  E.1.1 Online trained models

Figure 20: Training a GFlowNet (online and on-policy) on 4 different hypergrid
tasks. Corners, the most frequently used task in the hypergrid environment,
appears to be most difficult for learning $p(x)$. Figure 21: Training a
GFlowNet (online and on-policy) on 2 different sequence tasks. The edit
distance task appears to be significantly more difficult than the constant
task for learning $p(x)$. Given the auto-regressive nature of the sequence
environment, it is not unsurprising how well a GFlowNet performs on the
constant task (i.e. the sequence environment is not challenging on its own).
In contrast, we reify that the edit distance reward is indeed a challenging
task in this environment. Figure 22: MAE error for only test states on the
sequence task for online and on-policy. Note that for this experiment, the
test states are not necessarily unvisited by the model

####  E.1.2 Distilled flows and forward policy models

Figure 23: Training a MLP on 4 different tasks to regress to $F$ (left) and
$R$ (right) in the hypergrid environment. Task difficulty between supervised
models and online and on-policy trained GFlowNets appears consistent. Figure
24: Training a model to distill edge flows and policies on 4 different
hypergrid tasks for 90%-10% train-test split (top) and 60%-40% train-test
split (bottom). It generally seems that (a) doing so recovers the intended
distribution fairly well, and (b) modeling $P_{F}$ appears easier than
modeling $F$, insofar as it recovers $p(x)$ better, in the hypergrid
environment. Figure 25: Training a model to distill edge flows and policies on
2 different sequence tasks for 90%-10% train-test split. It generally seems
that (a) doing so recovers the intended distribution fairly well, and (b)
modeling $P_{F}$ appears easier than modeling $F$, insofar as it recovers
$p(x)$ better, in the sequence environment for the edit distance reward. For
the constant reward, there appears to be no apparent difference in difficulty
between learning $P_{F}$ or $F$. Figure 26: MAE error for only test states on
the sequence task for distilled model, regressed to $P_{F}$. Note that for
this experiment, the test states are unvisited by the model.

###  E.2 Reward transformation

Figure 27: Training models distilled to $P_{F}$ and GFlowNet models trained
online/on-policy for a range of monotonic skew values $\gamma$ on the sequence
task. Transforming the distribution of the edit distance reward to contain a
larger proportion of high reward values generally improves the performance for
approximating $p(x)$.

###  E.3 Memorization gap experiments

####  E.3.1 Memorization in hypergrids

Figure 28: Memorization gap training curves for branin, currin, and corners
tasks on the hypergrid environment. Maintaining _flow_ structure in the
learning problem (learning $P_{F}$ under shuffled $\tilde{R}$) generally
induces a smaller _memorization gap_ relative to the fully de-structured
setting. Note that for the branin reward, there is not apparent difference
between learning the true $R$ and the shuffled $\tilde{R}$, suggesting the
reward alongside the hypergrid environment induces a task with minimal
difficulty.

####  E.3.2 Memorization in sequences

Figure 29: Memorization gap training curves for the edit distance task on the
sequence environment. Maintaining _flow_ structure in the learning problem
(learning $P_{F}$ under shuffled $\tilde{R}$) generally induces a smaller
_memorization gap_ relative to the fully de-structured setting.

###  E.4 Offline and off-policy training

####  E.4.1 Offline and off-policy in hypergrids

Figure 30: Evaluation curves for offline and off-policy trained GFlowNets on
branin, currin, and corners hypergrid tasks for different choices of
${\mathbb{P}}_{{\mathcal{X}}}$ when training using a subset of the full
dataset (90%-10% train-test split). In this setting, agnostic to the choice of
${\mathbb{P}}_{{\mathcal{X}}}$ and hypergrid reward, we observe the GFlowNet
models converge when considering both the JS divergence and the MAE
distributional metrics for approximating $p(x)$. This differs from the graph
generation tasks where models struggle to converge on these metrics in this
setting. For the more difficult task (corners), sampling $x$ from sets that
include the proxy-policy $p(x;\theta)$ result in the best performance on JS
divergence while sampling $x\sim{\mathcal{U}}$ and $x\sim|\log
p(x;\theta)-\log p(x)|$ results in the best performance on the MAE metric.

####  E.4.2 Offline and off-policy in sequences

Figure 31: Evaluation curves for offline and off-policy trained GFlowNets on
edit distance sequence task for different choices of
${\mathbb{P}}_{{\mathcal{X}}}$ when training using a subset of the full
dataset (90%-10% train-test split). In this setting, agnostic to the choice of
${\mathbb{P}}_{{\mathcal{X}}}$, we observe the GFlowNet models converge when
considering JS divergence and MAE distributional metrics for approximating
$p(x)$. This differs from the graph generation tasks where models struggle to
converge on these metrics in this setting.