GFlowNet-EM for Learning Compositional Latent Variable Models

Edward J. HuNikolay MalkinMoksh JainKatie E. EverettAlexandros GraikosYoshua Bengio

article2023ICML50 citations

Proposes GFlowNet-EM, a framework that replaces the intractable expectation-maximization E-step with an amortized GFlowNet sampler to train expressive latent variable models over discrete compositional structures without imposing restrictive independence assumptions.

Listen

Real-world data such as natural language and visual scenes often rely on hidden, structured components—like parse trees in sentences or object arrangements in images. Latent variable models aim to capture these hidden structures; however, standard maximum-likelihood learning using the Expectation-Maximization (EM) framework becomes mathematically intractable when evaluating combinatorially large discrete spaces. Traditional approaches either enforce rigid independence assumptions, which severely limit the expressiveness of the model, or use slow Markov Chain Monte Carlo sampling, which scales poorly. Consequently, machine learning practitioners have struggled to train expressive models that accurately capture complex, interdependent latent structures without computational bottlenecks.

The article introduces and evaluates GFlowNet-EM, a novel training framework that incorporates Generative Flow Networks (GFlowNets) into the EM algorithm to learn expressive latent variable models with discrete compositional hidden states. The objective is to demonstrate that GFlowNets can act as flexible, amortized samplers for the intractable inference step (E-step) without relying on artificial conditional independence assumptions.

The authors assess the proposed approach across synthetic and real-world benchmark tasks, including hierarchical Gaussian mixtures, grammar induction on the Penn Treebank dataset (under standard context-free, energy-guided, and non-context-free grammars), and discrete image representation learning using the static MNIST dataset. To make joint optimization stable and prevent the learned posterior from collapsing into a single mode, the approach incorporates several stabilizing mechanisms, including adaptive thresholding to schedule updates, off-policy exploration, a sleep-phase hallucination step, and short refinement chains.

The empirical findings demonstrate clear performance gains over standard baselines. First, in grammar induction, GFlowNet-EM matches exact baseline methods on standard context-free grammars and enables tractable learning on more complex non-context-free grammars, achieving a superior test negative log-likelihood (at or below 5.46) where exact methods are computationally prohibitive. Second, when augmented with an energy-based prior to guide grammatical tree structures, GFlowNet-EM improves human-annotated constituent alignment from a 39.51 F1 score (standard marginalization) to 48.41. Third, in discrete image modeling, GFlowNet-EM consistently outperforms standard vector-quantized variational autoencoders (VQ-VAEs) across dictionary sizes, lowering test negative log-likelihood on MNIST from 80.84 to 70.74 for an eight-token dictionary while enabling simultaneous learning of the latent prior.

These findings indicate that machine learning models do not need to sacrifice representational richness to remain computationally viable. By eliminating restrictive factorization and conditional independence requirements, GFlowNet-EM unlocks the training of more sophisticated generative architectures for language and vision, offering better alignment with structured human knowledge and domain rules. Although the method reduces asymptotic theoretical complexity on complex grammars from quintic to quadratic in sequence length, it introduces practical computational overhead during training due to sequential autoregressive sampling and decoupled optimization alternating between network updates.

Organizations and research teams developing generative systems for highly structured data should consider GFlowNet-EM as an alternative to restrictive variational approximations when modeling complex discrete dependencies. When deploying this method, teams must factor in the trade-off between higher training time and enhanced expressiveness, and they should leverage stabilization techniques like sleep-phase exploration to prevent posterior collapse. Future work should focus on software and algorithmic optimizations to reduce training overhead and expand the framework to continuous or hybrid continuous-discrete latent representations.

CIFAR Google Massachusetts Institute of Technology Mila Stony Brook University Université de Montréal grammar induction inference inference algorithm GFlowNet-EM compositional latent variable models generative flow networks

arXiv: 2302.06576 GFNOrg/GFlowNet-EM

Paper: GFlowNet Foundations, Yoshua Bengio et al. (2023). Read this for the flow-conservation objectives and amortized-sampling theory that GFlowNet-EM uses to approximate its intractable E-step.
Paper: Generative Flow Networks for Discrete Probabilistic Modeling, Dinghuai Zhang et al. (2022). Its energy-based GFlowNet setup makes the discrete-space sampling machinery behind GFlowNet-EM concrete before the method combines it with EM.

No sufficiently relevant recommendations were found.

Cover for GFlowNet-EM for Learning Compositional Latent Variable Models

Abstract

Latent variable models (LVMs) with discrete compositional latents are an important but challenging setting due to a combinatorially large number of possible configurations of the latents. A key tradeoff in modeling the posteriors over latents is between expressivity and tractable optimization. For algorithms based on expectation-maximization (EM), the E-step is often intractable without restrictive approximations to the posterior. We propose the use of GFlowNets, algorithms for sampling from an unnormalized density by learning a stochastic policy for sequential construction of samples, for this intractable E-step. By training GFlowNets to sample from the posterior over latents, we take advantage of their strengths as amortized variational inference algorithms for complex distributions over discrete structures. Our approach, GFlowNet-EM, enables the training of expressive LVMs with discrete compositional latents, as shown by experiments on non-context-free grammar induction and on images using discrete variational autoencoders (VAEs) without conditional independence enforced in the encoder.

Topics

grammar induction inference inference algorithm latent variable models optimization stochastic variational autoencoders variational inference latent factor variational approximation

Keywords

GFlowNet-EM compositional latent variable models generative flow networks conditional GFlowNet probabilistic context-free grammars variational EM variational inference discrete compositional latents

1. Introduction
2. Background
2.1. Expectation-Maximization (EM)
2.2. GFlowNets
3. Motivating Example: Pitfalls of Factorization
4. GFlowNet-EM
4.1. GFlowNet-EM Optimization Techniques
4.2. Improving Posterior Estimation
5. Empirical Results
5.1. Hierarchical Mixture Revisited
5.2. Grammar Induction on Penn Tree Bank (PTB)
5.2.1. Context-Free Grammar
5.2.2. CFG with Energy-Based Model Guidance
5.2.3. Non-Context-Free Grammar
5.3. Discrete Variational Autoencoders
6. Related Work
6.1. GFlowNets
6.2. Latent Variable Models and EM
6.3. Applications
7. Conclusions
Acknowledgements
References
A. On GFlowNet optimization techniques
B. Hierarchical mixture
C. Grammar induction
C.1. Experiment setup
C.2. GFlowNet parametrization
C.3. Marginalizing preterminals
C.4. Training energy-based model prior
C.5. Non-context-free grammar parameterization
C.6. Time complexity analysis
C.7. Ablation studies
C.8. Sample parses from grammars learned by GFlowNet-EM
D. Discrete VAE
D.1. Experiment setup
D.2. GFlowNet-EM visualizations
E. Computation cost in practice

Knowls

Knowl 1 — GFlowNets provide an amortized E-step for compositional latent-variable models
model/method
GFlowNet-EM jointly trains a generative model $p_\theta(z,x)=p_\theta(z)p_\theta(x\mid z)$ and a conditional GFlowNet whose terminal distribution $q_\phi(z\mid x)$ approximates the posterior over latent structures. For each observed data point $x$ , the GFlowNet is trained with reward
$R_\theta(z\mid x)=p_\theta(z)p_\theta(x\mid z),$
so that an ideally trained GFlowNet terminates at $z$ with probability proportional to the posterior’s unnormalized density. The E-step updates the GFlowNet using trajectories conditioned on $x$ and this reward. The M-step samples $z\sim q_\phi(z\mid x)$ and updates the generative-model parameters to increase $\log p_\theta(z)p_\theta(x\mid z)$ . Thus, the M-step’s sampled gradient estimates the update under the approximate posterior, while the GFlowNet can represent dependencies among components of a discrete latent without imposing a factorized posterior.
Knowl 2 — Loss thresholding adapts the number of E-step updates
algorithm
GFlowNet-EM uses a moving average of the GFlowNet training loss as a heuristic for deciding when its posterior approximation is accurate enough to update the generative model. Given data, generative-model parameters $\theta$ , conditional-GFlowNet parameters $\phi$ , an exploratory sampling policy $P_F^\#(\cdot\mid x)$ , and a loss threshold $\alpha$ , the procedure is:
Input: Data, generative model parameters $\theta$ , GFlowNet parameters $\phi$ , exploratory policy $P_F^\#$ , threshold schedule $\alpha$
Repeat until a convergence condition is met:
    Sample an observed example $x$
    Sample a trajectory $\tau$ from $P_F^\#(\cdot\mid x)$ and let $z$ be its terminal state
    Compute the trajectory-balance loss using reward $p_\theta(z)p_\theta(x\mid z)$
    Update $\phi$ to reduce the trajectory-balance loss
    If the moving-average GFlowNet loss is below $\alpha$ :
        Sample $z$ from the current conditional GFlowNet $P_F(\cdot\mid x)$
        Update $\theta$ to increase $\log p_\theta(z)p_\theta(x\mid z)$
A lower threshold demands a more accurate posterior approximation before an M-step; a higher threshold triggers M-steps more frequently. For the grammar experiments, the threshold was linearly decreased from 6 to 3 over 10,000 updates. The authors report that thresholding is needed to prevent generative-model collapse, in which most grammar symbols receive almost no posterior mass; threshold choice trades off fitting speed and accuracy.
Knowl 3 — A trajectory-based sleep phase helps prevent posterior collapse
model/method
The sleep phase trains a conditional GFlowNet on examples drawn from the current generative model: sample a latent structure $z\sim p_\theta(z)$ and data $x\sim p_\theta(x\mid z)$ , then sample a trajectory $\tau$ from the backward policy leading to $z$ given $x$ . Update the forward policy to increase the probability of that trajectory conditioned on $x$ , equivalently minimizing $-\log P_F(\tau\mid x)$ . This avoids computing the generally intractable terminal-state probability, which sums the probabilities of all trajectories leading to $z$ . For fixed backward-policy parameters, the global optimum of this trajectory-likelihood objective with respect to the forward policy maximizes the terminal probability of recovering $z$ given $x$ . The authors use sleep-phase training to improve mode coverage and mitigate posterior collapse; they caution that it may be unhelpful when model-generated data is unlike real data.

Knowl 4 — Exploration, local credit assignment, and MCMC stabilize joint training

empirical result

To address the changing reward during joint learning, the grammar experiments combine off-policy exploration, a forward-looking subtrajectory-balance (SubTB) loss, a sleep phase, and short MCMC chains initialized from GFlowNet samples. Exploration uses either a tempered policy, with action probabilities proportional to $P_F(s'\mid s,x)^\beta$ for $\beta<1$ , or an $\epsilon$ -uniform mixture with the uniform action distribution. Forward-looking SubTB incorporates the accumulated partial log-reward into the state-flow estimate, giving local credit when reward factors accrue during sequential construction. The grammar experiments used 10 MCMC steps.

A five-seed ablation on joint CFG learning reported the following held-out negative log-likelihood (NLL) per word and sentence F1:

Configuration	NLL / word $\downarrow$	Sentence F1 $\uparrow$
GFlowNet-EM	$5.70 \pm 0.03$	$34.85 \pm 3.39$
Without MCMC	$6.02 \pm 0.01$	$28.56 \pm 0.55$
Without sleep phase	$5.91 \pm 0.04$	$28.13 \pm 0.43$
Without SubTB	$5.84 \pm 0.08$	$26.56 \pm 7.82$
Without exploration	$5.70 \pm 0.02$	$31.87 \pm 1.06$
Sleep phase only	$6.08 \pm 0.06$	$48.41 \pm 1.38$

The complete configuration gave the strongest combined likelihood and F1 results among the reported ablations. Removing MCMC, sleep, or SubTB worsened both metrics; removing exploration left NLL unchanged but lowered F1. Sleep-only training had the highest F1 in this ablation but the worst NLL, illustrating that the methods do not improve every metric uniformly.

Knowl 5 — A bottom-up GFlowNet constructs parse trees from ordered forests
model/method
For grammar induction, a GFlowNet state is an ordered forest whose trees represent partial constituents of a sentence. An action chooses two adjacent trees and joins them beneath a new parent; a sentence of length $n$ therefore requires $n-1$ joining actions to form a binary parse. The policy processes the roots of the forest with a six-layer Transformer of hidden dimension 512. A recursive MLP aggregator combines each node’s symbol embedding with embeddings of its children, so each root representation summarizes its whole tree. MLP heads on the Transformer representations parameterize the forward and backward policies and state-flow estimator; sum pooling gives the flow estimator a scalar output for forests of different sizes. Rather than sampling preterminal tags as extra actions, the method marginalizes over their possible values when evaluating tree rewards, which can be done in linear time for the grammars studied.
Knowl 6 — Parent-conditioned production rules yield a non-context-free grammar
model/method
The paper extends a context-free grammar by allowing a nonterminal’s expansion to depend on its parent nonterminal. If $P$ is the parent of nonterminal $X$ , and $L,R$ are the left and right child symbols in the production $X\to LR$ , the rule probability is proportional to a product of two factors:
$p_\theta(L,R\mid X,P)\propto f_1(L,R,X;\theta)f_2(L,R,P;\theta),$
with normalization over possible $(L,R)$ productions. Setting $f_2$ identically to 1 recovers the context-free case. Although a grammar of this form can be represented as a context-free grammar with $|NT|^2$ nonterminal symbols, directly training that enlarged grammar was considered prohibitive. Exact posterior sampling for the parent-conditioned grammar has quintic time complexity, motivating the use of an amortized GFlowNet posterior sampler.
Knowl 7 — GFlowNet-EM matches exact EM on a hierarchical mixture where mean field fails
empirical result
The hierarchical-mixture experiment has four supercluster means arranged on a grid, with four petal subclusters at fixed offsets around each mean. Each observation is generated by choosing a supercluster and petal uniformly and sampling from a fixed-variance normal centered at the corresponding component; the means are the parameters to learn. Because the posterior over supercluster and petal assignments is dependent and multimodal, a factorized approximation can make incorrect joint assignments.

Across 20 random seeds after 60 iterations, data log-likelihood per sample was $-5.79\pm0.74$ for exact EM, $-7.26\pm1.12$ for factorized variational EM, and $-5.77\pm0.48$ for GFlowNet-EM. The ground-truth means give $-5.62\pm0.01$ . Exact EM and GFlowNet-EM nearly recovered the true means, whereas factorized variational EM failed to learn them. In this small experiment, the GFlowNet used two sequential actions to select the supercluster and petal, and its reward could omit the uniform prior because it was constant across assignments.

Knowl 8 — Grammar induction results show benefits for non-CFGs and shape priors

data/table

The grammar comparison evaluates held-out NLL per word and sentence F1 against human-labeled constituent spans. Models were trained for 10,000 grammar updates on Penn Tree Bank sentences of at most 20 tokens; the table reports means and standard deviations over five seeds. For the non-context-free grammar, the reported NLL is a variational estimate, and exact marginalization or exact posterior sampling is intractable.

Grammar	Method	NLL / word $\downarrow$	Sentence F1 $\uparrow$
CFG	Marginalization	$5.61 \pm 0.01$	$39.51 \pm 7.01$
CFG	Exact-sampling EM	$5.74 \pm 0.05$	$31.17 \pm 6.06$
CFG	MC-EM	$5.88 \pm 0.01$	$22.31 \pm 1.04$
CFG	MC-EM + EBM prior	$5.91 \pm 0.02$	$23.81 \pm 1.41$
CFG	GFlowNet-EM	$5.70 \pm 0.03$	$34.85 \pm 3.39$
CFG	GFlowNet-EM + EBM prior	$5.79 \pm 0.03$	$48.41 \pm 1.38$
Non-CFG	MC-EM	$-$	$18.98 \pm 0.26$
Non-CFG	GFlowNet-EM	$\leq 5.46 \pm 0.07$	$38.68 \pm 1.90$

Marginalization performs best among the CFG methods on NLL and F1, while GFlowNet-EM improves over exact-sampling EM on both reported metrics. The CFG with an annealed energy-based prior (EBM) achieves the highest CFG sentence F1, showing that the prior can encourage tree shapes resembling human annotations. For the non-CFG grammar, GFlowNet-EM obtains substantially higher F1 than MC-EM and reports the variational NLL estimate; the experiment demonstrates learning with a grammar whose exact posterior methods are intractable.

Knowl 9 — An autoregressive GFlowNet serves as a discrete-image encoder
model/method
For discrete image latents, the encoder is a GFlowNet that constructs the latent representation one categorical entry at a time, conditioned on the image and entries already sampled. This autoregressive policy avoids imposing the conditional-independence assumptions of a factorized encoder while keeping the encoder more tractable than an unrestricted joint distribution. The E-step trains the encoder toward the posterior over latents; the M-step trains the decoder to reconstruct the image from an encoder sample. The approach does not pass approximate gradients through a continuous relaxation. Unlike the VQ-VAE comparison model, GFlowNet-EM can learn a latent prior jointly with the decoder; the experiments use a PixelCNN prior. On static MNIST, the latent grid is $4\times4$ , with dictionary size $K\in\{4,8,10\}$ and embedding dimension $D=1$ . Training alternates 400 E-step and 400 M-step updates and uses sleep-phase samples for exploration. Greedy encoder samples can be used to train the decoder for faster convergence, but bias the M-step objective.

Knowl 10 — GFlowNet-EM lowers static-MNIST test NLL at increased training cost

data/table

The following static-MNIST test NLLs are means and standard deviations over five runs; lower is better. GFlowNet-EM NLL is estimated with 5,000 importance-weighted samples. The comparison uses the same $4\times4$ latent representation for each method and varies the categorical dictionary size $K$ .

Method	$K=4$	$K=8$	$K=10$
VQ-VAE	$86.36 \pm 0.14$	$80.84 \pm 0.39$	$82.96 \pm 0.38$
GFlowNet-EM	$74.18 \pm 0.41$	$70.74 \pm 0.99$	$70.67 \pm 0.72$
GFlowNet-EM + Greedy Decoder Training (GD)	$76.22 \pm 0.58$	$72.03 \pm 0.98$	$72.69 \pm 1.56$
GFlowNet-EM + GD + jointly learned prior	$78.59 \pm 1.48$	$70.84 \pm 1.06$	$71.69 \pm 1.90$

Every listed GFlowNet-EM variant has lower test NLL than VQ-VAE at each codebook size. Greedy decoder training and joint prior learning do not uniformly improve NLL over base GFlowNet-EM, although joint prior learning enables unconditional sampling from a learned prior. The posterior’s expressiveness comes with higher compute cost: the paper reports about three hours for GFlowNet-EM training versus about 15 minutes for VQ-VAE in these experiments.

Coverage note — No substantial contributed material is omitted; detailed EBM-training hyperparameters and example parse visualizations are left out because they are implementation specifics or illustrations of the grammar results.

References

1.Agrawal, A. and Domke, J. Amortized variational inference for simple hierarchical models. Neural Information Processing Systems (NeurIPS), 2021.
2.Baevski, A., Zhou, Y., Mohamed, A., and Auli, M. wav2vec 2.0: A framework for self-supervised learning of speech representations. Neural Information Processing Systems (NeurIPS), 2020.
3.Baker, J. K. Trainable grammars for speech recognition. The Journal of the Acoustical Society of America, 65(S1): S132–S132, 1979.
4.Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
5.Bengio, Y., Leonard, N., and Courville, A. C. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint 1308.3432, 2013.
6.Bengio, Y., Lahlou, S., Deleu, T., Hu, E., Tiwari, M., and Bengio, E. GFlowNet foundations. Journal of Machine Learning Research (JMLR), 2023. To appear.
7.Bishop, C. M. Pattern Recognition and Machine Learning. Springer, 2006.
8.Bornschein, J. and Bengio, Y. Reweighted wake-sleep. International Conference on Learning Representations (ICLR), 2015.
9.Chomsky, N. Aspects of the Theory of Syntax. MIT Press, 1965.
10.Deleu, T., Gois, A., Emezue, C., Rankawat, M., Lacoste-Julien, S., Bauer, S., and Bengio, Y. Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
11.Dempster, A. P., Laird, N. M., and Rubin, D. B. Maximum likelihood from incomplete data via the EM algorithm. Journal of the Royal Statistical Society B, 39(1):1–38, 1977.
12.Deng, L. The MNIST database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141–142, 2012.
13.Dhariwal, P., Jun, H., Payne, C., Kim, J. W., Radford, A., and Sutskever, I. Jukebox: A generative model for music. arXiv preprint 2005.00341, 2020.
14.Dieng, A. B., Ruiz, F. J. R., and Blei, D. M. Topic modeling in embedding spaces. Transactions of the Association for Computational Linguistics, 8:439–453, 2020. doi: 10.1162/tacl_a_00325. URL https://aclanthology.org/2020.tacl-1.29.
15.Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. Computer Vision and Pattern Recognition (CVPR), 2021.
16.Frey, B. and Jojic, N. A comparison of algorithms for inference and learning in probabilistic graphical models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 27(9):1392–1416, 2005.
17.Ganchev, K., Graça, J., Gillenwater, J., and Taskar, B. Posterior regularization for structured latent variable models. Journal of Machine Learning Research (JMLR), 11:2001–2049, aug 2010.
18.Ghahramani, Z. Factorial learning and the EM algorithm. Neural Information Processing Systems (NIPS), 1994.
19.Goyal, A. and Bengio, Y. Inductive biases for deep learning of higher-level cognition. Proceedings of the Royal Society A, 478(2266):20210068, 2022.
20.Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. International Conference on Machine Learning (ICML), 2018.
21.Hewitt, L. B., Le, T. A., and Tenenbaum, J. B. Learning to learn generative programs with memoised wake-sleep. Uncertainty in Artificial Intelligence (UAI), 2020.
22.Hinton, G. E., Dayan, P., Frey, B. J., and Neal, R. M. The “wake-sleep” algorithm for unsupervised neural networks. Science, 268 5214:1158–61, 1995.
23.Jain, M., Bengio, E., Hernandez-Garcia, A., Rector-Brooks, J., Dossou, B. F., Ekbote, C., Fu, J., Zhang, T., Kilgour, M., Zhang, D., Simine, L., Das, P., and Bengio, Y. Biological sequence design with GFlowNets. International Conference on Machine Learning (ICML), 2022a.
24.Jain, M., Raparthy, S. C., Hernandez-Garcia, A., Rector-Brooks, J., Bengio, Y., Miret, S., and Bengio, E. Multi-objective GFlowNets. arXiv preprint 2210.12765, 2022b.
25.Jin, L., Oh, B.-D., and Schuler, W. Character-based PCFG induction for modeling the syntactic acquisition of morphologically rich languages. In Findings of the Association for Computational Linguistics: EMNLP 2021, pp. 4367–4378, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.findings-emnlp.371. URL https://aclanthology.org/2021.findings-emnlp.371.
26.Kim, Y., Dyer, C., and Rush, A. Compound probabilistic context-free grammars for grammar induction. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2369–2385, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-1228. URL https://aclanthology.org/P19-1228.
27.Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. International Conference on Learning Representations (ICLR), 2014.
28.Koller, D. and Friedman, N. Probabilistic graphical models: principles and techniques. MIT press, 2009.
29.Lari, K. and Young, S. The estimation of stochastic context-free grammars using the inside-outside algorithm. Computer Speech and Language, 4(1):35–56, 1990.
30.Le, T. A., Kosiorek, A. R., Siddharth, N., Teh, Y. W., and Wood, F. Revisiting reweighted wake-sleep for models with stochastic control flow. Neural Information Processing Systems (NeurIPS), 2019.
31.Lieck, R. and Rohrmeier, M. Recursive Bayesian networks: Generalising and unifying probabilistic context-free grammars and dynamic Bayesian networks. Neural Information Processing Systems (NeurIPS), 2021.
32.Liu, D., Jain, M., Dossou, B. F. P., Shen, Q., Lahlou, S., Goyal, A., Malkin, N., Emezue, C. C., Zhang, D., Hassen, N., Ji, X., Kawaguchi, K., and Bengio, Y. GFlowOut: Dropout with generative flow networks. arXiv preprint 2210.12928, 2022.
33.Madan, K., Rector-Brooks, J., Korablyov, M., Bengio, E., Jain, M., Nica, A., Bosc, T., Bengio, Y., and Malkin, N. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2023.
34.Malkin, N., Ortiz, A., and Jojic, N. Mining self-similarity: Label super-resolution with epitomic representations. European Conference on Computer Vision (ECCV), 2020.
35.Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Y. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems (NeurIPS), 2022.
36.Malkin, N., Lahlou, S., Deleu, T., Ji, X., Hu, E., Everett, K., Zhang, D., and Bengio, Y. GFlowNets and variational inference. International Conference on Learning Representations (ICLR), 2023.
37.Marcus, M. P., Santorini, B., Marcinkiewicz, M. A., and Taylor, A. Treebank-3. Linguistic Data Consortium, Philadelphia, 14, 1999.
38.Neal, R. M. and Hinton, G. E. A view of the em algorithm that justifies incremental, sparse, and other variants. In Learning in graphical models, pp. 355–368. Springer, 1998.
39.Neath, R. C. et al. On convergence properties of the Monte Carlo EM algorithm. Advances in modern statistical theory and applications: a Festschrift in Honor of Morris L. Eaton, pp. 43–62, 2013.
40.Nishikawa-Toomey, M., Deleu, T., Subramanian, J., Bengio, Y., and Charlin, L. Bayesian learning of causal structure and mechanisms with GFlowNets and variational bayes. arXiv preprint 2211.02763, 2022.
41.Pan, L., Malkin, N., Zhang, D., and Bengio, Y. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023.
42.Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., and Sutskever, I. Zero-shot text-to-image generation. International Conference on Machine Learning (ICML), 2021.
43.Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. International Conference on Machine Learning (ICML), 2014.
44.Rush, A. Torch-struct: Deep structured prediction library. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics: System Demonstrations, pp. 335–342, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.acl-demos.38. URL https://aclanthology.org/2020.acl-demos.38.
45.van den Oord, A., Kalchbrenner, N., and Kavukcuoglu, K. Pixel recurrent neural networks. International Conference on Machine Learning (ICML), 2016.
46.van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. Neural Information Processing Systems (NIPS), 2017.
47.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. Neural Information Processing Systems (NIPS), 2017.
48.Wang, Y., Blei, D. M., and Cunningham, J. P. Posterior collapse and latent variable non-identifiability. Neural Information Processing Systems (NeurIPS), 2021.
49.Yang, S., Zhao, Y., and Tu, K. PCFGs can do better: Inducing probabilistic context-free grammars with many symbols. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1487–1498, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.naacl-main.117. URL https://aclanthology.org/2021.naacl-main.117.
50.Zhang, D., Malkin, N., Liu, Z., Volokhova, A., Courville, A., and Bengio, Y. Generative flow networks for discrete probabilistic modeling. International Conference on Machine Learning (ICML), 2022.
51.Zhang, D., Chen, R. T. Q., Malkin, N., and Bengio, Y. Unifying generative models with GFlowNets and beyond. arXiv preprint 2209.02606v2, 2023a.
52.Zhang, D., Rainone, C., Peschl, M., and Bondesan, R. Robust scheduling with GFlowNets. International Conference on Learning Representations (ICLR), 2023b.
53.Zhao, Y. and Titov, I. An empirical study of compound PCFGs. In Proceedings of the Second Workshop on Domain Adaptation for NLP, pp. 166–171, Kyiv, Ukraine, April 2021. Association for Computational Linguistics. URL https://aclanthology.org/2021.adaptnlp-1.17.
54.Zimmermann, H., Lindsten, F., van de Meent, J.-W., and Naesseth, C. A. A variational perspective on generative flow networks. arXiv preprint 2210.07992, 2022.

Citation

MLA

Hu, E. J., et al. “GFlowNet-EM for Learning Compositional Latent Variable Models”. International Conference on Machine Learning, vol. 202, 2023, pp. 13528–49, https://proceedings.mlr.press/v202/hu23c.html.

APA

Hu, E. J., Malkin, N., Jain, M., Everett, K. E., Graikos, A., & Bengio, Y. (2023). GFlowNet-EM for Learning Compositional Latent Variable Models. International Conference on Machine Learning, 202, 13528–13549. https://proceedings.mlr.press/v202/hu23c.html

Chicago

Hu, E. J., N. Malkin, M. Jain, K. E. Everett, A. Graikos, and Y. Bengio. 2023. “GFlowNet-EM for Learning Compositional Latent Variable Models”. International Conference on Machine Learning 202: 13528–49. https://proceedings.mlr.press/v202/hu23c.html.

Harvard

Hu, E.J. et al. (2023) “GFlowNet-EM for Learning Compositional Latent Variable Models”, International Conference on Machine Learning. PMLR, pp. 13528–13549. Available at: https://proceedings.mlr.press/v202/hu23c.html.

Vancouver

1. Hu EJ, Malkin N, Jain M, Everett KE, Graikos A, Bengio Y (2023) GFlowNet-EM for Learning Compositional Latent Variable Models. In: International Conference on Machine Learning. PMLR, pp 13528–13549

BibTeX

@InProceedings{pmlr-v202-hu23c,
  title = 	 {{GF}low{N}et-{EM} for Learning Compositional Latent Variable Models},
  author =       {Hu, Edward J and Malkin, Nikolay and Jain, Moksh and Everett, Katie E and Graikos, Alexandros and Bengio, Yoshua},
  booktitle = 	 {Proceedings of the 40th International Conference on Machine Learning},
  pages = 	 {13528--13549},
  year = 	 {2023},
  editor = 	 {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan},
  volume = 	 {202},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {23--29 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v202/hu23c/hu23c.pdf},
  url = 	 {https://proceedings.mlr.press/v202/hu23c.html},
  abstract = 	 {Latent variable models (LVMs) with discrete compositional latents are an important but challenging setting due to a combinatorially large number of possible configurations of the latents. A key tradeoff in modeling the posteriors over latents is between expressivity and tractable optimization. For algorithms based on expectation-maximization (EM), the E-step is often intractable without restrictive approximations to the posterior. We propose the use of GFlowNets, algorithms for sampling from an unnormalized density by learning a stochastic policy for sequential construction of samples, for this intractable E-step. By training GFlowNets to sample from the posterior over latents, we take advantage of their strengths as amortized variational inference algorithms for complex distributions over discrete structures. Our approach, GFlowNet-EM, enables the training of expressive LVMs with discrete compositional latents, as shown by experiments on non-context-free grammar induction and on images using discrete variational autoencoders (VAEs) without conditional independence enforced in the encoder.}
}

Metadata:DOI registry

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

GFlowNet-EM for Learning Compositional Latent Variable Models

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — GFlowNets provide an amortized E-step for compositional latent-variable models

Knowl 2 — Loss thresholding adapts the number of E-step updates

Knowl 3 — A trajectory-based sleep phase helps prevent posterior collapse

Knowl 4 — Exploration, local credit assignment, and MCMC stabilize joint training

Knowl 5 — A bottom-up GFlowNet constructs parse trees from ordered forests

Knowl 6 — Parent-conditioned production rules yield a non-context-free grammar

Knowl 7 — GFlowNet-EM matches exact EM on a hierarchical mixture where mean field fails