GFlowNet-EM for learning compositional latent variable models

TL;DR

This paper introduces GFlowNet-EM, a novel method that leverages GFlowNets to perform efficient posterior sampling in discrete latent variable models, enabling more expressive modeling without restrictive assumptions.

Contribution

The paper proposes GFlowNet-EM, a new approach that uses GFlowNets for the intractable E-step in EM, allowing training of complex discrete latent variable models.

Findings

Effective sampling from complex posteriors demonstrated in grammar induction.

Successful application to discrete VAEs on image data.

Enables training of more expressive discrete latent 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.