Joint Learning of Energy-based Models and their Partition Function
This addresses a fundamental bottleneck in probabilistic modeling for machine learning, offering a novel approach to approximate EBM learning in combinatorially-large spaces, though it is incremental in advancing existing EBM frameworks.
The paper tackles the intractability of learning energy-based models (EBMs) in large discrete spaces by proposing a method to jointly learn an energy model and its log-partition function using neural networks, enabling tractable optimization without MCMC and achieving competitive results in multilabel classification and label ranking tasks.
Energy-based models (EBMs) offer a flexible framework for parameterizing probability distributions using neural networks. However, learning EBMs by exact maximum likelihood estimation (MLE) is generally intractable, due to the need to compute the partition function (normalization constant). In this paper, we propose a novel formulation for approximately learning probabilistic EBMs in combinatorially-large discrete spaces, such as sets or permutations. Our key idea is to jointly learn both an energy model and its log-partition, both parameterized as a neural network. Our approach not only provides a novel tractable objective criterion to learn EBMs by stochastic gradient descent (without relying on MCMC), but also a novel means to estimate the log-partition function on unseen data points. On the theoretical side, we show that our approach recovers the optimal MLE solution when optimizing in the space of continuous functions. Furthermore, we show that our approach naturally extends to the broader family of Fenchel-Young losses, allowing us to obtain the first tractable method for optimizing the sparsemax loss in combinatorially-large spaces. We demonstrate our approach on multilabel classification and label ranking.