Mitigating Out-of-Distribution Data Density Overestimation in Energy-Based Models

TL;DR

This paper identifies why short-run Langevin Monte Carlo causes energy-based models to overestimate out-of-distribution data density and proposes a novel USP scheme to improve OOD detection.

Contribution

It reveals the issues with heuristic LMC modifications and introduces USP, a new method to better estimate densities in EBMs for improved OOD detection.

Findings

USP significantly improves OOD detection on Fashion-MNIST.

Heuristic modifications to LMC are the main cause of density overestimation.

USP provides a more accurate support approximation for EBMs.

Abstract

Deep energy-based models (EBMs), which use deep neural networks (DNNs) as energy functions, are receiving increasing attention due to their ability to learn complex distributions. To train deep EBMs, the maximum likelihood estimation (MLE) with short-run Langevin Monte Carlo (LMC) is often used. While the MLE with short-run LMC is computationally efficient compared to an MLE with full Markov Chain Monte Carlo (MCMC), it often assigns high density to out-of-distribution (OOD) data. To address this issue, here we systematically investigate why the MLE with short-run LMC can converge to EBMs with wrong density estimates, and reveal that the heuristic modifications to LMC introduced by previous works were the main problem. We then propose a Uniform Support Partitioning (USP) scheme that optimizes a set of points to evenly partition the support of the EBM and then uses the resulting points to approximate the EBM-MLE loss gradient. We empirically demonstrate that USP avoids the pitfalls of short-run LMC, leading to significantly improved OOD data detection performance on Fashion-MNIST.