Delta-AI: Local objectives for amortized inference in sparse graphical models
Authors: Jean-Pierre René Falet
Hae Beom Lee
Nikolay Malkin
Chen Sun
Dragos Secrieru
Dinghuai Zhang
Guillaume Lajoie
Yoshua Bengio
Venue: ICLR 2024 poster (2024)
Keywords: amortized inference
variational inference
graphical models
Markov random fields
generative flow networks
GFlowNets
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: An objective for amortized samplers of sparse graphical models that achieves efficient credit assignment by matching local conditionals, applied to sampling energy models and training sparse latent variable models.

Abstract: We present a new algorithm for amortized inference in sparse probabilistic graphical models (PGMs), which we call $\Delta$-amortized inference ($\Delta$-AI). Our approach is based on the observation that when the sampling of variables in a PGM is seen as a sequence of actions taken by an agent, sparsity of the PGM enables local credit assignment in the agent's policy learning objective. This yields a local constraint that can be turned into a local loss in the style of generative flow networks (GFlowNets) that enables off-policy training but avoids the need to instantiate all the random variables for each parameter update, thus speeding up training considerably. The $\Delta$-AI objective matches the conditional distribution of a variable given its Markov blanket in a tractable learned sampler, which has the structure of a Bayesian network, with the same conditional distribution under the target PGM. As such, the trained sampler recovers marginals and conditional distributions of interest and enables inference of partial subsets of variables. We illustrate $\Delta$-AI's effectiveness for sampling from synthetic PGMs and training latent variable models with sparse factor structure. Code: https://github.com/GFNOrg/Delta-AI.

Published as a conference paper at ICLR 2024
DELTA-AI: LOCAL OBJECTIVES FOR AMORTIZED INFERENCE IN SPARSE GRAPHICAL MODELS
Jean-Pierre Falet", Hae Beom Lee", Nikolay Malkin", Chen Sun, Dragos Secrieru, Dinghuai Zhang, Guillaume Lajoie', Yoshua Bengio Mila Québec Al Institute, Université de Montréal
Montreal, Quebec, Canada (jean-pierre. falet, hae-beon lee, ni mikolay.mal kin]@mi la quebec
ABSTRACT
We present new algorithm for amortized inference in sparse probabilistic graph- ical models (PGMs), which we call A d-amortized inference (A-AI). Our approach is based on the observation that when the sampling of variables in PGM is seen as sequence of actions taken by an agent, sparsity of the PGM enables local credit assignment in the agent's policy learning objective. This yields local con- straint that can be turned into local loss in the style of generative flow networks (GFlowNets) that enables off- off-policy training but avoids the need to instantiate all the random variables for each parameter update, thus speeding up training con- siderably. The A-AI objective matches the conditional distribution of variable given its Markov blanket in tractable learned sampler. which has the structure of a Bayesian network, with the same conditional distribution under the target PGM. As such, the trained sampler recovers marginals and conditional distributions of interest and enables inference of partial subsets of variables. We illustrate A-Al's effectiveness for sampling from synthetic PGMs and training latent variable mod- els with sparse factor structure.
Code: https: / thub .cmm//FNOrrgDeelaa-..
INTRODUCTION
Probabilistic modeling in high-dimensional spaces is challenging due the combinatorially large number of possible modes. When the modes are not known a priori and are well-separated from each other, the convergence of local exploration algorithms such Markov chain Monte Carlo (MCMC) is prohibitively slow. In contrast, amortized inference methods, which train models to perform approximate sampling from the distribution of interest, potentially scalable to high- dimensional spaces and come with guarantees on generation time, but may also suffer from mode collapse issues. Generative flow networks (GFlowNets; Bengio al., 2021) are family of methods that have recently shown success in sampling distributions over high- high-dimensional discrete spaces. GFlowNets amortized sequential samplers that compatible with local exploration (Zhang et al., 2022) and allow stable off-policy training (Malkin et al., 2023), helping to overcome the obstacle of mode collapse.
A major limitation of GFlowNets is that they scale poorly in the length of the generative trajectory, which equals number of variables if they are sampled one time as in Zhang et al. (2022). This is because the (not necessarily normalized) target density the GFlowNet's 's reward function takes as input the instantiated values of all the variables. As result, the sampling cost scales linearly in the dimension. In addition, because the learning signal for all steps in sampling trajectory comes from the terminal reward, GFlowNets suffer from inefficient credit assignment when trajectories are long, unless the reward function can be meaningfully decomposed into partial rewards for partially instantiated variables (Pan et al., 2023).
Motivated by the GelowNet methodology and its limitations, we propose A-amortized inference (A- AI), an algorithm for cnergy-based probabilistic inference and training that scales well with respect to the dimension of variables. A-AI training recovers the conditional probability distributions in Bayesian network that defines the same distribution as a target graphical model. A-AI makes use of known graphical model structure to allow for local credit assignment: updating the parameters of the
Equul contribution. CCIAAR AI Chair. CCIAAR Senior Fellow.

---

Published as a conference paper at ICLR 2024
amortized sampler requires only instantiating single variable and its Markov blanket, in contrast to all variables. It can also be seen as a novel training algorithm for GFlowNets of a particular structure, where the GFlowNet policy performs amortized inference.
We summarize the advantages of A-AI as follows. First, A-AI enables significantly faster training than regular GFlowNets and other amortized inference methods because i) each training step can involve only small subset of variables, so the sampling cost is negligible, ii) the local training signal is much stronger than that in regular GFlowNet methods, as it uses the decomposition of the energy function into terms, each of which is directly matched with the corresponding condition- als and marginals. Second, the memory cost is very low because computing each gradient update only involves small subset of variables (for a sparse graphical model), Lastly, with Á-Al we can benefit from flexible probabilistic inference by amortizing many possible sampling orders into sin- gle sampler (by learning Bayesian networks with multiple Markov-equivas structures), allowing inference over partial subsets of variables.
The paper is structured as follows:
In $2, we introduce the necessary background on graphical models and present the GFlowNet approach to amortized inference as motivation.
In $3, we define and analyze the proposed A-AI algorithm.
In $4, we validate our idea on various synthetic energy- -based models, showing that A-AI provides faster wall-clock convergence time compared to algorithms from prior work, with A-Al's conver- gence time even comparable to or better than that of (training-free MCMC sampling, noting that the trade-off gets better if more samples are needed, as with amortized inference in general. In $5, validate A-AI on the task of image generation using latent variable model. We impose the inductive bias of a graphical model structure on the joint over observed and latent variables and use A-AI as the posterior sampler in amortized variational expectation- maniriiaaitiiiiiiioonnn (EM) procedure.
2
BACKGROUND
2.1 PROBABILISTIC GRAPHICAL MODELS
In this section, we introduce the relevant notation and review the background on probabilistic graph- ical models (PGMs) that is essential to understand our main methodology. Details and further back- ground can be found in $B.1 and in Koller & Friedman (2009). Additional related work is in $A.
We consider a collection of random variables = (X,)vev indexed by a set V. For concreteness, we will assume the X, are binary, X takes values E X {0. I) We also assume that X has full support, i.e., its density p X - R is strictly positive. However, the entirety of the following discussion applies any discrete spaces, and much of generalizes to continuous random variables.
Undirected graphical models (Markov networks / factor graphs). Let G (V, E) be an undi- rected graph, with each vertex v E V corresponding to a variable X,. X is a Markov network with respect to G if it satisfies the local Markov property:
X, (vi((v)uN((v)) (Na(v) VvVV,
where No (v) denotes the set of neighbours of v in G and, for S S V, Xs denotes the projection of X onto the variables indexed by S, Xs = (Xp)yes. Markov networks are a natural way to encode known conditional independences in the joint distribution of the variables X,.
A common way to specify of Markov networks is via a decomposition of the joint density into local factors, each depending on subset of the variables. If $1...., Øx is collection of functions taking positive real values, where &2 depends on Xs. for some S& S V (i.e., dx {0, 1S - Rso), then the normalized product these factors defines a density:
1 K Z
[ k=1
p(x)
du (XS,),
xs := (x,)vES-
This is an energy model with energy E(x)= = Σg log d2 (XSz). ). Note that the normalization constant Z equals the sum of the product of factors over all x X and can be intractable. Distributions of the form (1) are equivalent to Markov networks with respect to a certain graph (see $B.1 for review).
Directed graphical models (Bayesian networks). Directed graphical models are a different way of encoding conditional independences among random variables. If D AV be directed acyelic graph (DAG), where C VxV V is the set of directed edges, then X is a Bayesian network with graph
2

---

Published as a conference paper at ICLR 2024
structure D if its density p is the product of the conditionals of each variable given its parents:
((x)=
p(xv xPa(v)).
Pa(v) := {u : (u, v) A}.
veV
with the conditionals understood to be unconditional marginals if the parent set is empty. Note that each factor in (2) normalized, unlike in (1). This property makes the joint distribution easy to sample: one chooses any topological ordering VI, ...... of V consistent with the DAG D and sam- ples the Xy, in order, conditioning each Xv, the (previously sampled) parents of vy. In a Bayesian network, every variable conditionally independent of its non-descendants given its parents.
From Bayesian networks to Markoy networks and back. Any Bayesian network with respect to a DAG D is also Markov network, as every condi- tional in (2) can considered a factor in a model of the form (I). However, the graph structure of the correspond- ing Markov network (called the moral graph) may have more edges than the underlying undirected graph of D, un- less D has no immoralities (induced sub- graphs of the form
undirected graph G
I-map
chordal completion
(eg, min fill alg.)
undirected
P-map
immorality-free
chordal graph G
underlying undir. graph
DAG D
Figure Summary of the relationships between undirected graphs defining Markov networks (left) and DAGs defining Bayesian networks (right). Chordalization strictly relaxes the conditional independence constraints on a Markov network, while Markov networks with respect to a chordal graph and Bayesian networks with respect to its P-map equivalent.
Conversely, if is a Markov network with respect to graph G, then X is Bayesian network with respect to a directed graph with underlying undirected graph if G is chordal (f.e., has cycle of length greater than 3 induced subgraph). Any graph G can be made into chordal graph G by adding edges process called triangulation chordalization), which relaxes the conditional independence conditions on X. In turn, G has an immorality-free ori- entation D, and X is a Bayesian network with respect to D. D is called an I-map for G and P-map for G.
These relationships are summarized in Fig. and more details are in $B.1.
2.2 TRAINING AND AMORTIZATION IN PGMs
Daximum-llkkllhoon training of PGMs. The problem of generative modeling consists in opti- mizing the parameters of a PGM so as to maximize the likelihood of a given dataset of independent observations of the variables, D (x
For Bayesian networks, where the conditional distributions Pe (x, xxp(v)) depend on a parameter vector v the gradient of the log- ligeliiel is simply
Vy log P (D) Σ Ve log Pv
(x())
V. log pv (xx" xxp(v))
(3)
PEV
Generative modeling with energy- -baedd models of the form (1), where the factors dx depend parameter vector v. is more difficult. The gradient the tog-likelihood of a sample is:
W log P (x) = Σ v log du (Is.) V log Z. V log Z = =Bx-poxxi log (x)].
(4)
Optimizing log (D) thus consists of minimizing the energy energy for given samples x (0) E D (positive phase first term in (4)) and maximizing it for samples from the model itself, x ~ Pw (x) (negative phase second term). The latter is complicated by the intractability of sampling the model exactly, motivating methods that draw approximate samples using MCMC (Hinton, 2002; Tieleman, 2008; Du et al., 2021) or an amortized inference method, such as GFlowNet (Zhang et al., 2022).
One can also consider the case of generative modeling with latent variables, i.e., where some subset H c V of the variables is unobserved. The likelihodd gradient for partial observation x is
(5) Thus the objectives (3) and (4) can be applied to partial observations, as long as one can sample the conditional distribution over the latents x given the observed values xy This can be achieved using MCMC, importance sampling (Bornschein & Bengio, 2015; Le et al., 2019a), or amortized estimation of this conditional distribution, known as variational Oopectationmmxximizatie (EM;
Vy log Ps (xvY)) Bxx-Po (xH (x [V & log P (x)].
3

---

Published as a conference paper at ICLR 2024
Dempster al., 1977; Neal & Hinton, 1998; Koller & Friedman, 2009). A GFlowNet can also be used as the amortized estimator in variational EM (Hu et al., 2023).
Amortized inference. Complementary to the generative modeling problem PGMs that of amortized inference: given factorized density p of the form (1) with unknown normalization con- stant, one wishes to approximate p by a distribution that can be tractably sampled. One option is to approximate p by sequential sampler, which iteratively chooses the value of one variable at a time conditioned on the values of those previously sampled. Training this sampler, as an agent that proposes action sequences and receives the unnormalized density as reward, can be seen as problem of policy optimization.
Given density p with factors du and corresponding Markov network structure G (V, E), suppose that D = (V, A) is an I-map for G. so p is also a Bayesian network with respect to D. If qo is any Bayesian network with respect D) with its conditional distributions given by parametric estimators q qo(x, | Xra(v)), we are interested in fitting the conditionals of qo such that the distributions qo and p are equal. If this is achieved, the sampler qa can be used to perform ancestral sampling of the variables in a fixed topological ordering FI, .2....., V|V|-
2.3 AMORTIZED INFERENCE WITH GENERATIVE FLOW NETWORKS
We give some background on generative flow networks (GFlowNets), group of deep reinforcement learning (RL) methods whose limitations in non-local credit assignment are motivation for our proposed algorithms. Taking an RL view, the training of the conditionals in qe to sample from can be seen as policy optimization problem. Ancestral sampling from qe is viewed as a sequence of actions taken by sampling agent, which obtains a value of all the variables x and receives a reward of R(x)= [Tx &x (xS,)) ). The aim of training the policy is to make the likelihood of producing the sample x proportional to the reward: (x) x R(x).
GFlowNets, which are a specialized family of reinforcement learning algorithms that train agents to match a target distribution, have yielded the state of the art in discrete probabilistic modeling problems (Zhang et al., 2022). However, in Zhang et al. (2022), the target density is unstructured; the sampler does not take advantage of known factor structure and learns to sample the variables in an arbitrary order. Here we summarize, in simplified form, the GFlowNet objectives that can be used to train a Bayesian network as sequential sampler that approximates the target distribution.
We describe TB and DB, two losses that achieve this. (We also use SubTB, variance-reducing interpolation between TB and DB, in some experiments; see $B.2.)
Trajectory balance. To train qu as a GelowNet with the trajectory balance (TB) objective (Malkin et al., 2022), one optimizes its parameters jointly with an estimate Zo of the normalization constant. The TB objective at sample x is
LIB(x) (log Zo log (x) log R(x))²
If (6) is optimized to 0, so (x) = Z R(x), for all samples then qe is a perfect sampler for and Zo equals Z, the true normalization constant of p. The objective can be optimized by gradient descent using on-policy samples x taken from 40 itself, in which case it equivalent in expected gra- dient to DKL as those drawn from tempered policy or from a known dataset, as done by Zhang et al. (2022): Zim- mermann et al. (2023) in order to avoid mode collapse and thus better cover the target distribution.
p) (Malkin et al., 2023; Zimmermann et al., 2023), or at off-policy samples, such
Detailed balance. The detailed balance (DB) objective (Bengio et al., 2023) is an alternative loss that typically has lower variance but poorer performance than TB. The DB objective depends on an individual step in a sampling trajectory, rather than on the complete sample. Recall that qo. as an agent, samples in topological order: N. also trains an auxiliary object, a flow function Fo that outputs a scalar for any partially instantiated sample. The DB objective at the i-th step of the sampling trajectory is
To train qo with the DB objective, one
LDB (x))........).... xw) = (log Fo (I))......)) + log qo (x,, NPa(wi) - log Fo ())) (7) with the additional constraint that x is fully instantiated sample, then F (x) = R(x). If the DB loss is optimized to 0 for every such transition, then qa is perfect sampler for p. and the flow function at the initial state, F(xo), equals the true normalization constant Z. Just as with TB, one can flexibly choose the samples at which (7) is optimized by gradient descent.
Pan et al. (2023) found that DB performs strongly if the reward has multiplicative decomposi- tion that enables a form of reward shaping called the forward-looking parametrization of flows. In
4

---

Published as a conference paper at ICLR 2024
of interest, the reward R(x) has multiplicative decomposition into the factors (Xs,) Rather than making Fe a neural network taking partial sample as input, one can learn Fg as multiplicative correction to a partially accumulated reward ((x iee,
Fo (x(p).....))) NNo (F(v)....)))) R (x(V)....))))
(8)
where NN is a neural network. We discuss two ways to define the partially reward R in $B.2.
However, both the TB and DB losses are non-local: they require a fully instantiated sample x to receive a training signal.
LOCAL CONSTRAINTS FOR MATCHING MARKOV AND BAYESIAN NETWORKS
In this section, we introduce A-AI, which aims to address the limitations of GelowNets in locality of credit assigmment. We develop novel GllowNet-style objective, but whose loss only depends on a small subset of variables at time. This is hypothesized to greatly improve training convergence and decrease resource requirements. We derive A-AI by enforeing the equality of local conditional distributions in pair of graphical models on the same set of variables.
Setting. We use the notation introduced in $2.1. Let X = (X,)vev be a collection of discrete random variables whose density p has factor structure of the form (I), with factors &u depending on sets of variables Sz. Let = (V. F) be the graph with edges between any two variables that cooccur in factor, that X is Markov network with respect G.
Let G be a chordal completion of G and D = AV an immorality- free orientation of G. By the results in $2.1, X is Bayesian network with respect to D, so p has factorization over the conditionals specified by D.
A-AI constraint. The A-AI constraint expresses the equality the conditional distributions over one of the variables given its Markov blanket under the two factorizations, p determined by the Markov network and q by the Bayesian network. To precise, suppose that x,,' X are two settings of the variables that differ in exactly one variable i.e., Xu # x and xV x Using the factorization for Markov network (1) and that of Bayesian network (2), and combining them appropriately, we have:
du (XSe)
q (xv (Pa(v))
1
Ch(u) == (v. ().
(9)
kweS, Φ (2'S2) we (a)UCh(() qf (x, (x(a(v))
We remark that the left side depends only on ((u)LNo(u)> while the right side depends only on X())uNE(u) Thus these constraints are local with respect to the PGM structure. They are nonetheless sufficient to recover the conditional distributions in the Bayesian network. The correctness and sufficiency of this constraint are formalized by the following proposition.
Proposition 1. Suppose that X - Rs0 is the density of a Markov network with factors 6x and that q X R=0 is the density of a Bayesian network with respect to an I-map D (V, A). Then the following are equivalent: for all x,x' differing in a single variable X (9) holds; (2)
All proofs can be found in $D. We remark that the proof of Proposition 1, which uses the transitive closure of the single-variable mutation relation, is similar to the principle of concrete score matching (Meng et al., 2022), although the motivations are quite different.
From constraints losses. We return to the problem of fitting Bayesian network { I-) to given factorized model p, as in $2.2. Proposition gives sufficient local constraints for qo to equal p. We can thus turn these constraints squared log-ratio losses in the style of (6) and (7):
2
Φu (ISe)
Qo (x, (Pa(i))
Ls(x. x.) =
Σ log RIMESa
Σ log
(10)
ΦR (2's2)
(E(M)UCh(M)) go xx) xxpatv)
where x' is defined as the perturbation of x with x,, changed to the value xy Proposition 1 easily implies that La(x, xx) = for all x, u, x,,, then qo is an amortized sampler for p.
Training policy and exploration. Just as in GFlowNet training ($2.3), the A-AI objective Lx u, x.)) can be optimized at on-policy samples x 4o(x) or those drawn from a tem- pered/dithered policy 4, dataset. The variable which perturbed is sampled uniformly in our experiments, although future work should investigate the effect of other perturbation policies. The full training algorithm is summarized in Algorithm
5

---

Published as a conference paper at ICLR 2024
Local credit assignment. We emphasize that (10) de- pends only on the values of x in the G-neighbourhood of u. The conditionals in this small subset of variables is given strong local supervision from only the factors that involve the variable u. This also significantly reduces the memory cost, we do not have to maintain all vari- ables in the computation graph to compute gradients.
Algorithm 1 A-amortized inference (basic form)
Require: factors dy. DAG D, model qe. op-
timization/eeplloation hyperparameters
1:
repeat
2:
Sample ~ (7) (training policy)
3:
Sample u E V. new value x * Xu A - grad update w.r.t. Lx(x, M,
Masked autoencoder for amortizing all conditionals. The form of the parametric model qo (xv XPa(v)) is not specified in (10). In our experiments, we use a single neural network, masked autoencoder (MAE), to model all the conditionals specified by D simul- taneously, allowing to make use of generalizable structure in the conditionals. To be precise, in the case of binary (=1) data, XPa(y) is encoded as a |V|-thgtth vector with all units except those corresponding to variables in Pa(v) set O, and the logit of xy is read off from the output unit corresponding to variable
4:
5:
until converged
A stochastic loss. If the number of terms on either side of (10) is prohibitively large, e.g., due to u having too many children added during chordal completion, an unbiased stochastic estimator of the gradient can be used. This estimator requires sampling only one child of u at a time; see SE
Amortizing over multiple DAG orders. An ad- vantage of the A-AI formulation is that it allows to perform amortized inference in multiple I-maps (Bayesian network DAGs D) simultaneously. In- deed, the A-AI constraint (9) is valid for any I-map D for the Markov network G, and the same paramet- ric model qe can be used for multiple DAGs (Fig. 2). Amortization of DAG orders is especially beneficial when using on-policy or tempered-polee policy training. If the loss is optimized only at variables near the root of the DAG, then the downstream variables do not need to be sampled to obtain a loss gradient (a direct con- sequence of the locality of the credit assignment), which reduces the overhead of computing 'rollouts' of the model to obtain samples for training. If the DAG order is freely chosen at each training itera- tion, then a given conditional will be trained as long as it is near the root of some I-map for the Markov network. While this procedure increases the number of functions that qo must learn to express, in prac- tice, it does not substantially slow down convergence (see Fig. F.2), showing that the shared structure the conditionals of p may in fact aid learning.
V/
V4
Vo.
Vi
V/2
V4
V/
Ve
VG
V.
Vg
V.3
V2
Figure 2: Generating and amortizing multiple DAG orders. The conditionals present two I- maps (DAGs) for the same undirected model are different: for example, the conditional p(v1 V2, v2)v3) appears in the row, not in those in the first row. A-AI learns model q that matches the conditionals the tar- get distribution p. If q has a structure that allows taking varying subsets of variables as imput, then it can be trained to match the conditionals appearing in multiple DAG structures simultaneously, and the resulting sampler can then be used for sam- pling any of these DAGs.
DAGs in the second
Uses of A-AL A-AI allows to solve both problems described in $2.2. If the parameters of the factors in graphical model p. an amortized sampler qe can be trained using A-AI to match the conditional distributions of p. thus amortizing the (generally intractable) sampling from p. If the parameters & of the factors are unknown, then the parameters of Pa can be updated using maximum- likellhood gradient using samples from qe. while q can be trained to match the condi- tional distributions of Pv that evolves over the course of training, as we describe next.
training of PGMs. As discussed at the end of $2.1, amortized inference models can be used to obtain the log-likelihone gradient in generative modeling settings. We elaborate two such settings in which an amortized model trained using A-AI can participate in bilevel optimization loop with a generative model.
A-AI in maximum-likelihon
In the training of an energy-based model, such as factorized Markov network Pv. an amortized inference model qe can be used to draw the samples from Pw needed in the negative phase in (4). In the training of generative model pv on variables V with latent variables H C V, an amortized inference model qe can be used to obtain the samples of xи conditioned on x needed in (5). If Pv has known Markov network structure G. then the conditional distribution of xн given X is a Markov network with respect to the induced subgraph of G on H and is therefore also

---

Published as a conference paper at ICLR 2024
(a) bing Lattice. M4
[b) Ising Lattice, [V|-1024
(c) Factor Lattice, /V|=64
(d) Factor Lattice, IV -1024
0.70
0.70
0.875
DB
0.65
0.650
FLO8
2 0.625
0.50
0.63
0.400
0.55
0.55
0.33
0.550
0.50
0.90
10a
10'
107
101
10t
101
107
122
10f
107
200
10'
207
301
101
20"
122
10*
201
Wall-clock Time (Sec)
Woll-clock Time (Soc)
Wall-dock Time (sec)
Wall-clock Time isec)
Figure 4: Comparison of 4-AI and GFlowNets. A-Al converges to the target distribution fastest.
a Bayesian network with respect an I-map for that subgraph. Thus one train amortized inference model qo as a sampler latent variables XH using the A-AI objectives, where the model Qo amortizes the dependence the observed variables Xv by explicitly conditioning them.
In both cases, one alternately optimizes & with respect to the A-AI objectives and updates & using (3) or (4). The locality of credit assignment also benefits this parameter-learning scenario, which both the parameters of both qa and P can be optimized after sampling only single variable and its Markov blanket. The scheduling of the optimization steps in these bilevel optimization loops is an important design choice. Past work that used GFlowNet as the amortized inference model performed alternating gradient steps with respect 0 and ø (Zhang al., 2022) or used adaptive schedules based on the loss values (Hu et al., 2023). We describe our choices in the experiment sections below.
Continuous A-AI and score matching. A variant of A-AI with real-valued variables has connec- tions to score matching; see &C for discussion and experiments.
4 EXPERIMENTS: SYNTHETIC DATA
In this section, we demonstrate the efficacy and efficiency of A-AI on sparse synthetic models in the case where the parameters of Ps are known, and study the parameter-learning scenario in $5. We test the baselines and our method on the three graphical models shown in Fig. 3. See $F.1 for the detailed experimental setup.
A-AI vs. GelowNets. We compare A-AI to regular GFlowNet losses: trajectory bal- ance (TB) and detailed balance (DB). Both require all variables to be instantiated to compute a single reward, thus become in- efficient when there is an excessive number of variables. We also compare against the Norward-looking parameterization (FL-DB; see $2.3 and $B.2), which reparametrizes the flow function with intermediate ener- gies to accelerate learning. All algorithms amortize over multiple DAG orders. The evaluation metric is negative log likelihood (NLL) of ground- truth samples from the target energy function (long-run Gibbs) with respect to each learned sampler.
(a) Ising ladder (b) Ising lattice (c) Factor lattice Figure 3: Graphical models. (a) and (b) are UGMs for Ising models and (c) shows the factor graph model, where each factor is small randomly initialized MLP with four arguments. (a) is chordal and (b,c) are non-chordal.
Fig. 4 shows that our A-AI achieves significantly faster training convergence, because even with the same parametric model qo being learned, the baselines still need all the variables be instantiated to receive a reward, resulting in longer wall-clock time and poor credit assignment. On the other hand, the A-AI objective requires only small subset of variables to be instantiated, which results in much faster wall-clock time and stronger training signals for those small subsets. These results demonstrate the superiority of the proposed local objectives.
A-AI vs. unstructured amortized inference. In 1, we also compare A-AI with amortized sam- plers from past work that do not use the constraint of graphical model structure, finding that they tend to converge even slower than the structured variants (Fig. F.1).
A-AI vs. MCMC. We further compare A-AI against MCMC methods: Gibbs sampling and its variant Gibbs-With- With-Gradients (GWG; Grathwohl et al., 2021). The goal of this experiment to measure the amortization costs and benefits by comparing wall-clock time to generate high-quality samples. In Fig. 5, the time axis of MCMC baselines means the wall-clock time of running 10k independent Gibbs chains, and for A-AI it is the training time. Note that for this experiment, we do not amortize multiple DAG orders in A-AL. Also note that the energy functions have a peaky
7

---

Published as a conference paper at ICLR 2024
(4] lking Ladder, [V]-64
(b) Ning Ladder, [V =256
[c] Ising Lattice, [V|664
(4) (d)1sing Lattice, |V/-256
150
150
50
Gibbs
50
GWG
123
125
6-Ai
100
100
MMD
30
30
20
20
10
10
50
100
150
200
250
100
200
203
400
500
50
100
150
200
230
100
200
200 400 500
Wall-cloc Time (sec)
call Time (Sec)
Wall WalliclockTTRe Time (Sec)
Wallccook Time (sec)
Figure 5: Comparison against MCMC. (a, b) chordal graph, (c, d) non-chordal graph. A-Al provides a substantial amortization benefit, with training time smaller than the mixing time of MCMC chains.
landscape (see $F.I), so the MCMC baselines have difficulties in traversing all the modes. The evaluation metric is linear MMD) between ground-truth samples (obtained by long-run Gibbs) and the samples generated from each amortized sampler. We can see from Fig. 5 that after short period during which the amortized A-AI sampler performs poorly (as the amortization network is close to initialization), soon A-AI achieves significantly lower MMD, showing its superior effectiveness in traversing all the modes in the peaky distribution.
5 EXPERIMENTS: VARIATIONAL EM ON REAL DATA
To demonstrate how A-AI can be used to learn a PGM from par- tial observations, we revisit the problem of latent variable mod- eling for MNIST images (Deng, 2012). Specifically, we want to train a generative model Py Xy( where x are latent vari- ables and N are the observed pixel values. To train this model using variational EM, we learn an amortized variational posterior qo(xH Xv(H) using A-AI or another inference method (E-step) and learn Py by maximizing log p (VH)) on samples from qo (M-step). Following past work on discrete probabilistic mod- eling. the images are discretized by interpreting the pixel values in [0, 1] as parameters of independent Bernoulli random variables, from which we sample the pixel values.
Figure 6: Structure of the PGM used in $5. Grey layers above the pixel layer are com- posed of latent variables. Nodes within the boundaries of a blue square represent a clique with the node in the layer above (dashed lines), except in the pixel layer, where edges between pixels removed enforce conditional independence given the latents. Only three clique windows are shown for better visibility; in practice,windows windows are tiled across the layers with specitic win- dow size and stride. The total number of binary latent variables in the three layers 144 16 4 = 164. The graph contains only 8.9% of the edges possible in model with conditionally in- dependent pixels.
We assume a pyramid-shaper graph structure for the latent variable model, shown in Fig. 6, in which dependencies between variables are local and sparse, and the observed variables are conditionally independent given the latents (see $F2 for details). We define Pv as a Bayesian network with a DAG structure that is an I-map for this pyramid-shaped undirected graphical model. The amortized posterior qa. also Bayesian network, has the same structure as Pw over the subgraph consisting of latent variables, but has all edges from latent to hidden variables oriented upward, i.e., conditioning the generation of latents in H on observed variables in V \ H. This choice of DAG structure for qe ensures it is an I-map for Ps con- ditioned on V \ H. We amortize qe over multiple DAG orders with the same underlying undirected graph, as described in $3.
Given data XV,,, we train qa(xH |
xV:H) be proportional
Po Xv)),, using the A-AI objective (10) treating Pv as a factorized model, or using a GFlowNet objective, using Pw Xviu)) as the reward for xH. In both cases, use gently exploratory training policy: xH is sampled from a tempered qe with a low off-policy exploration rate (see $F.2).
We compare A-AI with the following algorithms for approximating the posterior:
(1) GFlowNets: the objectives from $2.3 (TB, DB, FL.-DB,, trained with reward pv (XH, XV.H).
(2) Mean-field variational EM: a fully-factorized variational posterior trained using on- on-policy TB with the same reward, equivalentt to minimizing KL between the amortized and true posteriors.
(3) Wake-sleep: maximizing log (x( I xVvH) using one of the following: (i) Sleep: Uncon- ditional samples x from Pw the 'sleep' phase of wake-sleep (Hinton et al., 1995); (ii) IW: Imporance-weighted samples лн from the posterior given a dataset example xV correspond- ing to the 'wake" phase of reweighted wake-sleep (Bornschein & Bengio, 2015). The samples polsn,xvi))
are taken from q and weighted by TelxH|Xvia normalized over a batch.
(4) Gibbs sampling: sampling the posterior using K-step Gibbs sampling with respect to ps-
8

---

Published as a conference paper at ICLR 2024
(a) Ulnconditional Samples
(b) NLL vs. Time
[c) Classifier Entropy vs. Time
Vanilla VAE
J 9
V 4
1.4
Gibbs
: 3
3
Cihos
FL-DB
0
f
FL
D6
T6
5
2
/
2 7 177002
4
Entraps
0.8
DB
Mean-Field EM
2-
0.6
Mean-feeld EM
Sloop
2
9
9 9
Sleap
IW
2 0
f
9
A-Al
P
5
Ground Truth /
2
0.2 0.5
10
0.2
0.5 2
10
Walk-Clokk Time (h)
Wall-Clock Time (h)
Figure 7: Results on latent-variable modeling of MNIST. (a) Unconditional samples from Par For compar- ison, samples from a vanilla VAE (NLL = 85, entropy = 0.24), with an encoder and decoder parametrized using the same neural network architectures as for 40 and pg. respectively, and latent dimension equal to the shown
number of latent variables in the graphical model, are shown the top, and ground truth samples the bottom. (b) NLL of held-out test data for amortized methods. (c) Mean prediction entropy on unconditional samples of pretrained MNIST classifier. All models are trained for 12 hours; mean = std over 5 runs shown.
Algorithms (1), (2), and (3) use the same model architecture as A-AI for the conditionals in qa. while (4) is not an amortized inference method and involves no learning in the E-step.
We use the NLL of test-set samples under Pv. estimated using the importance weighted variational lower bound (Burda al., 2016), as the main evaluation metric to assess convergence of the bilevel objective for the amortized inference algorithms. To assess the quality of generated samples, moti- vated by -idtr---iitiiittionn detection methods (Liu et al., 2023), also measure the mean predic- tion entropy of a standard pretrained MNIST classifier evaluated on unconditional samples from pø-
Unconditional samples from Pv and evaluation metrics are plotted in Fig. 7. As we move from a variational factorized posterior with learned structured prior (mean-field EM), to structured pos- terior and structured prior (A-AI and GFlowNet baselines), the quality of unconditional samples improves. Wall-clock convergence for A-AI is quicker than all other baselines, with the IW variant closer to A-AI than the sleep variant, which is expected to be slower due to the initially poor quality of samples from pv- While they performed well in this setting, wake-sleep variants are limited to settings with a normalized target density (unlike GFlowNets, MCMC, and A-AT) and scale poorly with the latent dimension. Gibbs sampling of the posterior, the standard non-amortized sampling algorithm in graphical models, the third-slowest to converge. Although FL-DB uses the factoriza- tion of the energy function as an inductive bias for the state flow, it showed the slowest wall-clock convergence due to the expensive computation of the energy of every partially instantiated state.
Additional results are provided in $F.2; notably, we show that amortizing over multiple orders has barely any impact on convergence time for A-AI, as opposed to learning the conditionals of a single DAG (Fig. F.2), which not only demonstrates the benefit of parameter-sharing in qa. but also A-Al's unique ability to do inference over partial subsets, which reduces sampling time. We emphasize that none of the baselines are capable of partial inference over subsets of variables. A-AI is therefore ex- pected to lead to even more considerable improvements in wall-clock convergence time and memory requirements the dimensionality of the graphical model is increased.
6 DISCUSSION
We have proposed an objective for amortized sampling in probabilistic graphical models that uses only local information learning signal. In this work, evaluated our method in settings where the PGM structure was assumed known; however, in practice, the structure may need to be learned jointly with training of the amortized sampler by A-AL We note that the related algorithm family of GFlowNets has successfully been applied to structure learning of Bayesian networks, both with and without joint inference of the parameters (Deleu et al., 2022; 2023; Nishikawa-Toome but not yet to structure learning of undirected graphical models.
al., 2022),
A limitation of A-AI is the scaling to graphs for which chordalization results in very large Markov blankets, making training less efficient. Future work should thus consider methods such as stochastic losses ($E) or parametrization of small joins conditionals, such as those recovered by belief propaga- tion in junction trees, rather than only the univariate conditionals specified by the Bayesian network. A related question is the choice of the subsets of variables to sample when learning from incomplete observations, which is necessary for learning on very large graphs where instantiating all variables is prohibitive. Heuristies for variable selection that maximize information gain estimation of which can also be amortized - can be considered. Finally, A-AI should be applied to continuous spaces and real-world data where we can impose the inductive bias of PGM structure.
9

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY
Code is available
ACKNOWLEDGMENTS
The authors grateful to Zhen Liu for discussions in the initial stages of this project, to Olexa Bilaniuk for help with code optimization, and to Thomas Jiralerspong for his help during the author response period.
JF acknowlegdes support from the FRQS/MSSS.
GL acknowledges funding from CIFAR, NSERC, Samsung, and a Canada Research Chair in Neural Computation and Interfacing.
YB acknowledges funding from CIFAR, NSERC, Intel, and Samsung.
The research enabled in part by computational resources provided by the Digital Research Alliance of Canada (https: http///allannn awlaaaccaaaannnncaa//m/ ca), Mila (https: //ti://milll htt.m/ma....uuee......mm. NVIDIA.
quebec), and
REFERENCES
Barry C. Arnold and D. V. Gokhale. Distributions most nearly compatible with given families of conditional distributions. Test, 7(2):377-330, 1998.
Matthew J. Beal. Variational algorithms for approximate Bayesian inference, 2003. URL https:
///cse.bffaee htt.r//a...etteetaatt////eee//pp /baa103. ..f.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow net-
work based generative models for non non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, (24):1-76, 2023.
Julian Besag. Spatial interaction and the statistical analysis of lattice systems. Journal of the Royal Statistical Society, Series B, 36(2):192-236, 1974.
Jörg Bornschein and Yoshua Bengio. Reweighted wake-sleep. International Conference on Learn- ing Representations (ICLR), 2015.
Yuri Burda, Roger Baker Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. International Conference Learning Reprexentations (ICI.R, 2016.
Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
Tristan Deleu, Mizu Nismeaywa-Toomey, Jithendaraa Subramanian, Nikolay Malkin, Laurent Char- lin, and Yoshua Bengio. Joint Bayesian inference of graphical structure and parameters with single generative flow network. Neural Information Processing Systems (NeurIPS), 2023.
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the EM algorithm. Journal of the Royal Statistical Society B, 39(1):1338, 1977.
Li Deng. The MNIST database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 291))14 2012.
Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. International Conference on Machine Learning (ICML), 2021.
Samuel J. Gershman and Noah D. Goodman. Amortize inference in probabilistic reasoning. Cog- nitive Science, 36, 2014.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris Maddison. Oops I took a gradient: Scalable sampling for discrete distributions. International Conference on Ma- chine Learning (ICML), 2021.
10

---

Published as a conference paper at ICLR 2024
John M. Hammersley and Peter E. Clifford. Markov random fields finite graphs and lattices. 1971.
Luke B. Hewitt, Tuan Anh Le, and Joshua B. Tenenbaum. Learning to learn generative programs with memoised wake-sleep. Uncertainty in Artificial Intelligence (UAI), 2020.
Geoffrey E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation, 14:1771-1800, 2002.
Geoffrey E. Hinton, Peter Dayan, Brendan J. Frey, and R M Neal. The "wake-sleep" algorithm for unsupervised neural networks. Science, 268 5214:1155 61, 1995.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Neural Infor- mation Processing Systems (NeurIPS), 2020.
Edward Hu, Nikolay Malkin, Moksh Jain, Katie Everett, Alexandros Graikos, and Yoshua Bengio. GFlowNet-EM for learning compositional latent variable models. International Conference on Machine Learning (ICML), 2023.
Michael I. Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning, 37:183-233, 2004.
Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. International Conference Learning Representations (ICLR), 2014.
Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. MIT press, 2009.
Tuan Anh Le, Adam R. Kosiorek, Siddharth, Yee Whye Teh, and Frank Wood. Revisiting reweighted wake-sleep for models with stochastic control flow. Neural Information Processing Systems (NeurIPS), 2019a.
Tuan Anh Le, Adam R. Kosiorek, N. Siddharth, Yee Whye Tch, and Frank Wood. Revisiting reweighted wake-sleep for models with stochastic control flow. Uncertainty Artificial Intel- ligence (UAI), 2019h.
Tuan Anh Le, Katherine M. Collins, Luke B. Hewitt, Kevin Ellis, N. Siddharth, Samuel J. Gershman, and Joshua B. Tenenbaum. Hybrid memoised wake-sleep: Approximate inference at the discrete- continuous interface. International Conference Learning Representations (ICLR), 2022.
Xixi Liu, Yaroslava Lochman, and Zach Chrsitopher. GEN: Pushing the limits of softmax-ban out-of-distribution detection. Computer Vision and Pattern Recognition (CVPR), 2023.
Matthew Loper, Naureen Mahmood, Javier Romero, Gerard Pons-Moll, and Michael Black. SMPL: A skinned multi-person linear model. ACM Trans. Graphics (Proc. SIGGRAPH Asia), 34(6):248:1-24:1166 16, October 2015.
Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bose, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2023.
Naureen Mahmood, Nima Ghorbani, Nikolaus F. Troje, Gerard Pons-Moll, and Michael J. Black. AMASS: Archive of motion capture surface shapes. International Conference on Computer Vision (ICCV), 2019.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems (NeurIPS), 2022.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. International Conference on Learning Representations (ICLR), 2023.
11

---

Published as a conference paper at ICLR 2024
Christian Mandery, Ömer Terlemez, Martin Do, Nikolaus Vahrenkamp, and Tamim Asfour. Uni- fying representations and large-scale whole-body motion databases for studying human motion. IEEE Transactions Robolics, 32(4))796-889, 2016.
Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Gen- eralized score matching for discrete data. Neural Information Processing Systems (NeurIPS), 2022.
Radford M Neal and Geoffrey E Hinton. A view of the EM algorithm that justifies incremental, sparse, and other variants. Learning in graphical models, pp. 355-368. Springer, 1998.
Mizu Nishikawa- Toomey, Tristan Deleu, Jithendaraa Subramanian, Yoshua Bengio, and Laurent Charlin. Bayesian learning of causal structure and mechanisms with GFlowNets and variational bayes. arXiv preprint arXiv:221 1.02763, 2022.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023.
Tom Rainforth, Adam R. Kosiorek, Tuan Anh Le, Chris J. Maddison, Maximilian Igl. Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. International Conference on Machine Learning (ICML), 2018
Lawrence K. Saul, T. Jaakkola, and Michael I. Jordan. MMan field theory for sigmoid belief net- works. Journal of Artificial Intelligence Research, 4:61-76, 1996.
Vaidotas Simkus, Benjamin Rhodes, and Michael U. Gutmann. Variational Gibbs inference for statistical model estimation from incomplete data. Journal of Muchine Learning Research, (24):
2023.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Neurul Information Processing Systems (NeurIPS), 2019.
Andreas Stuhlmüller, Jacob Taylor, and Noah Goodman. Learning stochastic inverses. Neural Information Processing Systems (NIPS), 2013.
Tijmen Tieleman. Training restricted Boltzmann machines using approximations to the likelihood gradient. International Conference on Machine Learning (ICMI.,, 2008
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. International Conference on Machine Learning (ICML), 2022.
Heiko Zimmermann, Fredrik Lindsten, Jan-Willem van de Meent, and Christian A. Naesseth. A variational perspective on generative flow networks. Transactions on Machine Learning Research (TMLR), 2023.
12

---

Published as a conference paper at ICLR 2024
A OTHER RELATED WORK
Amortized inference in graphical models has been proposed as a model for human reasoning (Ger- shman & Goodman, 2014). PGMs are the setting for many early variational inference methods (Saul et al., 1996; Beal. 2003; Jordan et al., 2004) and for related problems, e.g., recovering tractable joints from unary conditionals such as those appearing in the A-AI objective (Arnold & Gokhale, 1998)
and fitting conditionals over latent variables using approximate samples (Stuhlmüller et al., 2013;
Simkus et al., 2023). Amortized inference is an ingredient in the training of structured generative models using variational EM ($2.1) and wake-sleep algorithms (Hinton et al., 1995; Bornschein & Bengio, 2015; Le et al., 2019b; Hewitt et al., 2020; Le et al., 2022). Beyond PGMs, amortized infer- ence is used in weakly structured spaces, notably in variational autoencoders (Kingma & Welling, 2014; Burda et al., 2016; Rainforth et al., 2018 8); score-based models and diffusion models (Song & Ermon, 2019; Ho et al., 2020) can also be seen as performing amortized posterior inference.
B EXTENDED BACKGROUND
B.I GRAPUICAL MODELS
Equivalence of factor models and Markov networks. A distribution of the form (1) is a Markov network with respect to certain graph G E). This graph is defined by the condition that uv E E if u and v cooccur in some factor's set of arguments Sy (i.e., 3k [M, v) S Sx). Remarkably, Clifford, 1971; Besag,
the converse also true: the Hammersley-Cite 1974) states that if X is a Markov network with respect to G, then its density has a factorization of the form (I), with each factor 6x depending on set of variables that forms maximal clique (complete subgraph) in G.
theorem (Hammersley
Conditional independences in Bayesian networks. The kinds of conditional independences that can expressed by a Bayesian network structure are distinct from those that can be expressed by a Markov network structure. For example, consider the DAG D network with this structure has X1 and X3 marginally independent, but not necessarily conditionally independent given X2. No Markov network structure can encode these constraints: X, and X3 must be in different connected components (since they are marginally independent), implying their independence given X2 as well.
2 - Any Bayesian
We proceed review the conversion from Bayesian to Markov networks and vice versa, elucidating the relationships that were summarized in Fig. 1.
From Bayesian networks to Markov networks. If X is a Bayesian network with respect to the graph D = (V, A), then every conditional in (2) can be considered factor in model of the form (1). Thus X is also a Markov network with respect to the graph G = (V, E) that has an edge between two variables if they cooccur in some conditional (i.e., (u, v) E E if' (u, v) E A, (v. A, (u, w) E A and (v, w) E A for some w E V). Equivalently, G is constructed by forming a clique out of every variable together with its parents.
The underlying undirected graph of D equals G the above procedure introduces no extra edges if and only if has no immoralities (induced subgraphs of the form - < -). Under these conditions, the set of distributions that are Bayesian networks with respect to D coincides with the set of distributions that are Markov networks with respect to G.
From Markoy networks Bayesian networks. Conversely, suppose that X is Markov net- work with respect to G = E). Recall that an acyclic orientation of G is a DAG D = (V,A) whose underlying undirected graph is G. It can be shown that G has acyelic orientation with no immoralities if and only G is chordal.
Any graph G can be converted into chordal graph G. known as a chordal completion of G, by adding extra edges. Although the optimal such chordalization is in general intractable to compute, the min-fill heuristic can be used to find chordal completions.
The set of distributions that are Markov networks with respect to G is subset of those that are Markov networks with respect to G. The chordal completion G has an acyclic, immorality-fre -ree orientation D) which is called P-map for G and an I-map for G. Any Markov network with respect to T also Bayesian network with respect its P-map D. Consequently, any Markov network with respect to the original graph G is a Bayesian network with respect to its I-map D.
13

---

Published as a conference paper at ICLR 2024
Sampling P-map for a chordal factor graph. Let G be chordal graph (note that any graph can chordalized by inserting additional edges, for example, using the min-fill algorithm). We summarize how to find P-map, i.e., an orientation of the edges of G that has no immoralities.
First, the mrxer-l cliques are found using a maximum cardinality search algorithm. Let C be the resulting set of maximal cliques. For maximal cliques i, J. let Sij |C; n CjI be the number of shared vertices, and let J = S) be the weighted graph with node corresponding to each clique and an edge of weight Sıj between every pair of cliques and j. junction tree, clique tree, can be found by first sampling maximum spanning tree for J and then directing its edges away from any chosen root clique.
Finally, the orientation of J is used to derive an orientation of G, by traversing the tree J in topological order, visiting the unvisited nodes within each clique in an arbitrary order, and orienting edges of G towards any newly visited nodes.
The arbitrary choices made in this procedure in the choice of maximum spanning tree, the choice of root clique, and the choice of ordering the unvisited nodes within each clique - can be varied to produce multiple P-maps for the same chordal graph.
B.2 GENERATIVE FLOW NETWORKS
Subtrajectory balance. We describe an interpolation between the TB loss (6) and DB loss (7), known as SubTB (Madan et al., 2023), as it applies to our setting.
SubTB requires learning the same estimators as the DB loss: the flow function Fa. in addition to the conditionals. For a sequence of sampled values Aug........ the objective is decomposed over sub-ranges:
?
L (x)
log Fe (x(vI.....))) +
ge (Nva (xPavza) log Fe (x
k=/+1
LSubTB (x) Eoske/siv 2 LSUTB..r))
Eogi<jsiv\ 1-11
Notice that j j=i iit1 recovers DB and i 0, j IV| recovers TB. The hyperparameter A controls the tradeoff between the losses for large j-i (closer TB) and small j (closer to DB), although values of A close to I are typical.
The forward-looking parametrization (8) is also applicable to SubTB.
Defining partial energies. We experimented with two ways of defining the partial reward Ř for partially instantiated variables. A natural way define R accumulates the product of the factors all of whose arguments are present in the partial sample:
R (X()))))))
Φx (xs₂)
(11)
Ar (
The partial reward signal is accumulated gradually as the arguments of more factors du are 'filled in', providing learning signal for partial samples.
However, we found that this option performs poorly in practice, possibly due the parametrization of NN, in (8) as a head on the masked autoencoder (cf. $3) that also predicts the logits of condition- als in qe. Such model must learn be sensitive to whether any given factor is fully instantiated (all entries nonzero- - either +1 or - 1), and we attribute the failure to this dependence being difficult to learn.
An alternative option, which works better in practice and is tested in our experiments, takes advan- tage of the factors always MLPs or bilinear functions, in our experiments being able to take arbitrary real-valuee inputs, not just the discrete values in the sample space. Therefore, not-yet- sampled variables can simply be set to zero (the value corresponding to a masked input). We simply define
R (I(VI....)))
[ dx (x (11.....1)..
(12)
where x vi...... denotes the masked sample that is the input the masked autoencoder, with indi- cating absence of a value for all not-yet-sampled variables. Thus the factors for which not all values have been instantiated are computed with inputs of for the missing variables.
14

---

Published as a conference paper at ICLR 2024
C A-AI WITH CONTINUOUS VARIABLES
C.I CONTINUOUS A-AI AND SCORE MATCHING
Beyond the discrete spaces considered in this paper, it is interesting to consider the generalization of A-AI to continuous sample spaces, where the variables take real values. Recall that for two densities and on Rd, the Fisher divergence is defined as
i ( d l p(x) { dx,
(x) V log (xx) - V, (x))())] Ex-9(x)
(13)
Ox,
The objective of score-based generative models (e.g., Song & Ermon (2019)) optimizes (13) with respect to (or its score V logg p) against a fixed distribution convolved with noise, which can be tractably sampled if dataset is available. This is distinct from our motivation of matching q to an intractable distribution p. However, we have the following result.
typically the data distribution
Proposition 2 (informal). In a PGM over real-valued variables, the expected A-AI objective (10) x sampled from p. chosen uniformly at random, and additive perturbations x" Xu h approaches the score matching objective as h 0.
Prop. 2 has implications for training generative models over continuous-valuee data with known graphical model structure. It can motivate the use of scores in place of the perturbation ratios in (10) when working in continuous spaces: the scores share with the perturbation ratios the property of being independent of the variables outside the Markov blanket of u.
Figure C.I: Illustration of the joint positions in the human poses dataset, Each human pose is specified by 24 joints, as illustrated above, according to the SMPL body model (Loper et al., 2015). Each joint is specified by 3-dimensional joint angles of range [-x, x and connected to set of other adjacent joints. The overall adjaceney graph forms trec, such that the graph is already chordal.
Single DAG order
0.5 Amortizing multiple DAG orders
A-Al
A-Al
TB
TB
DB
DB
Test MMD
FL-DB
FL-D8
0.2
0.2
0.1
0.1
0.0
0.0
20
40
60
80
20
40
Wall-clock Time (min)
Wall-clock Time (min)
Figure C.2: Experimental results on the Poses dataset. We report mean and standard deviation over 3 runs.
C.2 EXPERIMENT ON POSES DATASET
We further validate A-AI on a real-world dataset with continuous variables. This a fully-observed setting with no latent variables.
Dataset. We use the AMASS dataset (Mahmood et al., 2019), which a large collection of record- ings of real human poses. For this experiment, we focus on the KIT subset (Mandery et al., 2016),
15

---

Published as a conference paper at ICLR 2024
and sample 100,000 instances from this subset. 20% is held-out test set. Each human pose is specified by 24 "joints", each of which represent a body part (neck, spine, left knee, right knee, etc.) that can take be articulated relative adjacent body parts. See Figure C. for the illustration. Each joint orientation is specified by 3-dimensional axis-angle representation in [-x, x]3 The overall graph forms tree, which is chordal.
Factor graph and amortized sampler. We assume that the data- generating factor graph has as ad- jacency matrix the structural connectivity of the joints in the dataset (i.e. two joints are connected together by a body-part, such as the left wrist and left clbow, then these two variables will share an edge in the factor graph structure). Each factor takes variables (2 joints 3 angles) as input, and each factor is parameterized by a separate 3-layer multilayer perceptron with ReLU activations. The amortized sampler outputs the parameters of Truncated Gaussian random variable (truncated the range 1-1, I] to model the conditional distributions in the factor graph. We found that optimizing truncated Gaussian distributions more stable than Beta. We jointly train the sampler approx- imate the factor graph distribution using A-Al and the factor parameters (the energy function) using (4).
Baselines and evaluation metric. We consider the same set of baselines as in Figure 4: TB, DB, and FL-DB. Note that for all the baselines, the sampling is done using the same sparse graphical model constraints (each variable is sampled conditioned only on their parents) for fair comparison. The evaluation metric linear MMD between the test data and the generated samples from the amortized sampler.
Results. Figure C.2 shows that A-AI outperforms all the baselines by significant margin in terms of wall-clock convergence on the MMD. This tendency holds regardless of whether learn to sample according to single DAG I-map of the factor graph (Figure C.2, Left), or we amortizing multiple DAG orders using the same neural network (Figure C.2, Right).
D PROOFS
Proposition 1. Suppose that p : X - B. is the density of a Markov network with factors 6x and that q X R=0 is the density of a Bayesian network with respect to an I-map (V. Then the following are equivalent: for all x,x' differing in a single variable X (9) holds; (2)
Proof of Proposition 1. Part 1: (1) = (2). Suppose that 4 = p, i.e., the conditional distributions specified by the Bayesian network are the conditionals of the distribution p.
As stated before, the A-AI constraint stems from an equivalence of the Markov and Bayesian net- work factorizations.
Suppose that x, X are two settings of the variables that differ in exactly one variable u, i.e. xu # x and Xy Xy for all v E V {u}. Using the Markov network factorization (1), we have
p(x) 1 [l tu (Ase)
Φ (xS.)
Φx (xs,).
[T
(14)
p(x')
11 Φx (x'sz)
Φx (''s)
kinESa &x (x'ss.
where the last equality uses that x and x' differ only at u, while the factor Φe is independent of x unless u E Sx.
On the other hand, using the Bayesian network factorization (2), we have P (x, (PP(v))
p(x)
live P (xp xPa(v))
I
(xv (xPavv))
(15)
p(x')
Πvvv p (xt x
veV p x xppcc)) (E(M)CCh(u) P (x4 (xpatv)
The last equality similarly uses that the numerator and denominator can differ only if the conditional of xy given its parents depends on Xu. which occurs and only if or v child of The A-AI constraint (9) expresses precisely the equality of the two expressions (14) and (15).
Part 2: (1)=(()) Conversely, we must show that if the equality holds for all x, x, differing in a single variable, then p qq. Because the left side equals of the equality equals ((x)) by (14), and
its right side equals g(x) प(ズ" by (15), the equality, together with positivity of all densities, implies that p(x) p(x)) [ix) q(x))
16

---

Published as a conference paper at ICLR 2024
The transitive closure of the relation on X defined by x x' if and x' differ single variable the trivial equivalence relation, i.e., any x, E X are lined by chain x = x1 ~ X2 - Xn such that each x,, Xi+1 differ in single variable. It follows that for any x, E X, ((x) g(x")
plx) p(x))
Therefore, the two densities are proportional, which implies that they are equal.
Proposition 2 (informal). In a PGM over real-valued variables, the expected A-AI objective (10) over x sampled from p. u chosen uniformly at random, and additive perturbations x Xy + approaches the score matching objective as h - O.
Proof of Proposition 2. We assume the log-density log is continuously differentiable. Let ej be the unit vector along coordinate i. The A-Al objective then
log p(x + hez). p(x)
log go(x+ + hei)
log ((x)
i Olog (x()) h
h + 0(1)) dxi à alog (o())
Bx-p(x) Bi
= Bx-p(x)
(x
0x;
Ölog
= 2?2
B. Bx-p(x)BB
+ +(()
Ox;
dx,
h²2
d (Bx-p(x) II log p(x)
V. logg
We see that the A-AI objective, normalized by 4. approaches the gradient of the Fisher information (13) as h O.
E STOCHASTIC LOSSES FOR SQUARED-SUM OBJECTIVES
Proposition Suppose 1. Let S Ei fi and {g S)?, where g and each fi are
functions of 0. Let
Li 2 (x
Lij 2
(n 1)fi +5²
with and fi indicating that gradients are blocked. i.e.,
dă 00
0 and
of
0. Then
aL
0L:
iLi,j
E,
Bizj
(16)
00
00
00
where B::l] is uniform average over indices in {1......) and Bitj [] is uniform average over all pairs of different indices in ( {1.....n} n).
Proof. The gradient of L is:
0L.
0g as
(g + S)
00
00 00
0g
S 0g
35
0f
8 a0
+
+
Σ
fi
Σ
fi
(17)
00
00
00
80
[#]
Let us now show that we recover the same terms from the gradient of the right side of (16):
0L;
ög de
B;
E,
(x
nfii
00
0g
0g
S
00
00
n(g (n 1) fn) 0fj
B
B
00
as = 00
0fj
0fi
Σ
fj
+
00
which exactly recovers the five terms in (17).
17

---

Published as a conference paper at ICLR 2024
ising Ladder, dimtV)=64
Ladder, dimtV)2256
Ladder, dim
0.64
0.75
0.58
0.70
0.57
0.60
0.58
0.55
0.54
0.54
0.55
0.53
10
107
10
10
10
10
10
107
10'
10
107
30"
10
10
10a
10'
107
107
10*
107
107 Ising Lattice, dimiV)
(e) lsing Lattice, dimiV)-256
(1) linng Lattice, LimVVi=1044
0.60
0.75
0.59
0.70
0.50
0.80
0.65
0.57
0.60
0.50
0.5
0.54
10"
10'
(g)
dimoVi=64
elm(V1=256
dimi) 1024
0.55
0.64
0.62
0.00
0.53
0.58
0.58
2 0.50
0.51
0.54
0.54
0.52
0.50
0.50
10a
10
10"
10a
10"
100
10c
10
10"
10"
15
10'
10
10
10'
10"
10P
Wall-cook Time (sec)
Wall-clokk Time (sec)
Wall-llokk Time (sec)
TB
D8
FL-D8
SubTa
FL-SubTB
--Al
Figure F.I: Wall-clock training convergence on synthetic experiments with unstructured GFlowNet samplers from Zhang et al. (2022).
F EXTENDED EXPERIMENT DETAILS AND RESULTS
F.1 SYNTHETIC EXPERIMENTS
Energy models. In the synthetic experiments, we use either Ising models or factor graph models.
The Ising model is a Markov random field, distribution over D-dimensional binary vectors. Each element in vector has a value of either or +1. The following energy model & specifies the distribution over those binary vectors x: p(x) exp(-E(x)),
8(x) o (-xTxxxxxxxx
(18)
where J is a symmetric interaction matrix and is a unary potential vector. or > O controls the sharpness of the energy function. Note that J is sparse according to assumption of sparse graphical model (see Fig. 3a and Fig. 3b). Each non-zero element of J and b has a value of either -1 or +1. In Figure 4 (a.b) we set o 0.2, and in Figure 5 we set σ 2 (when [V] 64) and o 1 (when [V] 256).
In the factor graph model, each factor is defined as a tiny MLP which has hidden layer whose dimension is 10. Each MLP takes 4 arguments an input and outputs single scalar (see Fig. 3c). The parameters of those MLPs are randomly initialized with the spherical Gaussian distribution o2 We use o = 0.5 in Figure 4 (c,d).
Generating ground-truth Gibbs samples. We use NLL and MMD as an evaluation metric, which are based on the ground-trutt samples from the target energy function. Therefore, it is crucial to have accurate ground-truth samples for correct evaluation. We generate those ground-truth samples by running 10k independent Gibbs chains for 10k steps. When the target energy functions are too peaky (in Fig. 5), we further anneal the energy function for the additional first 10k steps (i.e., initially let the energy function be smooth and gradually lower the temperature during the first 10k steps) to prevent those chains from being attracted by only few dominant modes.
18

---

Published as a conference paper at ICLR 2024
Model architecture and training details. For all the GFN baselines and our method, the amorti- zation network an MAE of the following form: (Linear(512) - LayerNorm - ReLU) x 3 - Linear((V)). The other training details are as follows.
A-AI vs. GFNs (Fig, 4). We consider [V] E [64, 1024). Each model is trained for total 200k iterations.
Baseline GFNs: Batchsize is set to 1k. We periodically sample new DAG (every 50 iteration) to amortize over random DAG orders. Learning rate of the parameters of the amortized sampler is set to 10-3 and that of the partition function estimator is set to 10 Those learning rates are step-wisely decayed by 0.1 at 40k, 80K, I 120k, 160k, and 180k-th iteration. For the training policy, we use c-uniform policy such that the sampling probability p is defined as p = (1 c) . + C- Puuifoomm. where Pon and Puniform are on- and uniform policy, respectively We use
0.1.
A-AI: We periodically sample set of sub-DAGs (every 50 iteration) to amortize over random DAG orders. For instance, when there are 64 variables (or nodes), we sample 16 sub-DAGs for each node (see Fig. 2), resulting in total 64 x 16 = 1024 batchsize. Learning rate of the parameters of the amortized sampler is set to 10-3, with the same learning rate scheduling as the baseline GFNs. For the training policy, we simply use the tempered off- with temperature to 2.
A-AI vs. MCMCs (Fig, 5). We consider [V] E {64, 256).
Baseline MCMCs: We run 10k parallel Gibbs chains in this experiment. At each evaluation step, we collect all those 10k samples and measure the MMD between the ground-truth samples.
- A-AI: In this experiment, we do not amortize DAG orders because the flexibility of inference not the focus of this experiment. Note that when the DAG order is fixed, we cannot use the efficient on-policy training as discussed in $3. However, we can simply only periodically sample the full dimensional on-policy samples to avoid the cost (e.g.. every 10 iteration). Due to the peaky energy landscape, we found that gradually annealing the energy function for the first 10k steps helps stabilize the training. Also, for the first 10k steps, we use high temperature (e.g.. 10) for better exploration and gradually lower the temperature to 1 (on-policy). We train total 10k steps, with the batchsize set 10k and learning rate set to 10-3, Learning rate is step-wisely decayed at 2k, 4k, 6k, 8k, and 9k-th step.
A-AI vs. any-order GFlowNets. In past work, GFlowNets have been used as amortized samplers for unstructured energy-bat based models (Zhang et al., 2022). In that work, the GFlowNet sampler selects values for the variables in arbitrary order: while the models in $2.3 simply train the con- ditionals in the Bayesian network needed to sample the variables in topological order, there, the generative policy selects both a variable and its value.
Fig. F.1 shows the training convergence on the synthetic experiments in the same way as Fig. 4. Among the baselines, TB shows the fastest and most stable convergence, whereas DB and SubTB show relatively unstable convergence, suggesting that the flow function harder to train than for- ward transition probabilities. As result, while FL variants (FL-DB and FL-SubTB) outperform DB and SubTB, they still underperform TB. A-AI shows much faster convergence.
All experiments are run with 5 random seeds, and the mean std regions are shown in the plots.
F.2 LATENT VARIABLE MODELING OF MNIST
Pyramid graph structure. We expand on the specific pyramid-shaped graph structure used for the latent variable model in Fig. 6, which is used for all experiments. The first hidden layer has dimensions 6 xx6 6, the second 2 x 2, and the third is a single variable node. Each latent variable node is composed of 4 bernoulli random variables which are fully-connected. Each variable node in the first hidden layer form a clique with a window of 8 x variable nodes in the pixel-layer below. This clique window reduces to 4 x 4, then x 2, for the second and third hidden layers respectively. All edges in V H/ are removed such that pixels are conditionally independent given xn-
The resulting graph has 4608 edges between X and H (3.6% of all possible edges) and 7990 edges in the induced subgraph on H (59.8% of all possible edges).
Model architecture and training details. The model architecture in all experiments that involves an amortized sampler qo (all except the Gibbs baseline) is a masked autoencoder (MAB), as de- scribed in $3. It consists of a shared trunk (3 fully-connected layers with ELU activations), and one output head (a linear layer) for each variable in V. Layer normalization and residual connections
19

---

Published as a conference paper at ICLR 2024
(b) NLL vs. Time
(c) Classifier Entropy vs. Time
Single DAG
1.25
Single DAG
200
Multiple DAGs
Multiple DAGS
1.00
175
Entropy
0.75
150
125
0.50
100
0.25
2
4
6
10
12 0
4
8
10
12
Wall-Clock Time (h)
Wall-Clock Time (h)
Figure F.2: Comparison of A-AI trained for full inference single full DAG partial inference multiple DAGs. All DAGs are I-maps for the underlying undirected pyramid-shaped graphical model (sec Appendix F.2).
Reconstructions
Vanilla VAE
7
#
Gibbs
7
FL-D8
4
9
D8
77
TB
27
Mean-Field EM
7
Sleep
7
IW
7
4
4-5
7
Ground Truth
7
4
Figure F3: Reconstructions from A-Al and all baselines in $ F.2. The samples compared to the ground- truth samples (bottom) that used input, and vanilla VAE (top).
used in the shared trunk. All layers are of dimension 512. We used this MAE amortize condi- tionals, and used a separate vector of learnable parameters of size [V] to infer the marginal for each variable.
We tuned the following hyperparameters:
Learning rate for go and Pw in [10-f, 30 for the conditionals, with the learning rate for the marginals fixed to the learning rate for the conditionals multiplied by factor of 100.
Temperature (only applicable to A-AI and the GFlowNet baselines, which use tempered sam- ples for greater exploration) in [1.5, 2, 4]. Tempered samples were drawn with probability e E [0.05,0 0.2, In practice, we found A-AI to benefit from very light use of tempered sampling. with e 0.05 and temperature The GFlowNet baselines did best with larger values of epsilon.
Training of the bilevel objective was done in an alternating fashion: qo for 100 iterations, followed by Ps for 100 iterations. For the Gibbs variant, the number of steps set 1; any larger number was prohibitively slow.
Mini-batch gradient descent was used along with the Adam optimizer. step-wise learning rate decay (reducing the learning rate by half) was scheduled at 40%, 70%, and 90% of the total number of training iterations.
Batch size was 1148 for all experiments (which the number of latents, 164, multiplied by 7). In the case of the A-AI variant that learns multiple orders by doing partial inference over subsets, an I- map for the sub-graph consisting of variable and its Markov blanket was sampled for each variable (totalling 164 1-maps), and 7 samples were drawn from each I-map. For all other experiments pertaining to full inference, a single I-map for the full graphical model was sampled, and 1148 full samples were drawn. In the case of the A-AI variant trained for full inference on a single DAG (Fig. F22), each of the 164 variables were perturbed 7 times in 7 different samples. For the experiments learning multiple DAG orders, a new DAG I-map (or set of DAG I-maps in the case of A-AI trained on partial instantiations) sampled every 10 iterations. For the importance-weighted
20

---

Published as a conference paper at ICLR 2024
variant, used 7 independent importance-weightee samples per ground-truth sample (and reduced batch-size by factor of to maintain the final batch-size equal to the other baselines).
We estimated NLL using the importance-weiighted variational lower bound, with independent importance-weigheee samples per ground-truth sample. A standard pretrained MNIST classifier was used compute the prediction entropy on unconditional samples from P*- This classifier is a convolutional neural network with two convolution layers (10 channels, kernel size of 5, stride followed by 20 channels, kernel size of 5, stride 1), each followed by a max max-pooling layer and ReLU activation, followed by two ReLU linear layers. Dropout (probability of 0.5) was applied after the convolutional block and between the two fully connected layers. Test accuracy was 0.9887.
21


Reviews:
[{"review_id": "76h1yyRM5K", "reviewer": "Reviewer_mXkv", "summary": "This work presents Delta-AI, an algorithm for the problem of amortized inference in PGMs. The main idea behind the paper is to leverage the sparsity of PGMs to pose a local constraint that can then be used as a novel loss function for GFlowNets. This local constraint makes GFlowNets more efficient both in terms of memory and time since only the few relevant variables can be used in each training step. The proposed inference algorithm is tested on a sytnetic data set as well as MNIST data and outperforms vanilla GFlowNets and other amortized inference algorithms.", "questions": "1. How will the proposed inference algorithm scale to natural images?\n\n\nP.S: Being an emergency reviewer I might have missed some specifics and thus am lowering my confidence. Looking forward to the rebuttal.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Interesting use of local constraints in GFlowNets for amortized inference within the realm of probabilitic graphical models. \n\n2. The proposed DELTA-AI seems sound and performs pretty well in the shown experimental settings.", "weaknesses": "1. The experiment settings are a bit weak in my opinion. Formalizing the idea and getting a a proof of concept is fine using a synthetic data. Also if there are no real data sets available that can also be considered but just considering MNIST as the real data is a bit limiting and does not show the full power of DELTA-AI. For example why not use the DELTA-AI loss within PixelVAE and try inference on natural images?\n\n2. The paper is based on taking advantage of local credit assignment -> local losses this part gets kind of obsfuscated as the paper goes along. (but this is a simple fix by rewriting a bit of portions in sections 3)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qa7XQ0tcHs", "reviewer": "Reviewer_iUke", "summary": "The authors propose a technique for amortized inference in sparse probabilistic models which they call \\Delta-AI (\\Delta-amortized inference) that takes advantage of sparse model structure (previously specified). This is done by matching the conditional distribution of a variable given its Markov blanket. The sparsity of the graphical model allows for local credit assignment in the policy learning objective. They experiment with synthetic PGMs and latent variable models with sparse factor graph structure to show the algorithms effectiveness.", "questions": "Not having direct background / experience with GFlowNets, I found myself wondering why have two representations of the model (the \nfactorized Markov network p and the inference network q) - is the assumption that $p$ is provided (not just the structure / factors (sets of nodes) but also the exact parameters - or are the parameters of both learnt?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The authors show how to take advantage of  known (or assumed) graphical model structure to allow for local credit assignment. Computationally, this lowers the memory requirement as parameter updates only requires instantiating a single variable and its Markov blanket.", "weaknesses": "The paper was hard to read - while there was a fair amount of discussion about the graphical model basics, there was not much about GFlowNets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qEgbdLwBoP", "reviewer": "Reviewer_1ZeM", "summary": "This paper presents Δ-amortized inference (Δ-AI), a new algorithm for efficient inference in sparse probabilistic graphical models (PGMs). Δ-AI leverages sparsity to assign credit locally in the agent's policy learning, enabling faster off-policy training without the need to instantiate all random variables. The trained sampler matches conditional distributions and recovers marginals, making it effective for inference in partial variable subsets. Experimental results demonstrate its efficiency for synthetic PGMs and latent variable models with sparse structures.", "questions": "- The paper hinges on the assumption that the factors of the distribution $\\phi_{k}$ are known. This seems like a stringent condition, and it is worth exploring how the framework and algorithm proposed in the paper can be extended to scenarios where these factors are unknown and need to be learned.\n\n-  The paper asserts that it accomplishes simultaneous amortization over multiple DAG structures. However, it would be beneficial to provide a more detailed explanation of how this simultaneous amortization is achieved.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written and organized. It presents a clear and compelling motivation for the problem at hand. The discussion of prior works is thorough and well-structured, and the paper offers promising avenues for future research.\n- The paper develops a novel algorithm namely, Δ-AI that offers substantially faster training when compared to regular GFlowNets and other amortized inference methods, since each training step involves only a small subset of variables, resulting in negligible sampling costs. Furthermore, the memory cost associated with Δ-AI is exceptionally low, since the algorithm leverages the sparsity of the probabilistic graphical model to compute the gradient update locally over a small set of variables. Lastly Δ-AI provides the advantage of flexible probabilistic inference by amortizing numerous potential sampling orders into a single sampler, by learning Bayesian networks with multiple Markov-equivalent structures.", "weaknesses": "As the current paper falls outside the scope of my research interests, I am unable to identify any significant weaknesses in the paper. Consequently, my confidence in assessing the paper is limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)