GFlowNet-EM for learning compositional latent variable models (2302.06576v2)

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.

• The paper introduces GFlowNet-EM, highlighting its novel ability to sample complex discrete latent structures without imposing limiting conditional independence.
• It combines GFlowNets for the E-step with traditional likelihood maximization in the M-step, enhancing posterior inference and EM optimization.
• Empirical results demonstrate superior performance in grammar induction, hierarchical mixture models, and discrete VAEs, improving model expressivity and reducing NLL.

GFlowNet-EM for Learning Compositional Latent Variable Models

The paper "GFlowNet-EM for Learning Compositional Latent Variable Models" by Edward J. Hu et al. introduces a new framework that leverages Generative Flow Networks (GFlowNets) for the expectation-maximization (EM) of latent variable models (LVMs). The research addresses the difficulty of representing and optimizing posteriors over discrete compositional latents, which traditionally present challenges due to their combinatorial complexity.

Overview

Latent variable models with compositional discrete latent structures, such as those found in language and visual domains, require a balance between model expressivity and tractable optimization. Traditional EM methods often default to simplifying assumptions or restrictive approximations for computational feasibility during the E-step of the EM algorithm. These limitations hinder the expressiveness of the LVMs, especially in domains requiring the modeling of complex hierarchical structures like trees or graphs.

GFlowNets, as proposed in the paper, provide an innovative approach by learning to sample from the intractable posterior distribution over discrete latent structures. This is achieved by training a stochastic policy to construct samples sequentially. Through their approach termed "GFlowNet-EM," the authors present a method for parameterizing and optimizing LVMs without imposing conditional independence, as exemplified by successful applications in grammar induction and image representation tasks.

Methodological Contributions

1. GFlowNet-EM Framework: This framework combines the strength of GFlowNets for the E-step in EM with traditional likelihood maximization techniques for the M-step, enabling more expressive posterior inference for LVMs.
2. Sampling Complex Structures: By leveraging GFlowNets, which can efficiently sample combinatorial objects like trees or graphs, the framework supports amortized variational inference over structured latent spaces that are challenging for conventional approximate inference methods.
3. Adaptive Optimization Techniques: The authors discuss enhancements like adaptive E-step strategies, exploration through tempered policy sampling, and a sleep phase inspired by the wake-sleep algorithm to maintain high posterior true mode coverage and counteract posterior collapse.

Empirical Results and Implications

The empirical section of the paper validates the proposed method across multiple domains:

• Hierarchical Mixture Models: The GFlowNet-EM solution matches exact EM's performance in model likelihood while outperforming factorized variational EM, emphasizing its efficacy in managing interactions between latent variables.
• Grammar Induction: On the Penn Tree Bank dataset, GFlowNet-EM learns grammars with comparable likelihoods and improved linguistic fidelity compared to baselines. This is particularly notable when integrating an energy-based model (EBM) as a prior, capturing tree shapes akin to those imposed by human grammar rules.
• Discrete VAEs: GFlowNet-EM demonstrates a superior ability to model discrete latent representations with enhanced expressivity and reduced NLL on benchmark datasets, outperforming traditional VQ-VAE models.

Future Perspectives

This research opens avenues for further exploration of GFlowNets in more diverse LVM scenarios, including continuous latents or hybrid latent spaces. Moreover, the ability to incorporate domain-specific priors through EBMs suggests potential extensions into unsupervised learning and reinforcement learning applications. Future work might also consider the expansion of GFlowNet-EM's applications in real-time systems and its integration with other machine learning paradigms that require compositional reasoning.

In conclusion, GFlowNet-EM provides a robust framework for enhancing the expressiveness of latent variable models with discrete compositional structures, contributing to both theoretical advancements and practical applications in machine learning.