# Joint Learning of Energy-based Models and their Partition Function

arXiv: https://arxiv.org/abs/2501.18528 · HTML: https://mlsanity.com/paper/2501.18528/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 44 references that are in the map.
- First posted on arXiv 30 Jan 2025.

## Builds on (most cited first, 21 of 21)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [How to Train Your Energy-Based Models](https://mlsanity.com/paper/2101.03288/) (2021)
- [Fast Differentiable Sorting and Ranking](https://mlsanity.com/paper/2002.08871/) (2020)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [Concrete Score Matching: Generalized Score Matching for Discrete Data](https://mlsanity.com/paper/2211.00802/) (2022)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [Learning with Fenchel-Young Losses](https://mlsanity.com/paper/1901.02324/) (2019)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Generalized Energy Based Models](https://mlsanity.com/paper/2003.05033/) (2020)
- [The Elements of Differentiable Programming](https://mlsanity.com/paper/2403.14606/) (2024)
- [Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective](https://mlsanity.com/paper/2302.01425/) (2023)
- [Sparse Continuous Distributions and Fenchel-Young Losses](https://mlsanity.com/paper/2108.01988/) (2021)
- [Efficient Computation of Expectations under Spanning Tree Distributions](https://mlsanity.com/paper/2008.12988/) (2020)
- [A Smoother Way to Train Structured Prediction Models](https://mlsanity.com/paper/1902.03228/) (2019)
- [Learning Energy Networks with Generalized Fenchel-Young Losses](https://mlsanity.com/paper/2205.09589/) (2022)
- [Geometric Losses for Distributional Learning](https://mlsanity.com/paper/1905.06005/) (2019)

## Built on by (most cited first, 3 of 3)

- [Complexity Analysis of Normalizing Constant Estimation: from Jarzynski Equality to Annealed Importance Sampling and beyond](https://mlsanity.com/paper/2502.04575/) (2025)
- [Autoregressive Language Models are Secretly Energy-Based Models: Insights into the Lookahead Capabilities of Next-Token Prediction](https://mlsanity.com/paper/2512.15605/) (2025)
- [Learning Energy-Based Models by Self-normalising the Likelihood](https://mlsanity.com/paper/2503.07021/) (2025)

## Abstract (arXiv)

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.