A Contrastive Objective for Training Continuous Generative Flow Networks

• Tiago da Silva FGV
• Diego Mesquita FGV

Resumo

Generative Flow Networks (GFlowNets) are a novel class of flexible amortized samplers for distributions supported on complex objects (e.g., graphs and sequences), achieving significant success in problems such as combinatorial optimization, drug discovery and natural language processing. Nonetheless, training of GFlowNets is challenging—partly because it relies on estimating high-dimensional integrals, including the log-partition function, via stochastic gradient descent (SGD). In particular, for distributions supported on non-discrete spaces, which have received far less attention from the recent literature, every previously proposed learning objective either depends on estimating a log-partition function or is restricted to on-policy training, which is susceptible to mode-collapse. In this context, inspired by the success of contrastive learning for variational inference, we propose the continuous contrastive loss (CCL) as the first objective function natively enabling off-policy training of continuous GFlowNets without reliance on the approximation of high-dimensional integrals via SGD, extending previous work based on discrete distributions. Additionally, we show that minimizing the CCL objective is empirically effective and often leads to faster training convergence than alternatives.