Generative Flow Networks for Discrete Probabilistic Modeling

Introduction

Probabilistic modeling in discrete spaces, especially those with compositional structure, is important due to the universality of applications of discrete data structures, such as in natural language processing (Tai et al., 2015) or in symbolic reasoning (Besold et al., 2017). However, distributions in high-dimensional discrete spaces are generally hard to model, as they may feature rapid combinatorial growth of modes. These modes can be well separated from each other, which poses a challenge for Markov Chain Monte Carlo (MCMC) methods (Salakhutdinov, 2009). Mixing between modes is generally slow without a priori knowledge of the specific latent structure of the distribution.

Furthermore, generative modeling methods such as energy-based models (EBMs; LeCun et al., 2006) also suffer from these mixing issues when generating negative samples with MCMC (Tieleman & Hinton, 2009). The incapability of MCMC to capture the energy landscape results in the spurious mode problem (Desjardins et al., 2010; Bengio et al., 2013), in which new modes that do not occur in the true data distribution appear in the learned energy distribution.

In this paper, we propose to take advantage of generative flow networks, or GFlowNets (Bengio et al., 2021a, b), instead of MCMC methods in order to simultaneously learn a sampler and train an energy function from data. GFlowNets are generative models for compositional objects, i.e., they learn a stochastic policy that iteratively constructs the sampled object through a sequence of simpler steps. In past work, GFlowNets were trained to query a given energy function, rather than from a dataset (like typical generative models). We propose to adapt the GFlowNet methodology to the problem of learning from data, and jointly train the GFlowNet sampler and the energy function.

Because a GFlowNet can learn from the low-energy configurations it has already encountered, it has a chance to guess and sample yet-unvisited modes if there are learnable regularities in the underlying data distribution. The compositional generative structure of GFlowNet further enhances its ability to discover regularities and thus jump between modes without having to travel long low-probability paths in between (Fig. 1). Instead of having to exploit a priori structure to design long jumps in the MCMC, GFlowNets can discover that structure when it is not explicitly known.

This paper makes the following key contributions:

We cast a non-autoregressive sequential generation model for high-dimensional discrete data as a generative flow network (§3.1).

We introduce a GFlowNet-based MCMC proposal enabling efficient large jumps with low probability of rejection, taking advantage of the compositional structure learned by the GFlowNet generative policy (§3.3).

We describe a procedure based on such proposals for jointly training the GFlowNet sampler with an energy model given a dataset and state conditions under which this estimates the true data log-likelihood gradient with respect to the energy function’s parameters (§3.2, §3.3).

We test the algorithm on a variety of synthetic and real tasks, achieving competitive results (§4).

Preliminaries

Generative flow networks, or GFlowNets for short (Bengio et al., 2021a, b), are trainable generative policies on which this paper is built. They model the generation process of objects ${\mathbf{x}}\in{\mathcal{X}}$ by a sequence of discrete actions that incrementally modify a partially constructed object (state). The space of possible action sequences is represented by a directed acyclic graph (DAG, see Fig. 2) $G=({\mathcal{S}},{\mathcal{A}})$, where the vertices in ${\mathcal{S}}$ are states and the edges in ${\mathcal{A}}$ are actions that modify one state to another. We use the terms parents and children, respectively, for the tails of incoming edges and the heads of outgoing edges of a state. Generation begins at a special initial state ${\mathbf{s}}_{0}$ and terminates upon a transition to any terminal state, which is a state with no outgoing actions. The set of terminal states is identified with the target space ${\mathcal{X}}$. Note that multiple possible action sequences may lead to the same terminal state.

In this case, $P_{F}({\mathbf{s}}_{t+1}|{\mathbf{s}}_{t})$ is called a forward policy, and can be used to sample complete trajectories from the density $F$ and thus also their terminal states, objects ${\mathbf{x}}\in{\mathcal{X}}$.It is helpful to think of the particle analogy. For a Markovian flow, the distribution over complete trajectories ($P_{F}(\tau)\propto F(\tau)$) satisfies a history-independence property: a particle’s choice of route after reaching a state ${\mathbf{s}}$ is independent of how it reached ${\mathbf{s}}$. We write $P_{T}({\mathbf{x}})$ for the probability that a trajectory sampled from $P_{F}$ terminates in ${\mathbf{x}}$.

By Eq. (1), one can specify a Markovian flow by a learned scalar $\log Z_{\bm{\theta}}$ and a neural net with parameters ${\bm{\theta}}$ that outputs $P_{F}(\cdot|{\mathbf{s}};{\bm{\theta}})$ for any input state ${\mathbf{s}}$. Algorithms for learning $P_{F}$ from maximum-entropy reinforcement learning (Haarnoja et al., 2017) would solve this problem if there was only one way to construct an object ${\mathbf{x}}$ (the DAG is a tree) while GFlowNets are applicable in the more general DAG setting (Bengio et al., 2021a). The solution originally proposed by Bengio et al. (2021a) recasts an unnormalized forward policy as a network flow on $G$ (in the classical sense of Ford & Fulkerson (1956)) and performs gradient descent on the error in a flow conservation constraint at states sampled from a training policy, amounting to a generalization of temporal difference objectives from Sutton (2005).

Alternative objectives proposed by Bengio et al. (2021b); Malkin et al. (2022) require a model to produce 3 outputs: the scalar $\log Z_{\bm{\theta}}$, a forward policy $P_{F}(\cdot|\cdot;{\bm{\theta}})$, and a backward policy that produces distributions $P_{B}(\cdot|{\mathbf{s}};{\bm{\theta}})$ over the parents of an input state ${\mathbf{s}}$. (When action sequences are sampled in reverse from $P_{B}$, we will use dashed arrows ${\mathbf{s}}^{\prime}{\dashrightarrow}{\mathbf{s}}$ to indicate that the action ${\mathbf{s}}{\rightarrow}{\mathbf{s}}^{\prime}$ has been sampled against the direction of the DAG edges. The $\bm{\theta}$ will be dropped from $P_{F}$ and $P_{B}$ notation when it does not cause ambiguity.)

In modeling tasks with a fixed reward (Bengio et al., 2021a; Malkin et al., 2022), only the forward policy $P_{F}$ is used for sampling from the GFlowNet, and the backward policy $P_{B}$ is simply an training artifact. However, our approach makes use of the backward policy to perform local exploration in a contrastive divergence-like algorithm (§3.3, Figure 1).

2 Energy-based models

Energy-based models (EBMs) (LeCun et al., 2006; Song & Kingma, 2021) are a popular approach for probabilistic inference and modeling. An EBM specifies a distribution over a space ${\mathcal{X}}$ by a density $p_{\bm{\phi}}({\mathbf{x}})=\frac{1}{Z_{\bm{\phi}}}\exp(-{\mathcal{E}}_{\bm{\phi}}({\mathbf{x}}))$, where ${\mathcal{E}}_{\bm{\phi}}$ is the energy function, with model parameter $\bm{\phi}$, and $Z_{\bm{\phi}}$ is a normalizing constant independent of ${\mathbf{x}}$. The normalizing factor is not explicitly parametrized, but, for finite spaces, can theoretically be computed as $Z_{\bm{\phi}}=\sum_{{\mathbf{x}}\in{\mathcal{X}}}\exp(-{\mathcal{E}}_{\bm{\phi}}({\mathbf{x}}))$. This summation is finite, and thus the energy function ${\mathcal{E}}_{\bm{\phi}}({\mathbf{x}})$ can take an arbitrary form. However, this flexibility comes with a price: calculating $Z_{\bm{\phi}}$ can involve an exponentially large summation, making evaluation of the exact probability intractable. To avoid such expensive calculations, it is desirable to amortize sampling from such models by training a generative model.

Evaluating or estimating the second term in (3) involves taking negative samples ${\mathbf{x}}^{\prime}$ from the EBM distribution, which can require an expensive search. The classical Contrastive Divergence (CD) algorithm (Hinton, 2002) approximates this gradient update by changing the energy function parameter with the following stochastic approximation:

where $q_{K}({\mathbf{x}}^{\prime}|{\mathbf{x}})$ is the distribution obtained by using a $K$-step MCMC initialized at ${\mathbf{x}}$ to approximately sample from $p_{\bm{\phi}}$. For example, ${\mathbf{x}}^{\prime}$ can be taken from a $K$-step Metropolis-Hastings chain starting at a true data sample ${\mathbf{x}}\sim p_{\rm data}({\mathbf{x}})$, and the parameters updated with $\nabla_{\bm{\phi}}({\mathcal{E}}_{\bm{\phi}}({\mathbf{x}})-{\mathcal{E}}_{\bm{\phi}}({\mathbf{x}}^{\prime}))$. As $K\to\infty$, assuming mixing of MCMC chains, the distribution over negative samples $q_{K}({\mathbf{x}}^{\prime}|{\mathbf{x}})$ converges to $p_{\bm{\phi}}({\mathbf{x}}^{\prime})$, and we recover, in expectation, the true gradient (3).

Tieleman (2008) later proposed persistent CD (PCD), where the MCMC chains that give negative samples do not restart at true data at every training step, but are initialized with a previous state. While CD-type algorithms are efficient in computation and can accelerate learning, their gradient estimation is biased and thus may not model the true data distribution faithfully (Nijkamp et al., 2020).

Methodology

In this work, we aim to model a target distribution over discrete data with a GFlowNet. In the domains we consider, the data is in $D$-dimensional binary space, i.e., ${\mathbf{x}}\in{\mathcal{X}}\triangleq\{0,1\}^{D}$. As an example, ${\mathbf{x}}$ could be an image with $D$ pixels taking binary values. We model the generation of vectors in ${\mathcal{X}}$ by a GFlowNet. The state space ${\mathcal{S}}$ of the GFlowNet consists of vectors of length $d$ with entries in $\{0,1,\oslash\}$, where the void symbol $\oslash$ represents a yet unspecified entry that may be turned to 0 or 1 by a future action. To be precise:

The DAG structure on ${\mathcal{S}}$ is the $D$-th Cartesian power of the DAG with states $\{\oslash,0,1\}$, where 0 and 1 are children of $\oslash$. Concretely, the children of a state ${\mathbf{s}}=({\mathbf{s}}^{1},\dots,{\mathbf{s}}^{D})$ are vectors that can be obtained from ${\mathbf{s}}$ by changing any one entry ${\mathbf{s}}^{d}$ from $\oslash$ to 0 or 1, and its parents are states that can be obtained by changing a single entry ${\mathbf{s}}^{d}\in\{0,1\}$ to $\oslash$.

We define $|{\mathbf{s}}|\triangleq\#\{{\mathbf{s}}^{d}\mid{\mathbf{s}}^{d}\in\{0,1\},d=1,\ldots,D\},$ the number of non-void entries in ${\mathbf{s}}$, so ${\mathcal{X}}$ is naturally identified with $\{{\mathbf{s}}\in{\mathcal{S}}:|{\mathbf{s}}|=D\}$. There is an initial state ${\mathbf{s}}_{0}\triangleq(\oslash,\oslash,\ldots,\oslash)$. Any trajectory from ${\mathbf{s}}_{0}$ to ${\mathbf{x}}\in{\mathcal{X}}$ has exactly $D$ actions. A choice of trajectory from ${\mathbf{s}}_{0}$ to ${\mathbf{x}}$ amounts to a choice of the order in which the entries of ${\mathbf{x}}$ are assigned.

In this setting, the forward policy $P_{F}(\cdot|{\mathbf{s}};{\bm{\theta}})$ of a GFlowNet, introduced in §2.1, is a distribution over all ways to select a position with a void entry in ${\mathbf{s}}$ and a value (0 or 1) to assign to this entry. Thus the action space for a state ${\mathbf{s}}$ has size $2(D-|{\mathbf{s}}|)$. Correspondingly, the backward policy $P_{B}(\cdot|{\mathbf{s}};{\bm{\theta}})$ is a distribution over the $|{\mathbf{s}}|$ ways to select a position with a nonvoid entry in ${\mathbf{s}}$. We illustrate the mechanism of the forward and backward policies in Figure 3.

In our experiments, we take $P_{F}(\cdot|{\mathbf{s}};\bm{\theta})$ and $P_{B}(\cdot|{\mathbf{s}};\bm{\theta})$ to be neural networks with a multilayer perceptron (MLP) architecture, where the input is a vector ${\mathbf{s}}\in\{\oslash,0,1\}^{D}$ encoded using a value of $-1$ for $\oslash$, and $P_{F}$ and $P_{B}$ share all weights except the final weight matrix that outputs logits for the forward and backward actions. The scalar $Z_{\bm{\theta}}$ is parametrized in the log domain, as suggested by Malkin et al. (2022). Specific implementation details are given in §4.

2 GFlowNet training towards a target distribution

To train the parameters ${\bm{\theta}}$ of the GFlowNet, we use the trajectory balance objective proposed by Malkin et al. (2022). Trajectory balance optimizes the following objective along complete trajectories $\tau=({\mathbf{s}}_{0}{\rightarrow}{\mathbf{s}}_{1}{\rightarrow}\dots\rightarrow\dots\rightarrow{\mathbf{s}}_{n})$:

Proposition 1 of Malkin et al. (2022) shows that if this objective is globally minimized (i.e., zeroed out) for all complete trajectories $\tau$, then $P_{T}({\mathbf{x}})\propto R({\mathbf{x}})$, i.e., the forward policy samples proportionally to the reward. Trajectory balance improves training of GFlowNets under various metrics and characteristics of the reward landscape (Malkin et al., 2022) relative to previously proposed objectives.

Training policy.

With the trajectory balance objective, we train the GFlowNet with stochastic gradient

with some training trajectory distribution $\pi_{\bm{\theta}}(\tau)$. Akin to on-policy RL settings, Malkin et al. (2022) took $\pi_{\bm{\theta}}$ to be the distribution over trajectories sampled from the current policy $P_{F}(\cdot\mid\cdot;\bm{\theta})$, or a perturbed / tempered version of it. That is, $\tau$ is sampled with ${\mathbf{s}}_{t+1}\sim P_{F}(\cdot|{\mathbf{s}}_{t};\bm{\theta})$ starting from ${\mathbf{s}}_{0}$, perhaps raised to a power or mixed with a uniform action policy to ensure $\pi_{\bm{\theta}}$ has full support, which is a condition for obtaining the desired distribution (Malkin et al., 2022).

In addition to this forward sampling approach, we propose a complementary strategy to benefit from the circumstances where we are given some terminating states (data examples ${\mathbf{x}}\in{\mathcal{X}}$). For a terminating state ${\mathbf{x}}$, one can sample a reverse trajectory $\tau=({\mathbf{x}}={\mathbf{s}}_{D}{\dashrightarrow}{\mathbf{s}}_{D-1}{\dashrightarrow}\dots\dashrightarrow{\mathbf{s}}_{0})$, where ${\mathbf{s}}_{t}\sim P_{B}(\cdot|{\mathbf{s}}_{t+1};\bm{\theta})$. Empirically, this backward trajectory sampling technique enables us to obtain a different trajectory distribution from the forward sampling distribution, as such backward trajectories visit regions of ${\mathcal{S}}$ near the true data samples that may be poorly explored by $P_{F}$, and could thus stabilize the optimization. Hereafter, we use $P_{F}(\tau)$ and $P_{B}(\tau|{\mathbf{x}})$ to denote the trajectory distributions that sample forward from ${\mathbf{s}}_{0}$ using $P_{F}$ and backward from ${\mathbf{x}}$ using $P_{B}$, respectively. In experiments, we take the training trajectory distribution $\pi_{\bm{\theta}}$ to be a mixture of these two sampling methods (see steps 3-9 of Algorithm 1).

Bengio et al. (2021b) noted that while there may be multiple Markovian flows satisfying (2), for any choice of a fixed backward policy $P_{B}$, there is a unique forward policy $P_{F}$ such that the corresponding $P_{T}({\mathbf{x}})$ is proportional to the reward. Malkin et al. (2022) suggested fixing $P_{B}(\cdot|{\mathbf{s}})$ to be uniform over the parents of every state ${\mathbf{s}}$ as a canonical choice. We find this to be beneficial in some domains. Furthermore, in our setting, this choice also enforces a maximum-entropy property on the forward policy, as the following proposition shows.

The entropy of a Markovian flow $F$, denoted ${\mathcal{H}}[F]$, is the expected total entropy of its forward policy distributions along a complete trajectory:

Suppose $G$ is the DAG defined in §3.1. Let $R$ be a nonnegative reward function on ${\mathcal{X}}$ and let $P_{B}^{\circ}$ be the uniform backward policy on $G$. Let $F^{\circ}$ be the Markovian flow uniquely determined by $P_{B}^{\circ}$ and $R$ subject to the reward matching constraint (2). Then $F^{\circ}$ has maximal entropy among all Markovian flows satisfying (2).

Estimating GFlowNet data likelihood.

The most commonly used metric in probabilistic modeling is the model’s likelihood on a test set. A well-trained model should assign a high likelihood to data from the same underlying distribution as the training data, that is, the terminating probability distribution $P_{T}({\mathbf{x}})=\sum_{\tau=({\mathbf{s}}_{0}{\rightarrow}\ldots{\rightarrow}{\mathbf{s}}_{D}),{\mathbf{s}}_{D}={\mathbf{x}}}P_{F}(\tau)$ would be close to the true data distribution. We overcome the intractability of the sum defining $P_{T}({\mathbf{x}})$ (the number of terms is factorial in $D$) by importance sampling:

where $\tau^{j}\sim P_{B}(\tau|{\mathbf{x}})$ are trajectories sampled backward from ${\mathbf{x}}$ using $P_{B}$. We can then use the average GFlowNet log likelihood on test set, estimated using (9) with $M$ large enough, as an evaluation metric.

3 Interleaved updates of GFlowNet and energy

The training of GFlowNets relies on a given function $R({\mathbf{x}})$ to provide reward signals. In generative modeling, we typically set $R({\mathbf{x}})$ to be the unnormalized target probability. However, in many settings, we do not have access to this exact quantity, but only to a collection of data samples from a target distribution.

To address this issue, we propose to introduce an energy-based model ${\mathcal{E}}_{\bm{\phi}}({\mathbf{x}})$ as an intermediate object between the data and the GFlowNet, to serve as the reward function with which the GFlowNet can be trained. The “final products” of training are then twofold: the trained energy model ${\mathcal{E}}_{\bm{\phi}}({\mathbf{x}})$ and the GFlowNet sampling policy $P_{F}$.

We train the EBM associated with the GFlowNet in an approximate MLE manner similar to Eq. (4), but using the GFlowNet policy to generate negative examples, ${\mathbf{x}}^{\prime}\sim P_{T}({\mathbf{x}}^{\prime})$. In the basic form of the energy function training procedure, updates to $\bm{\phi}$ are made proportionally to

where the GFlowNet terminating probability distribution $P_{T}({\mathbf{x}})$ corresponds to marginalizing the forward trajectory distribution $P_{F}(\tau)$ on its non-terminating states. If the GFlowNet is perfectly trained (with zero training loss), its terminating probability distribution will be equal to the energy distribution, and thus this approximate MLE objective is an unbiased estimate of the maximum likelihood training of EBM. In this way, we amortize the MCMC sampling computation into the GFlowNet training process.

In contrastive divergence training of EBMs (4), $K$-step MCMC is used to generate the negative samples. For each step of traditional MCMC methods, one first proposes a local random perturbation of the current sample ${\mathbf{x}}$ (seen as a state of the Markov chain) to some nearby point ${\mathbf{x}}^{\prime}$, and then decides whether to accept this transitionThis should be distinguished from the ‘transition’ (action) in GFlowNets, which happens between two states ${\mathbf{s}}_{t}$ and ${\mathbf{s}}_{t+1}$ in ${\mathcal{S}}$. according to the Metropolis-Hastings (MH) rejection rule (Hastings, 1970).

Here, we show that a GFlowNet can also be used for generating a proposal ${\mathbf{x}}^{\prime}$ from a given point ${\mathbf{x}}$ (Fig. 1). Given ${\mathbf{x}}$ and a fixed number of steps $K$ ($1\leq K\leq D$), we first sample a $K$-step trajectory with the backward policy $P_{B}$:

and then sample with the forward policy $P_{F}$, starting at ${\mathbf{s}}_{D-K}$ until a new terminal state ${\mathbf{x}}^{\prime}$ is reached:

where we have ${\mathbf{s}}^{\prime}_{D-K}={\mathbf{s}}_{D-K}$. For convenience, we denote this back-and-forth trajectory by $(\tau,\tau^{\prime})$ and the reverse trajectory by

Similar to $K$-step MCMC proposals, this GFlowNet proposal only changes the values of at most $K$ different entries.

With $\tau,\tau^{\prime}$ as above, we extend the definitions for complete trajectories (Eq. 1) to $P_{F}(\tau)=\prod_{t=D-K}^{D-1}P_{F}({\mathbf{s}}_{t+1}|{\mathbf{s}}_{t})$ and $P_{B}(\tau|{\mathbf{x}})=\prod_{t=D-K}^{D-1}P_{B}({\mathbf{s}}_{t}|{\mathbf{s}}_{t+1})$. The probability of a transition from ${\mathbf{x}}$ to ${\mathbf{x}}^{\prime}$ along the back-and-forth trajectory $(\tau,\tau^{\prime})$ is $P_{B}(\tau|{\mathbf{x}})P_{F}(\tau^{\prime})$. Similarly, the probability of going ${\mathbf{x}}^{\prime}$ to ${\mathbf{x}}$ along the reverse trajectory $(\tau_{-}^{\prime},\tau_{-})$ is $P_{B}(\tau^{\prime}|{\mathbf{x}}^{\prime})P_{F}(\tau)$. With the MH rule, if the move from ${\mathbf{x}}$ to ${\mathbf{x}}^{\prime}$ is accepted with probability

then the stationary distribution of an iterated application of such steps is equal to the desired reward distribution.

The following proposition shows that with a perfectly trained GFlowNet, we can cheaply obtain a high-dimensional form of Gibbs sampling, where $K$ variables are updated at each stepVanilla block Gibbs sampling would require computation exponential in $K$ in order to compute $2^{K}$ possible $K$-bit changes, with their energies and their normalizing constant.:

If a GFlowNet fits the reward perfectly, i.e., satisfies (2), then $A_{\tau,\tau^{\prime}}({\mathbf{x}}\to{\mathbf{x}}^{\prime})=1$, so the MH rejection step will always accept the proposal.

The proof is in §B. Notice that if $K=D$, the proposed transition is equivalent to directly sampling from ${\mathbf{s}}_{D-K}={\mathbf{s}}_{0}$ with $P_{F}$, which is independent of the backward policy and of the starting sample ${\mathbf{x}}$.

As a relaxation of (10) that can accelerate learning, we propose to generate negative samples using this back-and-forth GFlowNet transition proposal. Unlike MCMC methods, we do not iterate this kernel, but perform only one single step to generate negative samples ${\mathbf{x}}^{\prime}$ for each EBM parameter update, similarly to one-step contrastive divergence. As with MCMC-based contrastive divergence, it may be beneficial to begin with a small $K$ and gradually increase it over the course of training. In all of our experiments, we either use a constant $K=D$ (unconditional samples from the GFlowNet are used as negative examples, as in Eq. (10)) or gradually increase $K$ from 1 to $D$. We summarize our use of the GFlowNet proposal in EBM training in Algorithm 2.

Summary.

We propose a joint training framework (Algorithm 1), where the EBM and the GFlowNet are optimized alternately: the energy function serves as the (negative log-) reward function for the GFlowNet, which is trained with the trajectory balance objective to sample from the evolving energy model, while the energy function is trained with an approximate MLE gradient, where the GFlowNet provides negative samples ${\mathbf{x}}^{\prime}$ through an MCMC transition proposal that approximates block Gibbs sampling.

Experiments

We validate the EB-GFN algorithm on the Ising model (Ising, 1925). The Ising model is an elementary example of a Markov random field and is widely studied in mathematics and physics (see MacKay, 2003, §31). A $D$-particle Ising model is a distribution over $D$-dimensional binary vectors, with entries called spins. To keep with established conventions, we call the two possible values of each spin $\{+1,-1\}$ rather than $\{0,1\}$. The distribution is given by an energy model, where the energy is a quadratic form (with symmetric $D\times D$ matrix $J$) evaluated on the spin vector:

We consider $10\times 10$ grids with $\sigma=0.1,0.2,\dots,0.5$ and $9\times 9$ grids with $\sigma=-0.1,-0.2$.We use an odd grid size when $\sigma<0$ because such a model has many modes, each resembling a checkerboard with ‘seams’ where the checkerboard pattern is violated. For each setting of $\sigma$ and $N$, we use standard methods (Wang & Swendsen, 1990) to generate 2000 samples $\{{\mathbf{x}}_{i}\}_{i=1}^{2000}$ from the Ising model with energy ${\mathcal{E}}_{J}$. We then use the EB-GFN algorithm to jointly fit a symmetric matrix $J_{\bm{\phi}}$, giving an estimated Ising EBM ${\mathcal{E}}_{J_{\bm{\phi}}}$, and a GFlowNet that samples from this Ising EBM. Note that the EB-GFN algorithm does not have access to the true data-generating matrix $J$, only to the collection of samples $\{{\mathbf{x}}_{i}\}$. This is a simple test case for EB-GFN, since the energy is parametrized by a single matrix $J_{\bm{\phi}}$, not by a deep model. We evaluate the discrepancy (RMSE) between the true matrix $J$ and the learned matrix $J_{\bm{\phi}}$.

For simplicity, we set $K=D$ for the negative sampling step and use a training policy with $\alpha=1$ (no backward paths from training examples). Details can be found in §C.1. We compare EB-GFN with the Gibbs and Gibbs-With-Gradients (GWG) PCD algorithms (Grathwohl et al., 2021b). Table 1 shows the advantage of EB-GFN. Fig. 4 shows how faithful both the samples and the energy function obtained by the GFlowNet are, suggesting that the GFlowNet is able to discover generalizable structure from the data.

2 Synthetic tasks

We compare with PCD-10 and ALOE (Dai et al., 2020), two baselines for energy modeling in discrete spaces. We use the same energy function architecture and training protocol as Dai et al. (2020). The PCD-10 baseline utilizes Gibbs sampling with a replay buffer and the random restart mechanism (Du & Mordatch, 2019) for negative sample generation. ALOE learns three neural networks for negative sampling: a proposal model to provide initial samples and a pair of models (local search policy and stop policy) used to iteratively refine these samples. The initial proposal model could be either a simple multinomial distribution or a large autoregressive network. We name them ALOE and ALOE+ respectively, as the former has a similar total number of parameters to our GFlowNet, while ALOE+ is $32\times$ larger. For GFlowNet training, we use a mixed training policy with $\alpha=0.5$ and a schedule with linearly increasing $K$ for the back-and-forth proposal. See §C.2 for details.

For qualitative evaluation, we first visualize in Fig. 5 the heatmaps of the learned energy functions and some GFlowNet-generated samples by remapping the Gray code representations of samples back to 2-D space. As a comparison, we also visualize the energy model baselines in Fig. C.2. We observe that for multimodal tasks such as checkerboard and 8gaussians, GFlowNets are much better at capturing the modes and their structure than the baselines (as illustrated schematically in Fig. 1). We hypothesize that multimodality is easier to handle with GFlowNets than with MCMC if there are generalizable regularities that make it possible for the GFlowNet to guess new modes from those already visited and from which it has learned.

We quantitatively evaluate the algorithms in Table 2 by showing for each method the NLL of a large independent sample of ground truth data and the exponential Hamming MMD (Gretton et al., 2012) between ground truth data and generated samples as performance metrics. Our method outperforms the baselines on all datasets and metrics except MMD on 2spirals and 8gaussians, where it still exceeds comparable-size baselines. Of note, ALOE does not surpass PCD for many tasks without a large initial proposal network, which shows a potential weakness of its local search strategy. We defer more results and related details to §C.2.

3 Discrete image modeling

Here, we aim to generatively model previously studied image datasets in discrete high-dimensional spaces. These are generally hard problems as we lose the information of continuous pixel values, and prevalent scalable methods (Welling & Teh, 2011; Ma et al., 2015) are not applicable to training deep EBMs on discrete data. An MLP is utilized as the energy function, and the baseline training methods are PCD with Gibbs sampling and the Gibbs-With-Gradients sampling method (GWG; Grathwohl et al., 2021b). Experiments are performed on four different binary image datasets. Following the experimental settings in Grathwohl et al. (2021b), the EBM is trained via PCD-100 with different negative sampling methods, and a replay buffer is adopted as in Du & Mordatch (2019). The validation and evaluation protocol is also kept aligned with Grathwohl et al. (2021b), where the checkpoint with the best negative log-likelihood on the validation set is reported.

The test set NLLs are displayed in Table 3. The results indicates that EB-GFN reaches state-of-the-art energy modeling performance, surpassing GWG, on three of the four datasets. As a supplement, a visualization of Dynamic MNIST data is shown in Figure 6. More details are in §C.3.

Related Work

Energy-based models. EBM, one of the central methods in generative modeling, has proved effective with energy functions parametrized by deep nets (Hinton et al., 2006; Salakhutdinov & Hinton, 2009). To avoid costly MCMC simulation with deep models, contrastive divergence-type methods (Hinton, 2002; Tieleman, 2008; Xie et al., 2016; Du et al., 2021) were proposed to approximate the energy gradient. It has also been shown that better objectives beyond vanilla CD helps EBM training (Yu et al., 2020). Training methods have been proposed for better stability, shorter mixing time, faster training (Nijkamp et al., 2019; Du & Mordatch, 2019; Grathwohl et al., 2021a; Gao et al., 2021). Recent work shows that it can be beneficial to learn the sampler or the proposal distribution as well (Dai et al., 2019; Arbel et al., 2021), a finding that this work extends to discrete spaces.

Related methods. Autoregressive models (Sutskever et al., 2011; Graves, 2013; Germain et al., 2015), like our GFlowNets, generate each entry of a data vector sequentially, but in a fixed order. Autoregressive models can also be used for data without a natural order (Uria et al., 2016; van den Oord et al., 2016b, a; Meng et al., 2021). Some recent work (Emelianenko et al., 2019; Li et al., 2021) allows generation order to be learned. We are the first to interpret learning of generation order as a joint inference of forward (construction) and reverse (erasure) Markov processes.

Discrete inference. Probabilistic inference in discrete settings is generally harder than in continuous spaces. Many optimized sampling methods for discrete spaces have been developed (Titsias & Yau, 2017; Zanella, 2019; Han et al., 2020; Zhang et al., 2022). Most applicable to our experiment domains, Grathwohl et al. (2021b) uses a continuous relaxation to approximate the local energy landscape, while Dai et al. (2020) introduces a local search strategy in the variational distribution to initialize negative sample generation.

Conclusion

We have extended GFlowNets to the setting where one is given a dataset rather than a fixed energy function, and we learn both the GFlowNet sampler and the energy function. In doing so, we introduced a new proposal for MH MCMC that approaches block Gibbs sampling as the GFlowNet training converges. The main advantage of this proposal is that it can perform large jumps in the state space, unlike simple Gibbs sampling, taking advantage of the compositional structure that the GFlowNet may have uncovered that allows it to generalize across modes of the distribution and more easily jump between them (Fig. 1). Future work can consider iterating such proposals from a trained GFlowNet for efficient exploration of the space, rather than just sampling complete trajectories ($K=D$) from the GFlowNet.

The cost of such an approach is that in addition to training the energy function, we also have to train the sampler. However, this cost can be amortized if we intend to use the sampler later, since generating samples from the GFlowNet is often much cheaper than from a MCMC (if we want to make sure to cover the modes well). We hypothesize that the above advantages explain the good comparative results obtained here, and expect that the proposed approach may be extended to many other types of generative tasks.

Acknowledgement

We thank Yilun Du, Ricky T. Q. Chen, and Mila GFlowNet group for helpful discussion. Dinghuai Zhang thanks the never-ending snow storm in Montreal for preventing him from any form of outdoor activity :). Zhen Liu thanks miHoYo for the joy from their awesome games in the tough winters. Aaron Courville thanks the support of Microsoft Research, Hitachi and CIFAR. Yoshua Bengio acknowledges the funding from CIFAR, Samsung, IBM and Microsoft.

References

Appendix A Summary of GFlowNet notation

Appendix B Proofs of propositions

Suppose $G$ is the DAG defined in §3.1. Let $R$ be a nonnegative reward function on ${\mathcal{X}}$ and let $P_{B}^{\circ}$ be the uniform backward policy on $G$. Let $F^{\circ}$ be the Markovian flow uniquely determined by $P_{B}^{\circ}$ and $R$ subject to the reward matching constraint (2). Then $F^{\circ}$ has maximal entropy among all Markovian flows satisfying (2).

We use the following definitions from Bengio et al. (2021b):

For a trajectory flow $F$ and for any state ${\mathbf{s}}$, define the state flow $F({\mathbf{s}})=\sum_{{\mathbf{s}}\in\tau}F(\tau)$, and, for any edge ${\mathbf{s}}{\rightarrow}{\mathbf{s}}^{\prime}$, the edge flow

Notice that $Z=F(s_{0})$ immediately from (1).

If $F$ is a Markovian flow, then $P_{F}$ and $P_{B}$ can be computed in terms of state and edge flows:

The following computation shows that the entropy of the forward policy, defined by (8), equals a similar expression for the backward policy:

Because we have assumed $F(x)=R(x)$ for all $x$ terminal (condition (2)), and we have $Z=\sum_{{\mathbf{x}}\in{\mathcal{X}}}F({\mathbf{x}})$ clearly from the definitions, the quantity $\Delta$ is independent of the choice of Markovian flow. Therefore, maximizing ${\mathcal{H}}[F]$ is equivalent to maximizing the expected entropy of $P_{B}$.

Finally, notice that every complete trajectory $s_{0}{\rightarrow}\dots\rightarrow s_{n}$ passes through exactly one state ${\mathbf{s}}_{d}$ with ${\mathcal{H}}[P_{B}^{\circ}(\cdot|{\mathbf{s}}_{d})]=\log d$ for each $d=1,\dots,D$, and that

with equality if $P_{B}$ is uniform over the parents of ${\mathbf{s}}_{d}$. Thus ${\mathcal{H}}[F]$ is maximized when $P_{B}(\cdot|{\mathbf{s}})=P_{B}^{\circ}(\cdot|{\mathbf{s}})$ for all ${\mathbf{s}}$. ∎

If a GFlowNet fits the reward perfectly, i.e., satisfies (2), then $A_{\tau,\tau^{\prime}}({\mathbf{x}}\to{\mathbf{x}}^{\prime})=1$, so the MH rejection step will always accept the proposal.

Recall that the acceptance probability for a move from ${\mathbf{x}}$ to ${\mathbf{x}}^{\prime}$ along a reverse trajectory $\tau$ and a forward trajectory $\tau^{\prime}$ is given by

According to Eq. (21) of Malkin et al. (2022), for a GFlowNet satisfying the reward matching constraint (2) with respect to a reward function $R$, we have $R({\mathbf{x}})P_{B}(\tau|{\mathbf{x}})P_{F}(\tau^{\prime})=R({\mathbf{x}}^{\prime})P_{B}(\tau^{\prime}|{\mathbf{x}}^{\prime})P_{F}(\tau)$. Elementary algebraic manipulation, and substituting $R({\mathbf{x}})=e^{-{\mathcal{E}}({\mathbf{x}})}$, yields that $A_{\tau,\tau^{\prime}}({\mathbf{x}}\to{\mathbf{x}}^{\prime})=1$. ∎

Appendix C More about experiments

C.2 Synthetic tasks

This synthetic task is also adopted by previous works such as Grathwohl et al. (2019); Dai et al. (2020). We keep a consistent setting with Dai et al. (2020) unless specified. The data is first generated by an infinite data oracle as 2D floating-points values. They are then turned into 16-bit Gray code. This problem is challenging itself even without the existence of Gray code transformation, which is highly nonlinear. We use a $4$ layer MLP with $256$ hidden dimension and ELU activation (Clevert et al., 2016) as the energy function. The training of this energy function lasts $10^{5}$ steps. An Adam optimizer with $1\times 10^{-3}$ learning rate is used to update the EBM. The batch size is $128$. For PCD baseline, we use $10$ steps Gibbs sampling to generate negative samples, and choose the best results from three different replay buffer re-initialization rate: $\{0.05,0.1,1\}$. For ALOE and ALOE+, we keep the same configuration as in Dai et al. (2020), where the former uses a naive $32$-dim multinomial initial distribution proposal, and the latter use an autoregressive model which contains $32$ MLPs that each has three layers with $512$ hidden dimension. Both methods have an editor network and stop policy network, either of which is three layer MLP with $512$ hidden feature. For the EB-GFN algorithm, the policy network is a similar three layer MLP (the forward policy and backward policy share the same first layers and differ in the last layer), and the output dimension is $3\times 32=96$. The GFlowNet is optimized with an Adam optimizer, where the learning rate is $1\times 10^{-3}$. We use an equal mix of $P_{F}(\tau)$ and $P_{B}(\tau|{\mathbf{x}})$ to generate training trajectories for trajectory balance objective (i.e., set $\alpha=0.5$ in Algorithm 1). For the back-and-forth proposal, we set $K$ to linearly increase from $1$ to $D$ through the training process. The NLL computation of GFlowNet follows the method described in §3.2. We set the value of $M$ to be 100, which is large enough to converge (as a reference, for checkerboard dataset, the NLL is $20.695967$ when $M=10$, $20.695692$ when $M=50$, $20.695967$ when $M=100$, $20.695583$ when $M=500$, $20.695490$ when $M=1000$). The number of samples is set to $10^{5}$, which is also enough for convergence in a similar sense.

To help better understanding the oracle of this task, we visualize the ground truth samples in Fig. C.1. We can see that EB-GFN could generate samples very close to these true data. We also plot the visualization of the baselines’ energy function in Figure C.2. It demonstrates that ALOE actually has a hard time modeling multimode distribution without the help of a large initial proposal model, as ALOE+ does.

In Dai et al. (2020), the authors mentioned using Hamming kernel MMD, while in their public code linear MMD is adopted. Further, in their public implementation, the MMD result is calculated within a fixed group of 4000 samples. The variance of such a calculation results in many results in their experimental table being negative (note that MMD is a non-negative metric in theory). Based on these considerations, we choose to a more commonly adopted exponential Hamming kernel with $0.1$ bandwidth in Table 2. Besides, we report the average of $10$ repeat results, each with $4000$ samples. To make a fair comparison, we also report the results given by linear kernel MMD in Table C.1 which is also used in ALOE public code. We can see that our algorithm keeps being state-of-the-art, and surpasses both PCD and the basic ALOE method on all datasets except swissroll.

One interesting property of the proposed EB-GFN framework, is that we can get samples either from the resulting GFlowNet (by sampling with the forward policy) or the learned EBM (by sampling with MCMC). Theoretically, GFlowNet would benefit more from its inductive bias as we discussed in §1, but we ideally want both models to achieve good performance. To this end, we also track the performance of the learned EBM. We find the learned EBM shares similar performance with the GFlowNet. The comparison is shown in Table C.2. On average, the EBM expresses slightly worse NLL and MMD than the GFlowNet, but is still very competitive if compared with other baseline methods. This can also demonstrate the benefit of the GFlowNet prior. As a result, we hypothesize that the EB-GFN algorithm can achieve a good GFlowNet even with a not-so-good reward function. This is because GFlowNet only needs the reward function to be relatively accurate with respect the true target distribution. We point out an interesting analogy to this phenomenon in reinforcement learning: a policy can have good performance even when the agent has learned a not-so-good Q function (Sutton & Barto, 2005; Bengio et al., 2020).

Ablation study on synthetic tasks. For completeness, we conduct ablation study to understand the importance of two features in EB-GFN algorithms: (1) backward trajectory sampling in GFlowNets training distribution mentioned in §3.2, and (2) the back-and-forth proposal proposed in §3.3. We do experiments on checkerboard and moons tasks. We first remove the usage of backward training samples and only use $\tau\sim P_{F}(\tau)$ to train the GFlowNet. The GFlowNet NLL on moons becomes $19.746$ from $19.732$, and Hamming exponential MMD becomes $0.342$ from $0.121$ (in units of $1\times 10^{-4}$). For checkerboard, the NLL becomes $20.709$ from $20.696$ and the MMD becomes $2.648$ from $1.206$ (in units of $1\times 10^{-4}$). This shows removing the backward trajectory feature would only do little harm to the performance. For the second part, once we remove the back-and-forth proposal and always use $K=D$, the training loss of EB-GFN quickly diverges on both tasks. This indicates the suggested proposal trick is crucial to a reasonable optimization landscape.

Understanding the learned backward policy. In Fig. C.3, we give a visualization to show that the learned erasure policy $P_{B}(\cdot|\cdot;\bm{\theta})$ is meaningful. The GFlowNet-damaged samples have a clear visual structure that, interestingly, indicates that the several highest-magnitude bits in the Gray code are the first to be deleted by $P_{B}$ and, correspondingly, are the last to be generated by a forward policy that minimizes the trajectory balance loss jointly with this $P_{B}$.

C.3 Discrete image modeling

We explain the details of our discrete image modeling task here. Discrete image modeling with EBM is a hard problem, and pure PCD training would diverge if the number of MCMC steps is not large enough or if there is no replay buffer trick to help training. This is not the case with continuous circumstances (Nijkamp et al., 2020). In this part, we follow the settings of Grathwohl et al. (2021b), which are stated below. We use Adam optimizer with $1\times 10^{-4}$ learning rate to update the energy function. The batch size is set to be $100$ and the training lasts for $5\times 10^{4}$ steps. The energy function is an MLP with $256$ hidden units and three hidden layersWe do not use the ResNet-18 backbone for EBM, because we find that it takes 2 weeks for the training of Gibbs-With-Gradients to finish with the original code of GWG public repo.. We do not use exponential moving average for simplicity. The GFlowNet is optimized with an Adam optimizer, where the learning rate is $1\times 10^{-3}$. The validation of EBM likelihood is achieved with $300000$ step annealed importance sampling (AIS). GFlowNet is modelled as an MLP with three hidden layers and $512$ hidden units. In this part we take the canonical design of backward policy, as we find that it could stabilize the training process. We choose the checkpoint which has the best validation result, and report the corresponding test set performance. The same GFlowNet training techniques are utilized as in synthetic tasks: we use an equal mix of $P_{F}(\tau)$ and $P_{B}(\tau|{\mathbf{x}})$ to generate training trajectories for trajectory balance objective (i.e., set $\alpha=0.5$ in Algorithm 1). For the back-and-forth proposal, we set $K$ to linearly increase from $1$ to $D$ through the training process. We also do ablation study on these features in this task, and we get different results from §C.2. To be precise, in the static mnist experiment, we find that these two techniques are both important: EB-GFN would diverge without either trick. This partially reflects the difficulty of this task.

We also find that introducing LayerNorm (Ba et al., 2016) into the forward policy network architecture is of great benefit to the generative modeling performance. We add a LayerNorm after each linear layer except the last one. The negative likelihood of this ablation is presented in the following table.